From 4e09d2b5eae7891becb1056c6f57058261d7414c Mon Sep 17 00:00:00 2001 From: "mingholy.lmh" Date: Fri, 7 Nov 2025 19:32:43 +0800 Subject: [PATCH] feat: enhance zed integration with TodoWriteTool and TaskTool support - Implemented detection and handling for TodoWriteTool to route updates as plan entries instead of tool call events. - Added sub-agent tool tracking for TaskTool, allowing for event emission and cleanup. - Updated event listeners to manage sub-agent tool calls and approval requests effectively. --- .../cli/src/zed-integration/zedIntegration.ts | 437 ++++++++++++++++-- packages/core/src/index.ts | 2 + packages/core/src/subagents/index.ts | 3 +- 3 files changed, 413 insertions(+), 29 deletions(-) diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index 49e73991..4a01ed7e 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -12,6 +12,12 @@ import type { GeminiChat, ToolCallConfirmationDetails, ToolResult, + SubAgentEventEmitter, + SubAgentToolCallEvent, + SubAgentToolResultEvent, + SubAgentApprovalRequestEvent, + AnyDeclarativeTool, + AnyToolInvocation, } from '@qwen-code/qwen-code-core'; import { AuthType, @@ -28,6 +34,10 @@ import { getErrorStatus, isWithinRoot, isNodeError, + SubAgentEventType, + TaskTool, + Kind, + TodoWriteTool, } from '@qwen-code/qwen-code-core'; import * as acp from './acp.js'; import { AcpFileSystemService } from './fileSystemService.js'; @@ -403,9 +413,34 @@ class Session { ); } + // Detect TodoWriteTool early - route to plan updates instead of tool_call events + const isTodoWriteTool = + fc.name === TodoWriteTool.Name || tool.name === TodoWriteTool.Name; + + // Declare subAgentToolEventListeners outside try block for cleanup in catch + let subAgentToolEventListeners: Array<() => void> = []; + try { const invocation = tool.build(args); + // Detect TaskTool and set up sub-agent tool tracking + const isTaskTool = tool.name === TaskTool.Name; + + if (isTaskTool && 'eventEmitter' in invocation) { + // Access eventEmitter from TaskTool invocation + const taskEventEmitter = ( + invocation as { + eventEmitter: SubAgentEventEmitter; + } + ).eventEmitter; + + // Set up sub-agent tool tracking + subAgentToolEventListeners = this.setupSubAgentToolTracking( + taskEventEmitter, + abortSignal, + ); + } + const confirmationDetails = await invocation.shouldConfirmExecute(abortSignal); @@ -460,7 +495,8 @@ class Session { throw new Error(`Unexpected: ${resultOutcome}`); } } - } else { + } else if (!isTodoWriteTool) { + // Skip tool_call event for TodoWriteTool await this.sendUpdate({ sessionUpdate: 'tool_call', toolCallId: callId, @@ -473,14 +509,61 @@ class Session { } const toolResult: ToolResult = await invocation.execute(abortSignal); - const content = toToolCallContent(toolResult); - await this.sendUpdate({ - sessionUpdate: 'tool_call_update', - toolCallId: callId, - status: 'completed', - content: content ? [content] : [], - }); + // Clean up event listeners + subAgentToolEventListeners.forEach((cleanup) => cleanup()); + + // Handle TodoWriteTool: extract todos and send plan update + if (isTodoWriteTool) { + // Extract todos from args (initial state) + let todos: Array<{ + id: string; + content: string; + status: 'pending' | 'in_progress' | 'completed'; + }> = []; + + if (Array.isArray(args['todos'])) { + todos = args['todos'] as Array<{ + id: string; + content: string; + status: 'pending' | 'in_progress' | 'completed'; + }>; + } + + // If returnDisplay has todos (e.g., modified by user), use those instead + if ( + toolResult.returnDisplay && + typeof toolResult.returnDisplay === 'object' && + 'type' in toolResult.returnDisplay && + toolResult.returnDisplay.type === 'todo_list' && + 'todos' in toolResult.returnDisplay && + Array.isArray(toolResult.returnDisplay.todos) + ) { + todos = toolResult.returnDisplay.todos; + } + + // Convert todos to plan entries and send plan update + if (todos.length > 0 || Array.isArray(args['todos'])) { + const planEntries = convertTodosToPlanEntries(todos); + await this.sendUpdate({ + sessionUpdate: 'plan', + entries: planEntries, + }); + } + + // Skip tool_call_update event for TodoWriteTool + // Still log and return function response for LLM + } else { + // Normal tool handling: send tool_call_update + const content = toToolCallContent(toolResult); + + await this.sendUpdate({ + sessionUpdate: 'tool_call_update', + toolCallId: callId, + status: 'completed', + content: content ? [content] : [], + }); + } const durationMs = Date.now() - startTime; logToolCall(this.config, { @@ -500,6 +583,9 @@ class Session { return convertToFunctionResponse(fc.name, callId, toolResult.llmContent); } catch (e) { + // Ensure cleanup on error + subAgentToolEventListeners.forEach((cleanup) => cleanup()); + const error = e instanceof Error ? e : new Error(String(e)); await this.sendUpdate({ @@ -515,6 +601,300 @@ class Session { } } + /** + * Sets up event listeners to track sub-agent tool calls within a TaskTool execution. + * Converts subagent tool call events into zedIntegration session updates. + * + * @param eventEmitter - The SubAgentEventEmitter from TaskTool + * @param abortSignal - Signal to abort tracking if parent is cancelled + * @returns Array of cleanup functions to remove event listeners + */ + private setupSubAgentToolTracking( + eventEmitter: SubAgentEventEmitter, + abortSignal: AbortSignal, + ): Array<() => void> { + const cleanupFunctions: Array<() => void> = []; + const toolRegistry = this.config.getToolRegistry(); + + // Track subagent tool call states + const subAgentToolStates = new Map< + string, + { + tool?: AnyDeclarativeTool; + invocation?: AnyToolInvocation; + args?: Record; + } + >(); + + // Listen for tool call start + const onToolCall = (...args: unknown[]) => { + const event = args[0] as SubAgentToolCallEvent; + if (abortSignal.aborted) return; + + const subAgentTool = toolRegistry.getTool(event.name); + let subAgentInvocation: AnyToolInvocation | undefined; + let toolKind: acp.ToolKind = 'other'; + let locations: acp.ToolCallLocation[] = []; + + if (subAgentTool) { + try { + subAgentInvocation = subAgentTool.build(event.args); + toolKind = this.mapToolKind(subAgentTool.kind); + locations = subAgentInvocation.toolLocations().map((loc) => ({ + path: loc.path, + line: loc.line ?? null, + })); + } catch (e) { + // If building fails, continue with defaults + console.warn(`Failed to build subagent tool ${event.name}:`, e); + } + } + + // Save state for subsequent updates + subAgentToolStates.set(event.callId, { + tool: subAgentTool, + invocation: subAgentInvocation, + args: event.args, + }); + + // Check if this is TodoWriteTool - if so, skip sending tool_call event + // Plan update will be sent in onToolResult when we have the final state + if (event.name === TodoWriteTool.Name) { + return; + } + + // Send tool call start update with rawInput + void this.sendUpdate({ + sessionUpdate: 'tool_call', + toolCallId: event.callId, + status: 'in_progress', + title: event.description || event.name, + content: [], + locations, + kind: toolKind, + rawInput: event.args, + }); + }; + + // Listen for tool call result + const onToolResult = (...args: unknown[]) => { + const event = args[0] as SubAgentToolResultEvent; + if (abortSignal.aborted) return; + + const state = subAgentToolStates.get(event.callId); + + // Check if this is TodoWriteTool - if so, route to plan updates + if (event.name === TodoWriteTool.Name) { + let todos: + | Array<{ + id: string; + content: string; + status: 'pending' | 'in_progress' | 'completed'; + }> + | undefined; + + // Try to extract todos from resultDisplay first (final state) + if (event.resultDisplay) { + try { + // resultDisplay might be a JSON stringified object + const parsed = + typeof event.resultDisplay === 'string' + ? JSON.parse(event.resultDisplay) + : event.resultDisplay; + + if ( + typeof parsed === 'object' && + parsed !== null && + 'type' in parsed && + parsed.type === 'todo_list' && + 'todos' in parsed && + Array.isArray(parsed.todos) + ) { + todos = parsed.todos; + } + } catch { + // If parsing fails, ignore - resultDisplay might not be JSON + } + } + + // Fallback to args if resultDisplay doesn't have todos + if (!todos && state?.args && Array.isArray(state.args['todos'])) { + todos = state.args['todos'] as Array<{ + id: string; + content: string; + status: 'pending' | 'in_progress' | 'completed'; + }>; + } + + // Send plan update if we have todos + if (todos) { + const planEntries = convertTodosToPlanEntries(todos); + void this.sendUpdate({ + sessionUpdate: 'plan', + entries: planEntries, + }); + } + + // Skip sending tool_call_update event for TodoWriteTool + // Clean up state + subAgentToolStates.delete(event.callId); + return; + } + + let content: acp.ToolCallContent[] = []; + + // If there's a result display, try to convert to ToolCallContent + if (event.resultDisplay && state?.invocation) { + // resultDisplay is typically a string + if (typeof event.resultDisplay === 'string') { + content = [ + { + type: 'content', + content: { + type: 'text', + text: event.resultDisplay, + }, + }, + ]; + } + } + + // Send tool call completion update + void this.sendUpdate({ + sessionUpdate: 'tool_call_update', + toolCallId: event.callId, + status: event.success ? 'completed' : 'failed', + content: content.length > 0 ? content : [], + title: state?.invocation?.getDescription() ?? event.name, + kind: state?.tool ? this.mapToolKind(state.tool.kind) : null, + locations: + state?.invocation?.toolLocations().map((loc) => ({ + path: loc.path, + line: loc.line ?? null, + })) ?? null, + rawInput: state?.args, + }); + + // Clean up state + subAgentToolStates.delete(event.callId); + }; + + // Listen for permission requests + const onToolWaitingApproval = async (...args: unknown[]) => { + const event = args[0] as SubAgentApprovalRequestEvent; + if (abortSignal.aborted) return; + + const state = subAgentToolStates.get(event.callId); + const content: acp.ToolCallContent[] = []; + + // Handle different confirmation types + if (event.confirmationDetails.type === 'edit') { + const editDetails = event.confirmationDetails as unknown as { + type: 'edit'; + fileName: string; + originalContent: string | null; + newContent: string; + }; + content.push({ + type: 'diff', + path: editDetails.fileName, + oldText: editDetails.originalContent ?? '', + newText: editDetails.newContent, + }); + } + + // Build permission request options from confirmation details + // event.confirmationDetails already contains all fields except onConfirm, + // which we add here to satisfy the type requirement for toPermissionOptions + const fullConfirmationDetails = { + ...event.confirmationDetails, + onConfirm: async () => { + // This is a placeholder - the actual response is handled via event.respond + }, + } as unknown as ToolCallConfirmationDetails; + + const params: acp.RequestPermissionRequest = { + sessionId: this.id, + options: toPermissionOptions(fullConfirmationDetails), + toolCall: { + toolCallId: event.callId, + status: 'pending', + title: event.description || event.name, + content, + locations: + state?.invocation?.toolLocations().map((loc) => ({ + path: loc.path, + line: loc.line ?? null, + })) ?? [], + kind: state?.tool ? this.mapToolKind(state.tool.kind) : 'other', + rawInput: state?.args, + }, + }; + + try { + // Request permission from zed client + const output = await this.client.requestPermission(params); + const outcome = + output.outcome.outcome === 'cancelled' + ? ToolConfirmationOutcome.Cancel + : z + .nativeEnum(ToolConfirmationOutcome) + .parse(output.outcome.optionId); + + // Respond to subagent with the outcome + await event.respond(outcome); + } catch (error) { + // If permission request fails, cancel the tool call + console.error( + `Permission request failed for subagent tool ${event.name}:`, + error, + ); + await event.respond(ToolConfirmationOutcome.Cancel); + } + }; + + // Register event listeners + eventEmitter.on(SubAgentEventType.TOOL_CALL, onToolCall); + eventEmitter.on(SubAgentEventType.TOOL_RESULT, onToolResult); + eventEmitter.on( + SubAgentEventType.TOOL_WAITING_APPROVAL, + onToolWaitingApproval, + ); + + // Return cleanup functions + cleanupFunctions.push(() => { + eventEmitter.off(SubAgentEventType.TOOL_CALL, onToolCall); + eventEmitter.off(SubAgentEventType.TOOL_RESULT, onToolResult); + eventEmitter.off( + SubAgentEventType.TOOL_WAITING_APPROVAL, + onToolWaitingApproval, + ); + }); + + return cleanupFunctions; + } + + /** + * Maps core Tool Kind enum to ACP ToolKind string literals. + * + * @param kind - The core Kind enum value + * @returns The corresponding ACP ToolKind string literal + */ + private mapToolKind(kind: Kind): acp.ToolKind { + const kindMap: Record = { + [Kind.Read]: 'read', + [Kind.Edit]: 'edit', + [Kind.Delete]: 'delete', + [Kind.Move]: 'move', + [Kind.Search]: 'search', + [Kind.Execute]: 'execute', + [Kind.Think]: 'think', + [Kind.Fetch]: 'fetch', + [Kind.Other]: 'other', + }; + return kindMap[kind] ?? 'other'; + } + async #resolvePrompt( message: acp.ContentBlock[], abortSignal: AbortSignal, @@ -859,6 +1239,27 @@ class Session { } } +/** + * Converts todo items to plan entries format for zed integration. + * Maps todo status to plan status and assigns a default priority. + * + * @param todos - Array of todo items with id, content, and status + * @returns Array of plan entries with content, priority, and status + */ +function convertTodosToPlanEntries( + todos: Array<{ + id: string; + content: string; + status: 'pending' | 'in_progress' | 'completed'; + }>, +): acp.PlanEntry[] { + return todos.map((todo) => ({ + content: todo.content, + priority: 'medium' as const, // Default priority since todos don't have priority + status: todo.status, + })); +} + function toToolCallContent(toolResult: ToolResult): acp.ToolCallContent | null { if (toolResult.error?.message) { throw new Error(toolResult.error.message); @@ -870,26 +1271,6 @@ function toToolCallContent(toolResult: ToolResult): acp.ToolCallContent | null { type: 'content', content: { type: 'text', text: toolResult.returnDisplay }, }; - } else if ( - 'type' in toolResult.returnDisplay && - toolResult.returnDisplay.type === 'todo_list' - ) { - // Handle TodoResultDisplay - convert to text representation - const todoText = toolResult.returnDisplay.todos - .map((todo) => { - const statusIcon = { - pending: '○', - in_progress: '◐', - completed: '●', - }[todo.status]; - return `${statusIcon} ${todo.content}`; - }) - .join('\n'); - - return { - type: 'content', - content: { type: 'text', text: todoText }, - }; } else if ( 'type' in toolResult.returnDisplay && toolResult.returnDisplay.type === 'plan_summary' diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index dd675380..883fb114 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -102,6 +102,8 @@ export * from './tools/web-search/index.js'; export * from './tools/read-many-files.js'; export * from './tools/mcp-client.js'; export * from './tools/mcp-tool.js'; +export * from './tools/task.js'; +export * from './tools/todoWrite.js'; // MCP OAuth export { MCPOAuthProvider } from './mcp/oauth-provider.js'; diff --git a/packages/core/src/subagents/index.ts b/packages/core/src/subagents/index.ts index 2b0468a9..5560b4fd 100644 --- a/packages/core/src/subagents/index.ts +++ b/packages/core/src/subagents/index.ts @@ -62,9 +62,10 @@ export type { SubAgentToolResultEvent, SubAgentFinishEvent, SubAgentErrorEvent, + SubAgentApprovalRequestEvent, } from './subagent-events.js'; -export { SubAgentEventEmitter } from './subagent-events.js'; +export { SubAgentEventEmitter, SubAgentEventType } from './subagent-events.js'; // Statistics and formatting export type {