diff --git a/src/execution/AgentExecutor.test.ts b/src/execution/AgentExecutor.test.ts index 3ef52beed..cefb674d8 100644 --- a/src/execution/AgentExecutor.test.ts +++ b/src/execution/AgentExecutor.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, beforeEach } from 'vitest'; +import { describe, it, expect, beforeEach, vi } from 'vitest'; import { AgentExecutor, ExecutionEvent } from './AgentExecutor'; import { createMockProvider } from '../providers/mock'; import { ToolRegistry } from '../tools'; @@ -205,6 +205,518 @@ describe('AgentExecutor', () => { expect(systemMessage?.content).toBe('You are a pirate'); }); + it('should not invoke a tool flagged with needsApproval, evaluated with actual args', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + const needsApproval = vi.fn((args: any) => args.amount > 100); + + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { execute } as any, + needsApproval, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .build(); + + const toolCall = { + id: 'call-1', + type: 'function' as const, + function: { + name: 'chargeCard', + arguments: JSON.stringify({ amount: 500 }), + }, + }; + + const result = await (AgentExecutor as any).executeToolCall( + toolCall, + agent, + toolRegistry + ); + + expect(execute).not.toHaveBeenCalled(); + expect(needsApproval).toHaveBeenCalledWith({ amount: 500 }); + expect(result.requiresApproval).toBe(true); + expect(result.args).toEqual({ amount: 500 }); + }); + + it('should leave unflagged tools unchanged', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { execute } as any, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .build(); + + const toolCall = { + id: 'call-2', + type: 'function' as const, + function: { + name: 'chargeCard', + arguments: JSON.stringify({ amount: 5 }), + }, + }; + + const result = await (AgentExecutor as any).executeToolCall( + toolCall, + agent, + toolRegistry + ); + + expect(execute).toHaveBeenCalledWith({ amount: 5 }, {}); + expect(result.requiresApproval).toBeUndefined(); + expect(result.result).toEqual({ ok: true }); + }); + + it('should pause and persist a snapshot when a tool needs approval', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { + description: 'Charge a card', + parameters: {}, + execute, + } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + // MockLLMProvider simulates a tool call when the last user message + // mentions the tool's name. + const approvalProvider = createMockProvider({ + name: 'mock', + responses: ['Charging your card now'], + }); + + const save = vi.fn().mockResolvedValue(undefined); + const resolve = vi.fn().mockResolvedValue(null); + + const result = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider: approvalProvider, + toolRegistry, + approvalStore: { save, resolve }, + }); + + expect(execute).not.toHaveBeenCalled(); + expect(result.finishReason).toBe('awaiting-approval'); + expect(result.approvalId).toBeDefined(); + expect(save).toHaveBeenCalledTimes(1); + + const [pendingArg, snapshotArg] = save.mock.calls[0]; + expect(pendingArg.toolName).toBe('chargeCard'); + expect(pendingArg.args).toEqual({ input: 'mock input' }); + expect(snapshotArg.currentMessages).toEqual(result.messages); + expect(snapshotArg.currentMessages.some((m: any) => m.role === 'user')).toBe(true); + }); + + it('should throw a clear error when approval is needed but no approvalStore is provided', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { + description: 'Charge a card', + parameters: {}, + execute, + } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const approvalProvider = createMockProvider({ + name: 'mock', + responses: ['Charging your card now'], + }); + + await expect( + AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider: approvalProvider, + toolRegistry, + }) + ).rejects.toThrow(/requires approval/); + }); + + it('should save a checkpoint after each tool result when sessionId + checkpointStore are provided', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + toolRegistry.register('noop', { + displayName: 'Noop', + tool: { description: 'noop', parameters: {}, execute } as any, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('noop', { tool: 'noop', options: {} }) + .build(); + + // Custom provider: emits a tool call for the first 3 generations, + // then stops. + let call = 0; + const scriptedProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + call++; + if (call <= 3) { + return { + text: `step ${call}`, + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { + id: `call-${call}`, + type: 'function' as const, + function: { name: 'noop', arguments: '{}' }, + }, + ], + }; + } + return { + text: 'done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + }, + }; + + const save = vi.fn().mockResolvedValue(undefined); + const checkpointStore = { save, load: vi.fn(), delete: vi.fn() }; + + await AgentExecutor.execute({ + agent, + input: 'go', + provider: scriptedProvider as any, + toolRegistry, + sessionId: 'session-checkpoint-test', + checkpointStore, + }); + + expect(save).toHaveBeenCalledTimes(3); + const lengths = save.mock.calls.map(([, checkpoint]) => checkpoint.messages.length); + expect(lengths[1]).toBeGreaterThan(lengths[0]); + expect(lengths[2]).toBeGreaterThan(lengths[1]); + }); + + it('should resume from a checkpoint after a simulated crash + restart with zero message loss', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + toolRegistry.register('noop', { + displayName: 'Noop', + tool: { description: 'noop', parameters: {}, execute } as any, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('noop', { tool: 'noop', options: {} }) + .build(); + + function makeCheckpointStore() { + const records = new Map(); + return { + save: vi.fn(async (sessionId: string, checkpoint: any) => { + records.set(sessionId, checkpoint); + }), + load: vi.fn(async (sessionId: string) => records.get(sessionId) ?? null), + delete: vi.fn(async (sessionId: string) => { + records.delete(sessionId); + }), + }; + } + + function toolCallResponse(callNumber: number) { + return { + text: `step ${callNumber}`, + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { + id: `call-${callNumber}`, + type: 'function' as const, + function: { name: 'noop', arguments: '{}' }, + }, + ], + }; + } + + function stopResponse() { + return { + text: 'done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + } + + // --- "Uninterrupted" baseline run: all 4 tool calls in one go --- + const baselineStore = makeCheckpointStore(); + let baselineCall = 0; + const baselineProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + baselineCall++; + return baselineCall <= 4 ? toolCallResponse(baselineCall) : stopResponse(); + }, + }; + + const baselineResult = await AgentExecutor.execute({ + agent, + input: 'go', + provider: baselineProvider as any, + toolRegistry, + sessionId: 'baseline-session', + checkpointStore: baselineStore, + }); + + // --- Interrupted run: crashes after 2 of 4 tool calls --- + const sharedStore = makeCheckpointStore(); + let crashCall = 0; + const crashingProvider = { + name: 'scripted-crash', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted-crash'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + crashCall++; + if (crashCall <= 2) { + return toolCallResponse(crashCall); + } + throw new Error('simulated crash'); + }, + }; + + await expect( + AgentExecutor.execute({ + agent, + input: 'go', + provider: crashingProvider as any, + toolRegistry, + sessionId: 'resume-session', + checkpointStore: sharedStore, + }) + ).rejects.toThrow('simulated crash'); + + // "Restart": a brand new execute() call with the same sessionId/store, + // completing the remaining tool calls. + let resumeCall = 0; + const resumeProvider = { + name: 'scripted-resume', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted-resume'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + resumeCall++; + return resumeCall <= 2 ? toolCallResponse(2 + resumeCall) : stopResponse(); + }, + }; + + const resumedResult = await AgentExecutor.execute({ + agent, + input: 'go', + provider: resumeProvider as any, + toolRegistry, + sessionId: 'resume-session', + checkpointStore: sharedStore, + }); + + expect(sharedStore.load).toHaveBeenCalledWith('resume-session'); + expect(resumedResult.messages).toHaveLength(baselineResult.messages.length); + }); + + it('should build messages from scratch when no checkpoint exists for a fresh sessionId', async () => { + const checkpointStore = { + save: vi.fn().mockResolvedValue(undefined), + load: vi.fn().mockResolvedValue(null), + delete: vi.fn().mockResolvedValue(undefined), + }; + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .setPrompt('You are a pirate') + .build(); + + const result = await AgentExecutor.execute({ + agent, + input: 'Hello', + provider, + sessionId: 'fresh-session', + checkpointStore, + }); + + expect(checkpointStore.load).toHaveBeenCalledWith('fresh-session'); + const systemMessage = result.messages.find((m) => m.role === 'system'); + expect(systemMessage?.content).toBe('You are a pirate'); + const userMessage = result.messages.find((m) => m.role === 'user'); + expect(userMessage?.content).toBe('Hello'); + }); + + it('should delete the checkpoint once a run reaches a terminal finish reason', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + toolRegistry.register('noop', { + displayName: 'Noop', + tool: { description: 'noop', parameters: {}, execute } as any, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('noop', { tool: 'noop', options: {} }) + .build(); + + let call = 0; + const scriptedProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + call++; + if (call === 1) { + return { + text: 'calling noop', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { + id: 'call-1', + type: 'function' as const, + function: { name: 'noop', arguments: '{}' }, + }, + ], + }; + } + return { + text: 'done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + }, + }; + + function makeCheckpointStore() { + const records = new Map(); + return { + save: vi.fn(async (sessionId: string, checkpoint: any) => { + records.set(sessionId, checkpoint); + }), + load: vi.fn(async (sessionId: string) => records.get(sessionId) ?? null), + delete: vi.fn(async (sessionId: string) => { + records.delete(sessionId); + }), + }; + } + + const checkpointStore = makeCheckpointStore(); + + await AgentExecutor.execute({ + agent, + input: 'go', + provider: scriptedProvider as any, + toolRegistry, + sessionId: 'terminal-session', + checkpointStore, + }); + + expect(checkpointStore.delete).toHaveBeenCalledWith('terminal-session'); + await expect(checkpointStore.load('terminal-session')).resolves.toBeNull(); + }); + + it('should use fresh input (not stale stored messages) when execute() is called again with the same sessionId after completion', async () => { + function makeCheckpointStore() { + const records = new Map(); + return { + save: vi.fn(async (sessionId: string, checkpoint: any) => { + records.set(sessionId, checkpoint); + }), + load: vi.fn(async (sessionId: string) => records.get(sessionId) ?? null), + delete: vi.fn(async (sessionId: string) => { + records.delete(sessionId); + }), + }; + } + + const checkpointStore = makeCheckpointStore(); + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .build(); + + // First run: completes normally (no tool calls), reaching a terminal + // finish reason and clearing the checkpoint. + const firstProvider = createMockProvider({ + name: 'mock', + responses: ['first response'], + }); + + await AgentExecutor.execute({ + agent, + input: 'first input', + provider: firstProvider, + sessionId: 'reused-session', + checkpointStore, + }); + + // Second run: same sessionId, brand-new input. If the stale checkpoint + // were still around and rehydrated, this fresh input would be + // silently ignored. + const secondProvider = createMockProvider({ + name: 'mock', + responses: ['second response'], + }); + + const secondResult = await AgentExecutor.execute({ + agent, + input: 'second input, completely different', + provider: secondProvider, + sessionId: 'reused-session', + checkpointStore, + }); + + const userMessage = secondResult.messages.find((m) => m.role === 'user'); + expect(userMessage?.content).toBe('second input, completely different'); + expect(secondResult.text).toBe('second response'); + }); + it('should pass temperature and maxTokens', async () => { const agent = AgentBuilder.create() .setType(AgentType.SmartAssistant) diff --git a/src/execution/AgentExecutor.ts b/src/execution/AgentExecutor.ts index 75554863b..af331d71a 100644 --- a/src/execution/AgentExecutor.ts +++ b/src/execution/AgentExecutor.ts @@ -3,9 +3,12 @@ * Executes agents with streaming support and tool calling */ +import { nanoid } from 'nanoid'; import { LLMProvider, Message, ToolCall } from '../providers'; import { AgentConfig } from '../types'; import { ToolRegistry } from '../tools'; +import { ApprovalStore, ExecutionSnapshot, PendingApproval } from './ApprovalGate'; +import { Checkpoint, CheckpointStore } from './checkpoint'; /** * Execution event types @@ -58,6 +61,30 @@ export interface ExecuteOptions { temperature?: number; maxTokens?: number; onEvent?: (event: ExecutionEvent) => void; + approvalStore?: ApprovalStore; + sessionId?: string; + checkpointStore?: CheckpointStore; + /** + * When true, `input` is treated as a complete, ready-to-send message + * array that already includes any system prompt it needs (e.g. messages + * reconstructed from an ExecutionSnapshot by resume.ts). buildMessages() + * will not prepend a fresh system message built from `agent.prompt` in + * this case, avoiding a duplicate system message. Only relevant on the + * "build from scratch" fallback path (no checkpoint loaded); ignored + * when a checkpoint is rehydrated, since that path never re-injects a + * system message anyway. + */ + skipSystemPromptInjection?: boolean; + /** + * Starting value for the step counter (and therefore the maxSteps + * safety-limit budget), used when there is no checkpoint to rehydrate + * `steps` from but execution is still a continuation of prior work - e.g. + * resume.ts resuming a run that was paused for approval after already + * taking some steps. Ignored whenever a checkpoint is loaded, since + * `checkpoint.stepIndex` is the source of truth in that case. Defaults to + * 0 (a genuinely fresh run) when omitted. + */ + initialSteps?: number; } /** @@ -74,6 +101,7 @@ export interface ExecutionResult { }; finishReason: string; steps: number; + approvalId?: string; } /** @@ -93,6 +121,11 @@ export class AgentExecutor { temperature, maxTokens, onEvent, + approvalStore, + sessionId, + checkpointStore, + skipSystemPromptInjection, + initialSteps, } = options; // Emit start event @@ -103,20 +136,39 @@ export class AgentExecutor { agentName: agent.name, }); - // Build messages - const messages = this.buildMessages(agent, input); - // Build tools const tools = this.buildTools(agent, toolRegistry); - const currentMessages = [...messages]; - const allToolCalls: ToolCall[] = []; - const totalUsage = { - promptTokens: 0, - completionTokens: 0, - totalTokens: 0, - }; - let steps = 0; + // If a checkpoint exists for this sessionId, rehydrate state from it + // instead of building messages from scratch. + let checkpoint: Checkpoint | null = null; + if (sessionId && checkpointStore) { + checkpoint = await checkpointStore.load(sessionId); + } + + let currentMessages: Message[]; + let allToolCalls: ToolCall[]; + let totalUsage: { promptTokens: number; completionTokens: number; totalTokens: number }; + let steps: number; + + if (checkpoint) { + currentMessages = [...checkpoint.messages]; + allToolCalls = [...(checkpoint.toolCalls as ToolCall[])]; + totalUsage = { ...checkpoint.usage }; + steps = checkpoint.stepIndex; + } else { + // Build messages from scratch (fallback path) + const messages = this.buildMessages(agent, input, skipSystemPromptInjection); + currentMessages = [...messages]; + allToolCalls = []; + totalUsage = { + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + }; + steps = initialSteps ?? 0; + } + let finalText = ''; let finishReason = 'stop'; @@ -173,6 +225,49 @@ export class AgentExecutor { toolRegistry ); + if (toolResult.requiresApproval) { + if (!approvalStore) { + throw new Error( + `Tool '${toolResult.toolName}' requires approval but no approvalStore was provided to AgentExecutor.execute()` + ); + } + + const pending: PendingApproval = { + id: nanoid(), + toolCallId: toolCall.id, + toolName: toolCall.function.name, + args: toolResult.args || {}, + agentId: agent.id, + createdAt: new Date().toISOString(), + }; + const snapshot: ExecutionSnapshot = { + agent, + currentMessages, + pendingToolCall: pending, + steps, + sessionId, + }; + + await approvalStore.save(pending, snapshot); + + this.emitEvent(onEvent, { + type: 'finish', + timestamp: new Date(), + finishReason: 'awaiting-approval', + usage: totalUsage, + }); + + return { + text: '', + messages: currentMessages, + toolCalls: allToolCalls, + usage: totalUsage, + finishReason: 'awaiting-approval', + steps, + approvalId: pending.id, + }; + } + this.emitEvent(onEvent, { type: 'tool-result', timestamp: new Date(), @@ -186,6 +281,19 @@ export class AgentExecutor { toolCallId: toolCall.id, toolName: toolCall.function.name, }); + + if (sessionId && checkpointStore) { + const checkpoint: Checkpoint = { + agentId: agent.id || '', + sessionId, + stepIndex: steps, + messages: [...currentMessages], + toolCalls: [...allToolCalls], + usage: totalUsage, + finishReason, + }; + await checkpointStore.save(sessionId, checkpoint); + } } // Continue loop for next generation @@ -214,6 +322,17 @@ export class AgentExecutor { usage: totalUsage, }); + // The run has reached a terminal state (either the model stopped + // requesting tools, or maxSteps was exhausted) - as opposed to the + // 'awaiting-approval' early-return above, which is a mid-flight pause + // where the checkpoint must stay in place so it can still be resumed. + // Clear the checkpoint here so a later execute() call reusing this + // sessionId builds fresh messages from its own `input` instead of + // silently resuming from this now-finished run. + if (sessionId && checkpointStore) { + await checkpointStore.delete(sessionId); + } + return { text: finalText, messages: currentMessages, @@ -229,12 +348,14 @@ export class AgentExecutor { */ private static buildMessages( agent: AgentConfig, - input: string | Message[] + input: string | Message[], + skipSystemPromptInjection = false ): Message[] { const messages: Message[] = []; - // Add system prompt - if (agent.prompt) { + // Add system prompt, unless the caller has indicated `input` already + // includes one (e.g. resume.ts rebuilding from an ExecutionSnapshot). + if (agent.prompt && !skipSystemPromptInjection) { messages.push({ role: 'system', content: agent.prompt, @@ -297,6 +418,8 @@ export class AgentExecutor { toolName: string; result: any; error?: string; + requiresApproval?: boolean; + args?: Record; }> { if (!toolRegistry) { return { @@ -319,6 +442,22 @@ export class AgentExecutor { } const args = JSON.parse(toolCall.function.arguments); + + const needsApproval = + typeof toolDesc.needsApproval === 'function' + ? await toolDesc.needsApproval(args) + : !!toolDesc.needsApproval; + + if (needsApproval) { + return { + toolCallId: toolCall.id, + toolName: toolCall.function.name, + result: null, + requiresApproval: true, + args, + }; + } + // The 'ai' SDK tool.execute expects (args, context) const result = await toolDesc.tool.execute(args, {} as any); diff --git a/src/execution/ApprovalGate.test.ts b/src/execution/ApprovalGate.test.ts new file mode 100644 index 000000000..dbce3fd6b --- /dev/null +++ b/src/execution/ApprovalGate.test.ts @@ -0,0 +1,149 @@ +import { describe, it, expect } from 'vitest'; +import { ExecutionSnapshot, StorageServiceApprovalStore, PendingApproval } from './ApprovalGate'; +import { AgentType } from '../types'; +import { StorageService, FileSystemAdapter, PathAdapter } from '../storage/StorageService'; + +/** + * A minimal in-memory fake of the Node fs/path modules, good enough to + * exercise StorageService's real read/write/lock logic without touching + * disk. Mirrors the mocking style already used in StorageService.test.ts, + * but backs it with real state so round-trips are meaningful. + */ +function createFakeFs(): { fs: FileSystemAdapter; path: PathAdapter } { + const files = new Map(); + + const fs: FileSystemAdapter = { + existsSync: (p: string) => files.has(p), + mkdirSync: () => undefined, + writeFileSync: (p: string, data: string | Uint8Array) => { + files.set(p, typeof data === 'string' ? data : Buffer.from(data).toString('utf8')); + }, + unlinkSync: (p: string) => { + files.delete(p); + }, + readFileSync: (p: string) => files.get(p) as any, + rmSync: (p: string) => { + files.delete(p); + }, + }; + + const path: PathAdapter = { + join: (...parts: string[]) => parts.join('/'), + resolve: (...parts: string[]) => parts.join('/'), + }; + + return { fs, path }; +} + +function createApprovalStore(): StorageServiceApprovalStore { + const { fs, path } = createFakeFs(); + const storageService = new StorageService('test-db-hash', 'test-schema', fs, path, '/test/root'); + return new StorageServiceApprovalStore(storageService); +} + +function buildSnapshot(pending: PendingApproval): ExecutionSnapshot { + return { + agent: { + id: pending.agentId, + name: 'Test Agent', + agentType: AgentType.SmartAssistant, + }, + currentMessages: [{ role: 'user', content: 'do the thing' }], + pendingToolCall: pending, + steps: 1, + }; +} + +describe('Execution - ApprovalGate types', () => { + it('should round-trip an ExecutionSnapshot through JSON with no loss', () => { + const snapshot: ExecutionSnapshot = { + agent: { + id: 'agent-1', + name: 'Test Agent', + agentType: AgentType.SmartAssistant, + prompt: 'You are a helpful assistant', + }, + currentMessages: [ + { role: 'system', content: 'You are a helpful assistant' }, + { role: 'user', content: 'Please charge the card' }, + ], + pendingToolCall: { + id: 'approval-1', + toolCallId: 'call-1', + toolName: 'chargeCard', + args: { amount: 500 }, + agentId: 'agent-1', + createdAt: new Date().toISOString(), + }, + steps: 1, + }; + + const roundTripped = JSON.parse(JSON.stringify(snapshot)) as ExecutionSnapshot; + + expect(roundTripped).toEqual(snapshot); + }); +}); + +describe('Execution - StorageServiceApprovalStore', () => { + it('should round-trip save then resolve, keeping concurrent saves isolated', async () => { + const store = createApprovalStore(); + + const pendingA: PendingApproval = { + id: 'approval-a', + toolCallId: 'call-a', + toolName: 'chargeCard', + args: { amount: 100 }, + agentId: 'agent-a', + createdAt: new Date().toISOString(), + }; + const pendingB: PendingApproval = { + id: 'approval-b', + toolCallId: 'call-b', + toolName: 'sendEmail', + args: { to: 'user@example.com' }, + agentId: 'agent-b', + createdAt: new Date().toISOString(), + }; + + await Promise.all([ + store.save(pendingA, buildSnapshot(pendingA)), + store.save(pendingB, buildSnapshot(pendingB)), + ]); + + const [resolvedA, resolvedB] = await Promise.all([ + store.resolve('approval-a'), + store.resolve('approval-b'), + ]); + + expect(resolvedA?.pending).toEqual(pendingA); + expect(resolvedA?.snapshot).toEqual(buildSnapshot(pendingA)); + expect(resolvedB?.pending).toEqual(pendingB); + expect(resolvedB?.snapshot).toEqual(buildSnapshot(pendingB)); + }); + + it('should return null when resolving an id that was already resolved', async () => { + const store = createApprovalStore(); + + const pending: PendingApproval = { + id: 'approval-once', + toolCallId: 'call-once', + toolName: 'chargeCard', + args: { amount: 250 }, + createdAt: new Date().toISOString(), + }; + + await store.save(pending, buildSnapshot(pending)); + + const first = await store.resolve('approval-once'); + expect(first?.pending).toEqual(pending); + + const second = await store.resolve('approval-once'); + expect(second).toBeNull(); + }); + + it('should return null for an id that was never saved', async () => { + const store = createApprovalStore(); + const result = await store.resolve('never-saved'); + expect(result).toBeNull(); + }); +}); diff --git a/src/execution/ApprovalGate.ts b/src/execution/ApprovalGate.ts new file mode 100644 index 000000000..9ba0ebd3d --- /dev/null +++ b/src/execution/ApprovalGate.ts @@ -0,0 +1,114 @@ +/** + * Approval Gate + * Types and storage for human-in-the-loop tool approval. + */ + +import { Message } from '../providers'; +import { AgentConfig } from '../types'; +import { StorageService } from '../storage'; + +/** + * A tool call that is waiting on a human decision before it can execute. + */ +export interface PendingApproval { + id: string; + toolCallId: string; + toolName: string; + args: Record; + agentId?: string; + createdAt: string; +} + +/** + * A human decision resolving a PendingApproval. + */ +export interface ApprovalDecision { + id: string; + approved: boolean; + note?: string; +} + +/** + * Enough state to resume an AgentExecutor.execute() run once a + * PendingApproval has been resolved. + */ +export interface ExecutionSnapshot { + agent: AgentConfig; + currentMessages: Message[]; + pendingToolCall: PendingApproval; + steps: number; + /** + * The sessionId the paused AgentExecutor.execute() run was using for + * durable-execution checkpointing (LOU-C9/C10), if any. Recorded so that + * resume.ts can proactively invalidate any stale checkpoint left behind + * under this sessionId once the approval is resolved - a checkpoint + * saved before the pause is stale by construction, since a pause always + * happens before the paused tool's result (and thus the next checkpoint + * write) exists. Undefined when the paused run wasn't using durable + * execution at all. + */ + sessionId?: string; +} + +/** + * A resolved approval record: the original pending tool call plus the + * snapshot needed to resume execution. + */ +export interface ResolvedApproval { + pending: PendingApproval; + snapshot: ExecutionSnapshot; +} + +/** + * Persists PendingApprovals (and the ExecutionSnapshot needed to resume + * them) so a paused agent run can survive a process restart. + */ +export interface ApprovalStore { + save(pending: PendingApproval, snapshot: ExecutionSnapshot): Promise; + resolve(id: string): Promise; +} + +/** + * Default ApprovalStore backed by the SDK's StorageService. + * + * Note: MemoryManager.ts (referenced by the LOU-C3 ticket as the source of + * the "lock then read/write JSON" pattern) actually stores memories through + * a MemoryRepository/data-layer abstraction, not StorageService directly - + * there is no existing lock+JSON-against-StorageService usage elsewhere in + * this codebase to copy verbatim. This implementation instead follows the + * pattern StorageService itself documents and tests: acquireLock -> + * read/write*JSON*Attachment -> releaseLock. + */ +export class StorageServiceApprovalStore implements ApprovalStore { + constructor(private readonly storageService: StorageService) {} + + private getStorageKey(id: string): string { + return `approvals/${id}.json`; + } + + async save(pending: PendingApproval, snapshot: ExecutionSnapshot): Promise { + const storageKey = this.getStorageKey(pending.id); + await this.storageService.acquireLock(storageKey); + try { + const record: ResolvedApproval = { pending, snapshot }; + this.storageService.writePlainJSONAttachment(storageKey, record); + } finally { + this.storageService.releaseLock(storageKey); + } + } + + async resolve(id: string): Promise { + const storageKey = this.getStorageKey(id); + await this.storageService.acquireLock(storageKey); + try { + if (!this.storageService.fileExists(storageKey)) { + return null; + } + const record = this.storageService.readPlainJSONAttachment(storageKey); + this.storageService.deleteAttachment(storageKey); + return record; + } finally { + this.storageService.releaseLock(storageKey); + } + } +} diff --git a/src/execution/checkpoint.test.ts b/src/execution/checkpoint.test.ts new file mode 100644 index 000000000..cf0ada81f --- /dev/null +++ b/src/execution/checkpoint.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect, vi } from 'vitest'; +import { Checkpoint, LocalStorageCheckpointStore } from './checkpoint'; +import { StorageService, FileSystemAdapter, PathAdapter } from '../storage/StorageService'; + +/** + * A minimal in-memory fake of the Node fs/path modules, good enough to + * exercise StorageService's real read/write/lock logic without touching + * disk. Mirrors the mocking style already used in StorageService.test.ts + * and ApprovalGate.test.ts. + */ +function createFakeFs(): { fs: FileSystemAdapter; path: PathAdapter } { + const files = new Map(); + + const fs: FileSystemAdapter = { + existsSync: (p: string) => files.has(p), + mkdirSync: () => undefined, + writeFileSync: (p: string, data: string | Uint8Array) => { + files.set(p, typeof data === 'string' ? data : Buffer.from(data).toString('utf8')); + }, + unlinkSync: (p: string) => { + files.delete(p); + }, + readFileSync: (p: string) => files.get(p) as any, + rmSync: (p: string) => { + files.delete(p); + }, + }; + + const path: PathAdapter = { + join: (...parts: string[]) => parts.join('/'), + resolve: (...parts: string[]) => parts.join('/'), + }; + + return { fs, path }; +} + +function createCheckpointStore(): LocalStorageCheckpointStore { + const { fs, path } = createFakeFs(); + const storageService = new StorageService('test-db-hash', 'test-schema', fs, path, '/test/root'); + return new LocalStorageCheckpointStore(storageService); +} + +function buildCheckpoint(overrides: Partial = {}): Checkpoint { + return { + agentId: 'agent-1', + sessionId: 'session-1', + stepIndex: 2, + messages: [ + { role: 'user', content: 'Hi' }, + { role: 'assistant', content: 'Hello!' }, + ], + toolCalls: [{ id: 'call-1', name: 'noop' }], + usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 }, + finishReason: 'tool_calls', + ...overrides, + }; +} + +describe('Execution - LocalStorageCheckpointStore', () => { + it('should round-trip save then load deep-equal', async () => { + const store = createCheckpointStore(); + const checkpoint = buildCheckpoint(); + + await store.save('session-1', checkpoint); + const loaded = await store.load('session-1'); + + expect(loaded).toEqual(checkpoint); + }); + + it('should return null after delete then load', async () => { + const store = createCheckpointStore(); + const checkpoint = buildCheckpoint(); + + await store.save('session-1', checkpoint); + await store.delete('session-1'); + + const loaded = await store.load('session-1'); + expect(loaded).toBeNull(); + }); + + it('should return null (not throw) when loading a never-saved sessionId', async () => { + const store = createCheckpointStore(); + await expect(store.load('never-saved')).resolves.toBeNull(); + }); + + it('should not throw when deleting a sessionId that was never saved', async () => { + const store = createCheckpointStore(); + await expect(store.delete('never-saved')).resolves.toBeUndefined(); + }); + + it('should acquire and release the same storage-key lock in load() as save()/delete() do', async () => { + const { fs, path } = createFakeFs(); + const storageService = new StorageService('test-db-hash', 'test-schema', fs, path, '/test/root'); + const store = new LocalStorageCheckpointStore(storageService); + + const acquireSpy = vi.spyOn(storageService, 'acquireLock'); + const releaseSpy = vi.spyOn(storageService, 'releaseLock'); + + const checkpoint = buildCheckpoint(); + await store.save('session-lock', checkpoint); + acquireSpy.mockClear(); + releaseSpy.mockClear(); + + await store.load('session-lock'); + + expect(acquireSpy).toHaveBeenCalledWith('checkpoints/session-lock.json'); + expect(releaseSpy).toHaveBeenCalledWith('checkpoints/session-lock.json'); + // Lock must be released even though load() returns from inside the + // try block (finally-based release, matching save()/delete()). + expect(acquireSpy.mock.invocationCallOrder[0]).toBeLessThan( + releaseSpy.mock.invocationCallOrder[0] + ); + }); +}); diff --git a/src/execution/checkpoint.ts b/src/execution/checkpoint.ts new file mode 100644 index 000000000..a7ef12119 --- /dev/null +++ b/src/execution/checkpoint.ts @@ -0,0 +1,76 @@ +/** + * Checkpoint + * Backend-agnostic durable-execution checkpointing for AgentExecutor runs. + */ + +import { Message } from '../providers'; +import { StorageService } from '../storage'; + +/** + * A snapshot of an in-progress agent run, saved after each tool result so + * execution can resume from here (e.g. after a process restart). + */ +export interface Checkpoint { + agentId: string; + sessionId: string; + stepIndex: number; + messages: Message[]; + toolCalls: unknown[]; + usage: { + promptTokens: number; + completionTokens: number; + totalTokens: number; + }; + finishReason?: string; +} + +/** + * Storage-backend-agnostic interface for persisting/loading Checkpoints. + */ +export interface CheckpointStore { + save(sessionId: string, checkpoint: Checkpoint): Promise; + load(sessionId: string): Promise; + delete(sessionId: string): Promise; +} + +/** + * Default CheckpointStore backed by the SDK's StorageService, keyed by + * `checkpoints/{sessionId}.json`. + */ +export class LocalStorageCheckpointStore implements CheckpointStore { + constructor(private readonly storageService: StorageService) {} + + private getStorageKey(sessionId: string): string { + return `checkpoints/${sessionId}.json`; + } + + async save(sessionId: string, checkpoint: Checkpoint): Promise { + const storageKey = this.getStorageKey(sessionId); + await this.storageService.acquireLock(storageKey); + try { + this.storageService.writePlainJSONAttachment(storageKey, checkpoint); + } finally { + this.storageService.releaseLock(storageKey); + } + } + + async load(sessionId: string): Promise { + const storageKey = this.getStorageKey(sessionId); + await this.storageService.acquireLock(storageKey); + try { + if (!this.storageService.fileExists(storageKey)) { + return null; + } + return this.storageService.readPlainJSONAttachment(storageKey); + } finally { + this.storageService.releaseLock(storageKey); + } + } + + async delete(sessionId: string): Promise { + const storageKey = this.getStorageKey(sessionId); + // deleteAttachment already swallows "not found" (it no-ops if the file + // doesn't exist), so no extra try/catch is needed here. + this.storageService.deleteAttachment(storageKey); + } +} diff --git a/src/execution/index.ts b/src/execution/index.ts index 9252e401a..398fd7ffb 100644 --- a/src/execution/index.ts +++ b/src/execution/index.ts @@ -8,3 +8,6 @@ export * from './MemoryManager'; export * from './ContextBuilder'; export * from './errors'; export * from './retry'; +export * from './ApprovalGate'; +export * from './resume'; +export * from './checkpoint'; diff --git a/src/execution/resume.test.ts b/src/execution/resume.test.ts new file mode 100644 index 000000000..5b2fbbb62 --- /dev/null +++ b/src/execution/resume.test.ts @@ -0,0 +1,511 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { resumeAfterApproval } from './resume'; +import { ApprovalStore, PendingApproval, ExecutionSnapshot } from './ApprovalGate'; +import { AgentExecutor } from './AgentExecutor'; +import { Checkpoint, CheckpointStore } from './checkpoint'; +import { createMockProvider } from '../providers/mock'; +import { ToolRegistry } from '../tools'; +import { AgentBuilder } from '../core'; +import { AgentType } from '../types'; + +/** Simple in-memory ApprovalStore, good enough for resume tests. */ +function createInMemoryApprovalStore(): ApprovalStore { + const records = new Map(); + return { + async save(pending, snapshot) { + records.set(pending.id, { pending, snapshot }); + }, + async resolve(id) { + const record = records.get(id); + if (!record) { + return null; + } + records.delete(id); + return record; + }, + }; +} + +/** Simple in-memory CheckpointStore, good enough for resume tests. */ +function createInMemoryCheckpointStore(): CheckpointStore { + const checkpoints = new Map(); + return { + async save(sessionId, checkpoint) { + checkpoints.set(sessionId, checkpoint); + }, + async load(sessionId) { + return checkpoints.get(sessionId) ?? null; + }, + async delete(sessionId) { + checkpoints.delete(sessionId); + }, + }; +} + +describe('Execution - resumeAfterApproval', () => { + let toolRegistry: ToolRegistry; + + beforeEach(() => { + toolRegistry = new ToolRegistry(); + }); + + it('should run the deferred tool exactly once when approved', async () => { + const execute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const provider = createMockProvider({ + name: 'mock', + responses: ['Charging now', 'All done'], + }); + + const approvalStore = createInMemoryApprovalStore(); + + const paused = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider, + toolRegistry, + approvalStore, + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + expect(paused.approvalId).toBeDefined(); + expect(execute).not.toHaveBeenCalled(); + + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + provider + ); + + expect(execute).toHaveBeenCalledTimes(1); + expect(resumed.finishReason).toBe('stop'); + // pre-pause history (user + assistant) + 1 new tool-result message + expect(resumed.messages).toHaveLength(paused.messages.length + 1); + expect(resumed.messages.slice(0, paused.messages.length)).toEqual(paused.messages); + expect(resumed.messages[resumed.messages.length - 1].role).toBe('tool'); + }); + + it('should never invoke the tool and should produce a rejection message when rejected', async () => { + const execute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const provider = createMockProvider({ + name: 'mock', + responses: ['Charging now', 'Understood, cancelled'], + }); + + const approvalStore = createInMemoryApprovalStore(); + + const paused = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider, + toolRegistry, + approvalStore, + }); + + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: false, note: 'Not authorized' }, + approvalStore, + toolRegistry, + provider + ); + + expect(execute).not.toHaveBeenCalled(); + const rejectionMessage = resumed.messages[paused.messages.length]; + expect(rejectionMessage.role).toBe('tool'); + const parsed = JSON.parse(rejectionMessage.content); + expect(parsed.note).toBe('Not authorized'); + }); + + it('should not duplicate the system message when resuming an agent that has agent.prompt set', async () => { + const execute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const systemPrompt = 'You are a careful billing assistant.'; + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .setPrompt(systemPrompt) + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const provider = createMockProvider({ + name: 'mock', + responses: ['Charging now', 'All done'], + }); + + const approvalStore = createInMemoryApprovalStore(); + + const paused = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider, + toolRegistry, + approvalStore, + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + expect(paused.approvalId).toBeDefined(); + + // Sanity check: exactly one system message before the pause too. + const pausedSystemMessages = paused.messages.filter((m) => m.role === 'system'); + expect(pausedSystemMessages).toHaveLength(1); + expect(pausedSystemMessages[0].content).toBe(systemPrompt); + + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + provider + ); + + // Regression check for the resume duplicate-system-message bug: after + // resuming, there must still be exactly ONE system message, matching + // the original agent.prompt content - not two. + const resumedSystemMessages = resumed.messages.filter((m) => m.role === 'system'); + expect(resumedSystemMessages).toHaveLength(1); + expect(resumedSystemMessages[0].content).toBe(systemPrompt); + }); + + it('should resolve (not reject) with an error-shaped tool message when the deferred tool throws on resume', async () => { + const execute = vi.fn().mockRejectedValue(new Error('payment gateway timeout')); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const provider = createMockProvider({ + name: 'mock', + responses: ['Charging now', 'All done'], + }); + + const approvalStore = createInMemoryApprovalStore(); + + const paused = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider, + toolRegistry, + approvalStore, + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + + // Should resolve, not reject, even though the deferred tool throws. + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + provider + ); + + expect(execute).toHaveBeenCalledTimes(1); + const toolMessage = resumed.messages[paused.messages.length]; + expect(toolMessage.role).toBe('tool'); + const parsed = JSON.parse(toolMessage.content); + expect(parsed.error).toBe('payment gateway timeout'); + }); + + it('should continue step-count accounting from the pre-pause step count on resume', async () => { + const execute = vi.fn().mockResolvedValue({ ok: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + // Scripted provider: two plain (no tool call) generations first, to + // burn 2 steps, then a generation that triggers the approval-gated + // tool call, then a final stop response after resume. + let call = 0; + const scriptedProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + call++; + if (call <= 2) { + return { + text: `thinking step ${call}`, + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + } + if (call === 3) { + return { + text: 'calling chargeCard', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { + id: 'call-1', + type: 'function' as const, + function: { name: 'chargeCard', arguments: '{}' }, + }, + ], + }; + } + return { + text: 'done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + }, + }; + + const approvalStore = createInMemoryApprovalStore(); + + // Force exactly one step per execute() call so the 2 "thinking" steps + // are pre-pause history the *next* execute() call must pick up from. + const firstStep = await AgentExecutor.execute({ + agent, + input: 'go', + provider: scriptedProvider as any, + toolRegistry, + approvalStore, + maxSteps: 1, + }); + expect(firstStep.steps).toBe(1); + + const secondStep = await AgentExecutor.execute({ + agent, + input: firstStep.messages, + provider: scriptedProvider as any, + toolRegistry, + approvalStore, + maxSteps: firstStep.steps + 1, + skipSystemPromptInjection: true, + initialSteps: firstStep.steps, + } as any); + expect(secondStep.steps).toBe(2); + + // Third call actually triggers the approval-gated tool call, continuing + // from step 2. + const paused = await AgentExecutor.execute({ + agent, + input: secondStep.messages, + provider: scriptedProvider as any, + toolRegistry, + approvalStore, + skipSystemPromptInjection: true, + initialSteps: secondStep.steps, + } as any); + expect(paused.finishReason).toBe('awaiting-approval'); + expect(paused.steps).toBe(3); + + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + scriptedProvider as any + ); + + // Continuation from step 3 (not reset to 0/1): one more generation + // happens on resume, so steps should be 4. + expect(resumed.steps).toBe(4); + }); + + it('reviewer repro: a stale checkpoint under the paused session must not clobber the resume state, even when sessionId+checkpointStore are explicitly passed through resume executeOptions', async () => { + // Reproduces the LOU-C bug: combining sessionId+checkpointStore + // (durable execution, LOU-C9/C10) with resumeAfterApproval() (approval + // gate, LOU-C6) on the SAME session silently lost data, because + // AgentExecutor's checkpoint-rehydration branch would unconditionally + // override the resume's carefully-reconstructed messages/steps with + // whatever stale checkpoint existed under that sessionId from before + // the pause. + const chargeExecute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute: chargeExecute } as any, + needsApproval: true, + }); + + const lookupExecute = vi.fn().mockResolvedValue({ found: true }); + toolRegistry.register('lookup', { + displayName: 'Lookup', + tool: { description: 'Look something up', parameters: {}, execute: lookupExecute } as any, + needsApproval: false, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('lookup', { tool: 'lookup', options: {} }) + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + // Deterministic scripted provider: generation 1 calls the (unguarded) + // 'lookup' tool - this is the "durable-execution step taken BEFORE the + // pause" that gets checkpointed. Generation 2 calls the + // approval-gated 'chargeCard' tool, which pauses execution before the + // checkpoint for THIS step is ever written - so the checkpoint saved + // after generation 1 is stale by construction from that point on. + // Generation 3 (post-resume) stops with no further tool calls. + let call = 0; + const scriptedProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + call++; + if (call === 1) { + return { + text: 'looking up', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { id: 'call-lookup', type: 'function' as const, function: { name: 'lookup', arguments: '{}' } }, + ], + }; + } + if (call === 2) { + return { + text: 'charging', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { + id: 'call-charge', + type: 'function' as const, + function: { name: 'chargeCard', arguments: '{}' }, + }, + ], + }; + } + return { + text: 'all done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + }, + }; + + const approvalStore = createInMemoryApprovalStore(); + const checkpointStore = createInMemoryCheckpointStore(); + const sessionId = 'shared-session-1'; + + const paused = await AgentExecutor.execute({ + agent, + input: 'go', + provider: scriptedProvider as any, + toolRegistry, + approvalStore, + sessionId, + checkpointStore, + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + expect(lookupExecute).toHaveBeenCalledTimes(1); + expect(chargeExecute).not.toHaveBeenCalled(); + + // Sanity: a checkpoint really was written under this sessionId from + // the pre-pause 'lookup' step, and it is stale relative to the pause + // point (it doesn't know about the pending 'chargeCard' call at all). + const preResumeCheckpoint = await checkpointStore.load(sessionId); + expect(preResumeCheckpoint).not.toBeNull(); + expect(preResumeCheckpoint!.stepIndex).toBe(1); + expect(preResumeCheckpoint!.messages.some((m) => m.toolName === 'chargeCard')).toBe(false); + + // The reviewer's repro: explicitly pass sessionId+checkpointStore + // through resume's executeOptions (bypassing the ResumeExecuteOptions + // Omit via `as any`, exactly as a caller not using strict TS - or one + // that force-casts - legally could at runtime). + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + scriptedProvider as any, + { sessionId, checkpointStore } as any, + checkpointStore + ); + + expect(chargeExecute).toHaveBeenCalledTimes(1); + + // The deferred tool's result must NOT be lost: it must be present in + // the final message list, not silently dropped by falling back to the + // stale checkpoint's messages (which never saw the chargeCard call). + const chargeResultMessage = resumed.messages.find( + (m) => m.role === 'tool' && m.toolName === 'chargeCard' + ); + expect(chargeResultMessage).toBeDefined(); + expect(JSON.parse(chargeResultMessage!.content).charged).toBe(true); + + // The lookup call/result from before the pause must also survive - + // the stale checkpoint's own data isn't what's wrong here, silently + // substituting it in place of the resume's own reconstruction is. + expect(resumed.messages.some((m) => m.role === 'tool' && m.toolName === 'lookup')).toBe(true); + + // Full reconstructed history: user, assistant(lookup), tool(lookup), + // assistant(chargeCard), tool(chargeCard) - no 6th message is pushed + // since the final generation returns no further tool calls. + expect(resumed.messages).toHaveLength(5); + + // Steps must continue from the ExecutionSnapshot's own step count + // (snapshot.steps was 2 at pause, +1 for the post-resume generation = + // 3), NOT from the stale checkpoint's stepIndex (which was 1, and + // would yield 2 if the checkpoint were wrongly rehydrated instead). + expect(resumed.steps).toBe(3); + + // Defense in depth: the stale checkpoint must have been proactively + // cleared so a later, unrelated execute() call reusing this sessionId + // can't rehydrate this pre-pause state either. + expect(await checkpointStore.load(sessionId)).toBeNull(); + }); + + it('should throw a clear error for an unknown or already-resolved approval id', async () => { + const approvalStore = createInMemoryApprovalStore(); + const provider = createMockProvider({ name: 'mock' }); + + await expect( + resumeAfterApproval({ id: 'never-existed', approved: true }, approvalStore, toolRegistry, provider) + ).rejects.toThrow(/No pending approval/); + }); +}); diff --git a/src/execution/resume.ts b/src/execution/resume.ts new file mode 100644 index 000000000..cbf8191df --- /dev/null +++ b/src/execution/resume.ts @@ -0,0 +1,162 @@ +/** + * Resume execution of an agent after a human has decided on a pending + * tool-call approval. + */ + +import { LLMProvider, Message } from '../providers'; +import { ToolRegistry } from '../tools'; +import { ApprovalDecision, ApprovalStore } from './ApprovalGate'; +import { AgentExecutor, ExecuteOptions, ExecutionResult } from './AgentExecutor'; +import { CheckpointStore } from './checkpoint'; + +/** + * Options passed through to the underlying AgentExecutor.execute() call + * once the deferred tool result (or rejection) has been appended to the + * conversation. + * + * `sessionId`/`checkpointStore` are also omitted here (in addition to the + * fields already omitted below): the resume path always fully reconstructs + * `messages`/`steps` from the ApprovalGate's ExecutionSnapshot, so it must + * never delegate to AgentExecutor's checkpoint-rehydration mechanism. Any + * checkpoint that exists under a given sessionId at this point was written + * BEFORE the pause (checkpoints are only saved after a tool *result*, and a + * pause happens before the deferred tool runs), so it is stale by + * construction and would silently clobber the just-reconstructed messages + * (including the deferred tool's own result) and step count if allowed + * through. See resumeAfterApproval() below for the belt-and-suspenders + * runtime guard that backs this type-level omission. + */ +export type ResumeExecuteOptions = Omit< + ExecuteOptions, + | 'agent' + | 'input' + | 'provider' + | 'toolRegistry' + | 'skipSystemPromptInjection' + | 'initialSteps' + | 'sessionId' + | 'checkpointStore' +>; + +/** + * Resume a paused AgentExecutor run after a human approves or rejects the + * tool call it was waiting on. + * + * Deviation from the ticket snippet: AgentExecutor exposes only a static + * `execute()` method (there is no instantiable AgentExecutor to inject), + * so this takes the LLMProvider needed to keep generating instead of an + * `executor: AgentExecutor` instance. + * + * @param checkpointStore Optional durable-execution CheckpointStore. This + * is NOT threaded through to the follow-up execute() call (see + * ResumeExecuteOptions) - it is used only, when the paused run's + * ExecutionSnapshot carries a sessionId, to proactively delete the stale + * checkpoint left behind under that sessionId from before the pause. This + * closes the gap for a LATER, unrelated execute() call against the same + * sessionId+checkpointStore (one that doesn't go through resume at all) + * so it can't silently rehydrate that stale pre-pause state either. + */ +export async function resumeAfterApproval( + decision: ApprovalDecision, + approvalStore: ApprovalStore, + toolRegistry: ToolRegistry, + provider: LLMProvider, + executeOptions: ResumeExecuteOptions = {}, + checkpointStore?: CheckpointStore +): Promise { + const record = await approvalStore.resolve(decision.id); + if (!record) { + throw new Error(`No pending approval found for id '${decision.id}' (unknown or already resolved)`); + } + + const { pending, snapshot } = record; + const messages: Message[] = [...snapshot.currentMessages]; + + // Defense in depth: a checkpoint under this sessionId (if any) was + // written before this pause and is therefore stale by construction - + // clear it now so it can't be loaded by this resume or by some later, + // unrelated execute() call against the same sessionId+checkpointStore. + if (snapshot.sessionId && checkpointStore) { + await checkpointStore.delete(snapshot.sessionId); + } + + if (decision.approved) { + const toolDesc = toolRegistry.get(pending.toolName); + if (!toolDesc || !toolDesc.tool || !toolDesc.tool.execute) { + throw new Error(`Tool '${pending.toolName}' not found in registry`); + } + + try { + const result = await toolDesc.tool.execute(pending.args, {} as any); + + messages.push({ + role: 'tool', + content: JSON.stringify(result), + name: pending.toolName, + toolCallId: pending.toolCallId, + toolName: pending.toolName, + }); + } catch (error) { + // Mirror AgentExecutor.executeToolCall's handling of a thrown tool + // error: turn it into a graceful tool-result message instead of + // letting it reject this promise. By this point the pending-approval + // record has already been deleted (ApprovalGate.resolve() is + // delete-on-read), so failing to catch here would mean the whole + // resume just fails with no retry path. Uses the same + // `(error as Error).message` extraction AgentExecutor's catch block + // uses, wrapped in the `{error}`-shaped payload this file's own + // rejection branch (below) already uses for non-approved decisions. + messages.push({ + role: 'tool', + content: JSON.stringify({ error: (error as Error).message }), + name: pending.toolName, + toolCallId: pending.toolCallId, + toolName: pending.toolName, + }); + } + } else { + messages.push({ + role: 'tool', + content: JSON.stringify({ + error: 'Tool execution was rejected by the reviewer', + note: decision.note, + }), + name: pending.toolName, + toolCallId: pending.toolCallId, + toolName: pending.toolName, + }); + } + + return AgentExecutor.execute({ + ...executeOptions, + agent: snapshot.agent, + input: messages, + provider, + toolRegistry, + // Belt-and-suspenders runtime guard backing the ResumeExecuteOptions + // type-level Omit above: `Omit<...>` only stops *type-checked* callers + // from passing sessionId/checkpointStore through executeOptions - it + // does nothing to stop a caller who bypasses the type (e.g. an `as + // any` cast, or a plain-JS caller) from putting them on the object at + // runtime, where `...executeOptions` would otherwise spread them + // straight into this call. Explicitly forcing both to undefined here, + // after the spread, guarantees AgentExecutor.execute() can never fall + // into its checkpoint-rehydration branch for a resume, regardless of + // what executeOptions actually contains. + sessionId: undefined, + checkpointStore: undefined, + // `messages` was reconstructed from the ExecutionSnapshot's + // currentMessages, which already include the original system message + // (if any) that AgentExecutor.buildMessages() built the first time + // this agent ran. Since no checkpoint/sessionId is threaded through + // this call, execute() would otherwise fall into its "build from + // scratch" path and prepend a second, duplicate system message built + // fresh from agent.prompt. + skipSystemPromptInjection: true, + // Continue step-budget accounting from where the paused run left off, + // rather than silently resetting to a full fresh maxSteps allowance. + // snapshot.steps is the step count AgentExecutor.execute() had already + // reached (see ExecutionSnapshot) at the moment it paused for approval. + initialSteps: snapshot.steps, + }); +} diff --git a/src/types/tool.ts b/src/types/tool.ts index 0b6ae3c77..f31246911 100644 --- a/src/types/tool.ts +++ b/src/types/tool.ts @@ -34,6 +34,7 @@ export interface ToolConfiguration { export interface ToolDescriptor { displayName: string; tool: AITool; + needsApproval?: boolean | ((args: any) => boolean | Promise); injectStreamingController?: (controller: ReadableStreamDefaultController) => void; }