From 1bc95470599d5ddcb09c69712de7e3c8feae725f Mon Sep 17 00:00:00 2001 From: ONLY-yours <1349021570@qq.com> Date: Tue, 21 Apr 2026 14:44:50 +0800 Subject: [PATCH] fix: add the agent runtimes hooks test --- .../modules/AgentRuntime/RuntimeExecutors.ts | 33 ++- .../__tests__/RuntimeExecutors.test.ts | 227 ++++++++++++++++++ .../hooks/__tests__/HookDispatcher.test.ts | 144 +++++++++++ 3 files changed, 387 insertions(+), 17 deletions(-) diff --git a/src/server/modules/AgentRuntime/RuntimeExecutors.ts b/src/server/modules/AgentRuntime/RuntimeExecutors.ts index cc582a985d..6239e9c335 100644 --- a/src/server/modules/AgentRuntime/RuntimeExecutors.ts +++ b/src/server/modules/AgentRuntime/RuntimeExecutors.ts @@ -1404,12 +1404,24 @@ export const createRuntimeExecutors = ( type: 'tool_start', }); + // Extract before try so catch block can access for hook dispatch + const chatToolPayload: ChatToolPayload = payload.toolCalling; + const toolName = `${chatToolPayload.identifier}/${chatToolPayload.apiName}`; + + // Track tool call count for hooks (before try so catch can access) + const toolKey = `${chatToolPayload.identifier}/${chatToolPayload.apiName}`; + const callIndex = (toolCallCounts.get(toolKey) ?? 0) + 1; + toolCallCounts.set(toolKey, callIndex); + + let parsedArgs: Record = {}; try { - // payload is { parentMessageId, toolCalling: ChatToolPayload } - const chatToolPayload: ChatToolPayload = payload.toolCalling; - - const toolName = `${chatToolPayload.identifier}/${chatToolPayload.apiName}`; + parsedArgs = + typeof chatToolPayload.arguments === 'string' + ? JSON.parse(chatToolPayload.arguments) + : (chatToolPayload.arguments ?? {}); + } catch {} + try { // Check if this is a client-side function tool — pause instead of executing const toolSource = state.operationToolSet?.sourceMap?.[chatToolPayload.identifier] ?? @@ -1470,19 +1482,6 @@ export const createRuntimeExecutors = ( chatToolPayload.executor === 'client' && typeof streamManager.sendToolExecute === 'function'; - // Track tool call count for hooks - const toolKey = `${chatToolPayload.identifier}/${chatToolPayload.apiName}`; - const callIndex = (toolCallCounts.get(toolKey) ?? 0) + 1; - toolCallCounts.set(toolKey, callIndex); - - let parsedArgs: Record = {}; - try { - parsedArgs = - typeof chatToolPayload.arguments === 'string' - ? JSON.parse(chatToolPayload.arguments) - : (chatToolPayload.arguments ?? {}); - } catch {} - // Dispatch beforeToolCall hook (may return mock result) let toolCallMocked = false; let mockResult: { content: string } | null = null; diff --git a/src/server/modules/AgentRuntime/__tests__/RuntimeExecutors.test.ts b/src/server/modules/AgentRuntime/__tests__/RuntimeExecutors.test.ts index 4a1e58ffee..23c5936035 100644 --- a/src/server/modules/AgentRuntime/__tests__/RuntimeExecutors.test.ts +++ b/src/server/modules/AgentRuntime/__tests__/RuntimeExecutors.test.ts @@ -3524,4 +3524,231 @@ describe('RuntimeExecutors', () => { } }); }); + + describe('hooks integration', () => { + const createToolState = (overrides?: Partial): AgentState => ({ + cost: createMockCost(), + createdAt: new Date().toISOString(), + lastModified: new Date().toISOString(), + maxSteps: 100, + messages: [], + metadata: { agentId: 'agent-123', topicId: 'topic-123' }, + operationId: 'op-123', + status: 'running', + stepCount: 0, + toolManifestMap: {}, + usage: createMockUsage(), + ...overrides, + }); + + const createToolInstruction = (overrides?: any) => ({ + payload: { + parentMessageId: 'parent-msg', + toolCalling: { + apiName: 'search_tweets', + arguments: '{"query":"test"}', + id: 'tc-1', + identifier: 'twitter', + type: 'default' as const, + }, + ...overrides, + }, + type: 'call_tool' as const, + }); + + describe('call_tool hooks', () => { + it('should dispatch beforeToolCall and afterToolCall hooks', async () => { + const beforeHandler = vi.fn(); + const afterHandler = vi.fn(); + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue(null), + }; + + const ctxWithHooks = { ...ctx, hookDispatcher: mockDispatcher as any }; + const executors = createRuntimeExecutors(ctxWithHooks); + + await executors.call_tool!(createToolInstruction(), createToolState()); + + expect(mockDispatcher.dispatchBeforeToolCall).toHaveBeenCalledWith( + 'op-123', + expect.objectContaining({ + apiName: 'search_tweets', + callIndex: 1, + identifier: 'twitter', + }), + ); + + // afterToolCall dispatched via dispatch() + expect(mockDispatcher.dispatch).toHaveBeenCalledWith( + 'op-123', + 'afterToolCall', + expect.objectContaining({ + apiName: 'search_tweets', + identifier: 'twitter', + mocked: false, + success: true, + }), + ); + }); + + it('should skip real execution when beforeToolCall returns mock', async () => { + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue({ content: '{"mocked":true}' }), + }; + + const ctxWithHooks = { ...ctx, hookDispatcher: mockDispatcher as any }; + const executors = createRuntimeExecutors(ctxWithHooks); + + const result = await executors.call_tool!(createToolInstruction(), createToolState()); + + // Real tool should NOT have been called + expect(mockToolExecutionService.executeTool).not.toHaveBeenCalled(); + + // afterToolCall should report mocked: true + expect(mockDispatcher.dispatch).toHaveBeenCalledWith( + 'op-123', + 'afterToolCall', + expect.objectContaining({ mocked: true, success: true }), + ); + + // Tool message should be persisted with mock content + expect(mockMessageModel.create).toHaveBeenCalledWith( + expect.objectContaining({ + content: '{"mocked":true}', + role: 'tool', + }), + ); + }); + + it('should dispatch onToolCallError when tool throws', async () => { + mockToolExecutionService.executeTool.mockRejectedValue(new Error('Connection refused')); + + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue(null), + }; + + const ctxWithHooks = { ...ctx, hookDispatcher: mockDispatcher as any }; + const executors = createRuntimeExecutors(ctxWithHooks); + + await executors.call_tool!(createToolInstruction(), createToolState()); + + expect(mockDispatcher.dispatch).toHaveBeenCalledWith( + 'op-123', + 'onToolCallError', + expect.objectContaining({ + apiName: 'search_tweets', + error: 'Connection refused', + identifier: 'twitter', + }), + ); + }); + + it('should increment callIndex across multiple calls', async () => { + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue(null), + }; + + const ctxWithHooks = { ...ctx, hookDispatcher: mockDispatcher as any }; + const executors = createRuntimeExecutors(ctxWithHooks); + const state = createToolState(); + + await executors.call_tool!(createToolInstruction(), state); + await executors.call_tool!(createToolInstruction(), state); + + expect(mockDispatcher.dispatchBeforeToolCall).toHaveBeenNthCalledWith( + 1, + 'op-123', + expect.objectContaining({ callIndex: 1 }), + ); + expect(mockDispatcher.dispatchBeforeToolCall).toHaveBeenNthCalledWith( + 2, + 'op-123', + expect.objectContaining({ callIndex: 2 }), + ); + }); + + it('should work without hookDispatcher (backward compat)', async () => { + const executors = createRuntimeExecutors(ctx); // no hookDispatcher + const result = await executors.call_tool!(createToolInstruction(), createToolState()); + + expect(result).toBeDefined(); + expect(mockToolExecutionService.executeTool).toHaveBeenCalled(); + }); + }); + + describe('compress_context hooks', () => { + it('should dispatch beforeCompact and afterCompact hooks', async () => { + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue(null), + }; + + const ctxWithHooks = { + ...ctx, + hookDispatcher: mockDispatcher as any, + topicId: 'topic-123', + }; + const executors = createRuntimeExecutors(ctxWithHooks); + + const state = createToolState({ metadata: { agentId: 'agent-123', topicId: 'topic-123' } }); + + const instruction = { + payload: { + currentTokenCount: 5000, + messages: [ + { content: 'hello', id: 'msg-1', role: 'user' }, + { content: 'hi there', id: 'msg-2', role: 'assistant' }, + ], + }, + type: 'compress_context' as const, + }; + + await executors.compress_context!(instruction, state); + + expect(mockDispatcher.dispatch).toHaveBeenCalledWith( + 'op-123', + 'beforeCompact', + expect.objectContaining({ tokenCount: 5000 }), + ); + }); + }); + + describe('request_human_approve hooks', () => { + it('should dispatch beforeHumanIntervention hook', async () => { + const mockDispatcher = { + dispatch: vi.fn().mockResolvedValue(undefined), + dispatchBeforeToolCall: vi.fn().mockResolvedValue(null), + }; + + const ctxWithHooks = { ...ctx, hookDispatcher: mockDispatcher as any }; + const executors = createRuntimeExecutors(ctxWithHooks); + + const state = createToolState({ + messages: [{ content: '', id: 'asst-1', role: 'assistant' }], + status: 'running', + }); + + const instruction = { + pendingToolsCalling: [ + { apiName: 'post_tweet', id: 'tc-1', identifier: 'twitter', type: 'default' }, + ], + type: 'request_human_approve' as const, + }; + + await executors.request_human_approve!(instruction, state); + + expect(mockDispatcher.dispatch).toHaveBeenCalledWith( + 'op-123', + 'beforeHumanIntervention', + expect.objectContaining({ + pendingTools: [{ apiName: 'post_tweet', identifier: 'twitter' }], + }), + ); + }); + }); + }); }); diff --git a/src/server/services/agentRuntime/hooks/__tests__/HookDispatcher.test.ts b/src/server/services/agentRuntime/hooks/__tests__/HookDispatcher.test.ts index 60029a3b72..21089c7e4e 100644 --- a/src/server/services/agentRuntime/hooks/__tests__/HookDispatcher.test.ts +++ b/src/server/services/agentRuntime/hooks/__tests__/HookDispatcher.test.ts @@ -567,6 +567,150 @@ describe('HookDispatcher', () => { }); }); + describe('dispatchBeforeToolCall — edge cases', () => { + it('should use the last mock() call when multiple handlers call mock()', async () => { + dispatcher.register(operationId, [ + { + handler: async (event: any) => { + event.mock({ content: '{"first":true}' }); + }, + id: 'mock-1', + type: 'beforeToolCall', + }, + { + handler: async (event: any) => { + event.mock({ content: '{"second":true}' }); + }, + id: 'mock-2', + type: 'beforeToolCall', + }, + ]); + + const result = await dispatcher.dispatchBeforeToolCall(operationId, { + apiName: 'search', + args: {}, + callIndex: 1, + identifier: 'twitter', + stepIndex: 0, + }); + + expect(result).toEqual({ content: '{"second":true}' }); + }); + + it('should return mock when only one of multiple handlers calls mock()', async () => { + const observeHandler = vi.fn(); + dispatcher.register(operationId, [ + { handler: observeHandler, id: 'observe', type: 'beforeToolCall' }, + { + handler: async (event: any) => { + event.mock({ content: '{"mocked":true}' }); + }, + id: 'mocker', + type: 'beforeToolCall', + }, + ]); + + const result = await dispatcher.dispatchBeforeToolCall(operationId, { + apiName: 'search', + args: {}, + callIndex: 1, + identifier: 'twitter', + stepIndex: 0, + }); + + expect(observeHandler).toHaveBeenCalled(); + expect(result).toEqual({ content: '{"mocked":true}' }); + }); + + it('should only mock in local mode, not production mode', async () => { + vi.mocked(isQueueAgentRuntimeEnabled).mockReturnValue(true); + + dispatcher.register(operationId, [ + { + handler: async (event: any) => { + event.mock({ content: '{"mocked":true}' }); + }, + id: 'mock-hook', + type: 'beforeToolCall', + }, + ]); + + // dispatchBeforeToolCall only runs in local mode + const result = await dispatcher.dispatchBeforeToolCall(operationId, { + apiName: 'search', + args: {}, + callIndex: 1, + identifier: 'twitter', + stepIndex: 0, + }); + + // In local mode this would return the mock, but hooks are still in-memory + // so it should still work (dispatchBeforeToolCall doesn't check queue mode) + expect(result).toEqual({ content: '{"mocked":true}' }); + }); + + it('should not affect other hook types when beforeToolCall is registered', async () => { + const afterStepHandler = vi.fn(); + const onCompleteHandler = vi.fn(); + + dispatcher.register(operationId, [ + { + handler: async (event: any) => { + event.mock({ content: 'mock' }); + }, + id: 'tool-mock', + type: 'beforeToolCall', + }, + { handler: afterStepHandler, id: 'after-step', type: 'afterStep' }, + { handler: onCompleteHandler, id: 'complete', type: 'onComplete' }, + ]); + + // beforeToolCall should not trigger afterStep or onComplete + await dispatcher.dispatchBeforeToolCall(operationId, { + apiName: 'search', + args: {}, + callIndex: 1, + identifier: 'twitter', + stepIndex: 0, + }); + + expect(afterStepHandler).not.toHaveBeenCalled(); + expect(onCompleteHandler).not.toHaveBeenCalled(); + + // afterStep should still work independently + await dispatcher.dispatch(operationId, 'afterStep', makeEvent({ stepIndex: 0 })); + expect(afterStepHandler).toHaveBeenCalledTimes(1); + }); + + it('should call handlers even after a previous handler throws', async () => { + const mockHandler = vi.fn().mockImplementation(async (event: any) => { + event.mock({ content: '{"recovered":true}' }); + }); + + dispatcher.register(operationId, [ + { + handler: async () => { + throw new Error('first handler fails'); + }, + id: 'failing', + type: 'beforeToolCall', + }, + { handler: mockHandler, id: 'recovering', type: 'beforeToolCall' }, + ]); + + const result = await dispatcher.dispatchBeforeToolCall(operationId, { + apiName: 'search', + args: {}, + callIndex: 1, + identifier: 'twitter', + stepIndex: 0, + }); + + expect(mockHandler).toHaveBeenCalled(); + expect(result).toEqual({ content: '{"recovered":true}' }); + }); + }); + describe('callAgent hooks', () => { it('should dispatch beforeCallAgent hooks', async () => { const handler = vi.fn();