mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-27 01:54:53 +00:00
398 lines
19 KiB
TypeScript
398 lines
19 KiB
TypeScript
import type { ToolCall } from '@moonshot-ai/kosong';
|
|
import { describe, expect, it, vi } from 'vitest';
|
|
|
|
import { HookEngine } from '../../src/agent/hooks';
|
|
import type { SessionSubagentHost } from '../../src/session/subagent-host';
|
|
import { createFakeKaos } from '../tools/fixtures/fake-kaos';
|
|
import { createCommandKaos, testAgent } from './harness/agent';
|
|
import { executeTool } from '../tools/fixtures/execute-tool';
|
|
|
|
const signal = new AbortController().signal;
|
|
|
|
describe('Agent tools', () => {
|
|
it('blocks tools through PreToolUse before permission and emits PostToolUseFailure', async () => {
|
|
const execWithEnv = vi.fn().mockRejectedValue(new Error('Bash should not execute'));
|
|
const triggered: Array<[string, string, number]> = [];
|
|
const hookEngine = new HookEngine(
|
|
[
|
|
{
|
|
event: 'PreToolUse',
|
|
matcher: 'Bash',
|
|
command: "echo 'blocked by PreToolUse' >&2; exit 2",
|
|
},
|
|
{
|
|
event: 'PostToolUseFailure',
|
|
matcher: 'Bash',
|
|
command: 'exit 0',
|
|
},
|
|
],
|
|
{
|
|
onTriggered: (event, target, count) => {
|
|
triggered.push([event, target, count]);
|
|
},
|
|
},
|
|
);
|
|
const ctx = testAgent({
|
|
kaos: createFakeKaos({ execWithEnv }),
|
|
hookEngine,
|
|
});
|
|
ctx.configure({ tools: ['Bash'] });
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will run Bash.' }, bashCall());
|
|
ctx.mockNextResponse({ type: 'text', text: 'The hook blocked Bash.' });
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Try Bash' }] });
|
|
|
|
await ctx.untilTurnEnd();
|
|
|
|
expect(execWithEnv).not.toHaveBeenCalled();
|
|
expect(triggered).toEqual([
|
|
['PreToolUse', 'Bash', 1],
|
|
['PostToolUseFailure', 'Bash', 1],
|
|
]);
|
|
expect(JSON.stringify(ctx.agent.context.data().history)).toContain('blocked by PreToolUse');
|
|
});
|
|
|
|
it('emits PostToolUse after successful tools', async () => {
|
|
const triggered: Array<[string, string, number]> = [];
|
|
const hookEngine = new HookEngine(
|
|
[
|
|
{
|
|
event: 'PostToolUse',
|
|
matcher: 'Bash',
|
|
command: 'exit 0',
|
|
},
|
|
],
|
|
{
|
|
onTriggered: (event, target, count) => {
|
|
triggered.push([event, target, count]);
|
|
},
|
|
},
|
|
);
|
|
const ctx = testAgent({
|
|
kaos: createCommandKaos('ok'),
|
|
hookEngine,
|
|
});
|
|
ctx.configure({ tools: ['Bash'] });
|
|
await ctx.rpc.setPermission({ mode: 'auto' });
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will run Bash.' }, bashCall());
|
|
ctx.mockNextResponse({ type: 'text', text: 'Bash returned ok.' });
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Run Bash' }] });
|
|
|
|
await ctx.untilTurnEnd();
|
|
|
|
expect(triggered).toEqual([['PostToolUse', 'Bash', 1]]);
|
|
});
|
|
|
|
it('uses builtin descriptions on tool call start events', async () => {
|
|
const ctx = testAgent({
|
|
kaos: createCommandKaos('ok'),
|
|
});
|
|
ctx.configure({ tools: ['Bash'] });
|
|
await ctx.rpc.setPermission({ mode: 'yolo' });
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will run Bash.' }, bashCall());
|
|
ctx.mockNextResponse({ type: 'text', text: 'Bash returned ok.' });
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Run Bash' }] });
|
|
await ctx.untilTurnEnd();
|
|
|
|
const started = ctx.allEvents.find(
|
|
(event) => event.type === '[rpc]' && event.event === 'tool.call.started',
|
|
);
|
|
expect(started?.args).toMatchObject({
|
|
description: 'Running: printf hook-output',
|
|
});
|
|
});
|
|
|
|
it('continues after a foreground Agent tool returns a max_tokens failure', async () => {
|
|
const completion = Promise.reject(
|
|
new Error('Subagent turn failed before completing its final summary: reason=max_tokens.'),
|
|
);
|
|
void completion.catch(() => undefined);
|
|
const subagentHost = {
|
|
spawn: vi.fn().mockResolvedValue({
|
|
agentId: 'agent-child',
|
|
profileName: 'coder',
|
|
resumed: false,
|
|
completion,
|
|
}),
|
|
resume: vi.fn(),
|
|
} as unknown as SessionSubagentHost;
|
|
const ctx = testAgent({ subagentHost });
|
|
ctx.configure({ tools: ['Agent'] });
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will ask a subagent.' }, agentCall());
|
|
ctx.mockNextResponse({
|
|
type: 'text',
|
|
text: 'The subagent failed with reason=max_tokens, so I will continue in the parent turn.',
|
|
});
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Delegate and recover' }] });
|
|
await ctx.untilTurnEnd();
|
|
|
|
expect(subagentHost.spawn).toHaveBeenCalledWith(
|
|
'coder',
|
|
expect.objectContaining({
|
|
parentToolCallId: 'call_agent',
|
|
prompt: 'Investigate deeply',
|
|
description: 'Investigate deeply',
|
|
runInBackground: false,
|
|
}),
|
|
);
|
|
expect(ctx.llmCalls).toHaveLength(2);
|
|
expect(ctx.allEvents).toContainEqual(
|
|
expect.objectContaining({
|
|
type: '[rpc]',
|
|
event: 'tool.result',
|
|
args: expect.objectContaining({
|
|
toolCallId: 'call_agent',
|
|
isError: true,
|
|
output: expect.stringContaining('reason=max_tokens'),
|
|
}),
|
|
}),
|
|
);
|
|
expect(JSON.stringify(ctx.llmCalls[1]?.history)).toContain('reason=max_tokens');
|
|
});
|
|
|
|
it('passes text from content-part error outputs to PostToolUseFailure hooks', async () => {
|
|
const lookupCall: ToolCall = {
|
|
type: 'function',
|
|
id: 'call_lookup',
|
|
name: 'Lookup',
|
|
arguments: '{"query":"moon"}',
|
|
};
|
|
const resolved: Array<[string, string, string]> = [];
|
|
const hookEngine = new HookEngine(
|
|
[
|
|
{
|
|
event: 'PostToolUseFailure',
|
|
matcher: 'Lookup',
|
|
command: hookErrorMessageAssertCommand('rich failure text'),
|
|
},
|
|
],
|
|
{
|
|
onResolved: (event, target, action) => {
|
|
resolved.push([event, target, action]);
|
|
},
|
|
},
|
|
);
|
|
const ctx = testAgent({ hookEngine });
|
|
ctx.configure();
|
|
await ctx.rpc.setPermission({ mode: 'auto' });
|
|
await ctx.rpc.registerTool({
|
|
name: 'Lookup',
|
|
description: 'Look up a short test value.',
|
|
parameters: {
|
|
type: 'object',
|
|
properties: {
|
|
query: { type: 'string' },
|
|
},
|
|
required: ['query'],
|
|
additionalProperties: false,
|
|
},
|
|
});
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will look it up.' }, lookupCall);
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Look up moon' }] });
|
|
await ctx.untilToolCall({
|
|
isError: true,
|
|
output: [{ type: 'text', text: 'rich failure text' }],
|
|
});
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'The lookup failed.' });
|
|
await ctx.untilTurnEnd();
|
|
|
|
await vi.waitFor(() => {
|
|
expect(resolved).toEqual([['PostToolUseFailure', 'Lookup', 'allow']]);
|
|
});
|
|
});
|
|
|
|
it('uses the active builtin tool set as the LLM visible tools', async () => {
|
|
const ctx = testAgent();
|
|
ctx.configure({ tools: ['Write', 'Bash'] });
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'ready' });
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Which tools are active?' }] });
|
|
|
|
await ctx.untilTurnEnd();
|
|
expect(ctx.lastLlmInput()).toMatchInlineSnapshot(`
|
|
system: <system-prompt>
|
|
tools: Bash, Write
|
|
messages:
|
|
user: text "Which tools are active?"
|
|
`);
|
|
await ctx.expectResumeMatches();
|
|
});
|
|
|
|
it('disables Bash background mode unless task management tools are active', async () => {
|
|
const ctx = testAgent();
|
|
ctx.configure({ tools: ['Bash'] });
|
|
|
|
const bashOnly = ctx.agent.tools.loopTools.find((tool) => tool.name === 'Bash');
|
|
expect(bashOnly).toBeDefined();
|
|
expect(bashOnly!.description).toContain('Background execution is disabled for this agent.');
|
|
expect(bashOnly!.description).not.toContain('the command will be started as a background task');
|
|
await expect(
|
|
executeTool(bashOnly!, {
|
|
turnId: '0',
|
|
toolCallId: 'call_bash',
|
|
args: { command: 'sleep 10', run_in_background: true, description: 'watch' },
|
|
signal,
|
|
}),
|
|
).resolves.toMatchObject({
|
|
isError: true,
|
|
output:
|
|
'Background execution is not available for this agent because TaskOutput and TaskStop are not enabled.',
|
|
});
|
|
|
|
ctx.agent.tools.setActiveTools(['Bash', 'TaskList', 'TaskOutput', 'TaskStop']);
|
|
|
|
const managedBash = ctx.agent.tools.loopTools.find((tool) => tool.name === 'Bash');
|
|
expect(managedBash).toBeDefined();
|
|
expect(managedBash!.description).toContain('run_in_background=true');
|
|
});
|
|
|
|
it('routes registered user tools through tool.call request/response', async () => {
|
|
const lookupCall: ToolCall = {
|
|
type: 'function',
|
|
id: 'call_lookup',
|
|
name: 'Lookup',
|
|
arguments: '{"query":"moon"}',
|
|
};
|
|
const ctx = testAgent();
|
|
ctx.configure();
|
|
await ctx.rpc.setPermission({ mode: 'auto' });
|
|
await ctx.rpc.registerTool({
|
|
name: 'Lookup',
|
|
description: 'Look up a short test value.',
|
|
parameters: {
|
|
type: 'object',
|
|
properties: {
|
|
query: { type: 'string' },
|
|
},
|
|
required: ['query'],
|
|
additionalProperties: false,
|
|
},
|
|
});
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'I will look it up.' }, lookupCall);
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Look up moon' }] });
|
|
expect(
|
|
await ctx.untilToolCall({
|
|
content: 'moon-result',
|
|
output: 'moon-result',
|
|
}),
|
|
).toMatchInlineSnapshot(`
|
|
[wire] permission.set_mode { "mode": "auto", "time": "<time>" }
|
|
[emit] agent.status.updated { "model": "mock-model", "contextTokens": 0, "maxContextTokens": 1000000, "contextUsage": 0, "planMode": false, "permission": "auto" }
|
|
[wire] tools.register_user_tool { "name": "Lookup", "description": "Look up a short test value.", "parameters": { "type": "object", "properties": { "query": { "type": "string" } }, "required": [ "query" ], "additionalProperties": false }, "time": "<time>" }
|
|
[wire] turn.prompt { "input": [ { "type": "text", "text": "Look up moon" } ], "origin": { "kind": "user" }, "time": "<time>" }
|
|
[emit] turn.started { "turnId": 0, "origin": { "kind": "user" } }
|
|
[wire] context.append_message { "message": { "role": "user", "content": [ { "type": "text", "text": "Look up moon" } ], "toolCalls": [], "origin": { "kind": "user" } }, "time": "<time>" }
|
|
[wire] context.append_message { "message": { "role": "user", "content": [ { "type": "text", "text": "<auto-mode-enter-reminder>" } ], "toolCalls": [], "origin": { "kind": "injection", "variant": "permission_mode" } }, "time": "<time>" }
|
|
[wire] context.append_loop_event { "event": { "type": "step.begin", "uuid": "<uuid-1>", "turnId": "0", "step": 1 }, "time": "<time>" }
|
|
[emit] turn.step.started { "turnId": 0, "step": 1, "stepId": "<uuid-1>" }
|
|
[emit] assistant.delta { "turnId": 0, "delta": "I will look it up." }
|
|
[emit] tool.call.delta { "turnId": 0, "toolCallId": "call_lookup", "name": "Lookup", "argumentsPart": "{\\"query\\":\\"moon\\"}" }
|
|
[wire] context.append_loop_event { "event": { "type": "content.part", "uuid": "<uuid-2>", "turnId": "0", "step": 1, "stepUuid": "<uuid-1>", "part": { "type": "text", "text": "I will look it up." } }, "time": "<time>" }
|
|
[wire] context.append_loop_event { "event": { "type": "tool.call", "uuid": "call_lookup", "turnId": "0", "step": 1, "stepUuid": "<uuid-1>", "toolCallId": "call_lookup", "name": "Lookup", "args": { "query": "moon" } }, "time": "<time>" }
|
|
[emit] tool.call.started { "turnId": 0, "toolCallId": "call_lookup", "name": "Lookup", "args": { "query": "moon" } }
|
|
[emit] toolCall { "turnId": 0, "toolCallId": "call_lookup", "args": { "query": "moon" } }
|
|
`);
|
|
expect(ctx.lastLlmInput()).toMatchInlineSnapshot(`
|
|
system: <system-prompt>
|
|
tools: Lookup
|
|
messages:
|
|
user: text "Look up moon"
|
|
user: text <auto-mode-enter-reminder>
|
|
`);
|
|
|
|
ctx.mockNextResponse({ type: 'text', text: 'The lookup result is moon-result.' });
|
|
expect(await ctx.untilTurnEnd()).toMatchInlineSnapshot(`
|
|
[wire] context.append_loop_event { "event": { "type": "tool.result", "parentUuid": "call_lookup", "toolCallId": "call_lookup", "result": { "output": "moon-result" } }, "time": "<time>" }
|
|
[emit] tool.result { "turnId": 0, "toolCallId": "call_lookup", "output": "moon-result" }
|
|
[wire] context.append_loop_event { "event": { "type": "step.end", "uuid": "<uuid-1>", "turnId": "0", "step": 1, "usage": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "tool_use" }, "time": "<time>" }
|
|
[emit] turn.step.completed { "turnId": 0, "step": 1, "stepId": "<uuid-1>", "usage": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "tool_use" }
|
|
[wire] usage.record { "model": "mock-model", "usage": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 }, "usageScope": "turn", "time": "<time>" }
|
|
[emit] agent.status.updated { "model": "mock-model", "contextTokens": 104, "maxContextTokens": 1000000, "contextUsage": 0.000104, "planMode": false, "permission": "auto", "usage": { "byModel": { "mock-model": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 } }, "total": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 }, "currentTurn": { "inputOther": 88, "output": 16, "inputCacheRead": 0, "inputCacheCreation": 0 } } }
|
|
[wire] context.append_loop_event { "event": { "type": "step.begin", "uuid": "<uuid-3>", "turnId": "0", "step": 2 }, "time": "<time>" }
|
|
[emit] turn.step.started { "turnId": 0, "step": 2, "stepId": "<uuid-3>" }
|
|
[emit] assistant.delta { "turnId": 0, "delta": "The lookup result is moon-result." }
|
|
[wire] context.append_loop_event { "event": { "type": "content.part", "uuid": "<uuid-4>", "turnId": "0", "step": 2, "stepUuid": "<uuid-3>", "part": { "type": "text", "text": "The lookup result is moon-result." } }, "time": "<time>" }
|
|
[wire] context.append_loop_event { "event": { "type": "step.end", "uuid": "<uuid-3>", "turnId": "0", "step": 2, "usage": { "inputOther": 108, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn" }, "time": "<time>" }
|
|
[emit] turn.step.completed { "turnId": 0, "step": 2, "stepId": "<uuid-3>", "usage": { "inputOther": 108, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn" }
|
|
[wire] usage.record { "model": "mock-model", "usage": { "inputOther": 108, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "usageScope": "turn", "time": "<time>" }
|
|
[emit] agent.status.updated { "model": "mock-model", "contextTokens": 120, "maxContextTokens": 1000000, "contextUsage": 0.00012, "planMode": false, "permission": "auto", "usage": { "byModel": { "mock-model": { "inputOther": 196, "output": 28, "inputCacheRead": 0, "inputCacheCreation": 0 } }, "total": { "inputOther": 196, "output": 28, "inputCacheRead": 0, "inputCacheCreation": 0 }, "currentTurn": { "inputOther": 196, "output": 28, "inputCacheRead": 0, "inputCacheCreation": 0 } } }
|
|
[emit] turn.ended { "turnId": 0, "reason": "completed" }
|
|
`);
|
|
expect(ctx.lastLlmInput()).toMatchInlineSnapshot(`
|
|
messages:
|
|
<last>
|
|
assistant: text "I will look it up." calls call_lookup:Lookup { "query": "moon" }
|
|
tool[call_lookup]: text "moon-result"
|
|
`);
|
|
|
|
await ctx.rpc.unregisterTool({ name: 'Lookup' });
|
|
ctx.mockNextResponse({ type: 'text', text: 'No lookup tool is available.' });
|
|
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Can you still use Lookup?' }] });
|
|
|
|
expect(await ctx.untilTurnEnd()).toMatchInlineSnapshot(`
|
|
[wire] tools.unregister_user_tool { "name": "Lookup", "time": "<time>" }
|
|
[wire] turn.prompt { "input": [ { "type": "text", "text": "Can you still use Lookup?" } ], "origin": { "kind": "user" }, "time": "<time>" }
|
|
[emit] turn.started { "turnId": 1, "origin": { "kind": "user" } }
|
|
[wire] context.append_message { "message": { "role": "user", "content": [ { "type": "text", "text": "Can you still use Lookup?" } ], "toolCalls": [], "origin": { "kind": "user" } }, "time": "<time>" }
|
|
[wire] context.append_loop_event { "event": { "type": "step.begin", "uuid": "<uuid-5>", "turnId": "1", "step": 1 }, "time": "<time>" }
|
|
[emit] turn.step.started { "turnId": 1, "step": 1, "stepId": "<uuid-5>" }
|
|
[emit] assistant.delta { "turnId": 1, "delta": "No lookup tool is available." }
|
|
[wire] context.append_loop_event { "event": { "type": "content.part", "uuid": "<uuid-6>", "turnId": "1", "step": 1, "stepUuid": "<uuid-5>", "part": { "type": "text", "text": "No lookup tool is available." } }, "time": "<time>" }
|
|
[wire] context.append_loop_event { "event": { "type": "step.end", "uuid": "<uuid-5>", "turnId": "1", "step": 1, "usage": { "inputOther": 128, "output": 10, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn" }, "time": "<time>" }
|
|
[emit] turn.step.completed { "turnId": 1, "step": 1, "stepId": "<uuid-5>", "usage": { "inputOther": 128, "output": 10, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn" }
|
|
[wire] usage.record { "model": "mock-model", "usage": { "inputOther": 128, "output": 10, "inputCacheRead": 0, "inputCacheCreation": 0 }, "usageScope": "turn", "time": "<time>" }
|
|
[emit] agent.status.updated { "model": "mock-model", "contextTokens": 138, "maxContextTokens": 1000000, "contextUsage": 0.000138, "planMode": false, "permission": "auto", "usage": { "byModel": { "mock-model": { "inputOther": 324, "output": 38, "inputCacheRead": 0, "inputCacheCreation": 0 } }, "total": { "inputOther": 324, "output": 38, "inputCacheRead": 0, "inputCacheCreation": 0 }, "currentTurn": { "inputOther": 128, "output": 10, "inputCacheRead": 0, "inputCacheCreation": 0 } } }
|
|
[emit] turn.ended { "turnId": 1, "reason": "completed" }
|
|
`);
|
|
expect(ctx.lastLlmInput()).toMatchInlineSnapshot(`
|
|
tools: []
|
|
messages:
|
|
<last>
|
|
assistant: text "The lookup result is moon-result."
|
|
user: text "Can you still use Lookup?"
|
|
`);
|
|
await ctx.expectResumeMatches();
|
|
});
|
|
});
|
|
|
|
function bashCall(): ToolCall {
|
|
return {
|
|
type: 'function',
|
|
id: 'call_bash',
|
|
name: 'Bash',
|
|
arguments: '{"command":"printf hook-output","timeout":60}',
|
|
};
|
|
}
|
|
|
|
function agentCall(): ToolCall {
|
|
return {
|
|
type: 'function',
|
|
id: 'call_agent',
|
|
name: 'Agent',
|
|
arguments: JSON.stringify({
|
|
prompt: 'Investigate deeply',
|
|
description: 'Investigate deeply',
|
|
subagent_type: 'coder',
|
|
}),
|
|
};
|
|
}
|
|
|
|
function hookErrorMessageAssertCommand(expected: string): string {
|
|
const script = [
|
|
"let input = '';",
|
|
"process.stdin.on('data', (chunk) => { input += chunk; });",
|
|
"process.stdin.on('end', () => {",
|
|
' const payload = JSON.parse(input);',
|
|
` if (payload.error?.message === ${JSON.stringify(expected)}) process.exit(0);`,
|
|
" console.error(payload.error?.message ?? '<missing>');",
|
|
' process.exit(2);',
|
|
'});',
|
|
].join('');
|
|
return `node -e ${JSON.stringify(script)}`;
|
|
}
|