mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 22:21:22 -07:00
feat(hooks): Hook Execution Engine (#9092)
This commit is contained in:
729
packages/core/src/hooks/hookRunner.test.ts
Normal file
729
packages/core/src/hooks/hookRunner.test.ts
Normal file
@@ -0,0 +1,729 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { spawn, type ChildProcessWithoutNullStreams } from 'node:child_process';
|
||||
import { HookRunner } from './hookRunner.js';
|
||||
import { HookEventName, HookType } from './types.js';
|
||||
import type { HookConfig } from './types.js';
|
||||
import type { HookInput } from './types.js';
|
||||
import type { Readable, Writable } from 'node:stream';
|
||||
|
||||
// Mock type for the child_process spawn
|
||||
type MockChildProcessWithoutNullStreams = ChildProcessWithoutNullStreams & {
|
||||
mockStdoutOn: ReturnType<typeof vi.fn>;
|
||||
mockStderrOn: ReturnType<typeof vi.fn>;
|
||||
mockProcessOn: ReturnType<typeof vi.fn>;
|
||||
};
|
||||
|
||||
// Mock child_process with importOriginal for partial mocking
|
||||
vi.mock('node:child_process', async (importOriginal) => {
|
||||
const actual = (await importOriginal()) as object;
|
||||
return {
|
||||
...actual,
|
||||
spawn: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
// Mock console methods
|
||||
const mockConsole = {
|
||||
log: vi.fn(),
|
||||
warn: vi.fn(),
|
||||
error: vi.fn(),
|
||||
debug: vi.fn(),
|
||||
};
|
||||
|
||||
vi.stubGlobal('console', mockConsole);
|
||||
|
||||
describe('HookRunner', () => {
|
||||
let hookRunner: HookRunner;
|
||||
let mockSpawn: MockChildProcessWithoutNullStreams;
|
||||
|
||||
const mockInput: HookInput = {
|
||||
session_id: 'test-session',
|
||||
transcript_path: '/path/to/transcript',
|
||||
cwd: '/test/project',
|
||||
hook_event_name: 'BeforeTool',
|
||||
timestamp: '2025-01-01T00:00:00.000Z',
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
|
||||
hookRunner = new HookRunner();
|
||||
|
||||
// Mock spawn with accessible mock functions
|
||||
const mockStdoutOn = vi.fn();
|
||||
const mockStderrOn = vi.fn();
|
||||
const mockProcessOn = vi.fn();
|
||||
|
||||
mockSpawn = {
|
||||
stdin: {
|
||||
write: vi.fn(),
|
||||
end: vi.fn(),
|
||||
} as unknown as Writable,
|
||||
stdout: {
|
||||
on: mockStdoutOn,
|
||||
} as unknown as Readable,
|
||||
stderr: {
|
||||
on: mockStderrOn,
|
||||
} as unknown as Readable,
|
||||
on: mockProcessOn,
|
||||
kill: vi.fn(),
|
||||
killed: false,
|
||||
mockStdoutOn,
|
||||
mockStderrOn,
|
||||
mockProcessOn,
|
||||
} as unknown as MockChildProcessWithoutNullStreams;
|
||||
|
||||
vi.mocked(spawn).mockReturnValue(mockSpawn);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe('executeHook', () => {
|
||||
describe('command hooks', () => {
|
||||
const commandConfig: HookConfig = {
|
||||
type: HookType.Command,
|
||||
command: './hooks/test.sh',
|
||||
timeout: 5000,
|
||||
};
|
||||
|
||||
it('should execute command hook successfully', async () => {
|
||||
const mockOutput = { decision: 'allow', reason: 'All good' };
|
||||
|
||||
// Mock successful execution
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(
|
||||
() => callback(Buffer.from(JSON.stringify(mockOutput))),
|
||||
10,
|
||||
);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.output).toEqual(mockOutput);
|
||||
expect(result.exitCode).toBe(0);
|
||||
expect(mockSpawn.stdin.write).toHaveBeenCalledWith(
|
||||
JSON.stringify(mockInput),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle command hook failure', async () => {
|
||||
const errorMessage = 'Command failed';
|
||||
|
||||
mockSpawn.mockStderrOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(errorMessage)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(1), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(false);
|
||||
expect(result.exitCode).toBe(1);
|
||||
expect(result.stderr).toBe(errorMessage);
|
||||
});
|
||||
|
||||
it('should handle command hook timeout', async () => {
|
||||
const shortTimeoutConfig: HookConfig = {
|
||||
type: HookType.Command,
|
||||
command: './hooks/slow.sh',
|
||||
timeout: 50, // Very short timeout for testing
|
||||
};
|
||||
|
||||
let closeCallback: ((code: number) => void) | undefined;
|
||||
let killWasCalled = false;
|
||||
|
||||
// Mock a hanging process that registers the close handler but doesn't call it initially
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
closeCallback = callback; // Store the callback but don't call it yet
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
// Mock the kill method to simulate the process being killed
|
||||
mockSpawn.kill = vi.fn().mockImplementation((_signal: string) => {
|
||||
killWasCalled = true;
|
||||
// Simulate that killing the process triggers the close event
|
||||
if (closeCallback) {
|
||||
setTimeout(() => {
|
||||
closeCallback!(128); // Exit code 128 indicates process was killed by signal
|
||||
}, 5);
|
||||
}
|
||||
return true;
|
||||
});
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
shortTimeoutConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(false);
|
||||
expect(killWasCalled).toBe(true);
|
||||
expect(result.error?.message).toContain('timed out');
|
||||
expect(mockSpawn.kill).toHaveBeenCalledWith('SIGTERM');
|
||||
});
|
||||
|
||||
it('should expand environment variables in commands', async () => {
|
||||
const configWithEnvVar: HookConfig = {
|
||||
type: HookType.Command,
|
||||
command: '$GEMINI_PROJECT_DIR/hooks/test.sh',
|
||||
};
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
await hookRunner.executeHook(
|
||||
configWithEnvVar,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(spawn).toHaveBeenCalledWith(
|
||||
'/test/project/hooks/test.sh',
|
||||
expect.objectContaining({
|
||||
shell: true,
|
||||
env: expect.objectContaining({
|
||||
GEMINI_PROJECT_DIR: '/test/project',
|
||||
CLAUDE_PROJECT_DIR: '/test/project',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('executeHooksParallel', () => {
|
||||
it('should execute multiple hooks in parallel', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
// Mock both commands to succeed
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksParallel(
|
||||
configs,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results.every((r) => r.success)).toBe(true);
|
||||
expect(spawn).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should handle mixed success and failure', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
let callCount = 0;
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
const exitCode = callCount++ === 0 ? 0 : 1; // First succeeds, second fails
|
||||
setTimeout(() => callback(exitCode), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksParallel(
|
||||
configs,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].success).toBe(true);
|
||||
expect(results[1].success).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('executeHooksSequential', () => {
|
||||
it('should execute multiple hooks in sequence', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
const executionOrder: string[] = [];
|
||||
|
||||
// Mock both commands to succeed
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
const command =
|
||||
vi.mocked(spawn).mock.calls[executionOrder.length][0];
|
||||
executionOrder.push(command);
|
||||
setTimeout(() => callback(0), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksSequential(
|
||||
configs,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results.every((r) => r.success)).toBe(true);
|
||||
expect(spawn).toHaveBeenCalledTimes(2);
|
||||
// Verify they were called sequentially
|
||||
expect(executionOrder).toEqual(['./hook1.sh', './hook2.sh']);
|
||||
});
|
||||
|
||||
it('should continue execution even if a hook fails', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
{ type: HookType.Command, command: './hook3.sh' },
|
||||
];
|
||||
|
||||
let callCount = 0;
|
||||
mockSpawn.mockStderrOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data' && callCount === 1) {
|
||||
// Second hook fails
|
||||
setTimeout(() => callback(Buffer.from('Hook 2 failed')), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
const exitCode = callCount++ === 1 ? 1 : 0; // Second fails, others succeed
|
||||
setTimeout(() => callback(exitCode), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksSequential(
|
||||
configs,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(3);
|
||||
expect(results[0].success).toBe(true);
|
||||
expect(results[1].success).toBe(false);
|
||||
expect(results[2].success).toBe(true);
|
||||
expect(spawn).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('should pass modified input from one hook to the next for BeforeAgent', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
const mockBeforeAgentInput = {
|
||||
...mockInput,
|
||||
prompt: 'Original prompt',
|
||||
};
|
||||
|
||||
const mockOutput1 = {
|
||||
decision: 'allow' as const,
|
||||
hookSpecificOutput: {
|
||||
additionalContext: 'Context from hook 1',
|
||||
},
|
||||
};
|
||||
|
||||
let hookCallCount = 0;
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
if (hookCallCount === 0) {
|
||||
setTimeout(
|
||||
() => callback(Buffer.from(JSON.stringify(mockOutput1))),
|
||||
10,
|
||||
);
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
hookCallCount++;
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksSequential(
|
||||
configs,
|
||||
HookEventName.BeforeAgent,
|
||||
mockBeforeAgentInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].success).toBe(true);
|
||||
expect(results[0].output).toEqual(mockOutput1);
|
||||
|
||||
// Verify that the second hook received modified input
|
||||
const secondHookInput = JSON.parse(
|
||||
vi.mocked(mockSpawn.stdin.write).mock.calls[1][0],
|
||||
);
|
||||
expect(secondHookInput.prompt).toContain('Original prompt');
|
||||
expect(secondHookInput.prompt).toContain('Context from hook 1');
|
||||
});
|
||||
|
||||
it('should pass modified LLM request from one hook to the next for BeforeModel', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
const mockBeforeModelInput = {
|
||||
...mockInput,
|
||||
llm_request: {
|
||||
model: 'gemini-1.5-pro',
|
||||
messages: [{ role: 'user', content: 'Hello' }],
|
||||
},
|
||||
};
|
||||
|
||||
const mockOutput1 = {
|
||||
decision: 'allow' as const,
|
||||
hookSpecificOutput: {
|
||||
llm_request: {
|
||||
temperature: 0.7,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
let hookCallCount = 0;
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
if (hookCallCount === 0) {
|
||||
setTimeout(
|
||||
() => callback(Buffer.from(JSON.stringify(mockOutput1))),
|
||||
10,
|
||||
);
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
hookCallCount++;
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksSequential(
|
||||
configs,
|
||||
HookEventName.BeforeModel,
|
||||
mockBeforeModelInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].success).toBe(true);
|
||||
|
||||
// Verify that the second hook received modified input
|
||||
const secondHookInput = JSON.parse(
|
||||
vi.mocked(mockSpawn.stdin.write).mock.calls[1][0],
|
||||
);
|
||||
expect(secondHookInput.llm_request.model).toBe('gemini-1.5-pro');
|
||||
expect(secondHookInput.llm_request.temperature).toBe(0.7);
|
||||
});
|
||||
|
||||
it('should not modify input if hook fails', async () => {
|
||||
const configs: HookConfig[] = [
|
||||
{ type: HookType.Command, command: './hook1.sh' },
|
||||
{ type: HookType.Command, command: './hook2.sh' },
|
||||
];
|
||||
|
||||
mockSpawn.mockStderrOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from('Hook failed')), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(1), 20); // All hooks fail
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const results = await hookRunner.executeHooksSequential(
|
||||
configs,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results.every((r) => !r.success)).toBe(true);
|
||||
|
||||
// Verify that both hooks received the same original input
|
||||
const firstHookInput = JSON.parse(
|
||||
vi.mocked(mockSpawn.stdin.write).mock.calls[0][0],
|
||||
);
|
||||
const secondHookInput = JSON.parse(
|
||||
vi.mocked(mockSpawn.stdin.write).mock.calls[1][0],
|
||||
);
|
||||
expect(firstHookInput).toEqual(secondHookInput);
|
||||
});
|
||||
});
|
||||
|
||||
describe('invalid JSON handling', () => {
|
||||
const commandConfig: HookConfig = {
|
||||
type: HookType.Command,
|
||||
command: './hooks/test.sh',
|
||||
};
|
||||
|
||||
it('should handle invalid JSON output gracefully', async () => {
|
||||
const invalidJson = '{ "decision": "allow", incomplete';
|
||||
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(invalidJson)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.exitCode).toBe(0);
|
||||
// Should convert plain text to structured output
|
||||
expect(result.output).toEqual({
|
||||
decision: 'allow',
|
||||
systemMessage: invalidJson,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle malformed JSON with exit code 0', async () => {
|
||||
const malformedJson = 'not json at all';
|
||||
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(malformedJson)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.output).toEqual({
|
||||
decision: 'allow',
|
||||
systemMessage: malformedJson,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle invalid JSON with exit code 1 (non-blocking error)', async () => {
|
||||
const invalidJson = '{ broken json';
|
||||
|
||||
mockSpawn.mockStderrOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(invalidJson)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(1), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(false);
|
||||
expect(result.exitCode).toBe(1);
|
||||
expect(result.output).toEqual({
|
||||
decision: 'allow',
|
||||
systemMessage: `Warning: ${invalidJson}`,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle invalid JSON with exit code 2 (blocking error)', async () => {
|
||||
const invalidJson = '{ "error": incomplete';
|
||||
|
||||
mockSpawn.mockStderrOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(invalidJson)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(2), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(false);
|
||||
expect(result.exitCode).toBe(2);
|
||||
expect(result.output).toEqual({
|
||||
decision: 'deny',
|
||||
reason: invalidJson,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle empty JSON output', async () => {
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from('')), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.exitCode).toBe(0);
|
||||
expect(result.output).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle double-encoded JSON string', async () => {
|
||||
const mockOutput = { decision: 'allow', reason: 'All good' };
|
||||
const doubleEncodedJson = JSON.stringify(JSON.stringify(mockOutput));
|
||||
|
||||
mockSpawn.mockStdoutOn.mockImplementation(
|
||||
(event: string, callback: (data: Buffer) => void) => {
|
||||
if (event === 'data') {
|
||||
setTimeout(() => callback(Buffer.from(doubleEncodedJson)), 10);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
mockSpawn.mockProcessOn.mockImplementation(
|
||||
(event: string, callback: (code: number) => void) => {
|
||||
if (event === 'close') {
|
||||
setTimeout(() => callback(0), 20);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await hookRunner.executeHook(
|
||||
commandConfig,
|
||||
HookEventName.BeforeTool,
|
||||
mockInput,
|
||||
);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.output).toEqual(mockOutput);
|
||||
});
|
||||
});
|
||||
});
|
||||
358
packages/core/src/hooks/hookRunner.ts
Normal file
358
packages/core/src/hooks/hookRunner.ts
Normal file
@@ -0,0 +1,358 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { spawn } from 'node:child_process';
|
||||
import type { HookConfig } from './types.js';
|
||||
import { HookEventName } from './types.js';
|
||||
import type {
|
||||
HookInput,
|
||||
HookOutput,
|
||||
HookExecutionResult,
|
||||
BeforeAgentInput,
|
||||
BeforeModelInput,
|
||||
BeforeModelOutput,
|
||||
} from './types.js';
|
||||
import type { LLMRequest } from './hookTranslator.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
|
||||
/**
|
||||
* Default timeout for hook execution (60 seconds)
|
||||
*/
|
||||
const DEFAULT_HOOK_TIMEOUT = 60000;
|
||||
|
||||
/**
|
||||
* Exit code constants for hook execution
|
||||
*/
|
||||
const EXIT_CODE_SUCCESS = 0;
|
||||
const EXIT_CODE_BLOCKING_ERROR = 2;
|
||||
const EXIT_CODE_NON_BLOCKING_ERROR = 1;
|
||||
|
||||
/**
|
||||
* Hook runner that executes command hooks
|
||||
*/
|
||||
export class HookRunner {
|
||||
constructor() {}
|
||||
|
||||
/**
|
||||
* Execute a single hook
|
||||
*/
|
||||
async executeHook(
|
||||
hookConfig: HookConfig,
|
||||
eventName: HookEventName,
|
||||
input: HookInput,
|
||||
): Promise<HookExecutionResult> {
|
||||
const startTime = Date.now();
|
||||
|
||||
try {
|
||||
return await this.executeCommandHook(
|
||||
hookConfig,
|
||||
eventName,
|
||||
input,
|
||||
startTime,
|
||||
);
|
||||
} catch (error) {
|
||||
const duration = Date.now() - startTime;
|
||||
const hookSource = hookConfig.command || 'unknown';
|
||||
const errorMessage = `Hook execution failed for event '${eventName}' (source: ${hookSource}): ${error}`;
|
||||
debugLogger.warn(`Hook execution error (non-fatal): ${errorMessage}`);
|
||||
|
||||
return {
|
||||
hookConfig,
|
||||
eventName,
|
||||
success: false,
|
||||
error: error instanceof Error ? error : new Error(errorMessage),
|
||||
duration,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute multiple hooks in parallel
|
||||
*/
|
||||
async executeHooksParallel(
|
||||
hookConfigs: HookConfig[],
|
||||
eventName: HookEventName,
|
||||
input: HookInput,
|
||||
): Promise<HookExecutionResult[]> {
|
||||
const promises = hookConfigs.map((config) =>
|
||||
this.executeHook(config, eventName, input),
|
||||
);
|
||||
|
||||
return await Promise.all(promises);
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute multiple hooks sequentially
|
||||
*/
|
||||
async executeHooksSequential(
|
||||
hookConfigs: HookConfig[],
|
||||
eventName: HookEventName,
|
||||
input: HookInput,
|
||||
): Promise<HookExecutionResult[]> {
|
||||
const results: HookExecutionResult[] = [];
|
||||
let currentInput = input;
|
||||
|
||||
for (const config of hookConfigs) {
|
||||
const result = await this.executeHook(config, eventName, currentInput);
|
||||
results.push(result);
|
||||
|
||||
// If the hook succeeded and has output, use it to modify the input for the next hook
|
||||
if (result.success && result.output) {
|
||||
currentInput = this.applyHookOutputToInput(
|
||||
currentInput,
|
||||
result.output,
|
||||
eventName,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply hook output to modify input for the next hook in sequential execution
|
||||
*/
|
||||
private applyHookOutputToInput(
|
||||
originalInput: HookInput,
|
||||
hookOutput: HookOutput,
|
||||
eventName: HookEventName,
|
||||
): HookInput {
|
||||
// Create a copy of the original input
|
||||
const modifiedInput = { ...originalInput };
|
||||
|
||||
// Apply modifications based on hook output and event type
|
||||
if (hookOutput.hookSpecificOutput) {
|
||||
switch (eventName) {
|
||||
case HookEventName.BeforeAgent:
|
||||
if ('additionalContext' in hookOutput.hookSpecificOutput) {
|
||||
// For BeforeAgent, we could modify the prompt with additional context
|
||||
const additionalContext =
|
||||
hookOutput.hookSpecificOutput['additionalContext'];
|
||||
if (
|
||||
typeof additionalContext === 'string' &&
|
||||
'prompt' in modifiedInput
|
||||
) {
|
||||
(modifiedInput as BeforeAgentInput).prompt +=
|
||||
'\n\n' + additionalContext;
|
||||
}
|
||||
}
|
||||
break;
|
||||
|
||||
case HookEventName.BeforeModel:
|
||||
if ('llm_request' in hookOutput.hookSpecificOutput) {
|
||||
// For BeforeModel, we update the LLM request
|
||||
const hookBeforeModelOutput = hookOutput as BeforeModelOutput;
|
||||
if (
|
||||
hookBeforeModelOutput.hookSpecificOutput?.llm_request &&
|
||||
'llm_request' in modifiedInput
|
||||
) {
|
||||
// Merge the partial request with the existing request
|
||||
const currentRequest = (modifiedInput as BeforeModelInput)
|
||||
.llm_request;
|
||||
const partialRequest =
|
||||
hookBeforeModelOutput.hookSpecificOutput.llm_request;
|
||||
(modifiedInput as BeforeModelInput).llm_request = {
|
||||
...currentRequest,
|
||||
...partialRequest,
|
||||
} as LLMRequest;
|
||||
}
|
||||
}
|
||||
break;
|
||||
|
||||
default:
|
||||
// For other events, no special input modification is needed
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return modifiedInput;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a command hook
|
||||
*/
|
||||
private async executeCommandHook(
|
||||
hookConfig: HookConfig,
|
||||
eventName: HookEventName,
|
||||
input: HookInput,
|
||||
startTime: number,
|
||||
): Promise<HookExecutionResult> {
|
||||
const timeout = hookConfig.timeout ?? DEFAULT_HOOK_TIMEOUT;
|
||||
|
||||
return new Promise((resolve) => {
|
||||
if (!hookConfig.command) {
|
||||
const errorMessage = 'Command hook missing command';
|
||||
debugLogger.warn(
|
||||
`Hook configuration error (non-fatal): ${errorMessage}`,
|
||||
);
|
||||
resolve({
|
||||
hookConfig,
|
||||
eventName,
|
||||
success: false,
|
||||
error: new Error(errorMessage),
|
||||
duration: Date.now() - startTime,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
let stdout = '';
|
||||
let stderr = '';
|
||||
let timedOut = false;
|
||||
const command = this.expandCommand(hookConfig.command, input);
|
||||
|
||||
// Set up environment variables
|
||||
const env = {
|
||||
...process.env,
|
||||
GEMINI_PROJECT_DIR: input.cwd,
|
||||
CLAUDE_PROJECT_DIR: input.cwd, // For compatibility
|
||||
};
|
||||
|
||||
const child = spawn(command, {
|
||||
env,
|
||||
cwd: input.cwd,
|
||||
stdio: ['pipe', 'pipe', 'pipe'],
|
||||
shell: true,
|
||||
});
|
||||
|
||||
// Set up timeout
|
||||
const timeoutHandle = setTimeout(() => {
|
||||
timedOut = true;
|
||||
child.kill('SIGTERM');
|
||||
|
||||
// Force kill after 5 seconds
|
||||
setTimeout(() => {
|
||||
if (!child.killed) {
|
||||
child.kill('SIGKILL');
|
||||
}
|
||||
}, 5000);
|
||||
}, timeout);
|
||||
|
||||
// Send input to stdin
|
||||
if (child.stdin) {
|
||||
child.stdin.write(JSON.stringify(input));
|
||||
child.stdin.end();
|
||||
}
|
||||
|
||||
// Collect stdout
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
stdout += data.toString();
|
||||
});
|
||||
|
||||
// Collect stderr
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
stderr += data.toString();
|
||||
});
|
||||
|
||||
// Handle process exit
|
||||
child.on('close', (exitCode) => {
|
||||
clearTimeout(timeoutHandle);
|
||||
const duration = Date.now() - startTime;
|
||||
|
||||
if (timedOut) {
|
||||
resolve({
|
||||
hookConfig,
|
||||
eventName,
|
||||
success: false,
|
||||
error: new Error(`Hook timed out after ${timeout}ms`),
|
||||
stdout,
|
||||
stderr,
|
||||
duration,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// Parse output
|
||||
let output: HookOutput | undefined;
|
||||
if (exitCode === EXIT_CODE_SUCCESS && stdout.trim()) {
|
||||
try {
|
||||
let parsed = JSON.parse(stdout.trim());
|
||||
if (typeof parsed === 'string') {
|
||||
// If the output is a string, parse it in case
|
||||
// it's double-encoded JSON string.
|
||||
parsed = JSON.parse(parsed);
|
||||
}
|
||||
if (parsed) {
|
||||
output = parsed as HookOutput;
|
||||
}
|
||||
} catch {
|
||||
// Not JSON, convert plain text to structured output
|
||||
output = this.convertPlainTextToHookOutput(stdout.trim(), exitCode);
|
||||
}
|
||||
} else if (exitCode !== EXIT_CODE_SUCCESS && stderr.trim()) {
|
||||
// Convert error output to structured format
|
||||
output = this.convertPlainTextToHookOutput(
|
||||
stderr.trim(),
|
||||
exitCode || EXIT_CODE_NON_BLOCKING_ERROR,
|
||||
);
|
||||
}
|
||||
|
||||
resolve({
|
||||
hookConfig,
|
||||
eventName,
|
||||
success: exitCode === EXIT_CODE_SUCCESS,
|
||||
output,
|
||||
stdout,
|
||||
stderr,
|
||||
exitCode: exitCode || EXIT_CODE_SUCCESS,
|
||||
duration,
|
||||
});
|
||||
});
|
||||
|
||||
// Handle process errors
|
||||
child.on('error', (error) => {
|
||||
clearTimeout(timeoutHandle);
|
||||
const duration = Date.now() - startTime;
|
||||
|
||||
resolve({
|
||||
hookConfig,
|
||||
eventName,
|
||||
success: false,
|
||||
error,
|
||||
stdout,
|
||||
stderr,
|
||||
duration,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Expand command with environment variables and input context
|
||||
*/
|
||||
private expandCommand(command: string, input: HookInput): string {
|
||||
return command
|
||||
.replace(/\$GEMINI_PROJECT_DIR/g, input.cwd)
|
||||
.replace(/\$CLAUDE_PROJECT_DIR/g, input.cwd); // For compatibility
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert plain text output to structured HookOutput
|
||||
*/
|
||||
private convertPlainTextToHookOutput(
|
||||
text: string,
|
||||
exitCode: number,
|
||||
): HookOutput {
|
||||
if (exitCode === EXIT_CODE_SUCCESS) {
|
||||
// Success - treat as system message or additional context
|
||||
return {
|
||||
decision: 'allow',
|
||||
systemMessage: text,
|
||||
};
|
||||
} else if (exitCode === EXIT_CODE_BLOCKING_ERROR) {
|
||||
// Blocking error
|
||||
return {
|
||||
decision: 'deny',
|
||||
reason: text,
|
||||
};
|
||||
} else {
|
||||
// Non-blocking error (EXIT_CODE_NON_BLOCKING_ERROR or any other code)
|
||||
return {
|
||||
decision: 'allow',
|
||||
systemMessage: `Warning: ${text}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user