wave-agent-sdk 0.5.0 → 0.6.2

package/src/managers/aiManager.ts

@@ -3,7 +3,6 @@ import {
   compressMessages,
   type CallAgentOptions,
 } from "../services/aiService.js";
-import { getMessagesToCompress } from "../utils/messageOperations.js";
 import { convertMessagesForAPI } from "../utils/convertMessagesForAPI.js";
 import { calculateComprehensiveTotalTokens } from "../utils/tokenCalculation.js";
 import * as fs from "node:fs/promises";
@@ -21,10 +20,7 @@ import { ChatCompletionMessageFunctionToolCall } from "openai/resources.js";
 import type { HookManager } from "./hookManager.js";
 import type { ExtendedHookExecutionContext } from "../types/hooks.js";
 import type { PermissionManager } from "./permissionManager.js";
-import {
-  DEFAULT_SYSTEM_PROMPT,
-  buildSystemPrompt,
-} from "../constants/prompts.js";
+import { buildSystemPrompt } from "../prompts/index.js";
 
 export interface AIManagerCallbacks {
   onCompressionStateChange?: (isCompressing: boolean) => void;
@@ -34,6 +30,7 @@ export interface AIManagerCallbacks {
 export interface AIManagerOptions {
   messageManager: MessageManager;
   toolManager: ToolManager;
+  taskManager: import("../services/taskManager.js").TaskManager;
   logger?: Logger;
   backgroundTaskManager?: BackgroundTaskManager;
   hookManager?: HookManager;
@@ -60,6 +57,7 @@ export class AIManager {
   private logger?: Logger;
   private toolManager: ToolManager;
   private messageManager: MessageManager;
+  private taskManager: import("../services/taskManager.js").TaskManager;
   private backgroundTaskManager?: BackgroundTaskManager;
   private hookManager?: HookManager;
   private reversionManager?: import("./reversionManager.js").ReversionManager;
@@ -79,6 +77,7 @@ export class AIManager {
   constructor(options: AIManagerOptions) {
     this.messageManager = options.messageManager;
     this.toolManager = options.toolManager;
+    this.taskManager = options.taskManager;
     this.backgroundTaskManager = options.backgroundTaskManager;
     this.hookManager = options.hookManager;
     this.reversionManager = options.reversionManager;
@@ -159,15 +158,6 @@ export class AIManager {
     this.setIsLoading(false);
   }
 
-  /**
-   * Abort the AI recursion loop immediately.
-   * This is used when a tool is backgrounded via Ctrl-B, even if no foreground task was active.
-   */
-  public abortRecursion(): void {
-    this.logger?.info("Aborting AI recursion loop");
-    this.abortAIMessage();
-  }
-
   // Helper method to generate compactParams
   private generateCompactParams(
     toolName: string,
@@ -180,6 +170,7 @@ export class AIManager {
       if (toolPlugin?.formatCompactParams) {
         const context: ToolContext = {
           workdir: this.workdir,
+          taskManager: this.taskManager,
         };
         return toolPlugin.formatCompactParams(toolArgs, context);
       }
@@ -213,9 +204,7 @@ export class AIManager {
       );
 
       // Check if messages need compression
-      const { messagesToCompress, insertIndex } = getMessagesToCompress(
-        this.messageManager.getMessages(),
-      );
+      const messagesToCompress = this.messageManager.getMessages();
 
       // If there are messages to compress, perform compression
       if (messagesToCompress.length > 0) {
@@ -248,7 +237,6 @@ export class AIManager {
 
           // Execute message reconstruction and sessionId update after compression
           this.messageManager.compressMessagesAndUpdateSession(
-            insertIndex,
             compressionResult.content,
             compressionUsage,
           );
@@ -305,6 +293,9 @@ export class AIManager {
       return;
     }
 
+    // Save session in each recursion to ensure message persistence
+    await this.messageManager.saveSession();
+
     // Only create new AbortControllers for the initial call (recursionDepth === 0)
     // For recursive calls, reuse existing controllers to maintain abort signal
     let abortController: AbortController;
@@ -319,14 +310,10 @@ export class AIManager {
       this.toolAbortController = toolAbortController;
     } else {
       // Reuse existing controllers for recursive calls
-      // Fallback to new controllers if they were cleared (should not happen in normal flow but good for tests)
-      abortController = this.abortController || new AbortController();
-      toolAbortController = this.toolAbortController || new AbortController();
+      abortController = this.abortController!;
+      toolAbortController = this.toolAbortController!;
     }
 
-    // Save session in each recursion to ensure message persistence
-    await this.messageManager.saveSession();
-
     // Only set loading state for the initial call
     if (recursionDepth === 0) {
       this.setIsLoading(true);
@@ -354,17 +341,14 @@ export class AIManager {
         this.getModelConfig().permissionMode,
       );
       const toolsConfig = this.getFilteredToolsConfig(tools);
-      let effectiveSystemPrompt = buildSystemPrompt(
-        this.systemPrompt || DEFAULT_SYSTEM_PROMPT,
-        toolsConfig,
-      );
+      const toolNames = new Set(toolsConfig.map((t) => t.function.name));
+      const filteredToolPlugins = this.toolManager
+        .getTools()
+        .filter((t) => toolNames.has(t.name));
 
-      // Inject language prompt if configured
-      const language = this.getLanguage();
-      if (language) {
-        const languagePrompt = `\n\n# Language\nAlways respond in ${language}. Technical terms (e.g., code, tool names, file paths) should remain in their original language or English where appropriate.`;
-        effectiveSystemPrompt = (effectiveSystemPrompt || "") + languagePrompt;
-      }
+      let planModeOptions:
+        | { planFilePath: string; planExists: boolean }
+        | undefined;
 
       if (currentMode === "plan") {
         const planFilePath = this.permissionManager?.getPlanFilePath();
@@ -376,10 +360,7 @@ export class AIManager {
           } catch {
             planExists = false;
           }
-
-          const reminder = `\n\nPlan mode is active. The user indicated that they do not want you to execute yet -- you MUST NOT make any edits (with the exception of the plan file mentioned below), run any non-readonly tools (including changing configs or making commits), or otherwise make any changes to the system. This supercedes any other instructions you have received.\n\n## Plan File Info:\n${planExists ? `A plan file already exists at ${planFilePath}. You can read it and make incremental edits using the Edit tool if you need to.` : `No plan file exists yet. You should create your plan at ${planFilePath} using the Write tool if you need to.`}\nYou should build your plan incrementally by writing to or editing this file. NOTE that this is the only file you are allowed to edit - other than this you are only allowed to take READ-ONLY actions. You may also use the AskUserQuestion tool to gather requirements or clarify intent before finalizing your plan.`;
-
-          effectiveSystemPrompt = (effectiveSystemPrompt || "") + reminder;
+          planModeOptions = { planFilePath, planExists };
         }
       }
 
@@ -390,11 +371,20 @@ export class AIManager {
         messages: recentMessages,
         sessionId: this.messageManager.getSessionId(),
         abortSignal: abortController.signal,
-        memory: combinedMemory, // Pass combined memory content
         workdir: this.workdir, // Pass working directory
         tools: toolsConfig, // Pass filtered tool configuration
         model: model, // Use passed model
-        systemPrompt: effectiveSystemPrompt, // Pass custom system prompt
+        systemPrompt: buildSystemPrompt(
+          this.systemPrompt,
+          filteredToolPlugins,
+          {
+            workdir: this.workdir,
+            memory: combinedMemory,
+            language: this.getLanguage(),
+            isSubagent: !!this.subagentType,
+            planMode: planModeOptions,
+          },
+        ), // Pass custom system prompt
         maxTokens: maxTokens, // Pass max tokens override
       };
 
@@ -628,6 +618,15 @@ export class AIManager {
                 backgroundTaskManager: this.backgroundTaskManager,
                 workdir: this.workdir,
                 messageId: this.messageManager.getMessages().slice(-1)[0]?.id,
+                sessionId: this.messageManager.getSessionId(),
+                taskManager: this.taskManager,
+                onShortResultUpdate: (shortResult: string) => {
+                  this.messageManager.updateToolBlock({
+                    id: toolId,
+                    shortResult,
+                    stage: "running", // Keep it in running stage while updating shortResult
+                  });
+                },
               };
 
               // Execute tool
@@ -637,23 +636,6 @@ export class AIManager {
                 context,
               );
 
-              // Check if the tool was backgrounded via Ctrl-B
-              // If it was backgrounded, we should abort the AI recursion
-              if (
-                toolResult.success &&
-                toolResult.content.includes(
-                  "Command was manually backgrounded by user",
-                )
-              ) {
-                this.logger?.info(
-                  `Tool ${toolName} was backgrounded via Ctrl-B, aborting AI recursion`,
-                );
-                // Use abortAIMessage directly instead of abortRecursion to avoid double logging
-                // and ensure we don't trigger the "Request was aborted" error block
-                this.abortAIMessage();
-                return;
-              }
-
               // Update message state - tool execution completed
               this.messageManager.updateToolBlock({
                 id: toolId,
@@ -666,6 +648,7 @@ export class AIManager {
                 stage: "end",
                 name: toolName,
                 shortResult: toolResult.shortResult,
+                isManuallyBackgrounded: toolResult.isManuallyBackgrounded,
               });
 
               // Execute PostToolUse hooks after successful tool completion
@@ -690,6 +673,7 @@ export class AIManager {
                 stage: "end",
                 name: toolName,
                 compactParams,
+                isManuallyBackgrounded: false,
               });
             }
           },
@@ -721,7 +705,25 @@ export class AIManager {
         const isCurrentlyAborted =
           abortController.signal.aborted || toolAbortController.signal.aborted;
 
-        if (!isCurrentlyAborted) {
+        // Check if all tools were manually backgrounded
+        const lastMessage =
+          this.messageManager.getMessages()[
+            this.messageManager.getMessages().length - 1
+          ];
+        const toolBlocks =
+          lastMessage?.blocks.filter(
+            (block): block is import("../types/messaging.js").ToolBlock =>
+              block.type === "tool",
+          ) || [];
+        const hasBackgrounded =
+          toolBlocks.length > 0 &&
+          toolBlocks.some((block) => block.isManuallyBackgrounded);
+
+        if (hasBackgrounded) {
+          this.logger?.info(
+            "Some tools were manually backgrounded, stopping recursion.",
+          );
+        } else if (!isCurrentlyAborted) {
           // Recursively call AI service, increment recursion depth, and pass same configuration
           await this.sendAIMessage({
             recursionDepth: recursionDepth + 1,
@@ -733,18 +735,9 @@ export class AIManager {
         }
       }
     } catch (error) {
-      // Check if the error is an abort error
-      // Use the local variables to avoid null reference if this.abortController was cleared
-      const isCurrentlyAborted =
-        abortController.signal.aborted || toolAbortController.signal.aborted;
-
-      if (isCurrentlyAborted) {
-        this.logger?.info("AI message processing was aborted");
-      } else {
-        this.messageManager.addErrorBlock(
-          error instanceof Error ? error.message : "Unknown error occurred",
-        );
-      }
+      this.messageManager.addErrorBlock(
+        error instanceof Error ? error.message : "Unknown error occurred",
+      );
     } finally {
       // Only execute cleanup and hooks for the initial call
       if (recursionDepth === 0) {

package/src/managers/backgroundTaskManager.ts

@@ -69,6 +69,26 @@ export class BackgroundTaskManager {
       status: "running",
       stdout: "",
       stderr: "",
+      onStop: () => {
+        try {
+          if (child.pid) {
+            process.kill(-child.pid, "SIGTERM");
+            setTimeout(() => {
+              if (child.pid && !child.killed) {
+                try {
+                  process.kill(-child.pid, "SIGKILL");
+                } catch (error) {
+                  logger.error("Failed to force kill process:", error);
+                }
+              }
+            }, 1000);
+          } else {
+            child.kill("SIGTERM");
+          }
+        } catch {
+          child.kill("SIGTERM");
+        }
+      },
     };
 
     this.tasks.set(id, shell);
@@ -155,6 +175,26 @@ export class BackgroundTaskManager {
       status: "running",
       stdout: initialStdout,
       stderr: initialStderr,
+      onStop: () => {
+        try {
+          if (child.pid) {
+            process.kill(-child.pid, "SIGTERM");
+            setTimeout(() => {
+              if (child.pid && !child.killed) {
+                try {
+                  process.kill(-child.pid, "SIGKILL");
+                } catch (error) {
+                  logger.error("Failed to force kill process:", error);
+                }
+              }
+            }, 1000);
+          } else {
+            child.kill("SIGTERM");
+          }
+        } catch {
+          child.kill("SIGTERM");
+        }
+      },
     };
 
     this.tasks.set(id, shell);
@@ -232,65 +272,29 @@ export class BackgroundTaskManager {
       return false;
     }
 
-    if (task.type === "shell") {
-      const shell = task as BackgroundShell;
+    if (task.onStop) {
       try {
-        // Try to kill process group first
-        if (shell.process.pid) {
-          process.kill(-shell.process.pid, "SIGTERM");
-
-          // Force kill after timeout
-          setTimeout(() => {
-            if (
-              shell.status === "running" &&
-              shell.process.pid &&
-              !shell.process.killed
-            ) {
-              try {
-                process.kill(-shell.process.pid, "SIGKILL");
-              } catch (error) {
-                logger.error("Failed to force kill process:", error);
-              }
-            }
-          }, 1000);
-        }
-
-        shell.status = "killed";
-        shell.endTime = Date.now();
-        shell.runtime = shell.endTime - shell.startTime;
-        this.notifyTasksChange();
-        return true;
-      } catch {
-        // Fallback to direct process kill
-        try {
-          shell.process.kill("SIGTERM");
-          setTimeout(() => {
-            if (!shell.process.killed) {
-              shell.process.kill("SIGKILL");
-            }
-          }, 1000);
-          shell.status = "killed";
-          shell.endTime = Date.now();
-          shell.runtime = shell.endTime - shell.startTime;
-          this.notifyTasksChange();
-          return true;
-        } catch (directKillError) {
-          logger.error("Failed to kill child process:", directKillError);
-          return false;
+        const result = task.onStop();
+        if (result instanceof Promise) {
+          result.catch((error) => {
+            logger.error("Error in background task onStop callback:", error);
+          });
         }
+      } catch (error) {
+        logger.error("Error in background task onStop callback:", error);
       }
-    } else if (task.type === "subagent") {
-      // Subagent termination logic will be handled by aborting the AI loop
-      // which is already managed by the SubagentManager and AIManager.
-      // Here we just update the status.
-      task.status = "killed";
-      task.endTime = Date.now();
-      task.runtime = task.endTime - task.startTime;
-      this.notifyTasksChange();
-      return true;
     }
 
-    return false;
+    // If it's a subagent task, we should also notify the subagent manager to cleanup
+    // However, to avoid circular dependency, we rely on the onStop callback
+    // which is already set to instance.aiManager.abortAIMessage()
+    // The subagentManager.cleanupInstance will be called by the tool or by status change.
+
+    task.status = "killed";
+    task.endTime = Date.now();
+    task.runtime = task.endTime - task.startTime;
+    this.notifyTasksChange();
+    return true;
   }
 
   public cleanup(): void {

package/src/managers/foregroundTaskManager.ts

@@ -14,8 +14,9 @@ export class ForegroundTaskManager implements IForegroundTaskManager {
   }
 
   public async backgroundCurrentTask(): Promise<void> {
-    const task = this.activeForegroundTasks.pop();
-    if (task) {
+    const tasks = [...this.activeForegroundTasks].reverse();
+    this.activeForegroundTasks = [];
+    for (const task of tasks) {
       await task.backgroundHandler();
     }
   }

package/src/managers/mcpManager.ts

@@ -3,7 +3,7 @@ import { join } from "path";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { ChatCompletionFunctionTool } from "openai/resources.js";
-import { createMcpToolPlugin, findToolServer } from "@/utils/mcpUtils.js";
+import { createMcpToolPlugin, findToolServer } from "../utils/mcpUtils.js";
 import type { ToolPlugin, ToolResult, ToolContext } from "../tools/types.js";
 import type {
   Logger,
@@ -476,8 +476,11 @@ export class McpManager {
       const server = findToolServer(tool.name, servers);
 
       if (server) {
-        const plugin = createMcpToolPlugin(tool, server.name, (name, args) =>
-          this.executeMcpTool(name, args),
+        const plugin = createMcpToolPlugin(
+          tool,
+          server.name,
+          (name: string, args: Record<string, unknown>) =>
+            this.executeMcpTool(name, args),
         );
         mcpTools.set(plugin.name, plugin);
       }