wave-agent-sdk 0.5.0 → 0.6.2

package/src/managers/subagentManager.ts

@@ -8,20 +8,19 @@ import type {
   ModelConfig,
   Usage,
 } from "../types/index.js";
-import type { SessionData } from "../services/session.js";
 import { AIManager } from "./aiManager.js";
 import { MessageManager } from "./messageManager.js";
 import { ToolManager } from "./toolManager.js";
 import { HookManager } from "./hookManager.js";
-import {
-  UserMessageParams,
-  type AgentToolBlockUpdateParams,
-} from "../utils/messageOperations.js";
 import {
   addConsolidatedAbortListener,
   createAbortPromise,
 } from "../utils/abortUtils.js";
 import { BackgroundTaskManager } from "./backgroundTaskManager.js";
+import {
+  UserMessageParams,
+  type AgentToolBlockUpdateParams,
+} from "../utils/messageOperations.js";
 
 export interface SubagentManagerCallbacks {
   // Granular subagent message callbacks (015-subagent-message-callbacks)
@@ -51,6 +50,11 @@ export interface SubagentManagerCallbacks {
   ) => void;
   /** Triggered when subagent messages change */
   onSubagentMessagesChange?: (subagentId: string, messages: Message[]) => void;
+  /** Triggered when subagent latest total tokens change */
+  onSubagentLatestTotalTokensChange?: (
+    subagentId: string,
+    tokens: number,
+  ) => void;
 }
 
 export interface SubagentInstance {
@@ -61,14 +65,16 @@ export interface SubagentInstance {
   toolManager: ToolManager;
   status: "initializing" | "active" | "completed" | "error" | "aborted";
   messages: Message[];
+  lastTools: string[]; // Track last two tool names
   subagentType: string; // Store the subagent type for hook context
   backgroundTaskId?: string; // ID of the background task if transitioned
+  onUpdate?: () => void; // Optional callback for real-time updates
 }
 
 export interface SubagentManagerOptions {
   workdir: string;
   parentToolManager: ToolManager;
-  parentMessageManager: MessageManager;
+  taskManager: import("../services/taskManager.js").TaskManager;
   callbacks?: SubagentManagerCallbacks; // Use SubagentManagerCallbacks instead of parentCallbacks
   logger?: Logger;
   getGatewayConfig: () => GatewayConfig;
@@ -87,7 +93,7 @@ export class SubagentManager {
 
   private workdir: string;
   private parentToolManager: ToolManager;
-  private parentMessageManager: MessageManager;
+  private taskManager: import("../services/taskManager.js").TaskManager;
   private callbacks?: SubagentManagerCallbacks; // Use SubagentManagerCallbacks instead of parentCallbacks
   private logger?: Logger;
   private getGatewayConfig: () => GatewayConfig;
@@ -102,7 +108,7 @@ export class SubagentManager {
   constructor(options: SubagentManagerOptions) {
     this.workdir = options.workdir;
     this.parentToolManager = options.parentToolManager;
-    this.parentMessageManager = options.parentMessageManager;
+    this.taskManager = options.taskManager;
     this.callbacks = options.callbacks; // Store SubagentManagerCallbacks
     this.logger = options.logger;
     this.getGatewayConfig = options.getGatewayConfig;
@@ -168,6 +174,7 @@ export class SubagentManager {
       subagent_type: string;
     },
     runInBackground?: boolean,
+    onUpdate?: () => void,
   ): Promise<SubagentInstance> {
     if (!this.parentToolManager) {
       throw new Error(
@@ -196,6 +203,7 @@ export class SubagentManager {
     const aiManager = new AIManager({
       messageManager,
       toolManager,
+      taskManager: this.taskManager,
       logger: this.logger,
       workdir: this.workdir,
       systemPrompt: configuration.systemPrompt,
@@ -239,22 +247,13 @@ export class SubagentManager {
       toolManager,
       status: "initializing",
       messages: [],
+      lastTools: [], // Initialize lastTools
       subagentType: parameters.subagent_type, // Store the subagent type
+      onUpdate,
     };
 
     this.instances.set(subagentId, instance);
 
-    // Create subagent block in parent message manager
-    this.parentMessageManager.addSubagentBlock(
-      subagentId,
-      configuration.name,
-      messageManager.getSessionId(),
-      configuration,
-      "active",
-      parameters,
-      runInBackground,
-    );
-
     return instance;
   }
 
@@ -278,9 +277,6 @@ export class SubagentManager {
 
       // Set status to active and update parent
       this.updateInstanceStatus(instance.subagentId, "active");
-      this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
-        status: "active",
-      });
 
       if (runInBackground && this.backgroundTaskManager) {
         const taskId = this.backgroundTaskManager.generateId();
@@ -294,6 +290,11 @@ export class SubagentManager {
           description: instance.configuration.description,
           stdout: "",
           stderr: "",
+          subagentId: instance.subagentId,
+          onStop: () => {
+            instance.aiManager.abortAIMessage();
+            this.cleanupInstance(instance.subagentId);
+          },
         });
 
         instance.backgroundTaskId = taskId;
@@ -331,9 +332,6 @@ export class SubagentManager {
       return await this.internalExecute(instance, prompt, abortSignal);
     } catch (error) {
       this.updateInstanceStatus(instance.subagentId, "error");
-      this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
-        status: "error",
-      });
       throw error;
     }
   }
@@ -359,15 +357,15 @@ export class SubagentManager {
       description: instance.configuration.description,
       stdout: "",
       stderr: "",
+      subagentId: instance.subagentId,
+      onStop: () => {
+        instance.aiManager.abortAIMessage();
+        this.cleanupInstance(instance.subagentId);
+      },
     });
 
     instance.backgroundTaskId = taskId;
 
-    // Update parent message manager to reflect background status
-    this.parentMessageManager.updateSubagentBlock(subagentId, {
-      runInBackground: true,
-    });
-
     return taskId;
   }
 
@@ -378,24 +376,18 @@ export class SubagentManager {
   ): Promise<string> {
     // Set up consolidated abort handler to prevent listener accumulation
     let abortCleanup: (() => void) | undefined;
-    if (abortSignal) {
+    // Only link to parent abort signal if NOT running in background
+    if (abortSignal && !instance.backgroundTaskId) {
       abortCleanup = addConsolidatedAbortListener(abortSignal, [
         () => {
           // Update status to aborted
-          // Only update status if it's NOT a background task
-          if (!instance.backgroundTaskId) {
-            this.updateInstanceStatus(instance.subagentId, "aborted");
-            this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
-              status: "aborted",
-            });
-          }
+          this.updateInstanceStatus(instance.subagentId, "aborted");
+          // Cleanup instance immediately on abort
+          this.cleanupInstance(instance.subagentId);
         },
         () => {
           // Abort the AI execution
-          // Only abort if it's NOT a background task
-          if (!instance.backgroundTaskId) {
-            instance.aiManager.abortAIMessage();
-          }
+          instance.aiManager.abortAIMessage();
         },
       ]);
     }
@@ -441,8 +433,6 @@ export class SubagentManager {
       });
 
       // If we have an abort signal, race against it using utilities to prevent listener accumulation
-      // BUT: If this is a background task, we DON'T want to race against the abort signal
-      // because the abort signal (Esc) should only stop the tool watching the task, not the task itself.
       if (abortSignal && !instance.backgroundTaskId) {
         await Promise.race([
           executeAI,
@@ -470,9 +460,6 @@ export class SubagentManager {
 
       // Update status to completed and update parent
       this.updateInstanceStatus(instance.subagentId, "completed");
-      this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
-        status: "completed",
-      });
 
       // If this was transitioned to background, update the background task
       if (instance.backgroundTaskId && this.backgroundTaskManager) {
@@ -576,98 +563,6 @@ export class SubagentManager {
     this.instances.clear();
   }
 
-  /**
-   * Restore subagent instances from saved session data
-   * This method is called during agent initialization to restore previous subagent states
-   */
-  async restoreSubagentSessions(
-    subagentSessions: Array<{
-      sessionData: SessionData;
-      subagentId: string;
-      configuration: SubagentConfiguration;
-    }>,
-  ): Promise<void> {
-    for (const { sessionData, subagentId, configuration } of subagentSessions) {
-      try {
-        // Use the configuration from the SubagentBlock
-
-        // Create the subagent instance without executing - just restore the state
-
-        // Create MessageManager for the restored subagent
-        const subagentCallbacks = this.createSubagentCallbacks(subagentId);
-        const messageManager = new MessageManager({
-          callbacks: subagentCallbacks,
-          workdir: this.workdir,
-          logger: this.logger,
-          sessionType: "subagent",
-          subagentType: configuration.name, // Use configuration name for restored sessions
-          memoryRuleManager: this.memoryRuleManager,
-        });
-
-        // Use the parent tool manager
-        const toolManager = this.parentToolManager;
-
-        // Determine model to use
-        let modelToUse: string;
-        const parentModelConfig = this.getModelConfig();
-
-        if (!configuration.model || configuration.model === "inherit") {
-          modelToUse = parentModelConfig.agentModel;
-        } else if (configuration.model === "fastModel") {
-          modelToUse = parentModelConfig.fastModel;
-        } else {
-          modelToUse = configuration.model;
-        }
-
-        // Create AIManager for the restored subagent
-        const aiManager = new AIManager({
-          messageManager,
-          toolManager,
-          logger: this.logger,
-          workdir: this.workdir,
-          systemPrompt: configuration.systemPrompt,
-          subagentType: configuration.name, // Use configuration name as subagent type for restored instances
-          hookManager: this.hookManager,
-          permissionManager: this.parentToolManager.getPermissionManager(),
-          getGatewayConfig: this.getGatewayConfig,
-          getModelConfig: () => ({
-            ...parentModelConfig,
-            agentModel: modelToUse,
-          }),
-          getMaxInputTokens: this.getMaxInputTokens,
-          getLanguage: this.getLanguage,
-          callbacks: {
-            onUsageAdded: this.onUsageAdded,
-          },
-        });
-
-        // Create restored instance
-        const instance: SubagentInstance = {
-          subagentId,
-          configuration,
-          aiManager,
-          messageManager,
-          toolManager,
-          status: "completed", // Restored sessions are considered completed
-          messages: sessionData.messages,
-          subagentType: configuration.name, // Use configuration name as subagent type for restored instances
-        };
-
-        // IMPORTANT: Store instance in map BEFORE calling setMessages
-        // This ensures the callback can find the instance
-        this.instances.set(subagentId, instance);
-
-        messageManager.initializeFromSession(sessionData);
-      } catch (error) {
-        this.logger?.warn(
-          `Failed to restore subagent session ${subagentId}:`,
-          error,
-        );
-        // Continue with other sessions even if one fails
-      }
-    }
-  }
-
   /**
    * Create subagent callbacks for a specific subagent ID
    * Extracted to reuse in both create and restore flows
@@ -710,6 +605,20 @@ export class SubagentManager {
       },
 
       onToolBlockUpdated: (params: AgentToolBlockUpdateParams) => {
+        // Track last two tool names when a tool starts running
+        if (params.stage === "running" && params.name) {
+          const instance = this.instances.get(subagentId);
+          if (instance) {
+            // Add to lastTools if it's different from the last one or we want to show duplicates
+            // Based on "ToolA, ToolB" requirement, we'll just keep the last two
+            instance.lastTools.push(params.name);
+            if (instance.lastTools.length > 2) {
+              instance.lastTools.shift();
+            }
+            instance.onUpdate?.();
+          }
+        }
+
         // Forward tool block updates to parent via SubagentManager callbacks
         if (this.callbacks?.onSubagentToolBlockUpdated) {
           this.callbacks.onSubagentToolBlockUpdated(subagentId, params);
@@ -721,6 +630,8 @@ export class SubagentManager {
         const instance = this.instances.get(subagentId);
         if (instance) {
           instance.messages = messages;
+          // Trigger the onUpdate callback if provided
+          instance.onUpdate?.();
           // Forward subagent message changes to parent via callbacks
           if (this.callbacks?.onSubagentMessagesChange) {
             this.callbacks.onSubagentMessagesChange(subagentId, messages);
@@ -728,11 +639,16 @@ export class SubagentManager {
         }
       },
 
-      onSessionIdChange: (newSessionId: string) => {
-        // Update the subagent block with the new session ID
-        this.parentMessageManager.updateSubagentBlock(subagentId, {
-          sessionId: newSessionId,
-        });
+      onLatestTotalTokensChange: (tokens: number) => {
+        const instance = this.instances.get(subagentId);
+        if (instance) {
+          // Trigger the onUpdate callback if provided
+          instance.onUpdate?.();
+        }
+        // Forward latest total tokens to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentLatestTotalTokensChange) {
+          this.callbacks.onSubagentLatestTotalTokensChange(subagentId, tokens);
+        }
       },
     };
   }

package/src/managers/toolManager.ts

@@ -13,10 +13,15 @@ import { globTool } from "../tools/globTool.js";
 import { grepTool } from "../tools/grepTool.js";
 import { lsTool } from "../tools/lsTool.js";
 import { readTool } from "../tools/readTool.js";
-import { todoWriteTool } from "../tools/todoWriteTool.js";
 import { lspTool } from "../tools/lspTool.js";
 import { createTaskTool } from "../tools/taskTool.js";
 import { createSkillTool } from "../tools/skillTool.js";
+import {
+  taskCreateTool,
+  taskGetTool,
+  taskUpdateTool,
+  taskListTool,
+} from "../tools/taskManagementTools.js";
 import { McpManager } from "./mcpManager.js";
 import { PermissionManager } from "./permissionManager.js";
 import { ChatCompletionFunctionTool } from "openai/resources.js";
@@ -39,6 +44,8 @@ export interface ToolManagerOptions {
   permissionManager?: PermissionManager;
   /** Foreground task manager for backgrounding tasks */
   foregroundTaskManager?: import("../types/processes.js").IForegroundTaskManager;
+  /** Task manager for task management */
+  taskManager?: import("../services/taskManager.js").TaskManager;
   /** Reversion manager for file snapshots */
   reversionManager?: ReversionManager;
   /** Background task manager for background execution */
@@ -63,6 +70,7 @@ class ToolManager {
   private permissionManager?: PermissionManager;
   private foregroundTaskManager?: import("../types/processes.js").IForegroundTaskManager;
   private reversionManager?: ReversionManager;
+  private taskManager?: import("../services/taskManager.js").TaskManager;
   private backgroundTaskManager?: import("./backgroundTaskManager.js").BackgroundTaskManager;
   private permissionMode?: PermissionMode;
   private canUseToolCallback?: PermissionCallback;
@@ -72,6 +80,7 @@ class ToolManager {
     this.lspManager = options.lspManager;
     this.logger = options.logger;
     this.permissionManager = options.permissionManager;
+    this.taskManager = options.taskManager;
     this.foregroundTaskManager = options.foregroundTaskManager;
     this.reversionManager = options.reversionManager;
     this.backgroundTaskManager = options.backgroundTaskManager;
@@ -128,8 +137,11 @@ class ToolManager {
       grepTool,
       lsTool,
       readTool,
-      todoWriteTool,
       lspTool,
+      taskCreateTool,
+      taskGetTool,
+      taskUpdateTool,
+      taskListTool,
     ];
 
     for (const tool of builtInTools) {
@@ -177,11 +189,13 @@ class ToolManager {
       permissionMode: effectivePermissionMode,
       canUseToolCallback: this.canUseToolCallback,
       permissionManager: this.permissionManager,
+      taskManager: this.taskManager!,
       reversionManager: this.reversionManager,
       backgroundTaskManager: this.backgroundTaskManager,
       foregroundTaskManager: this.foregroundTaskManager,
       mcpManager: this.mcpManager,
       lspManager: this.lspManager,
+      sessionId: context.sessionId,
     };
 
     this.logger?.debug("Executing tool with enhanced context", {
@@ -252,6 +266,13 @@ class ToolManager {
     return [...builtInToolsConfig, ...mcpToolsConfig];
   }
 
+  /**
+   * Get the list of registered tool plugins
+   */
+  public getTools(): ToolPlugin[] {
+    return Array.from(this.tools.values());
+  }
+
   /**
    * Get the current permission mode
    */
@@ -278,6 +299,15 @@ class ToolManager {
   public getPermissionManager(): PermissionManager | undefined {
     return this.permissionManager;
   }
+
+  /**
+   * Get the task manager
+   */
+  public getTaskManager():
+    | import("../services/taskManager.js").TaskManager
+    | undefined {
+    return this.taskManager;
+  }
 }
 
 // Export tool registry class and types