import React, { useRef, useState } from 'react'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { act, create, type ReactTestRenderer } from 'react-test-renderer'; import { useStore } from '../../store'; import type { AIToolCall } from '../../types'; import { useAIChatLocalTools } from './useAIChatLocalTools'; const dispatchAIChatPayloadMock = vi.hoisted(() => vi.fn(async (_options: any) => 'stream')); const executeLocalAIToolCallMock = vi.hoisted(() => vi.fn(async ({ toolCall }: { toolCall: AIToolCall }) => ({ content: `result:${toolCall.function.name}`, success: true, toolName: toolCall.function.name, }))); vi.mock('./aiChatPayloadDispatch', () => ({ dispatchAIChatPayload: dispatchAIChatPayloadMock, })); vi.mock('./aiLocalToolExecutor', () => ({ executeLocalAIToolCall: executeLocalAIToolCallMock, buildToolResultMessage: ({ id, timestamp, toolCall, execution }: any) => ({ id, role: 'tool', content: execution.content, timestamp, tool_call_id: toolCall.id, tool_name: execution.toolName, success: execution.success, }), })); const SESSION_ID = 'session-local-tools'; const buildToolCall = (name: string): AIToolCall => ({ id: `call-${name}`, type: 'function', function: { name, arguments: '{}', }, }); const updateMessage = ( sessionId: string, messageId: string, patch: Parameters['updateAIChatMessage']>[2], ) => useStore.getState().updateAIChatMessage(sessionId, messageId, patch); let latestHook: ReturnType | undefined; const LocalToolsHarness = () => { const [sending, setSending] = useState(false); const pendingJVMPlanContextRef = useRef(undefined); const pendingJVMDiagnosticPlanContextRef = useRef(undefined); latestHook = useAIChatLocalTools({ sid: SESSION_ID, activeProviderModel: 'gpt-5', availableTools: [{ type: 'function', function: { name: 'inspect_active_tab', description: 'inspect tab', parameters: { type: 'object', properties: {} }, }, }], buildSystemContextMessages: async () => [{ role: 'system', content: 'system-context' }], dynamicModels: ['gpt-5'], mcpTools: [], nextMessageId: () => `generated-${Math.random().toString(36).slice(2, 6)}`, pendingJVMPlanContextRef, pendingJVMDiagnosticPlanContextRef, setSending, skills: [], updateAIChatMessage: updateMessage, userPromptSettings: { global: '', database: '', jvm: '', jvmDiagnostic: '', }, }); return ; }; describe('useAIChatLocalTools', () => { beforeEach(() => { vi.useFakeTimers(); dispatchAIChatPayloadMock.mockClear(); executeLocalAIToolCallMock.mockClear(); latestHook = undefined; useStore.setState({ activeContext: { connectionId: 'conn-1', dbName: 'crm' }, aiChatHistory: { [SESSION_ID]: [ { id: 'user-1', role: 'user', content: '查一下当前页签', timestamp: 1 }, { id: 'assistant-1', role: 'assistant', content: '', timestamp: 2, loading: true, phase: 'tool_calling', tool_calls: [buildToolCall('inspect_active_tab')], }, ], }, aiChatSessions: [{ id: SESSION_ID, title: '查一下当前页签', updatedAt: 1 }], aiActiveSessionId: SESSION_ID, connections: [{ id: 'conn-1', name: '主库', config: { type: 'mysql', host: '127.0.0.1', port: 3306, user: 'root', }, }], tabs: [{ id: 'tab-1', title: '订单查询', type: 'query', connectionId: 'conn-1', dbName: 'crm', query: 'select * from orders', }], activeTabId: 'tab-1', aiContexts: {}, sqlLogs: [], savedQueries: [], sqlSnippets: [], externalSQLDirectories: [], }); }); afterEach(() => { vi.useRealTimers(); useStore.setState({ activeContext: null, aiChatHistory: {}, aiChatSessions: [], aiActiveSessionId: null, tabs: [], activeTabId: null, aiContexts: {}, sqlLogs: [], savedQueries: [], sqlSnippets: [], externalSQLDirectories: [], }); }); it('writes tool results, closes the tool-calling message, and excludes connecting placeholders from the chained request', async () => { let renderer: ReactTestRenderer | undefined; await act(async () => { renderer = create(); }); expect(latestHook).toBeDefined(); const run = latestHook!.executeLocalTools([buildToolCall('inspect_active_tab')], 'assistant-1'); await act(async () => { await vi.advanceTimersByTimeAsync(150); await run; }); const messages = useStore.getState().aiChatHistory[SESSION_ID] || []; const assistant = messages.find((message) => message.id === 'assistant-1'); const toolResult = messages.find((message) => message.role === 'tool'); const connecting = messages.find((message) => message.phase === 'connecting'); expect(executeLocalAIToolCallMock).toHaveBeenCalledTimes(1); expect(assistant).toMatchObject({ loading: false, phase: 'idle' }); expect(toolResult).toMatchObject({ content: 'result:inspect_active_tab', success: true, tool_name: 'inspect_active_tab', }); expect(connecting).toMatchObject({ content: '汇总探针执行结果中', loading: true }); expect(dispatchAIChatPayloadMock).toHaveBeenCalledTimes(1); const dispatchArgs = dispatchAIChatPayloadMock.mock.calls[0][0] as any; expect(dispatchArgs.messages[0]).toEqual({ role: 'system', content: 'system-context' }); expect(JSON.stringify(dispatchArgs.messages)).toContain('result:inspect_active_tab'); expect(JSON.stringify(dispatchArgs.messages)).not.toContain('汇总探针执行结果中'); expect(dispatchArgs.tools).toHaveLength(1); await act(async () => { renderer?.unmount(); }); }); });