feat(agent): add iterative tool use loop with max iterations
Rewrites NativeAgent.process() from single-turn to an iterative tool loop. When toolRegistry and toolExecutor are provided, the agent calls the model, executes any requested tool calls, feeds results back, and loops until the model returns a text response or max iterations hit. - Backward compatible: works exactly as before without tools - Supports onToolUse callback for frontend status display - Max iterations (default 10) prevents infinite loops - Handles multiple tool calls per model response - 5 new tests (8 total)
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { NativeAgent } from './agent.js';
|
||||
import type { ModelClient, ChatResponse } from '../../models/types.js';
|
||||
import { ToolRegistry, ToolExecutor } from '../../tools/index.js';
|
||||
import { HookEngine } from '../../hooks/index.js';
|
||||
import type { Tool, ToolResult } from '../../tools/index.js';
|
||||
|
||||
describe('NativeAgent', () => {
|
||||
const createMockClient = (): ModelClient => ({
|
||||
@@ -67,3 +70,191 @@ describe('NativeAgent', () => {
|
||||
expect(mockSession.addMessage).toHaveBeenNthCalledWith(2, { role: 'assistant', content: 'Hello!' });
|
||||
});
|
||||
});
|
||||
|
||||
// Simple test tool
|
||||
const echoTool: Tool = {
|
||||
name: 'test.echo',
|
||||
description: 'Echo',
|
||||
inputSchema: { type: 'object', properties: { text: { type: 'string' } }, required: ['text'] },
|
||||
execute: async (args) => ({ success: true, output: (args as { text: string }).text }),
|
||||
};
|
||||
|
||||
describe('NativeAgent tool loop', () => {
|
||||
it('executes tool calls and feeds results back', async () => {
|
||||
let callCount = 0;
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) {
|
||||
// First call: model requests tool use
|
||||
return {
|
||||
content: '',
|
||||
stopReason: 'tool_use',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
toolCalls: [{ id: 'call_1', name: 'test.echo', args: { text: 'hello' } }],
|
||||
};
|
||||
}
|
||||
// Second call: model gives final text response
|
||||
return {
|
||||
content: 'The tool returned: hello',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 15, outputTokens: 10 },
|
||||
};
|
||||
}),
|
||||
};
|
||||
|
||||
const registry = new ToolRegistry();
|
||||
registry.register(echoTool);
|
||||
const hooks = new HookEngine({ confirm: [], log: [], silent: [] });
|
||||
const executor = new ToolExecutor(registry, hooks);
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
toolRegistry: registry,
|
||||
toolExecutor: executor,
|
||||
});
|
||||
|
||||
const response = await agent.process('echo hello');
|
||||
expect(response).toBe('The tool returned: hello');
|
||||
expect(mockClient.chat).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('respects max iterations', async () => {
|
||||
// Model always returns tool_use
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: '',
|
||||
stopReason: 'tool_use',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
toolCalls: [{ id: 'call_1', name: 'test.echo', args: { text: 'loop' } }],
|
||||
}),
|
||||
};
|
||||
|
||||
const registry = new ToolRegistry();
|
||||
registry.register(echoTool);
|
||||
const hooks = new HookEngine({ confirm: [], log: [], silent: [] });
|
||||
const executor = new ToolExecutor(registry, hooks);
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
toolRegistry: registry,
|
||||
toolExecutor: executor,
|
||||
maxIterations: 3,
|
||||
});
|
||||
|
||||
const response = await agent.process('loop forever');
|
||||
expect(response).toContain('max iterations');
|
||||
expect(mockClient.chat).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('works without tools (backward compatible)', async () => {
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: 'Hello!',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
}),
|
||||
};
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
});
|
||||
|
||||
const response = await agent.process('Hi');
|
||||
expect(response).toBe('Hello!');
|
||||
});
|
||||
|
||||
it('calls onToolUse callback on start and end', async () => {
|
||||
let callCount = 0;
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) {
|
||||
return {
|
||||
content: '',
|
||||
stopReason: 'tool_use',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
toolCalls: [{ id: 'call_1', name: 'test.echo', args: { text: 'hi' } }],
|
||||
};
|
||||
}
|
||||
return {
|
||||
content: 'Done',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 15, outputTokens: 10 },
|
||||
};
|
||||
}),
|
||||
};
|
||||
|
||||
const registry = new ToolRegistry();
|
||||
registry.register(echoTool);
|
||||
const hooks = new HookEngine({ confirm: [], log: [], silent: [] });
|
||||
const executor = new ToolExecutor(registry, hooks);
|
||||
const onToolUse = vi.fn();
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
toolRegistry: registry,
|
||||
toolExecutor: executor,
|
||||
onToolUse,
|
||||
});
|
||||
|
||||
await agent.process('echo hi');
|
||||
|
||||
expect(onToolUse).toHaveBeenCalledTimes(2);
|
||||
expect(onToolUse).toHaveBeenNthCalledWith(1, expect.objectContaining({
|
||||
type: 'start',
|
||||
tool: 'test.echo',
|
||||
args: { text: 'hi' },
|
||||
}));
|
||||
expect(onToolUse).toHaveBeenNthCalledWith(2, expect.objectContaining({
|
||||
type: 'end',
|
||||
tool: 'test.echo',
|
||||
result: expect.objectContaining({ success: true, output: 'hi' }),
|
||||
}));
|
||||
});
|
||||
|
||||
it('handles multiple tool calls in single response', async () => {
|
||||
let callCount = 0;
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) {
|
||||
return {
|
||||
content: '',
|
||||
stopReason: 'tool_use',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
toolCalls: [
|
||||
{ id: 'call_1', name: 'test.echo', args: { text: 'first' } },
|
||||
{ id: 'call_2', name: 'test.echo', args: { text: 'second' } },
|
||||
],
|
||||
};
|
||||
}
|
||||
return {
|
||||
content: 'Got both results',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 15, outputTokens: 10 },
|
||||
};
|
||||
}),
|
||||
};
|
||||
|
||||
const registry = new ToolRegistry();
|
||||
registry.register(echoTool);
|
||||
const hooks = new HookEngine({ confirm: [], log: [], silent: [] });
|
||||
const executor = new ToolExecutor(registry, hooks);
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
toolRegistry: registry,
|
||||
toolExecutor: executor,
|
||||
});
|
||||
|
||||
const response = await agent.process('echo both');
|
||||
expect(response).toBe('Got both results');
|
||||
expect(mockClient.chat).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user