chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,338 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { TestRig, normalizePath } from './test-helper.js';
|
||||
import { join } from 'node:path';
|
||||
import { writeFileSync } from 'node:fs';
|
||||
|
||||
describe('Hooks Agent Flow', () => {
|
||||
let rig: TestRig;
|
||||
|
||||
beforeEach(() => {
|
||||
rig = new TestRig();
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
if (rig) {
|
||||
await rig.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
describe('BeforeAgent Hooks', () => {
|
||||
it('should inject additional context via BeforeAgent hook', async () => {
|
||||
await rig.setup('should inject additional context via BeforeAgent hook', {
|
||||
fakeResponsesPath: join(
|
||||
import.meta.dirname,
|
||||
'hooks-agent-flow.responses',
|
||||
),
|
||||
});
|
||||
|
||||
const hookScript = `
|
||||
try {
|
||||
const output = {
|
||||
decision: "allow",
|
||||
hookSpecificOutput: {
|
||||
hookEventName: "BeforeAgent",
|
||||
additionalContext: "SYSTEM INSTRUCTION: This is injected context."
|
||||
}
|
||||
};
|
||||
process.stdout.write(JSON.stringify(output));
|
||||
} catch (e) {
|
||||
console.error('Failed to write stdout:', e);
|
||||
process.exit(1);
|
||||
}
|
||||
console.error('DEBUG: BeforeAgent hook executed');
|
||||
`;
|
||||
|
||||
const scriptPath = join(rig.testDir!, 'before_agent_context.cjs');
|
||||
writeFileSync(scriptPath, hookScript);
|
||||
|
||||
await rig.setup('should inject additional context via BeforeAgent hook', {
|
||||
settings: {
|
||||
hooksConfig: {
|
||||
enabled: true,
|
||||
},
|
||||
hooks: {
|
||||
BeforeAgent: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: `node "${scriptPath}"`,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await rig.run({ args: 'Hello test' });
|
||||
|
||||
// Verify hook execution and telemetry
|
||||
const hookTelemetryFound = await rig.waitForTelemetryEvent('hook_call');
|
||||
expect(hookTelemetryFound).toBeTruthy();
|
||||
|
||||
const hookLogs = rig.readHookLogs();
|
||||
const beforeAgentLog = hookLogs.find(
|
||||
(log) => log.hookCall.hook_event_name === 'BeforeAgent',
|
||||
);
|
||||
|
||||
expect(beforeAgentLog).toBeDefined();
|
||||
expect(beforeAgentLog?.hookCall.stdout).toContain('injected context');
|
||||
expect(beforeAgentLog?.hookCall.stdout).toContain('"decision":"allow"');
|
||||
expect(beforeAgentLog?.hookCall.stdout).toContain(
|
||||
'SYSTEM INSTRUCTION: This is injected context.',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('AfterAgent Hooks', () => {
|
||||
it('should receive prompt and response in AfterAgent hook', async () => {
|
||||
await rig.setup('should receive prompt and response in AfterAgent hook', {
|
||||
fakeResponsesPath: join(
|
||||
import.meta.dirname,
|
||||
'hooks-agent-flow.responses',
|
||||
),
|
||||
});
|
||||
|
||||
const hookScript = `
|
||||
const fs = require('fs');
|
||||
try {
|
||||
const input = fs.readFileSync(0, 'utf-8');
|
||||
console.error('DEBUG: AfterAgent hook input received');
|
||||
process.stdout.write("Received Input: " + input);
|
||||
} catch (err) {
|
||||
console.error('Hook Failed:', err);
|
||||
process.exit(1);
|
||||
}
|
||||
`;
|
||||
|
||||
const scriptPath = rig.createScript('after_agent_verify.cjs', hookScript);
|
||||
|
||||
rig.setup('should receive prompt and response in AfterAgent hook', {
|
||||
settings: {
|
||||
hooksConfig: {
|
||||
enabled: true,
|
||||
},
|
||||
hooks: {
|
||||
AfterAgent: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: normalizePath(`node "${scriptPath}"`)!,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await rig.run({ args: 'Hello validation' });
|
||||
|
||||
const hookTelemetryFound = await rig.waitForTelemetryEvent('hook_call');
|
||||
expect(hookTelemetryFound).toBeTruthy();
|
||||
|
||||
const hookLogs = rig.readHookLogs();
|
||||
const afterAgentLog = hookLogs.find(
|
||||
(log) => log.hookCall.hook_event_name === 'AfterAgent',
|
||||
);
|
||||
|
||||
expect(afterAgentLog).toBeDefined();
|
||||
// Verify the hook stdout contains the input we echoed which proves the
|
||||
// hook received the prompt and response
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('Received Input');
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('Hello validation');
|
||||
// The fake response contains "Hello World"
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('Hello World');
|
||||
});
|
||||
|
||||
it('should process clearContext in AfterAgent hook output', async () => {
|
||||
rig.setup('should process clearContext in AfterAgent hook output', {
|
||||
fakeResponsesPath: join(
|
||||
import.meta.dirname,
|
||||
'hooks-system.after-agent.responses',
|
||||
),
|
||||
});
|
||||
|
||||
// BeforeModel hook to track message counts across LLM calls
|
||||
const messageCountFile = join(rig.testDir!, 'message-counts.json');
|
||||
const escapedPath = JSON.stringify(messageCountFile);
|
||||
const beforeModelScript = `
|
||||
const fs = require('fs');
|
||||
const input = JSON.parse(fs.readFileSync(0, 'utf-8'));
|
||||
const messageCount = input.llm_request?.contents?.length || 0;
|
||||
let counts = [];
|
||||
try { counts = JSON.parse(fs.readFileSync(${escapedPath}, 'utf-8')); } catch (e) {}
|
||||
counts.push(messageCount);
|
||||
fs.writeFileSync(${escapedPath}, JSON.stringify(counts));
|
||||
console.log(JSON.stringify({ decision: 'allow' }));
|
||||
`;
|
||||
const beforeModelScriptPath = rig.createScript(
|
||||
'before_model_counter.cjs',
|
||||
beforeModelScript,
|
||||
);
|
||||
|
||||
const afterAgentScript = `
|
||||
const fs = require('fs');
|
||||
const input = JSON.parse(fs.readFileSync(0, 'utf-8'));
|
||||
if (input.stop_hook_active) {
|
||||
// Retry turn: allow execution to proceed (breaks the loop)
|
||||
console.log(JSON.stringify({ decision: 'allow' }));
|
||||
} else {
|
||||
// First call: block and clear context to trigger the retry
|
||||
console.log(JSON.stringify({
|
||||
decision: 'block',
|
||||
reason: 'Security policy triggered',
|
||||
hookSpecificOutput: {
|
||||
hookEventName: 'AfterAgent',
|
||||
clearContext: true
|
||||
}
|
||||
}));
|
||||
}
|
||||
`;
|
||||
const afterAgentScriptPath = rig.createScript(
|
||||
'after_agent_clear.cjs',
|
||||
afterAgentScript,
|
||||
);
|
||||
|
||||
rig.setup('should process clearContext in AfterAgent hook output', {
|
||||
settings: {
|
||||
hooksConfig: {
|
||||
enabled: true,
|
||||
},
|
||||
hooks: {
|
||||
BeforeModel: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: normalizePath(`node "${beforeModelScriptPath}"`)!,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
AfterAgent: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: normalizePath(`node "${afterAgentScriptPath}"`)!,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await rig.run({ args: 'Hello test' });
|
||||
|
||||
const hookTelemetryFound = await rig.waitForTelemetryEvent('hook_call');
|
||||
expect(hookTelemetryFound).toBeTruthy();
|
||||
|
||||
const hookLogs = rig.readHookLogs();
|
||||
const afterAgentLog = hookLogs.find(
|
||||
(log) => log.hookCall.hook_event_name === 'AfterAgent',
|
||||
);
|
||||
|
||||
expect(afterAgentLog).toBeDefined();
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('clearContext');
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('true');
|
||||
expect(result).toContain('Security policy triggered');
|
||||
|
||||
// Verify context was cleared: second call should not have more messages than first
|
||||
const countsRaw = rig.readFile('message-counts.json');
|
||||
const counts = JSON.parse(countsRaw) as number[];
|
||||
expect(counts.length).toBeGreaterThanOrEqual(2);
|
||||
expect(counts[1]).toBeLessThanOrEqual(counts[0]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Multi-step Loops', () => {
|
||||
it('should fire BeforeAgent and AfterAgent exactly once per turn despite tool calls', async () => {
|
||||
await rig.setup(
|
||||
'should fire BeforeAgent and AfterAgent exactly once per turn despite tool calls',
|
||||
{
|
||||
fakeResponsesPath: join(
|
||||
import.meta.dirname,
|
||||
'hooks-agent-flow-multistep.responses',
|
||||
),
|
||||
},
|
||||
);
|
||||
|
||||
// Create script files for hooks
|
||||
const baPath = rig.createScript(
|
||||
'ba_fired.cjs',
|
||||
"console.log('BeforeAgent Fired');",
|
||||
);
|
||||
const aaPath = rig.createScript(
|
||||
'aa_fired.cjs',
|
||||
"console.log('AfterAgent Fired');",
|
||||
);
|
||||
|
||||
await rig.setup(
|
||||
'should fire BeforeAgent and AfterAgent exactly once per turn despite tool calls',
|
||||
{
|
||||
settings: {
|
||||
hooksConfig: {
|
||||
enabled: true,
|
||||
},
|
||||
hooks: {
|
||||
BeforeAgent: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: normalizePath(`node "${baPath}"`)!,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
AfterAgent: [
|
||||
{
|
||||
hooks: [
|
||||
{
|
||||
type: 'command',
|
||||
command: normalizePath(`node "${aaPath}"`)!,
|
||||
timeout: 5000,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
await rig.run({ args: 'Do a multi-step task' });
|
||||
|
||||
const hookLogs = rig.readHookLogs();
|
||||
const beforeAgentLogs = hookLogs.filter(
|
||||
(log) => log.hookCall.hook_event_name === 'BeforeAgent',
|
||||
);
|
||||
const afterAgentLogs = hookLogs.filter(
|
||||
(log) => log.hookCall.hook_event_name === 'AfterAgent',
|
||||
);
|
||||
|
||||
expect(beforeAgentLogs).toHaveLength(1);
|
||||
|
||||
expect(afterAgentLogs).toHaveLength(1);
|
||||
|
||||
const afterAgentLog = afterAgentLogs[0];
|
||||
expect(afterAgentLog).toBeDefined();
|
||||
expect(afterAgentLog?.hookCall.stdout).toContain('AfterAgent Fired');
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user