From 2d58fd8c62ec9e89fb368e978c7fdd58722b0582 Mon Sep 17 00:00:00 2001 From: digua Date: Tue, 18 Aug 2026 17:18:59 +0800 Subject: [PATCH] fix(ai): track parallel tool lifecycles Keep the shared runtime in tool_running until every concurrent call finishes. Route renderer progress and completion by toolCallId so one result cannot hide another active tool. --- .../ai/agent/__tests__/event-handler.test.ts | 28 +++++ .../src/ai/agent/event-handler.ts | 11 +- src/stores/aiChat.ts | 100 ++++++++++-------- src/stores/aiToolLifecycle.test.ts | 28 +++++ src/stores/aiToolLifecycle.ts | 80 ++++++++++++++ 5 files changed, 202 insertions(+), 45 deletions(-) create mode 100644 src/stores/aiToolLifecycle.test.ts create mode 100644 src/stores/aiToolLifecycle.ts diff --git a/packages/node-runtime/src/ai/agent/__tests__/event-handler.test.ts b/packages/node-runtime/src/ai/agent/__tests__/event-handler.test.ts index 32377387..8363e381 100644 --- a/packages/node-runtime/src/ai/agent/__tests__/event-handler.test.ts +++ b/packages/node-runtime/src/ai/agent/__tests__/event-handler.test.ts @@ -45,6 +45,34 @@ describe('AgentEventHandler', () => { assert.equal(startChunk?.toolCallId, 'call_1') }) + it('keeps tool_running status until every parallel tool call finishes', () => { + const chunks: AgentStreamChunk[] = [] + const handler = new AgentEventHandler({ + onChunk: (chunk) => chunks.push(chunk), + context: {}, + systemPrompt: 'test', + }) + + handler.handleCoreEvent({ type: 'tool_start', toolCallId: 'slow-call', toolName: 'slow_tool', toolParams: {} }, []) + handler.handleCoreEvent({ type: 'tool_start', toolCallId: 'fast-call', toolName: 'fast_tool', toolParams: {} }, []) + handler.handleCoreEvent( + { type: 'tool_end', toolCallId: 'fast-call', toolName: 'fast_tool', toolResult: null, isError: false }, + [] + ) + + const statusesAfterFirstResult = chunks.filter((chunk) => chunk.type === 'status') + assert.equal(statusesAfterFirstResult.at(-1)?.status?.phase, 'tool_running') + assert.equal(statusesAfterFirstResult.at(-1)?.status?.currentTool, 'slow_tool') + + handler.handleCoreEvent( + { type: 'tool_end', toolCallId: 'slow-call', toolName: 'slow_tool', toolResult: null, isError: false }, + [] + ) + + const finalStatuses = chunks.filter((chunk) => chunk.type === 'status') + assert.equal(finalStatuses.at(-1)?.status?.phase, 'thinking') + }) + it('forwards tool progress without changing the result payload', () => { const chunks: AgentStreamChunk[] = [] const handler = new AgentEventHandler({ diff --git a/packages/node-runtime/src/ai/agent/event-handler.ts b/packages/node-runtime/src/ai/agent/event-handler.ts index 9471eaa2..36253068 100644 --- a/packages/node-runtime/src/ai/agent/event-handler.ts +++ b/packages/node-runtime/src/ai/agent/event-handler.ts @@ -159,6 +159,7 @@ export class AgentEventHandler { private readonly onChunk: (chunk: AgentStreamChunk) => void private readonly context: EventHandlerContext private readonly systemPrompt: string + private readonly activeTools = new Map() constructor(config: EventHandlerConfig) { this.onChunk = config.onChunk @@ -186,6 +187,7 @@ export class AgentEventHandler { case 'tool_start': { const params = this.normalizeToolParams(event.toolName, event.toolParams) this.toolsUsed.push(event.toolName) + this.activeTools.set(event.toolCallId, event.toolName) this.onChunk({ type: 'tool_start', toolCallId: event.toolCallId, toolName: event.toolName, toolParams: params }) this.emitStatus('tool_running', messages, { currentTool: event.toolName, force: true }) break @@ -199,6 +201,7 @@ export class AgentEventHandler { }) break case 'tool_end': + this.activeTools.delete(event.toolCallId) this.onChunk({ type: 'tool_result', toolCallId: event.toolCallId, @@ -206,9 +209,15 @@ export class AgentEventHandler { toolResult: event.toolResult, toolIsError: event.isError, }) - this.emitStatus('thinking', messages, { force: true }) + if (this.activeTools.size > 0) { + const currentTool = Array.from(this.activeTools.values()).at(-1) + this.emitStatus('tool_running', messages, { currentTool, force: true }) + } else { + this.emitStatus('thinking', messages, { force: true }) + } break case 'turn_end': + this.activeTools.clear() this.toolRounds = event.round this.emitStatus('thinking', messages, { force: true }) break diff --git a/src/stores/aiChat.ts b/src/stores/aiChat.ts index a614a287..80b765b2 100644 --- a/src/stores/aiChat.ts +++ b/src/stores/aiChat.ts @@ -20,7 +20,7 @@ import { useAgentStreamService } from '@/services/ai-stream/service' import { buildSerializablePreprocessConfig, shouldEnsureDesensitizeRulesBeforeSerialize } from './aiPreprocessConfig' import type { ChartPayload, ChatEvidencePayload } from '@openchatlab/core' import { extractToolResultText, truncateToolResultText } from '@openchatlab/core' -import { getDefaultGeneralAssistantId, type ToolProgress } from '@openchatlab/shared-types' +import { getDefaultGeneralAssistantId } from '@openchatlab/shared-types' import { createRenderOnlyToolPendingBlock, extractChartPayloads, @@ -39,6 +39,9 @@ import { type PlanDraftContentBlock, } from '@/services/ai/planBlocks' import { toSerializableContentBlocks } from './aiChatContentBlocks' +import { createToolLifecycleTracker, type ToolStatus } from './aiToolLifecycle' + +export type { ToolStatus } from './aiToolLifecycle' // 工具调用记录 export interface ToolCallRecord { @@ -124,15 +127,6 @@ export interface SourceMessage { type: number } -// 工具状态类型 -export interface ToolStatus { - name: string - displayName: string - status: 'running' | 'done' | 'error' - progress?: ToolProgress - result?: unknown -} - interface OwnerInfo { platformId: string displayName: string @@ -951,6 +945,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { targetBuffer.messages.push(aiMessage) let aiMessageIndex = targetBuffer.messages.length - 1 let hasStreamError = false + const toolLifecycle = createToolLifecycleTracker() const { updateAIMessage, @@ -1035,7 +1030,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { switch (chunk.type) { case 'content': if (chunk.content) { - state.currentToolStatus = null + state.currentToolStatus = toolLifecycle.current() appendTextToBlocks(chunk.content) } break @@ -1051,11 +1046,10 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { case 'tool_start': if (chunk.toolName) { const toolParams = chunk.toolParams as Record | undefined - state.currentToolStatus = { + state.currentToolStatus = toolLifecycle.start({ name: chunk.toolName, - displayName: chunk.toolName, - status: 'running', - } + toolCallId: chunk.toolCallId, + }) state.toolsUsedInCurrentRound.push(chunk.toolName) if (isRenderOnlyTool(chunk.toolName)) { addRenderOnlyToolPendingBlock(chunk.toolName, toolParams, chunk.toolCallId) @@ -1066,13 +1060,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { break case 'tool_update': - if ( - chunk.toolName && - chunk.toolProgress && - state.currentToolStatus?.name === chunk.toolName && - state.currentToolStatus.status === 'running' - ) { - state.currentToolStatus = { ...state.currentToolStatus, progress: chunk.toolProgress } + if (chunk.toolName && chunk.toolProgress) { + state.currentToolStatus = toolLifecycle.update({ + name: chunk.toolName, + toolCallId: chunk.toolCallId, + progress: chunk.toolProgress, + }) } break @@ -1110,6 +1103,11 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { isError: chunk.toolIsError, }) } + state.currentToolStatus = toolLifecycle.finish({ + name: chunk.toolName, + toolCallId: chunk.toolCallId, + status: toolFailed ? 'error' : 'done', + }) } state.isLoadingSource = false break @@ -1152,6 +1150,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { break case 'done': + toolLifecycle.clear() state.currentToolStatus = null if (!hasStreamError) updatePlanBlockStatus('done') if (chunk.usage) { @@ -1175,8 +1174,11 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { ...state.currentToolStatus, status: 'error', } - updateToolBlockStatus(state.currentToolStatus.name, 'error') + updateToolBlockStatus(state.currentToolStatus.name, 'error', { + toolCallId: state.currentToolStatus.toolCallId, + }) } + toolLifecycle.clear() if (!hasStreamError) { hasStreamError = true const blocks = targetBuffer.messages[aiMessageIndex].contentBlocks || [] @@ -1502,6 +1504,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { const thisRequestId = generateId('req') let lastDoneUsage: TokenUsage | undefined let hasStreamError = false + const toolLifecycle = createToolLifecycleTracker() setActiveTaskMeta(chatKey, content, thisRequestId, state.currentAIChatId) applySessionAssistantSelection(chatKey) @@ -1635,7 +1638,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { if (state.isAborted || thisRequestId !== state.currentRequestId) return switch (chunk.type) { case 'content': - state.currentToolStatus = null + state.currentToolStatus = toolLifecycle.current() appendTextToBlocks(chunk.content || '') break case 'think': @@ -1646,7 +1649,10 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { case 'tool_start': if (chunk.toolName) { const toolParams = chunk.toolParams as Record | undefined - state.currentToolStatus = { name: chunk.toolName, displayName: chunk.toolName, status: 'running' } + state.currentToolStatus = toolLifecycle.start({ + name: chunk.toolName, + toolCallId: chunk.toolCallId, + }) state.toolsUsedInCurrentRound.push(chunk.toolName) if (isRenderOnlyTool(chunk.toolName)) { addRenderOnlyToolPendingBlock(chunk.toolName, toolParams, chunk.toolCallId) @@ -1656,19 +1662,19 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { } break case 'tool_update': - if ( - chunk.toolName && - chunk.toolProgress && - state.currentToolStatus?.name === chunk.toolName && - state.currentToolStatus.status === 'running' - ) { - state.currentToolStatus = { ...state.currentToolStatus, progress: chunk.toolProgress } + if (chunk.toolName && chunk.toolProgress) { + state.currentToolStatus = toolLifecycle.update({ + name: chunk.toolName, + toolCallId: chunk.toolCallId, + progress: chunk.toolProgress, + }) } break case 'tool_result': if (chunk.toolName) { const charts = extractChartPayloads(chunk.toolResult) const renderOnlyError = toRenderOnlyToolErrorBlock(chunk.toolName, chunk.toolResult) + const toolFailed = renderOnlyError !== null || chunk.toolIsError === true if (isRenderOnlyTool(chunk.toolName)) { updateRenderOnlyToolResult(chunk.toolName, chunk.toolCallId, charts, renderOnlyError) } else { @@ -1685,19 +1691,19 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { } if (!isRenderOnlyTool(chunk.toolName)) { const resultText = extractToolResultText(chunk.toolResult) - updateToolBlockStatus( - chunk.toolName, - renderOnlyError !== null || chunk.toolIsError === true ? 'error' : 'done', - { - toolCallId: chunk.toolCallId, - result: truncateToolResultText(resultText), - displayResult: resultText, - isError: chunk.toolIsError, - } - ) + updateToolBlockStatus(chunk.toolName, toolFailed ? 'error' : 'done', { + toolCallId: chunk.toolCallId, + result: truncateToolResultText(resultText), + displayResult: resultText, + isError: chunk.toolIsError, + }) } + state.currentToolStatus = toolLifecycle.finish({ + name: chunk.toolName, + toolCallId: chunk.toolCallId, + status: toolFailed ? 'error' : 'done', + }) } - state.currentToolStatus = null state.isLoadingSource = false break case 'plan_delta': @@ -1718,6 +1724,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { } break case 'done': + toolLifecycle.clear() state.currentToolStatus = null if (!hasStreamError) updatePlanBlockStatus('done') if (chunk.usage) lastDoneUsage = { ...chunk.usage } @@ -1727,7 +1734,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { removePlanDraftsFromBlocks() updatePlanBlockStatus('skipped') hasStreamError = true - if (state.currentToolStatus) updateToolBlockStatus(state.currentToolStatus.name, 'error') + if (state.currentToolStatus) { + updateToolBlockStatus(state.currentToolStatus.name, 'error', { + toolCallId: state.currentToolStatus.toolCallId, + }) + } + toolLifecycle.clear() const blocks = targetBuffer.messages[aiMessageIndex].contentBlocks || [] blocks.push({ type: 'error', error: normalizeSerializedError(chunk.error) }) updateAIMessage({ contentBlocks: [...blocks], isStreaming: false }) diff --git a/src/stores/aiToolLifecycle.test.ts b/src/stores/aiToolLifecycle.test.ts new file mode 100644 index 00000000..291b2386 --- /dev/null +++ b/src/stores/aiToolLifecycle.test.ts @@ -0,0 +1,28 @@ +import { describe, it } from 'node:test' +import assert from 'node:assert/strict' +import { createToolLifecycleTracker } from './aiToolLifecycle' + +describe('createToolLifecycleTracker', () => { + it('keeps remaining parallel calls active and routes progress by toolCallId', () => { + const tracker = createToolLifecycleTracker() + + tracker.start({ name: 'lookup', toolCallId: 'slow-call' }) + tracker.start({ name: 'lookup', toolCallId: 'fast-call' }) + + const afterFastCall = tracker.finish({ name: 'lookup', toolCallId: 'fast-call', status: 'done' }) + assert.equal(afterFastCall?.toolCallId, 'slow-call') + assert.equal(afterFastCall?.status, 'running') + + const progressing = tracker.update({ + name: 'lookup', + toolCallId: 'slow-call', + progress: { phase: 'searching' }, + }) + assert.deepEqual(progressing?.progress, { phase: 'searching' }) + + const finished = tracker.finish({ name: 'lookup', toolCallId: 'slow-call', status: 'done' }) + assert.equal(finished?.toolCallId, 'slow-call') + assert.equal(finished?.status, 'done') + assert.equal(tracker.current(), null) + }) +}) diff --git a/src/stores/aiToolLifecycle.ts b/src/stores/aiToolLifecycle.ts new file mode 100644 index 00000000..018987a9 --- /dev/null +++ b/src/stores/aiToolLifecycle.ts @@ -0,0 +1,80 @@ +import type { ToolProgress } from '@openchatlab/shared-types' + +export interface ToolStatus { + name: string + displayName: string + status: 'running' | 'done' | 'error' + toolCallId?: string + progress?: ToolProgress + result?: unknown +} + +interface ToolIdentity { + name: string + toolCallId?: string +} + +function getToolKey(tool: ToolIdentity): string { + return tool.toolCallId ? `id:${tool.toolCallId}` : `name:${tool.name}` +} + +export function createToolLifecycleTracker() { + const activeTools = new Map() + let currentKey: string | null = null + + const getCurrentActive = (): ToolStatus | null => { + if (currentKey) { + const current = activeTools.get(currentKey) + if (current) return current + } + const fallback = Array.from(activeTools.entries()).at(-1) + currentKey = fallback?.[0] ?? null + return fallback?.[1] ?? null + } + + return { + start(tool: ToolIdentity): ToolStatus { + const key = getToolKey(tool) + const status: ToolStatus = { + name: tool.name, + displayName: tool.name, + status: 'running', + toolCallId: tool.toolCallId, + } + activeTools.set(key, status) + currentKey = key + return status + }, + + update(tool: ToolIdentity & { progress: ToolProgress }): ToolStatus | null { + const key = getToolKey(tool) + const current = activeTools.get(key) + if (!current) return getCurrentActive() + + const updated = { ...current, progress: tool.progress } + activeTools.set(key, updated) + currentKey = key + return updated + }, + + finish(tool: ToolIdentity & { status: 'done' | 'error' }): ToolStatus | null { + const key = getToolKey(tool) + const finished = activeTools.get(key) + activeTools.delete(key) + + if (currentKey === key) currentKey = null + const remaining = getCurrentActive() + if (remaining) return remaining + return finished ? { ...finished, status: tool.status } : null + }, + + current(): ToolStatus | null { + return getCurrentActive() + }, + + clear(): void { + activeTools.clear() + currentKey = null + }, + } +}