@defai.digital/ax-cli 3.7.2 → 3.8.2

package/dist/agent/llm-agent.js

@@ -20,6 +20,8 @@ import { SubagentOrchestrator } from "./subagent-orchestrator.js";
 import { getTaskPlanner, isComplexRequest, } from "../planner/index.js";
 import { PLANNER_CONFIG } from "../constants.js";
 import { resolveMCPReferences, extractMCPReferences } from "../mcp/resources.js";
+import { SDKError, SDKErrorCode } from "../sdk/errors.js";
+import { getStatusReporter } from "./status-reporter.js";
 export class LLMAgent extends EventEmitter {
     llmClient;
     textEditor;
@@ -51,6 +53,11 @@ export class LLMAgent extends EventEmitter {
     thinkingConfig;
     /** Track if agent has been disposed */
     disposed = false;
+    /** Tool approval system for VSCode integration */
+    requireToolApproval = false;
+    toolApprovalCallbacks = new Map();
+    /** BUG FIX: Track approval timeouts for cleanup to prevent memory leaks */
+    toolApprovalTimeouts = new Map();
     constructor(apiKey, baseURL, model, maxToolRounds) {
         super();
         const manager = getSettingsManager();
@@ -76,10 +83,15 @@ export class LLMAgent extends EventEmitter {
         // Load sampling configuration from settings (supports env vars, project, and user settings)
         this.samplingConfig = manager.getSamplingSettings();
         // Wire up checkpoint callback for automatic checkpoint creation
+        // CRITICAL FIX: Deep clone chatHistory to prevent race conditions
+        // The checkpoint creation is async and chatHistory can be modified during the operation
         this.textEditor.setCheckpointCallback(async (files, description) => {
+            // Create immutable snapshot of chat history at callback time
+            // This prevents inconsistencies if messages are added during checkpoint creation
+            const chatHistorySnapshot = JSON.parse(JSON.stringify(this.chatHistory));
             await this.checkpointManager.createCheckpoint({
                 files,
-                conversationState: this.chatHistory,
+                conversationState: chatHistorySnapshot,
                 description,
                 metadata: {
                     model: this.llmClient.getCurrentModel(),
@@ -110,6 +122,17 @@ export class LLMAgent extends EventEmitter {
             role: "system",
             content: `Current working directory: ${process.cwd()}\nTimestamp: ${new Date().toISOString().split('T')[0]}`,
         });
+        // NEW: Listen for context pruning to generate summaries
+        // CRITICAL FIX: Wrap async callback to prevent uncaught promise rejections
+        // Event listeners don't handle async errors automatically, so we must catch them
+        this.contextManager.on('before_prune', (data) => {
+            this.handleContextOverflow(data).catch((error) => {
+                const errorMsg = extractErrorMessage(error);
+                console.error('Error handling context overflow:', errorMsg);
+                // Emit error event for monitoring
+                this.emit('error', error);
+            });
+        });
     }
     initializeCheckpointManager() {
         // Initialize checkpoint manager in the background
@@ -130,7 +153,6 @@ export class LLMAgent extends EventEmitter {
     }
     async initializeMCP() {
         // Initialize MCP in the background without blocking
-        // Single error handler - no redundant catch needed since inner try-catch handles all errors
         Promise.resolve().then(async () => {
             try {
                 const config = loadMCPConfig();
@@ -144,6 +166,9 @@ export class LLMAgent extends EventEmitter {
                 console.warn("MCP initialization failed:", errorMsg);
                 this.emit('system', `MCP initialization failed: ${errorMsg}`);
             }
+        }).catch((error) => {
+            // Catch any errors from emit() or other unexpected failures
+            console.error("Unexpected MCP initialization error:", error);
         });
     }
     /**
@@ -188,32 +213,134 @@ export class LLMAgent extends EventEmitter {
     getSamplingConfig() {
         return this.samplingConfig;
     }
+    /**
+     * Enable or disable tool approval requirement
+     * When enabled, text_editor operations will emit 'tool:approval_required' events
+     * and wait for approval before executing
+     *
+     * This is used by VSCode extension to show diff previews
+     *
+     * @param enabled - Whether to require approval for text_editor operations
+     */
+    setRequireToolApproval(enabled) {
+        this.requireToolApproval = enabled;
+    }
+    /**
+     * Approve or reject a pending tool call
+     * Called by external integrations (e.g., VSCode extension) in response to
+     * 'tool:approval_required' events
+     *
+     * @param toolCallId - The ID of the tool call to approve/reject
+     * @param approved - true to execute the tool, false to reject it
+     */
+    approveToolCall(toolCallId, approved) {
+        const callback = this.toolApprovalCallbacks.get(toolCallId);
+        if (callback) {
+            // BUG FIX: Clear the timeout when approval is received (prevents memory leak)
+            const timeout = this.toolApprovalTimeouts.get(toolCallId);
+            if (timeout) {
+                clearTimeout(timeout);
+                this.toolApprovalTimeouts.delete(toolCallId);
+            }
+            callback(approved);
+            this.toolApprovalCallbacks.delete(toolCallId);
+        }
+    }
+    /**
+     * Wait for external approval of a tool call
+     * Emits 'tool:approval_required' event and waits for approveToolCall() to be called
+     *
+     * @param toolCall - The tool call awaiting approval
+     * @returns Promise<boolean> - true if approved, false if rejected or timeout
+     */
+    waitForToolApproval(toolCall) {
+        return new Promise((resolve) => {
+            // Emit event so external integrations can show diff preview
+            this.emit('tool:approval_required', toolCall);
+            // Store callback
+            this.toolApprovalCallbacks.set(toolCall.id, resolve);
+            // BUG FIX: Track the timeout so it can be cleared on approval/disposal
+            // This prevents memory leaks from dangling timers
+            const timeoutId = setTimeout(() => {
+                // Clean up both the callback and timeout tracking
+                this.toolApprovalTimeouts.delete(toolCall.id);
+                if (this.toolApprovalCallbacks.has(toolCall.id)) {
+                    this.toolApprovalCallbacks.delete(toolCall.id);
+                    resolve(false); // Auto-reject on timeout
+                }
+            }, 5 * 60 * 1000);
+            // Track the timeout for cleanup
+            this.toolApprovalTimeouts.set(toolCall.id, timeoutId);
+        });
+    }
+    /**
+     * Handle context overflow by generating a summary
+     * Called when context manager is about to prune messages
+     */
+    async handleContextOverflow(data) {
+        try {
+            const reporter = getStatusReporter();
+            const summary = await reporter.generateContextSummary(data.messages, this.chatHistory, 'context_overflow', data.tokenCount);
+            // Log for debugging
+            if (process.env.DEBUG) {
+                console.log(`[Context Overflow] Summary generated: ${summary.path}`);
+            }
+            // Add a chat entry to inform user (non-blocking)
+            const summaryEntry = {
+                type: 'assistant',
+                content: `⚠️ Context window approaching limit (${data.tokenCount.toLocaleString()} tokens). Summary saved to:\n\`${summary.path}\``,
+                timestamp: new Date(),
+            };
+            this.chatHistory.push(summaryEntry);
+            // Emit event for UI/logging
+            this.emit('context:summary', summary);
+        }
+        catch (error) {
+            // Summary generation failure should not block execution
+            const errorMsg = extractErrorMessage(error);
+            console.warn('Failed to generate context summary:', errorMsg);
+        }
+    }
     /**
      * Apply context pruning to both messages and chatHistory
      * BUGFIX: Prevents chatHistory from growing unbounded
      */
     applyContextPruning() {
+        // Prune LLM messages if needed
         if (this.contextManager.shouldPrune(this.messages, this.tokenCounter)) {
-            // Prune LLM messages
             this.messages = this.contextManager.pruneMessages(this.messages, this.tokenCounter);
-            // Also prune chatHistory to prevent unlimited growth
-            // Keep last 200 entries which is more than enough for UI display
-            const MAX_CHAT_HISTORY_ENTRIES = 200;
-            if (this.chatHistory.length > MAX_CHAT_HISTORY_ENTRIES) {
-                const entriesToRemove = this.chatHistory.length - MAX_CHAT_HISTORY_ENTRIES;
-                this.chatHistory = this.chatHistory.slice(entriesToRemove);
-                // Update tool call index map after pruning
-                // Clear and rebuild only for remaining entries
-                this.toolCallIndexMap.clear();
-                this.chatHistory.forEach((entry, index) => {
-                    if (entry.type === "tool_call" && entry.toolCall?.id) {
-                        this.toolCallIndexMap.set(entry.toolCall.id, index);
-                    }
-                    else if (entry.type === "tool_result" && entry.toolCall?.id) {
-                        this.toolCallIndexMap.set(entry.toolCall.id, index);
-                    }
-                });
-            }
+        }
+        // CRITICAL FIX: Always check and prune chatHistory to prevent unbounded growth
+        // This must happen UNCONDITIONALLY, even if context pruning is disabled
+        // Keep last 200 entries which is more than enough for UI display
+        const MAX_CHAT_HISTORY_ENTRIES = 200;
+        if (this.chatHistory.length > MAX_CHAT_HISTORY_ENTRIES) {
+            const entriesToRemove = this.chatHistory.length - MAX_CHAT_HISTORY_ENTRIES;
+            this.chatHistory = this.chatHistory.slice(entriesToRemove);
+            // Update tool call index map after pruning
+            // Clear and rebuild only for remaining entries
+            this.toolCallIndexMap.clear();
+            this.chatHistory.forEach((entry, index) => {
+                if (entry.type === "tool_call" && entry.toolCall?.id) {
+                    this.toolCallIndexMap.set(entry.toolCall.id, index);
+                }
+                else if (entry.type === "tool_result" && entry.toolCall?.id) {
+                    this.toolCallIndexMap.set(entry.toolCall.id, index);
+                }
+            });
+        }
+        // CRITICAL FIX: Add hard limit for messages array as safety backstop
+        // In case contextManager.shouldPrune() always returns false
+        const MAX_MESSAGES = 500;
+        if (this.messages.length > MAX_MESSAGES) {
+            // Keep system message (if exists) + last N messages
+            const systemMessages = this.messages.filter(m => m.role === 'system');
+            const nonSystemMessages = this.messages.filter(m => m.role !== 'system');
+            const keepMessages = Math.min(nonSystemMessages.length, MAX_MESSAGES - systemMessages.length);
+            this.messages = [
+                ...systemMessages,
+                ...nonSystemMessages.slice(-keepMessages)
+            ];
         }
     }
     /**
@@ -235,13 +362,16 @@ export class LLMAgent extends EventEmitter {
         try {
             const args = JSON.parse(toolCall.function.arguments || '{}');
             this.toolCallArgsCache.set(toolCall.id, args);
-            // Prevent unbounded memory growth - limit cache size
+            // CRITICAL FIX: Prevent unbounded memory growth with proper cache eviction
+            // When cache exceeds limit, reduce to 80% capacity (not just remove 100 entries)
             if (this.toolCallArgsCache.size > 500) {
+                const targetSize = 400; // 80% of max capacity
+                const toRemove = this.toolCallArgsCache.size - targetSize;
                 let deleted = 0;
                 for (const key of this.toolCallArgsCache.keys()) {
                     this.toolCallArgsCache.delete(key);
                     deleted++;
-                    if (deleted >= 100)
+                    if (deleted >= toRemove)
                         break;
                 }
             }
@@ -727,6 +857,29 @@ export class LLMAgent extends EventEmitter {
             }
             // Emit plan completed event
             this.emit("plan:completed", { plan, result: planResult });
+            // Generate status report on plan completion
+            try {
+                const reporter = getStatusReporter();
+                const tokenCount = this.tokenCounter.countMessageTokens(this.messages);
+                const statusReport = await reporter.generateStatusReport({
+                    messages: this.messages,
+                    chatHistory: this.chatHistory,
+                    tokenCount,
+                    plan,
+                });
+                // Notify user of status report
+                yield {
+                    type: "content",
+                    content: `\n📊 Status report saved to: \`${statusReport.path}\`\n`,
+                };
+                // Emit event for UI/logging
+                this.emit("plan:report", statusReport);
+            }
+            catch (error) {
+                // Status report generation failure should not block execution
+                const errorMsg = extractErrorMessage(error);
+                console.warn("Failed to generate status report:", errorMsg);
+            }
             this.currentPlan = null;
         }
         catch (error) {
@@ -1167,85 +1320,102 @@ export class LLMAgent extends EventEmitter {
         let accumulatedContent = "";
         let toolCallsYielded = false;
         let usageData = null;
-        for await (const chunk of stream) {
-            // Check for cancellation in the streaming loop
-            if (this.isCancelled()) {
-                yield* this.yieldCancellation();
-                // Return empty state after cancellation to avoid processing partial results
-                return { accumulated: {}, content: "", yielded: false };
-            }
-            if (!chunk.choices?.[0])
-                continue;
-            // Capture usage data from chunks (usually in the final chunk)
-            if (chunk.usage) {
-                usageData = chunk.usage;
-            }
-            // Accumulate the message using reducer
-            accumulatedMessage = this.messageReducer(accumulatedMessage, chunk);
-            // Check for tool calls - yield when we have complete tool calls with function names
-            const toolCalls = accumulatedMessage.tool_calls;
-            if (!toolCallsYielded && toolCalls && Array.isArray(toolCalls) && toolCalls.length > 0) {
-                const hasCompleteTool = toolCalls.some((tc) => tc.function?.name);
-                if (hasCompleteTool) {
+        // CRITICAL FIX: Ensure stream is properly closed on cancellation or error
+        // Without this, HTTP connections and buffers remain in memory
+        try {
+            for await (const chunk of stream) {
+                // Check for cancellation in the streaming loop
+                if (this.isCancelled()) {
+                    yield* this.yieldCancellation();
+                    // Return empty state after cancellation to avoid processing partial results
+                    return { accumulated: {}, content: "", yielded: false };
+                }
+                if (!chunk.choices?.[0])
+                    continue;
+                // Capture usage data from chunks (usually in the final chunk)
+                if (chunk.usage) {
+                    usageData = chunk.usage;
+                }
+                // Accumulate the message using reducer
+                accumulatedMessage = this.messageReducer(accumulatedMessage, chunk);
+                // Check for tool calls - yield when we have complete tool calls with function names
+                const toolCalls = accumulatedMessage.tool_calls;
+                if (!toolCallsYielded && toolCalls && Array.isArray(toolCalls) && toolCalls.length > 0) {
+                    const hasCompleteTool = toolCalls.some((tc) => tc.function?.name);
+                    if (hasCompleteTool) {
+                        yield {
+                            type: "tool_calls",
+                            toolCalls: toolCalls,
+                        };
+                        toolCallsYielded = true;
+                    }
+                }
+                // Stream reasoning content (GLM-4.6 thinking mode)
+                // Safety check: ensure choices[0] exists before accessing
+                if (chunk.choices[0]?.delta?.reasoning_content) {
                     yield {
-                        type: "tool_calls",
-                        toolCalls: toolCalls,
+                        type: "reasoning",
+                        reasoningContent: chunk.choices[0].delta.reasoning_content,
                     };
-                    toolCallsYielded = true;
+                }
+                // Stream content as it comes
+                if (chunk.choices[0]?.delta?.content) {
+                    accumulatedContent += chunk.choices[0].delta.content;
+                    yield {
+                        type: "content",
+                        content: chunk.choices[0].delta.content,
+                    };
+                    // Emit token count update (throttled and optimized)
+                    const now = Date.now();
+                    if (now - lastTokenUpdate.value > 1000) { // Increased throttle to 1s for better performance
+                        lastTokenUpdate.value = now;
+                        // Use fast estimation during streaming (4 chars ≈ 1 token)
+                        // This is ~70% faster than tiktoken encoding
+                        const estimatedOutputTokens = Math.floor(accumulatedContent.length / 4) +
+                            (accumulatedMessage.tool_calls
+                                ? Math.floor(JSON.stringify(accumulatedMessage.tool_calls).length / 4)
+                                : 0);
+                        totalOutputTokens.value = estimatedOutputTokens;
+                        yield {
+                            type: "token_count",
+                            tokenCount: inputTokens + estimatedOutputTokens,
+                        };
+                    }
                 }
             }
-            // Stream reasoning content (GLM-4.6 thinking mode)
-            // Safety check: ensure choices[0] exists before accessing
-            if (chunk.choices[0]?.delta?.reasoning_content) {
-                yield {
-                    type: "reasoning",
-                    reasoningContent: chunk.choices[0].delta.reasoning_content,
-                };
-            }
-            // Stream content as it comes
-            if (chunk.choices[0]?.delta?.content) {
-                accumulatedContent += chunk.choices[0].delta.content;
-                yield {
-                    type: "content",
-                    content: chunk.choices[0].delta.content,
-                };
-                // Emit token count update (throttled and optimized)
-                const now = Date.now();
-                if (now - lastTokenUpdate.value > 1000) { // Increased throttle to 1s for better performance
-                    lastTokenUpdate.value = now;
-                    // Use fast estimation during streaming (4 chars ≈ 1 token)
-                    // This is ~70% faster than tiktoken encoding
-                    const estimatedOutputTokens = Math.floor(accumulatedContent.length / 4) +
-                        (accumulatedMessage.tool_calls
-                            ? Math.floor(JSON.stringify(accumulatedMessage.tool_calls).length / 4)
-                            : 0);
-                    totalOutputTokens.value = estimatedOutputTokens;
+            // Track usage if available and emit accurate final token count
+            if (usageData) {
+                const tracker = getUsageTracker();
+                tracker.trackUsage(this.llmClient.getCurrentModel(), usageData);
+                // Emit accurate token count from API usage data (replaces estimation)
+                const totalTokens = usageData.total_tokens;
+                const completionTokens = usageData.completion_tokens;
+                if (totalTokens) {
+                    totalOutputTokens.value = completionTokens || 0;
                     yield {
                         type: "token_count",
-                        tokenCount: inputTokens + estimatedOutputTokens,
+                        tokenCount: totalTokens,
                     };
                 }
             }
+            // CRITICAL: Yield the accumulated result so the main loop can access it!
+            const result = { accumulated: accumulatedMessage, content: accumulatedContent, yielded: toolCallsYielded };
+            yield result;
+            return result;
         }
-        // Track usage if available and emit accurate final token count
-        if (usageData) {
-            const tracker = getUsageTracker();
-            tracker.trackUsage(this.llmClient.getCurrentModel(), usageData);
-            // Emit accurate token count from API usage data (replaces estimation)
-            const totalTokens = usageData.total_tokens;
-            const completionTokens = usageData.completion_tokens;
-            if (totalTokens) {
-                totalOutputTokens.value = completionTokens || 0;
-                yield {
-                    type: "token_count",
-                    tokenCount: totalTokens,
-                };
+        finally {
+            // CRITICAL FIX: Properly close the async iterator to release HTTP connections and buffers
+            // This prevents socket leaks when streams are cancelled or errors occur
+            if (typeof stream.return === 'function') {
+                try {
+                    await stream.return();
+                }
+                catch (cleanupError) {
+                    // Log but don't throw - cleanup errors shouldn't break the flow
+                    console.warn('Stream cleanup warning:', cleanupError);
+                }
             }
         }
-        // CRITICAL: Yield the accumulated result so the main loop can access it!
-        const result = { accumulated: accumulatedMessage, content: accumulatedContent, yielded: toolCallsYielded };
-        yield result;
-        return result;
     }
     /**
      * Add assistant message to history and conversation
@@ -1485,6 +1655,27 @@ export class LLMAgent extends EventEmitter {
                 return { success: false, error: parseResult.error };
             }
             const args = parseResult.args;
+            // Check if tool approval is required (for VSCode integration)
+            if (this.requireToolApproval) {
+                // Only require approval for file modification operations
+                const needsApproval = toolCall.function.name === "create_file" ||
+                    toolCall.function.name === "str_replace_editor" ||
+                    toolCall.function.name === "insert_text";
+                if (needsApproval) {
+                    // Emit event and wait for approval
+                    const approved = await this.waitForToolApproval(toolCall);
+                    if (!approved) {
+                        // User rejected the change
+                        this.emit('tool:rejected', toolCall);
+                        return {
+                            success: false,
+                            error: 'Change rejected by user'
+                        };
+                    }
+                    // User approved
+                    this.emit('tool:approved', toolCall);
+                }
+            }
             // Helper to safely get string argument with validation
             const getString = (key, required = true) => {
                 const value = args[key];
@@ -1724,6 +1915,9 @@ export class LLMAgent extends EventEmitter {
             // Safely preserve system message if it exists
             const systemMessage = this.messages.length > 0 ? this.messages[0] : null;
             this.messages = systemMessage ? [systemMessage] : [];
+            // CRITICAL FIX: Track tool calls to validate tool results
+            // Prevents API errors from orphaned tool results without corresponding tool calls
+            const toolCallIds = new Set();
             for (const entry of conversationState) {
                 if (entry.type === 'user') {
                     this.messages.push({
@@ -1732,6 +1926,14 @@ export class LLMAgent extends EventEmitter {
                     });
                 }
                 else if (entry.type === 'assistant') {
+                    // Track tool call IDs from assistant messages
+                    if (entry.toolCalls && Array.isArray(entry.toolCalls)) {
+                        for (const toolCall of entry.toolCalls) {
+                            if (toolCall?.id) {
+                                toolCallIds.add(toolCall.id);
+                            }
+                        }
+                    }
                     this.messages.push({
                         role: 'assistant',
                         content: entry.content,
@@ -1739,11 +1941,18 @@ export class LLMAgent extends EventEmitter {
                     });
                 }
                 else if (entry.type === 'tool_result' && entry.toolCall) {
-                    this.messages.push({
-                        role: 'tool',
-                        content: entry.content,
-                        tool_call_id: entry.toolCall.id,
-                    });
+                    // CRITICAL FIX: Only add tool result if corresponding tool call exists
+                    // This prevents "tool message without corresponding tool call" API errors
+                    if (toolCallIds.has(entry.toolCall.id)) {
+                        this.messages.push({
+                            role: 'tool',
+                            content: entry.content,
+                            tool_call_id: entry.toolCall.id,
+                        });
+                    }
+                    else {
+                        console.warn(`Skipping orphaned tool result for tool_call_id: ${entry.toolCall.id}`);
+                    }
                 }
             }
             this.emit('system', `Conversation rewound to checkpoint ${checkpointId}`);
@@ -1890,7 +2099,6 @@ export class LLMAgent extends EventEmitter {
      */
     checkDisposed() {
         if (this.disposed) {
-            const { SDKError, SDKErrorCode } = require('../sdk/errors.js');
             throw new SDKError(SDKErrorCode.AGENT_DISPOSED, 'Agent has been disposed and cannot be used. Create a new agent instance.');
         }
     }
@@ -1927,17 +2135,29 @@ export class LLMAgent extends EventEmitter {
         this.disposed = true;
         // Remove all event listeners to prevent memory leaks
         this.removeAllListeners();
+        // CRITICAL FIX: Remove event listeners from contextManager to prevent memory leak
+        // The 'before_prune' listener was registered in constructor (line 188) but never removed
+        this.contextManager.removeAllListeners('before_prune');
         // Dispose tools that have cleanup methods
         this.bash.dispose();
         // Clear in-memory caches
         this.recentToolCalls.clear();
         this.toolCallIndexMap.clear();
         this.toolCallArgsCache.clear();
+        // BUG FIX: Clear all pending tool approval timeouts to prevent memory leaks
+        // These timers would otherwise keep running for up to 5 minutes after dispose
+        for (const timeout of this.toolApprovalTimeouts.values()) {
+            clearTimeout(timeout);
+        }
+        this.toolApprovalTimeouts.clear();
+        this.toolApprovalCallbacks.clear();
         // Clear conversation history to free memory
         this.chatHistory = [];
         this.messages = [];
-        // Dispose token counter and context manager
-        this.tokenCounter.dispose();
+        // Dispose context manager (tokenCounter is a singleton, don't dispose)
+        // CRITICAL FIX: tokenCounter is obtained via getTokenCounter() which returns
+        // a shared singleton instance. Disposing it would break other agent instances
+        // using the same model. The singleton manages its own lifecycle.
         this.contextManager.dispose();
         // Abort any in-flight requests
         if (this.abortController) {