wave-agent-sdk 0.5.0 → 0.6.0

package/src/managers/aiManager.ts

@@ -1,11 +1,13 @@
+import * as os from "node:os";
+import * as path from "node:path";
 import {
   callAgent,
   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 fsSync from "node:fs";
 import * as fs from "node:fs/promises";
 import type {
   Logger,
@@ -21,10 +23,24 @@ 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 "../constants/prompts.js";
+
+function isGitRepository(dirPath: string): string {
+  try {
+    // Check if .git directory exists in current directory or any parent directory
+    let currentPath = path.resolve(dirPath);
+    while (currentPath !== path.dirname(currentPath)) {
+      const gitPath = path.join(currentPath, ".git");
+      if (fsSync.existsSync(gitPath)) {
+        return "Yes";
+      }
+      currentPath = path.dirname(currentPath);
+    }
+    return "No";
+  } catch {
+    return "No";
+  }
+}
 
 export interface AIManagerCallbacks {
   onCompressionStateChange?: (isCompressing: boolean) => void;
@@ -34,6 +50,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 +77,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 +97,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 +178,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 +190,7 @@ export class AIManager {
       if (toolPlugin?.formatCompactParams) {
         const context: ToolContext = {
           workdir: this.workdir,
+          taskManager: this.taskManager,
         };
         return toolPlugin.formatCompactParams(toolArgs, context);
       }
@@ -213,9 +224,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 +257,6 @@ export class AIManager {
 
           // Execute message reconstruction and sessionId update after compression
           this.messageManager.compressMessagesAndUpdateSession(
-            insertIndex,
             compressionResult.content,
             compressionUsage,
           );
@@ -305,6 +313,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 +330,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 +361,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 +380,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 +391,23 @@ 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,
+            isGitRepo: isGitRepository(this.workdir),
+            platform: os.platform(),
+            osVersion: `${os.type()} ${os.release()}`,
+            today: new Date().toISOString().split("T")[0],
+            memory: combinedMemory,
+            language: this.getLanguage(),
+            planMode: planModeOptions,
+          },
+        ), // Pass custom system prompt
         maxTokens: maxTokens, // Pass max tokens override
       };
 
@@ -628,6 +641,8 @@ export class AIManager {
                 backgroundTaskManager: this.backgroundTaskManager,
                 workdir: this.workdir,
                 messageId: this.messageManager.getMessages().slice(-1)[0]?.id,
+                sessionId: this.messageManager.getSessionId(),
+                taskManager: this.taskManager,
               };
 
               // Execute tool
@@ -637,23 +652,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 +664,7 @@ export class AIManager {
                 stage: "end",
                 name: toolName,
                 shortResult: toolResult.shortResult,
+                isManuallyBackgrounded: toolResult.isManuallyBackgrounded,
               });
 
               // Execute PostToolUse hooks after successful tool completion
@@ -690,6 +689,7 @@ export class AIManager {
                 stage: "end",
                 name: toolName,
                 compactParams,
+                isManuallyBackgrounded: false,
               });
             }
           },
@@ -721,7 +721,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 +751,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,24 @@ 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;
+    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);
       }

package/src/managers/messageManager.ts

@@ -45,7 +45,7 @@ export interface MessageManagerCallbacks {
   onAssistantReasoningUpdated?: (chunk: string, accumulated: string) => void;
   onToolBlockUpdated?: (params: AgentToolBlockUpdateParams) => void;
   onErrorBlockAdded?: (error: string) => void;
-  onCompressBlockAdded?: (insertIndex: number, content: string) => void;
+  onCompressBlockAdded?: (content: string) => void;
   onCompressionStateChange?: (isCompressing: boolean) => void;
   onMemoryBlockAdded?: (
     content: string,
@@ -74,6 +74,10 @@ export interface MessageManagerCallbacks {
     subagentId: string,
     status: "active" | "completed" | "error" | "aborted",
   ) => void;
+  onFileHistoryBlockAdded?: (
+    snapshots: import("../types/reversion.js").FileSnapshot[],
+  ) => void;
+  onSubagentTaskStopRequested?: (subagentId: string) => void;
 }
 
 export interface MessageManagerOptions {
@@ -88,6 +92,7 @@ export interface MessageManagerOptions {
 export class MessageManager {
   // Private state properties
   private sessionId: string;
+  private rootSessionId: string;
   private messages: Message[];
   private latestTotalTokens: number;
   private userInputHistory: string[];
@@ -104,6 +109,7 @@ export class MessageManager {
 
   constructor(options: MessageManagerOptions) {
     this.sessionId = generateSessionId();
+    this.rootSessionId = this.sessionId;
     this.messages = [];
     this.latestTotalTokens = 0;
     this.userInputHistory = [];
@@ -125,6 +131,10 @@ export class MessageManager {
     return this.sessionId;
   }
 
+  public getRootSessionId(): string {
+    return this.rootSessionId;
+  }
+
   public getMessages(): Message[] {
     return [...this.messages];
   }
@@ -247,6 +257,7 @@ export class MessageManager {
         unsavedMessages, // Only append new messages
         this.workdir,
         this.sessionType,
+        this.rootSessionId,
       );
 
       // Update the saved message count
@@ -297,6 +308,7 @@ export class MessageManager {
   // Initialize state from session data
   public initializeFromSession(sessionData: SessionData): void {
     this.setSessionId(sessionData.id);
+    this.rootSessionId = sessionData.rootSessionId || sessionData.id;
     this.setMessages([...sessionData.messages]);
     this.updateFilesInContext(sessionData.messages);
     this.setlatestTotalTokens(sessionData.metadata.latestTotalTokens);
@@ -402,17 +414,7 @@ export class MessageManager {
   public updateToolBlock(params: AgentToolBlockUpdateParams): void {
     const newMessages = updateToolBlockInMessage({
       messages: this.messages,
-      id: params.id,
-      parameters: params.parameters,
-      result: params.result,
-      success: params.success,
-      error: params.error,
-      stage: params.stage,
-      name: params.name,
-      shortResult: params.shortResult,
-      images: params.images,
-      compactParams: params.compactParams,
-      parametersChunk: params.parametersChunk,
+      ...params,
     });
     this.setMessages(newMessages);
     this.callbacks.onToolBlockUpdated?.(params);
@@ -442,14 +444,14 @@ export class MessageManager {
   }
 
   /**
-   * Compress messages and update session, delete compressed messages, only keep compressed messages and subsequent messages
+   * Compress messages and update session, delete compressed messages, only keep compressed messages and last 3 messages
    */
   public compressMessagesAndUpdateSession(
-    insertIndex: number,
     compressedContent: string,
     usage?: Usage,
   ): void {
-    const currentMessages = this.messages;
+    // Get last 3 messages to preserve
+    const lastThreeMessages = this.messages.slice(-3);
 
     // Create compressed message
     const compressMessage: Message = {
@@ -464,24 +466,22 @@ export class MessageManager {
       ...(usage && { usage }),
     };
 
-    // Convert negative index to positive index
-    const actualIndex =
-      insertIndex < 0 ? currentMessages.length + insertIndex : insertIndex;
-
-    // Build new message array: keep compressed message and all messages from actualIndex onwards
-    const newMessages: Message[] = [
-      compressMessage,
-      ...currentMessages.slice(actualIndex),
-    ];
+    // Build new message array: keep the compressed message and last 3 messages
+    const newMessages: Message[] = [compressMessage, ...lastThreeMessages];
 
     // Update sessionId
     this.setSessionId(generateSessionId());
 
+    // Trigger task list update if this is the main session to ensure continuity
+    if (this.sessionType === "main") {
+      this.callbacks.onSessionIdChange?.(this.sessionId);
+    }
+
     // Set new message list
     this.setMessages(newMessages);
 
-    // Trigger compression callback, insertIndex remains unchanged
-    this.callbacks.onCompressBlockAdded?.(insertIndex, compressedContent);
+    // Trigger compression callback
+    this.callbacks.onCompressBlockAdded?.(compressedContent);
   }
 
   public addFileHistoryBlock(
@@ -496,6 +496,7 @@ export class MessageManager {
         snapshots,
       } as unknown as import("../types/index.js").MessageBlock);
       this.setMessages([...this.messages]);
+      this.callbacks.onFileHistoryBlockAdded?.(snapshots);
     }
   }
 
@@ -729,6 +730,16 @@ export class MessageManager {
 
     // Truncate messages in memory
     const newMessages = this.messages.slice(0, index);
+
+    // Identify subagent tasks to stop
+    for (const message of messagesToRemove) {
+      for (const block of message.blocks) {
+        if (block.type === "subagent" && block.subagentId) {
+          this.callbacks.onSubagentTaskStopRequested?.(block.subagentId);
+        }
+      }
+    }
+
     this.setMessages(newMessages);
 
     // Update persistence: rewrite the session file