From f0afff68c40b1707d11686d5892097d8c52039aa Mon Sep 17 00:00:00 2001 From: Syngnat Date: Tue, 9 Jun 2026 21:18:39 +0800 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(ai-chat):=20?= =?UTF-8?q?=E6=8B=86=E5=88=86=E6=9C=AC=E5=9C=B0=E5=B7=A5=E5=85=B7=E8=B0=83?= =?UTF-8?q?=E7=94=A8=E9=93=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 抽出 useAIChatLocalTools 承载工具执行、熔断和回灌模型逻辑 - 补齐重试消息的工具上下文依赖，避免配置变更后使用旧闭包 - 增加 hook 行为测试并同步 MCP 指南断言 --- .../AIChatPanel.message-boundary.test.tsx | 20 +- frontend/src/components/AIChatPanel.tsx | 238 ++++------------- .../components/ai/aiLocalToolExecutor.test.ts | 3 +- .../ai/useAIChatLocalTools.test.tsx | 194 ++++++++++++++ .../src/components/ai/useAIChatLocalTools.ts | 250 ++++++++++++++++++ 5 files changed, 505 insertions(+), 200 deletions(-) create mode 100644 frontend/src/components/ai/useAIChatLocalTools.test.tsx create mode 100644 frontend/src/components/ai/useAIChatLocalTools.ts diff --git a/frontend/src/components/AIChatPanel.message-boundary.test.tsx b/frontend/src/components/AIChatPanel.message-boundary.test.tsx index 9e42bd3..5ff0c9e 100644 --- a/frontend/src/components/AIChatPanel.message-boundary.test.tsx +++ b/frontend/src/components/AIChatPanel.message-boundary.test.tsx @@ -12,6 +12,7 @@ const resizeSource = readFileSync(new URL('./ai/useAIChatPanelResize.ts', import const runtimeResourcesSource = readFileSync(new URL('./ai/useAIChatRuntimeResources.ts', import.meta.url), 'utf8'); const sessionStateSource = readFileSync(new URL('./ai/useAIChatSessionState.ts', import.meta.url), 'utf8'); const titleGeneratorSource = readFileSync(new URL('./ai/useAIChatSessionTitleGenerator.ts', import.meta.url), 'utf8'); +const localToolsSource = readFileSync(new URL('./ai/useAIChatLocalTools.ts', import.meta.url), 'utf8'); const streamSubscriptionSource = readFileSync(new URL('./ai/useAIChatStreamSubscription.ts', import.meta.url), 'utf8'); const inspectionGuidanceSource = readFileSync(new URL('./ai/aiSystemInspectionGuidance.ts', import.meta.url), 'utf8'); const systemContextSource = readFileSync(new URL('./ai/aiSystemContextMessages.ts', import.meta.url), 'utf8'); @@ -41,7 +42,8 @@ describe('AIChatPanel message render isolation', () => { it('loads MCP tools and skills into the runtime tool chain', () => { expect(runtimeResourcesSource).toContain('AIListMCPTools'); expect(runtimeResourcesSource).toContain('AIGetSkills'); - expect(source).toContain('executeLocalAIToolCall'); + expect(source).toContain("import { useAIChatLocalTools } from './ai/useAIChatLocalTools';"); + expect(localToolsSource).toContain('executeLocalAIToolCall'); expect(systemContextSource).toContain('以下是当前启用的 Skill'); expect(source).toContain('buildAvailableAIChatTools'); }); @@ -53,11 +55,11 @@ describe('AIChatPanel message render isolation', () => { expect(inspectionGuidanceSource).toContain('inspect_current_connection'); expect(inspectionGuidanceSource).toContain('inspect_external_sql_directories'); expect(inspectionGuidanceSource).toContain('inspect_external_sql_file'); - expect(source).toContain('tabs: useStore.getState().tabs'); - expect(source).toContain('activeTabId: useStore.getState().activeTabId'); - expect(source).toContain('externalSQLDirectories: useStore.getState().externalSQLDirectories'); - expect(source).toContain('toolContextMap: toolContextMapRef.current'); - expect(source).toContain('buildToolResultMessage'); + expect(localToolsSource).toContain('tabs: currentState.tabs'); + expect(localToolsSource).toContain('activeTabId: currentState.activeTabId'); + expect(localToolsSource).toContain('externalSQLDirectories: currentState.externalSQLDirectories'); + expect(localToolsSource).toContain('toolContextMap: toolContextMapRef.current'); + expect(localToolsSource).toContain('buildToolResultMessage'); }); it('extracts chat runtime helpers so context compression and error cleanup stay out of the panel file', () => { @@ -65,11 +67,14 @@ describe('AIChatPanel message render isolation', () => { expect(source).toContain("import { useAIChatStreamSubscription } from './ai/useAIChatStreamSubscription';"); expect(source).toContain('compressContextIfNeeded, getDynamicMaxContextChars'); expect(source).toContain('useAIChatStreamSubscription({'); + expect(source).toContain('useAIChatLocalTools({'); expect(runtimeSource).toContain('export const getDynamicMaxContextChars'); expect(runtimeSource).toContain('export const compressContextIfNeeded'); expect(runtimeSource).toContain('export const sanitizeErrorMsg'); expect(payloadDispatchSource).toContain('export const dispatchAIChatPayload'); expect(payloadDispatchSource).toContain('sanitizeErrorMsg'); + expect(localToolsSource).toContain('compressContextIfNeeded'); + expect(localToolsSource).toContain('dispatchAIChatPayload'); expect(streamSubscriptionSource).toContain('EventsOn(eventName, handler);'); expect(streamSubscriptionSource).toContain('请直接使用 function call 调用工具执行操作'); expect(streamSubscriptionSource).toContain('executeLocalTools(existing.tool_calls!, doneAssistantId)'); @@ -92,6 +97,7 @@ describe('AIChatPanel message render isolation', () => { expect(source).toContain("import { useAIChatAutoContext } from './ai/useAIChatAutoContext';"); expect(source).toContain("import { useAIChatSessionTitleGenerator } from './ai/useAIChatSessionTitleGenerator';"); expect(source).toContain("import { useAIChatPanelResize } from './ai/useAIChatPanelResize';"); + expect(source).toContain("import { useAIChatLocalTools } from './ai/useAIChatLocalTools';"); expect(planContextSource).toContain('export const useAIChatPlanContexts'); expect(planContextSource).toContain('pendingJVMPlanContextRef'); expect(autoContextSource).toContain('export const useAIChatAutoContext'); @@ -100,5 +106,7 @@ describe('AIChatPanel message render isolation', () => { expect(titleGeneratorSource).toContain('Failed to auto-generate title'); expect(resizeSource).toContain('export const useAIChatPanelResize'); expect(resizeSource).toContain('document.body.style.pointerEvents = \'none\''); + expect(localToolsSource).toContain('export const useAIChatLocalTools'); + expect(localToolsSource).toContain('MAX_TOOL_CALL_ROUNDS'); }); }); diff --git a/frontend/src/components/AIChatPanel.tsx b/frontend/src/components/AIChatPanel.tsx index 834e6d5..caf39d4 100644 --- a/frontend/src/components/AIChatPanel.tsx +++ b/frontend/src/components/AIChatPanel.tsx @@ -4,7 +4,6 @@ import { useStore } from '../store'; import type { OverlayWorkbenchTheme } from '../utils/overlayWorkbenchTheme'; import type { AIChatMessage, - AIToolCall, JVMAIPlanContext, JVMDiagnosticPlanContext, } from '../types'; @@ -27,11 +26,6 @@ import { compressContextIfNeeded, getDynamicMaxContextChars } from '../utils/aiC import { getShortcutPlatform, resolveShortcutBinding } from '../utils/shortcuts'; import { isMacLikePlatform } from '../utils/appearance'; import { buildAvailableAIChatTools } from '../utils/aiToolRegistry'; -import { - buildToolResultMessage, - executeLocalAIToolCall, - type AIToolContextEntry, -} from './ai/aiLocalToolExecutor'; import { buildAIChatInlineHistorySessions, buildAIChatInsights, @@ -49,6 +43,7 @@ import { useAIChatPanelResize } from './ai/useAIChatPanelResize'; import { useAIChatPlanContexts } from './ai/useAIChatPlanContexts'; import { useAIChatSessionState } from './ai/useAIChatSessionState'; import { useAIChatSessionTitleGenerator } from './ai/useAIChatSessionTitleGenerator'; +import { useAIChatLocalTools } from './ai/useAIChatLocalTools'; interface AIChatPanelProps { width?: number; @@ -88,8 +83,6 @@ export const AIChatPanel: React.FC = ({ const messagesEndRef = useRef(null); const textareaRef = useRef(null); - const toolCallRoundRef = useRef(0); // 连续失败轮次计数 - const totalToolRoundRef = useRef(0); // 全局工具调用总轮次计数（防止无限循环） const nudgeCountRef = useRef(0); // 催促模型使用 function call 的次数 const { getCurrentJVMPlanContext, @@ -219,6 +212,45 @@ export const AIChatPanel: React.FC = ({ setTimeout(() => textareaRef.current?.focus(), 50); }, [sid, truncateAIChatMessages, deleteAIChatMessage]); + const buildSystemContextMessages = useCallback(( + overrideJVMPlanContext?: JVMAIPlanContext, + overrideJVMDiagnosticPlanContext?: JVMDiagnosticPlanContext, + ) => { + const { activeContext, aiContexts, connections, tabs, activeTabId } = useStore.getState(); + return buildAISystemContextMessages({ + activeContext, + aiContexts, + connections, + tabs, + activeTabId, + availableToolNames: availableTools.map((tool) => tool.function.name), + skills, + userPromptSettings, + overrideJVMPlanContext, + overrideJVMDiagnosticPlanContext, + }); + }, [availableTools, skills, userPromptSettings]); + + const { + executeLocalTools, + resetToolCallState, + toolContextMapRef, + } = useAIChatLocalTools({ + sid, + activeProviderModel: activeProvider?.model, + availableTools, + buildSystemContextMessages, + dynamicModels, + mcpTools, + nextMessageId: genId, + pendingJVMPlanContextRef, + pendingJVMDiagnosticPlanContextRef, + setSending, + skills, + updateAIChatMessage, + userPromptSettings, + }); + const handleRetryMessage = useCallback(async (msg: AIChatMessage) => { const historyLocal = useStore.getState().aiChatHistory[sid] || []; const aiIndex = historyLocal.findIndex(m => m.id === msg.id); @@ -236,9 +268,7 @@ export const AIChatPanel: React.FC = ({ const userMsg = historyLocal[lastUserMsgIndex]; truncateAIChatMessages(sid, userMsg.id); - // 重置计数器（与 handleSend 保持一致） - toolCallRoundRef.current = 0; - totalToolRoundRef.current = 0; + resetToolCallState(); nudgeCountRef.current = 0; const retryJVMPlanContext = msg.jvmPlanContext || getCurrentJVMPlanContext(); const retryJVMDiagnosticPlanContext = @@ -282,192 +312,15 @@ export const AIChatPanel: React.FC = ({ } }, [ sid, + availableTools, + buildSystemContextMessages, truncateAIChatMessages, addAIChatMessage, getCurrentJVMPlanContext, getCurrentJVMDiagnosticPlanContext, + resetToolCallState, ]); - const buildSystemContextMessages = useCallback(( - overrideJVMPlanContext?: JVMAIPlanContext, - overrideJVMDiagnosticPlanContext?: JVMDiagnosticPlanContext, - ) => { - const { activeContext, aiContexts, connections, tabs, activeTabId } = useStore.getState(); - return buildAISystemContextMessages({ - activeContext, - aiContexts, - connections, - tabs, - activeTabId, - availableToolNames: availableTools.map((tool) => tool.function.name), - skills, - userPromptSettings, - overrideJVMPlanContext, - overrideJVMDiagnosticPlanContext, - }); - }, [availableTools, skills, userPromptSettings]); - - // 记录所有成功的 get_tables 调用结果，用于表级精确匹配 - const toolContextMapRef = useRef>(new Map()); - - const executeLocalTools = useCallback(async (toolCalls: AIToolCall[], currentAsstMsgId: string) => { - const currentAsstMsg = (useStore.getState().aiChatHistory[sid] || []).find(m => m.id === currentAsstMsgId); - const inheritedJVMPlanContext = currentAsstMsg?.jvmPlanContext || pendingJVMPlanContextRef.current; - const inheritedJVMDiagnosticPlanContext = - currentAsstMsg?.jvmDiagnosticPlanContext || pendingJVMDiagnosticPlanContextRef.current; - pendingJVMPlanContextRef.current = inheritedJVMPlanContext; - pendingJVMDiagnosticPlanContextRef.current = inheritedJVMDiagnosticPlanContext; - - // 【全局轮次熔断】防止模型（如 DeepSeek）在已生成答案后仍无限循环调用工具 - const MAX_TOOL_CALL_ROUNDS = 15; - totalToolRoundRef.current += 1; - if (totalToolRoundRef.current > MAX_TOOL_CALL_ROUNDS) { - updateAIChatMessage(sid, currentAsstMsgId, { loading: false, phase: 'idle' }); - useStore.getState().addAIChatMessage(sid, { - id: genId(), role: 'assistant', - content: `⚠️ 工具调用已达 ${MAX_TOOL_CALL_ROUNDS} 轮上限，自动终止循环。如需继续探索，请发送新的消息。`, - timestamp: Date.now(), - jvmPlanContext: inheritedJVMPlanContext, - jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, - }); - setSending(false); - return; - } - - const results: AIChatMessage[] = []; - const currentConnections = useStore.getState().connections; - // 【串行逐条执行 + 实时写入 store】 - for (const tc of toolCalls) { - const execution = await executeLocalAIToolCall({ - toolCall: tc, - connections: currentConnections, - activeContext: useStore.getState().activeContext, - aiContexts: useStore.getState().aiContexts, - aiChatHistory: useStore.getState().aiChatHistory, - aiChatSessions: useStore.getState().aiChatSessions, - activeSessionId: sid, - tabs: useStore.getState().tabs, - activeTabId: useStore.getState().activeTabId, - mcpTools, - toolContextMap: toolContextMapRef.current, - sqlLogs: useStore.getState().sqlLogs, - savedQueries: useStore.getState().savedQueries, - sqlSnippets: useStore.getState().sqlSnippets, - externalSQLDirectories: useStore.getState().externalSQLDirectories, - skills, - userPromptSettings, - dynamicModels, - }); - const toolResultMsg: AIChatMessage = buildToolResultMessage({ - id: genId(), - timestamp: Date.now(), - toolCall: tc, - execution, - }); - results.push(toolResultMsg); - - // 【实时写入】每执行完一条立即写入 store，让 UI 能实时看到进度打勾 - useStore.getState().addAIChatMessage(sid, toolResultMsg); - - // 延迟 150ms，给 UI 渲染时间，创造“逐个完成”的视觉节奏 - await new Promise(resolve => setTimeout(resolve, 150)); - } - - // 智能熔断：只计连续失败轮次，成功则重置 - const anySuccess = results.some(r => r.success === true); - if (anySuccess) { - toolCallRoundRef.current = 0; - } else { - toolCallRoundRef.current += 1; - if (toolCallRoundRef.current >= 3) { - useStore.getState().addAIChatMessage(sid, { - id: genId(), role: 'assistant', - content: '⚠️ 探针连续 3 轮执行失败，自动终止。请检查连接状态后重试。', - timestamp: Date.now(), - jvmPlanContext: inheritedJVMPlanContext, - jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, - }); - setSending(false); - return; - } - } - try { - // 【过渡状态】工具执行完毕，将上一条消息的 loading 关闭（消除闪烁光标） - updateAIChatMessage(sid, currentAsstMsgId, { loading: false, phase: 'idle' }); - - // 插入过渡气泡 - const chainConnectingMsg: AIChatMessage = { - id: genId(), role: 'assistant', phase: 'connecting', - content: '汇总探针执行结果中', - timestamp: Date.now(), loading: true, - jvmPlanContext: inheritedJVMPlanContext, - jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, - }; - useStore.getState().addAIChatMessage(sid, chainConnectingMsg); - - // 模拟人类视角的平滑多段过渡 - const safeUpdateTransition = (text: string) => { - const currentMsg = useStore.getState().aiChatHistory[sid]?.find(m => m.id === chainConnectingMsg.id); - // 只有当消息仍然处于连接过渡态时才允许修改文本；如果模型已经开始吐出思考、正文、工具或结束，直接退出 - if (currentMsg && currentMsg.phase === 'connecting' && currentMsg.loading) { - updateAIChatMessage(sid, chainConnectingMsg.id, { content: text }); - } - }; - - setTimeout(() => safeUpdateTransition('向模型回传运行时数据'), 200); - setTimeout(() => safeUpdateTransition('模型大脑深度推理中'), 500); - setTimeout(() => safeUpdateTransition('等待下发操作指令'), 1200); - setTimeout(() => safeUpdateTransition('正在深度思考链路与逻辑'), 3000); - - setSending(true); - const currentHistory = useStore.getState().aiChatHistory[sid] || []; - // 过滤掉 connecting 占位消息，不发给模型 - const messagesPayload = currentHistory.filter(m => m.phase !== 'connecting').map(toAIRequestMessage); - const sysMessages = await buildSystemContextMessages( - inheritedJVMPlanContext, - inheritedJVMDiagnosticPlanContext, - ); - - let finalMessagesPayload = messagesPayload; - // 在这里加入长度检查和自动摘要（带上动态限额） - const dynamicMaxLimit = getDynamicMaxContextChars(activeProvider?.model); - const summary = await compressContextIfNeeded(sid, messagesPayload, dynamicMaxLimit); - if (summary) { - const compressedMsg: AIChatMessage = { - id: genId(), role: 'assistant', content: `【自动记忆重塑】已将超长历史探针数据和对话压缩为摘要：\n\n${summary}`, timestamp: Date.now() - 1000 - }; - const continueMsg: AIChatMessage = { - id: genId(), role: 'user', content: '请根据上述最新状态与探索结果，继续完成你先前未竟的分析或执行下一步。', timestamp: Date.now() - 500 - }; - useStore.getState().replaceAIChatHistory(sid, [compressedMsg, continueMsg, chainConnectingMsg]); - finalMessagesPayload = [ - { role: 'assistant', content: compressedMsg.content }, - { role: 'user', content: continueMsg.content } - ]; - } - - const allMessages = [...sysMessages, ...finalMessagesPayload]; - - // 【软收敛】超过 10 轮工具调用后，不再传递 tools 参数，从物理层面强制模型只能用文本回答 - const SOFT_LIMIT_ROUNDS = 10; - const chainTools = totalToolRoundRef.current >= SOFT_LIMIT_ROUNDS ? [] : availableTools; - - await dispatchAIChatPayload({ - sid, - messages: allMessages, - tools: chainTools, - addAIChatMessage: (sessionId, message) => useStore.getState().addAIChatMessage(sessionId, message), - setSending, - nextMessageId: genId, - jvmPlanContext: inheritedJVMPlanContext, - jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, - }); - } catch (e) { - console.error('Failed to chain tool call', e); - setSending(false); - } - }, [availableTools, buildSystemContextMessages, dynamicModels, mcpTools, sid, skills]); - useAIChatStreamSubscription({ sid, sending, @@ -512,8 +365,7 @@ export const AIChatPanel: React.FC = ({ } setComposerNotice(null); - toolCallRoundRef.current = 0; // 重置工具调用轮次计数 - totalToolRoundRef.current = 0; // 重置总轮次计数 + resetToolCallState(); nudgeCountRef.current = 0; // 重置催促计数 const currentJVMPlanContext = getCurrentJVMPlanContext(); const currentJVMDiagnosticPlanContext = getCurrentJVMDiagnosticPlanContext(); diff --git a/frontend/src/components/ai/aiLocalToolExecutor.test.ts b/frontend/src/components/ai/aiLocalToolExecutor.test.ts index 8d55977..8586be9 100644 --- a/frontend/src/components/ai/aiLocalToolExecutor.test.ts +++ b/frontend/src/components/ai/aiLocalToolExecutor.test.ts @@ -445,9 +445,10 @@ describe('aiLocalToolExecutor', () => { expect(result.success).toBe(true); expect(result.content).toContain('"supportsWholeCommandAutoSplit":true'); - expect(result.content).toContain('"fullCommandPasteExample":"OPENAI_API_KEY=... uvx mcp-server-fetch --stdio"'); + expect(result.content).toContain('"fullCommandPasteExample":"$env:GITHUB_TOKEN=...; uvx mcp-server-github --stdio"'); expect(result.content).toContain('"title":"启动命令"'); expect(result.content).toContain('"example":"node / uvx / python"'); + expect(result.content).toContain('PowerShell $env:KEY=VALUE;'); expect(result.content).toContain('"title":"uvx 工具"'); expect(result.content).toContain('"exampleLaunchPreview":"uvx some-mcp-server"'); }); diff --git a/frontend/src/components/ai/useAIChatLocalTools.test.tsx b/frontend/src/components/ai/useAIChatLocalTools.test.tsx new file mode 100644 index 0000000..ef69374 --- /dev/null +++ b/frontend/src/components/ai/useAIChatLocalTools.test.tsx @@ -0,0 +1,194 @@ +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(); + }); + }); +}); diff --git a/frontend/src/components/ai/useAIChatLocalTools.ts b/frontend/src/components/ai/useAIChatLocalTools.ts new file mode 100644 index 0000000..926c59b --- /dev/null +++ b/frontend/src/components/ai/useAIChatLocalTools.ts @@ -0,0 +1,250 @@ +import { useCallback, useRef } from 'react'; +import type { MutableRefObject } from 'react'; + +import { useStore } from '../../store'; +import type { + AIChatMessage, + AIMCPToolDescriptor, + AISkillConfig, + AIToolCall, + AIUserPromptSettings, + JVMAIPlanContext, + JVMDiagnosticPlanContext, +} from '../../types'; +import { compressContextIfNeeded, getDynamicMaxContextChars } from '../../utils/aiChatRuntime'; +import { toAIRequestMessage } from '../../utils/aiMessagePayload'; +import type { AIChatToolDefinition } from '../../utils/aiToolRegistry'; +import { dispatchAIChatPayload } from './aiChatPayloadDispatch'; +import { + buildToolResultMessage, + executeLocalAIToolCall, + type AIToolContextEntry, +} from './aiLocalToolExecutor'; + +interface UseAIChatLocalToolsOptions { + sid: string; + activeProviderModel?: string; + availableTools: AIChatToolDefinition[]; + buildSystemContextMessages: ( + overrideJVMPlanContext?: JVMAIPlanContext, + overrideJVMDiagnosticPlanContext?: JVMDiagnosticPlanContext, + ) => any[] | Promise; + dynamicModels: string[]; + mcpTools: AIMCPToolDescriptor[]; + nextMessageId: () => string; + pendingJVMPlanContextRef: MutableRefObject; + pendingJVMDiagnosticPlanContextRef: MutableRefObject; + setSending: (sending: boolean) => void; + skills: AISkillConfig[]; + updateAIChatMessage: ( + sid: string, + messageId: string, + patch: Partial, + ) => void; + userPromptSettings: AIUserPromptSettings; +} + +const MAX_TOOL_CALL_ROUNDS = 15; +const SOFT_LIMIT_ROUNDS = 10; + +export const useAIChatLocalTools = ({ + sid, + activeProviderModel, + availableTools, + buildSystemContextMessages, + dynamicModels, + mcpTools, + nextMessageId, + pendingJVMPlanContextRef, + pendingJVMDiagnosticPlanContextRef, + setSending, + skills, + updateAIChatMessage, + userPromptSettings, +}: UseAIChatLocalToolsOptions) => { + const toolCallRoundRef = useRef(0); + const totalToolRoundRef = useRef(0); + const toolContextMapRef = useRef>(new Map()); + + const resetToolCallState = useCallback(() => { + toolCallRoundRef.current = 0; + totalToolRoundRef.current = 0; + }, []); + + const executeLocalTools = useCallback(async (toolCalls: AIToolCall[], currentAsstMsgId: string) => { + const store = useStore.getState(); + const currentAsstMsg = (store.aiChatHistory[sid] || []).find((message) => message.id === currentAsstMsgId); + const inheritedJVMPlanContext = currentAsstMsg?.jvmPlanContext || pendingJVMPlanContextRef.current; + const inheritedJVMDiagnosticPlanContext = + currentAsstMsg?.jvmDiagnosticPlanContext || pendingJVMDiagnosticPlanContextRef.current; + pendingJVMPlanContextRef.current = inheritedJVMPlanContext; + pendingJVMDiagnosticPlanContextRef.current = inheritedJVMDiagnosticPlanContext; + + totalToolRoundRef.current += 1; + if (totalToolRoundRef.current > MAX_TOOL_CALL_ROUNDS) { + updateAIChatMessage(sid, currentAsstMsgId, { loading: false, phase: 'idle' }); + useStore.getState().addAIChatMessage(sid, { + id: nextMessageId(), + role: 'assistant', + content: `⚠️ 工具调用已达 ${MAX_TOOL_CALL_ROUNDS} 轮上限，自动终止循环。如需继续探索，请发送新的消息。`, + timestamp: Date.now(), + jvmPlanContext: inheritedJVMPlanContext, + jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, + }); + setSending(false); + return; + } + + const results: AIChatMessage[] = []; + const currentConnections = useStore.getState().connections; + for (const toolCall of toolCalls) { + const currentState = useStore.getState(); + const execution = await executeLocalAIToolCall({ + toolCall, + connections: currentConnections, + activeContext: currentState.activeContext, + aiContexts: currentState.aiContexts, + aiChatHistory: currentState.aiChatHistory, + aiChatSessions: currentState.aiChatSessions, + activeSessionId: sid, + tabs: currentState.tabs, + activeTabId: currentState.activeTabId, + mcpTools, + toolContextMap: toolContextMapRef.current, + sqlLogs: currentState.sqlLogs, + savedQueries: currentState.savedQueries, + sqlSnippets: currentState.sqlSnippets, + externalSQLDirectories: currentState.externalSQLDirectories, + skills, + userPromptSettings, + dynamicModels, + }); + const toolResultMsg: AIChatMessage = buildToolResultMessage({ + id: nextMessageId(), + timestamp: Date.now(), + toolCall, + execution, + }); + results.push(toolResultMsg); + useStore.getState().addAIChatMessage(sid, toolResultMsg); + await new Promise((resolve) => setTimeout(resolve, 150)); + } + + const anySuccess = results.some((message) => message.success === true); + if (anySuccess) { + toolCallRoundRef.current = 0; + } else { + toolCallRoundRef.current += 1; + if (toolCallRoundRef.current >= 3) { + useStore.getState().addAIChatMessage(sid, { + id: nextMessageId(), + role: 'assistant', + content: '⚠️ 探针连续 3 轮执行失败，自动终止。请检查连接状态后重试。', + timestamp: Date.now(), + jvmPlanContext: inheritedJVMPlanContext, + jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, + }); + setSending(false); + return; + } + } + + try { + updateAIChatMessage(sid, currentAsstMsgId, { loading: false, phase: 'idle' }); + + const chainConnectingMsg: AIChatMessage = { + id: nextMessageId(), + role: 'assistant', + phase: 'connecting', + content: '汇总探针执行结果中', + timestamp: Date.now(), + loading: true, + jvmPlanContext: inheritedJVMPlanContext, + jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, + }; + useStore.getState().addAIChatMessage(sid, chainConnectingMsg); + + const safeUpdateTransition = (text: string) => { + const currentMsg = useStore.getState().aiChatHistory[sid]?.find((message) => message.id === chainConnectingMsg.id); + if (currentMsg && currentMsg.phase === 'connecting' && currentMsg.loading) { + updateAIChatMessage(sid, chainConnectingMsg.id, { content: text }); + } + }; + + setTimeout(() => safeUpdateTransition('向模型回传运行时数据'), 200); + setTimeout(() => safeUpdateTransition('模型大脑深度推理中'), 500); + setTimeout(() => safeUpdateTransition('等待下发操作指令'), 1200); + setTimeout(() => safeUpdateTransition('正在深度思考链路与逻辑'), 3000); + + setSending(true); + const currentHistory = useStore.getState().aiChatHistory[sid] || []; + const messagesPayload = currentHistory + .filter((message) => message.phase !== 'connecting') + .map(toAIRequestMessage); + const sysMessages = await buildSystemContextMessages( + inheritedJVMPlanContext, + inheritedJVMDiagnosticPlanContext, + ); + + let finalMessagesPayload = messagesPayload; + const dynamicMaxLimit = getDynamicMaxContextChars(activeProviderModel); + const summary = await compressContextIfNeeded(sid, messagesPayload, dynamicMaxLimit); + if (summary) { + const compressedMsg: AIChatMessage = { + id: nextMessageId(), + role: 'assistant', + content: `【自动记忆重塑】已将超长历史探针数据和对话压缩为摘要：\n\n${summary}`, + timestamp: Date.now() - 1000, + }; + const continueMsg: AIChatMessage = { + id: nextMessageId(), + role: 'user', + content: '请根据上述最新状态与探索结果，继续完成你先前未竟的分析或执行下一步。', + timestamp: Date.now() - 500, + }; + useStore.getState().replaceAIChatHistory(sid, [compressedMsg, continueMsg, chainConnectingMsg]); + finalMessagesPayload = [ + { role: 'assistant', content: compressedMsg.content }, + { role: 'user', content: continueMsg.content }, + ]; + } + + const allMessages = [...sysMessages, ...finalMessagesPayload]; + const chainTools = totalToolRoundRef.current >= SOFT_LIMIT_ROUNDS ? [] : availableTools; + + await dispatchAIChatPayload({ + sid, + messages: allMessages, + tools: chainTools, + addAIChatMessage: (sessionId, message) => useStore.getState().addAIChatMessage(sessionId, message), + setSending, + nextMessageId, + jvmPlanContext: inheritedJVMPlanContext, + jvmDiagnosticPlanContext: inheritedJVMDiagnosticPlanContext, + }); + } catch (error) { + console.error('Failed to chain tool call', error); + setSending(false); + } + }, [ + activeProviderModel, + availableTools, + buildSystemContextMessages, + dynamicModels, + mcpTools, + nextMessageId, + pendingJVMDiagnosticPlanContextRef, + pendingJVMPlanContextRef, + setSending, + sid, + skills, + updateAIChatMessage, + userPromptSettings, + ]); + + return { + executeLocalTools, + resetToolCallState, + toolContextMapRef, + }; +};