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.
2025-12-19 09:33:53 +00:00 · 2025-11-07 19:32:43 +08:00
parent d17c37af7d
commit 4e09d2b5ea
3 changed files with 413 additions and 29 deletions
--- 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,6 +509,52 @@ class Session {
      }

      const toolResult: ToolResult = await invocation.execute(abortSignal);
+
+      // 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({
@@ -481,6 +563,7 @@ class Session {
          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<string, unknown>;
+      }
+    >();
+
+    // 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, acp.ToolKind> = {
+      [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'
--- 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';
--- 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 {