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.
This commit is contained in:
digua committed 2026-08-18 17:27:48 +08:00
1 parent c2382b77f2
commit 2d58fd8c62
5 files changed
+202 -45

No files matched your search

@@ -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({
@@ -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<string, string>()
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
+56 -44
View File
@@ -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<string, unknown> | 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<string, unknown> | 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 })
+28
View File
@@ -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)
})
})
+80
View File
@@ -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<string, ToolStatus>()
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
},
}
}