diff --git a/desktop/src/components/chat/MessageList.test.tsx b/desktop/src/components/chat/MessageList.test.tsx index 9ad7523c..74bcfe25 100644 --- a/desktop/src/components/chat/MessageList.test.tsx +++ b/desktop/src/components/chat/MessageList.test.tsx @@ -1,6 +1,6 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { fireEvent, render, screen, waitFor, within } from '@testing-library/react' -import { MessageList, buildRenderItems } from './MessageList' +import { MessageList, buildRenderModel } from './MessageList' import { useChatStore } from '../../stores/chatStore' import { useTabStore } from '../../stores/tabStore' import type { UIMessage } from '../../types/chat' @@ -117,14 +117,66 @@ describe('MessageList nested tool calls', () => { }, ] - const toolUseIds = new Set(messages.filter((message) => message.type === 'tool_use').map((message) => message.toolUseId)) - const renderItems = buildRenderItems(messages, toolUseIds) + const { renderItems } = buildRenderModel(messages) const toolGroups = renderItems.filter((item) => item.kind === 'tool_group') expect(toolGroups).toHaveLength(2) expect(toolGroups.map((item) => item.toolCalls[0]?.toolUseId)).toEqual(['agent-1', 'write-1']) }) + it('keeps later nested tool calls after an interleaved user message', () => { + const messages: UIMessage[] = [ + { + id: 'tool-agent', + type: 'tool_use', + toolName: 'Agent', + toolUseId: 'agent-1', + input: { description: 'Inspect src/components' }, + timestamp: 1, + }, + { + id: 'tool-read', + type: 'tool_use', + toolName: 'Read', + toolUseId: 'read-1', + input: { file_path: '/tmp/example.ts' }, + timestamp: 2, + parentToolUseId: 'agent-1', + }, + { + id: 'user-follow-up', + type: 'user_text', + content: '顺便把刚才的问题也处理掉', + timestamp: 3, + }, + { + id: 'tool-write', + type: 'tool_use', + toolName: 'Write', + toolUseId: 'write-1', + input: { file_path: '/tmp/out.ts', content: 'export const value = 1' }, + timestamp: 4, + parentToolUseId: 'agent-1', + }, + ] + + const { renderItems, childToolCallsByParent } = buildRenderModel(messages) + const renderedKinds = renderItems.map((item) => + item.kind === 'tool_group' + ? `tool:${item.toolCalls[0]?.toolUseId}` + : `message:${item.message.id}`, + ) + + expect(renderedKinds).toEqual([ + 'tool:agent-1', + 'message:user-follow-up', + 'tool:write-1', + ]) + expect( + (childToolCallsByParent.get('agent-1') ?? []).map((toolCall) => toolCall.toolUseId), + ).toEqual(['read-1']) + }) + it('shows failed agent status and compact unavailable summary for Explore launch errors', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/components/chat/MessageList.tsx b/desktop/src/components/chat/MessageList.tsx index d98468e4..f97c6793 100644 --- a/desktop/src/components/chat/MessageList.tsx +++ b/desktop/src/components/chat/MessageList.tsx @@ -22,19 +22,58 @@ type RenderItem = | { kind: 'tool_group'; toolCalls: ToolCall[]; id: string } | { kind: 'message'; message: UIMessage } -export function buildRenderItems(messages: UIMessage[], toolUseIds: Set): RenderItem[] { - const items: RenderItem[] = [] - let pendingToolCalls: ToolCall[] = [] +type RenderModel = { + renderItems: RenderItem[] + toolResultMap: Map + childToolCallsByParent: Map +} - const flushGroup = () => { +function appendChildToolCall( + childToolCallsByParent: Map, + parentToolUseId: string, + toolCall: ToolCall, +) { + const siblings = childToolCallsByParent.get(parentToolUseId) + if (siblings) { + siblings.push(toolCall) + } else { + childToolCallsByParent.set(parentToolUseId, [toolCall]) + } +} + +export function buildRenderModel(messages: UIMessage[]): RenderModel { + const items: RenderItem[] = [] + const toolResultMap = new Map() + const childToolCallsByParent = new Map() + const toolUseIds = new Set() + let pendingToolCalls: ToolCall[] = [] + const inlineParentToolUseIds = new Set() + + const flushGroup = (resetInlineParents = false) => { if (pendingToolCalls.length > 0) { items.push({ kind: 'tool_group', toolCalls: [...pendingToolCalls], id: `group-${pendingToolCalls[0]!.id}`, }) + for (const toolCall of pendingToolCalls) { + inlineParentToolUseIds.add(toolCall.toolUseId) + } pendingToolCalls = [] } + + if (resetInlineParents) { + inlineParentToolUseIds.clear() + } + } + + for (const msg of messages) { + if (msg.type === 'tool_use') { + toolUseIds.add(msg.toolUseId) + } + if (msg.type === 'tool_result') { + toolResultMap.set(msg.toolUseId, msg) + } } for (const msg of messages) { @@ -43,24 +82,30 @@ export function buildRenderItems(messages: UIMessage[], toolUseIds: Set) } if (msg.type === 'tool_use') { - if (msg.parentToolUseId) { + const parentIsPending = msg.parentToolUseId + ? pendingToolCalls.some((toolCall) => toolCall.toolUseId === msg.parentToolUseId) + : false + + if (msg.parentToolUseId && (inlineParentToolUseIds.has(msg.parentToolUseId) || parentIsPending)) { flushGroup() + appendChildToolCall(childToolCallsByParent, msg.parentToolUseId, msg) + inlineParentToolUseIds.add(msg.toolUseId) continue } if (msg.toolName === 'AskUserQuestion') { - flushGroup() + flushGroup(true) items.push({ kind: 'message', message: msg }) } else { pendingToolCalls.push(msg) } } else { - flushGroup() + flushGroup(true) items.push({ kind: 'message', message: msg }) } } flushGroup() - return items + return { renderItems: items, toolResultMap, childToolCallsByParent } } export function MessageList() { @@ -77,31 +122,10 @@ export function MessageList() { bottomRef.current?.scrollIntoView?.({ behavior: 'smooth' }) }, [messages.length, streamingText]) - const { toolResultMap, childToolCallsByParent, renderItems } = useMemo(() => { - const toolUseIds = new Set() - const toolResultMap = new Map() - const childToolCallsByParent = new Map() - - for (const msg of messages) { - if (msg.type === 'tool_use') { - toolUseIds.add(msg.toolUseId) - if (msg.parentToolUseId) { - const siblings = childToolCallsByParent.get(msg.parentToolUseId) - if (siblings) { - siblings.push(msg) - } else { - childToolCallsByParent.set(msg.parentToolUseId, [msg]) - } - } - } - if (msg.type === 'tool_result' && msg.toolUseId) { - toolResultMap.set(msg.toolUseId, msg) - } - } - - const renderItems = buildRenderItems(messages, toolUseIds) - return { toolUseIds, toolResultMap, childToolCallsByParent, renderItems } - }, [messages]) + const { toolResultMap, childToolCallsByParent, renderItems } = useMemo( + () => buildRenderModel(messages), + [messages], + ) return (