@umsai/ums-code 0.0.14-v2 → 0.3.0-v1

package/dist/src/zed-integration/zedIntegration.js

@@ -3,16 +3,46 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { AuthType, clearCachedCredentialFile, convertToFunctionResponse, DiscoveredMCPTool, getErrorMessage, getErrorStatus, isNodeError, isWithinRoot, logToolCall, MCPServerConfig, StreamEventType, ToolConfirmationOutcome, } from '@umsai/ums-code-core';
+import { AuthType, clearCachedCredentialFile, convertToFunctionResponse, DiscoveredMCPTool, StreamEventType, DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL_AUTO, DEFAULT_GEMINI_FLASH_MODEL, MCPServerConfig, ToolConfirmationOutcome, logToolCall, logUserPrompt, getErrorStatus, isWithinRoot, isNodeError, SubAgentEventType, TaskTool, Kind, TodoWriteTool, UserPromptEvent, } from '@umsai/ums-code-core';
+import * as acp from './acp.js';
+import { AcpFileSystemService } from './fileSystemService.js';
+import { Readable, Writable } from 'node:stream';
+import { SettingScope } from '../config/settings.js';
 import * as fs from 'node:fs/promises';
 import * as path from 'node:path';
-import { Readable, Writable } from 'node:stream';
 import { z } from 'zod';
-import { SettingScope } from '../config/settings.js';
-import * as acp from './acp.js';
-import { AcpFileSystemService } from './fileSystemService.js';
 import { randomUUID } from 'node:crypto';
+import { getErrorMessage } from '../utils/errors.js';
+import { ExtensionStorage } from '../config/extension.js';
 import { loadCliConfig } from '../config/config.js';
+import { ExtensionEnablementManager } from '../config/extensions/extensionEnablement.js';
+import { handleSlashCommand, getAvailableCommands, } from '../nonInteractiveCliCommands.js';
+import { isSlashCommand } from '../ui/utils/commandUtils.js';
+/**
+ * Built-in commands that are allowed in ACP integration mode.
+ * Only these commands will be available when using handleSlashCommand
+ * or getAvailableCommands in ACP integration.
+ *
+ * Currently, only "init" is supported because `handleSlashCommand` in
+ * nonInteractiveCliCommands.ts only supports handling results where
+ * result.type is "submit_prompt". Other result types are either coupled
+ * to the UI or cannot send notifications to the client via ACP.
+ *
+ * If you have a good idea to add support for more commands, PRs are welcome!
+ */
+const ALLOWED_BUILTIN_COMMANDS_FOR_ACP = ['init'];
+/**
+ * Resolves the model to use based on the current configuration.
+ *
+ * If the model is set to "auto", it will use the flash model if in fallback
+ * mode, otherwise it will use the default model.
+ */
+export function resolveModel(model, isInFallbackMode) {
+    if (model === DEFAULT_GEMINI_MODEL_AUTO) {
+        return isInFallbackMode ? DEFAULT_GEMINI_FLASH_MODEL : DEFAULT_GEMINI_MODEL;
+    }
+    return model;
+}
 export async function runZedIntegration(config, settings, extensions, argv) {
     const stdout = Writable.toWeb(process.stdout);
     const stdin = Readable.toWeb(process.stdin);
@@ -41,21 +71,6 @@ class GeminiAgent {
     async initialize(args) {
         this.clientCapabilities = args.clientCapabilities;
         const authMethods = [
-            {
-                id: AuthType.LOGIN_WITH_GOOGLE,
-                name: 'Log in with Google',
-                description: null,
-            },
-            {
-                id: AuthType.USE_GEMINI,
-                name: 'Use Gemini API key',
-                description: 'Requires setting the `GEMINI_API_KEY` environment variable',
-            },
-            {
-                id: AuthType.USE_VERTEX_AI,
-                name: 'Vertex AI',
-                description: null,
-            },
             {
                 id: AuthType.USE_OPENAI,
                 name: 'Use OpenAI API key',
@@ -87,7 +102,7 @@ class GeminiAgent {
         this.settings.setValue(SettingScope.User, 'security.auth.selectedType', method);
     }
     async newSession({ cwd, mcpServers, }) {
-        const sessionId = randomUUID();
+        const sessionId = this.config.getSessionId() || randomUUID();
         const config = await this.newSessionConfig(sessionId, cwd, mcpServers);
         let isAuthenticated = false;
         if (this.settings.merged.security?.auth?.selectedType) {
@@ -108,8 +123,12 @@ class GeminiAgent {
         }
         const geminiClient = config.getGeminiClient();
         const chat = await geminiClient.startChat();
-        const session = new Session(sessionId, chat, config, this.client);
+        const session = new Session(sessionId, chat, config, this.client, this.settings);
         this.sessions.set(sessionId, session);
+        // Send available commands update as the first session update
+        setTimeout(async () => {
+            await session.sendAvailableCommandsUpdate();
+        }, 0);
         return {
             sessionId,
         };
@@ -124,7 +143,7 @@ class GeminiAgent {
             mergedMcpServers[name] = new MCPServerConfig(command, args, env, cwd);
         }
         const settings = { ...this.settings.merged, mcpServers: mergedMcpServers };
-        const config = await loadCliConfig(settings, this.extensions, sessionId, this.argv, cwd);
+        const config = await loadCliConfig(settings, this.extensions, new ExtensionEnablementManager(ExtensionStorage.getUserExtensionsDir(), this.argv.extensions), sessionId, this.argv, cwd);
         await config.initialize();
         return config;
     }
@@ -148,12 +167,15 @@ class Session {
     chat;
     config;
     client;
+    settings;
     pendingPrompt = null;
-    constructor(id, chat, config, client) {
+    turn = 0;
+    constructor(id, chat, config, client, settings) {
         this.id = id;
         this.chat = chat;
         this.config = config;
         this.client = client;
+        this.settings = settings;
     }
     async cancelPendingPrompt() {
         if (!this.pendingPrompt) {
@@ -166,9 +188,38 @@ class Session {
         this.pendingPrompt?.abort();
         const pendingSend = new AbortController();
         this.pendingPrompt = pendingSend;
-        const promptId = Math.random().toString(16).slice(2);
+        // Increment turn counter for each user prompt
+        this.turn += 1;
         const chat = this.chat;
-        const parts = await this.#resolvePrompt(params.prompt, pendingSend.signal);
+        const promptId = this.config.getSessionId() + '########' + this.turn;
+        // Extract text from all text blocks to construct the full prompt text for logging
+        const promptText = params.prompt
+            .filter((block) => block.type === 'text')
+            .map((block) => (block.type === 'text' ? block.text : ''))
+            .join(' ');
+        // Log user prompt
+        logUserPrompt(this.config, new UserPromptEvent(promptText.length, promptId, this.config.getContentGeneratorConfig()?.authType, promptText));
+        // Check if the input contains a slash command
+        // Extract text from the first text block if present
+        const firstTextBlock = params.prompt.find((block) => block.type === 'text');
+        const inputText = firstTextBlock?.text || '';
+        let parts;
+        if (isSlashCommand(inputText)) {
+            // Handle slash command - allow specific built-in commands for ACP integration
+            const slashCommandResult = await handleSlashCommand(inputText, pendingSend, this.config, this.settings, ALLOWED_BUILTIN_COMMANDS_FOR_ACP);
+            if (slashCommandResult) {
+                // Use the result from the slash command
+                parts = slashCommandResult;
+            }
+            else {
+                // Slash command didn't return a prompt, continue with normal processing
+                parts = await this.#resolvePrompt(params.prompt, pendingSend.signal);
+            }
+        }
+        else {
+            // Normal processing for non-slash commands
+            parts = await this.#resolvePrompt(params.prompt, pendingSend.signal);
+        }
         let nextMessage = { role: 'user', parts };
         while (nextMessage !== null) {
             if (pendingSend.signal.aborted) {
@@ -177,7 +228,7 @@ class Session {
             }
             const functionCalls = [];
             try {
-                const responseStream = await chat.sendMessageStream({
+                const responseStream = await chat.sendMessageStream(resolveModel(this.config.getModel(), this.config.isInFallbackMode()), {
                     message: nextMessage?.parts ?? [],
                     config: {
                         abortSignal: pendingSend.signal,
@@ -237,6 +288,27 @@ class Session {
         };
         await this.client.sessionUpdate(params);
     }
+    async sendAvailableCommandsUpdate() {
+        const abortController = new AbortController();
+        try {
+            const slashCommands = await getAvailableCommands(this.config, this.settings, abortController.signal, ALLOWED_BUILTIN_COMMANDS_FOR_ACP);
+            // Convert SlashCommand[] to AvailableCommand[] format for ACP protocol
+            const availableCommands = slashCommands.map((cmd) => ({
+                name: cmd.name,
+                description: cmd.description,
+                input: null,
+            }));
+            const update = {
+                sessionUpdate: 'available_commands_update',
+                availableCommands,
+            };
+            await this.sendUpdate(update);
+        }
+        catch (error) {
+            // Log error but don't fail session creation
+            console.error('Error sending available commands update:', error);
+        }
+    }
     async runTool(abortSignal, promptId, fc) {
         const callId = fc.id ?? `${fc.name}-${Date.now()}`;
         const args = (fc.args ?? {});
@@ -250,6 +322,7 @@ class Session {
                 function_name: fc.name ?? '',
                 function_args: args,
                 duration_ms: durationMs,
+                status: 'error',
                 success: false,
                 error: error.message,
                 tool_type: typeof tool !== 'undefined' && tool instanceof DiscoveredMCPTool
@@ -274,8 +347,20 @@ class Session {
         if (!tool) {
             return errorResponse(new Error(`Tool "${fc.name}" not found in registry.`));
         }
+        // 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 = [];
         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.eventEmitter;
+                // Set up sub-agent tool tracking
+                subAgentToolEventListeners = this.setupSubAgentToolTracking(taskEventEmitter, abortSignal);
+            }
             const confirmationDetails = await invocation.shouldConfirmExecute(abortSignal);
             if (confirmationDetails) {
                 const content = [];
@@ -321,7 +406,8 @@ class Session {
                     }
                 }
             }
-            else {
+            else if (!isTodoWriteTool) {
+                // Skip tool_call event for TodoWriteTool
                 await this.sendUpdate({
                     sessionUpdate: 'tool_call',
                     toolCallId: callId,
@@ -333,13 +419,45 @@ class Session {
                 });
             }
             const 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 = [];
+                if (Array.isArray(args['todos'])) {
+                    todos = args['todos'];
+                }
+                // 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, {
                 'event.name': 'tool_call',
@@ -347,6 +465,7 @@ class Session {
                 function_name: fc.name,
                 function_args: args,
                 duration_ms: durationMs,
+                status: 'success',
                 success: true,
                 prompt_id: promptId,
                 tool_type: typeof tool !== 'undefined' && tool instanceof DiscoveredMCPTool
@@ -356,6 +475,8 @@ 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({
                 sessionUpdate: 'tool_call_update',
@@ -368,6 +489,235 @@ class Session {
             return errorResponse(error);
         }
     }
+    /**
+     * 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
+     */
+    setupSubAgentToolTracking(eventEmitter, abortSignal) {
+        const cleanupFunctions = [];
+        const toolRegistry = this.config.getToolRegistry();
+        // Track subagent tool call states
+        const subAgentToolStates = new Map();
+        // Listen for tool call start
+        const onToolCall = (...args) => {
+            const event = args[0];
+            if (abortSignal.aborted)
+                return;
+            const subAgentTool = toolRegistry.getTool(event.name);
+            let subAgentInvocation;
+            let toolKind = 'other';
+            let locations = [];
+            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) => {
+            const event = args[0];
+            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;
+                // 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'];
+                }
+                // 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 = [];
+            // 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) => {
+            const event = args[0];
+            if (abortSignal.aborted)
+                return;
+            const state = subAgentToolStates.get(event.callId);
+            const content = [];
+            // Handle different confirmation types
+            if (event.confirmationDetails.type === 'edit') {
+                const editDetails = event.confirmationDetails;
+                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
+                },
+            };
+            const params = {
+                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
+     */
+    mapToolKind(kind) {
+        const kindMap = {
+            [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, abortSignal) {
         const FILE_URI_SCHEME = 'file://';
         const embeddedContext = [];
@@ -664,6 +1014,20 @@ 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) {
+    return todos.map((todo) => ({
+        content: todo.content,
+        priority: 'medium', // Default priority since todos don't have priority
+        status: todo.status,
+    }));
+}
 function toToolCallContent(toolResult) {
     if (toolResult.error?.message) {
         throw new Error(toolResult.error.message);
@@ -675,24 +1039,6 @@ function toToolCallContent(toolResult) {
                 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') {
             const planDisplay = toolResult.returnDisplay;
@@ -702,14 +1048,16 @@ function toToolCallContent(toolResult) {
                 content: { type: 'text', text: planText },
             };
         }
-        else if ('fileDiff' in toolResult.returnDisplay) {
-            // Handle FileDiff
-            return {
-                type: 'diff',
-                path: toolResult.returnDisplay.fileName,
-                oldText: toolResult.returnDisplay.originalContent,
-                newText: toolResult.returnDisplay.newContent,
-            };
+        else {
+            if ('fileName' in toolResult.returnDisplay) {
+                return {
+                    type: 'diff',
+                    path: toolResult.returnDisplay.fileName,
+                    oldText: toolResult.returnDisplay.originalContent,
+                    newText: toolResult.returnDisplay.newContent,
+                };
+            }
+            return null;
         }
     }
     return null;