wave-agent-sdk 0.0.7 → 0.0.10

package/src/managers/subagentManager.ts

@@ -7,12 +7,49 @@ import type {
   ModelConfig,
   Usage,
 } from "../types/index.js";
+import type { SessionData } from "../services/session.js";
 import { AIManager } from "./aiManager.js";
-import {
-  MessageManager,
-  type MessageManagerCallbacks,
-} from "./messageManager.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";
+
+export interface SubagentManagerCallbacks {
+  // Granular subagent message callbacks (015-subagent-message-callbacks)
+  /** Triggered when subagent adds user message */
+  onSubagentUserMessageAdded?: (
+    subagentId: string,
+    params: UserMessageParams,
+  ) => void;
+  /** Triggered when subagent creates assistant message */
+  onSubagentAssistantMessageAdded?: (subagentId: string) => void;
+  /** Triggered during subagent content streaming updates */
+  onSubagentAssistantContentUpdated?: (
+    subagentId: string,
+    chunk: string,
+    accumulated: string,
+  ) => void;
+  /** Triggered during subagent reasoning streaming updates */
+  onSubagentAssistantReasoningUpdated?: (
+    subagentId: string,
+    chunk: string,
+    accumulated: string,
+  ) => void;
+  /** Triggered when subagent tool block is updated */
+  onSubagentToolBlockUpdated?: (
+    subagentId: string,
+    params: AgentToolBlockUpdateParams,
+  ) => void;
+  /** Triggered when subagent messages change */
+  onSubagentMessagesChange?: (subagentId: string, messages: Message[]) => void;
+}
 
 export interface SubagentInstance {
   subagentId: string;
@@ -22,16 +59,19 @@ export interface SubagentInstance {
   toolManager: ToolManager;
   status: "initializing" | "active" | "completed" | "error" | "aborted";
   messages: Message[];
+  subagentType: string; // Store the subagent type for hook context
 }
 
 export interface SubagentManagerOptions {
   workdir: string;
   parentToolManager: ToolManager;
   parentMessageManager: MessageManager;
+  callbacks?: SubagentManagerCallbacks; // Use SubagentManagerCallbacks instead of parentCallbacks
   logger?: Logger;
-  gatewayConfig: GatewayConfig;
-  modelConfig: ModelConfig;
-  tokenLimit: number;
+  getGatewayConfig: () => GatewayConfig;
+  getModelConfig: () => ModelConfig;
+  getTokenLimit: () => number;
+  hookManager?: HookManager;
   onUsageAdded?: (usage: Usage) => void;
 }
 
@@ -42,20 +82,24 @@ export class SubagentManager {
   private workdir: string;
   private parentToolManager: ToolManager;
   private parentMessageManager: MessageManager;
+  private callbacks?: SubagentManagerCallbacks; // Use SubagentManagerCallbacks instead of parentCallbacks
   private logger?: Logger;
-  private gatewayConfig: GatewayConfig;
-  private modelConfig: ModelConfig;
-  private tokenLimit: number;
+  private getGatewayConfig: () => GatewayConfig;
+  private getModelConfig: () => ModelConfig;
+  private getTokenLimit: () => number;
+  private hookManager?: HookManager;
   private onUsageAdded?: (usage: Usage) => void;
 
   constructor(options: SubagentManagerOptions) {
     this.workdir = options.workdir;
     this.parentToolManager = options.parentToolManager;
     this.parentMessageManager = options.parentMessageManager;
+    this.callbacks = options.callbacks; // Store SubagentManagerCallbacks
     this.logger = options.logger;
-    this.gatewayConfig = options.gatewayConfig;
-    this.modelConfig = options.modelConfig;
-    this.tokenLimit = options.tokenLimit;
+    this.getGatewayConfig = options.getGatewayConfig;
+    this.getModelConfig = options.getModelConfig;
+    this.getTokenLimit = options.getTokenLimit;
+    this.hookManager = options.hookManager;
     this.onUsageAdded = options.onUsageAdded;
   }
 
@@ -112,12 +156,7 @@ export class SubagentManager {
       subagent_type: string;
     },
   ): Promise<SubagentInstance> {
-    if (
-      !this.parentToolManager ||
-      !this.gatewayConfig ||
-      !this.modelConfig ||
-      !this.tokenLimit
-    ) {
+    if (!this.parentToolManager) {
       throw new Error(
         "SubagentManager not properly initialized - call initialize() first",
       );
@@ -126,35 +165,19 @@ export class SubagentManager {
     const subagentId = randomUUID();
 
     // Create isolated MessageManager for the subagent
-    const subagentCallbacks: MessageManagerCallbacks = {
-      // These callbacks will be handled by the parent agent
-      onMessagesChange: (messages: Message[]) => {
-        const instance = this.instances.get(subagentId);
-        if (instance) {
-          instance.messages = messages;
-          // Update parent's subagent block with latest messages
-          this.parentMessageManager.updateSubagentBlock(subagentId, {
-            messages: messages,
-          });
-        }
-      },
-    };
+    const subagentCallbacks = this.createSubagentCallbacks(subagentId);
 
     const messageManager = new MessageManager({
       callbacks: subagentCallbacks,
       workdir: this.workdir,
       logger: this.logger,
+      sessionType: "subagent",
+      subagentType: parameters.subagent_type,
     });
 
     // Use the parent tool manager directly - tool restrictions will be handled by allowedTools parameter
     const toolManager = this.parentToolManager;
 
-    // Determine model to use
-    const modelToUse =
-      configuration.model && configuration.model !== "inherit"
-        ? configuration.model
-        : this.modelConfig.agentModel;
-
     // Create isolated AIManager for the subagent
     const aiManager = new AIManager({
       messageManager,
@@ -162,12 +185,31 @@ export class SubagentManager {
       logger: this.logger,
       workdir: this.workdir,
       systemPrompt: configuration.systemPrompt,
-      gatewayConfig: this.gatewayConfig,
-      modelConfig: {
-        ...this.modelConfig,
-        agentModel: modelToUse,
+      subagentType: parameters.subagent_type, // Pass subagent type for hook context
+      hookManager: this.hookManager,
+      getGatewayConfig: this.getGatewayConfig,
+      getModelConfig: () => {
+        // Determine model dynamically each time
+        const parentModelConfig = this.getModelConfig();
+        let modelToUse: string;
+
+        if (!configuration.model || configuration.model === "inherit") {
+          // Use parent's agentModel for "inherit" or undefined
+          modelToUse = parentModelConfig.agentModel;
+        } else if (configuration.model === "fastModel") {
+          // Use parent's fastModel for special "fastModel" value
+          modelToUse = parentModelConfig.fastModel;
+        } else {
+          // Use specific model name
+          modelToUse = configuration.model;
+        }
+
+        return {
+          ...parentModelConfig,
+          agentModel: modelToUse,
+        };
       },
-      tokenLimit: this.tokenLimit,
+      getTokenLimit: this.getTokenLimit,
       callbacks: {
         onUsageAdded: this.onUsageAdded,
       },
@@ -181,6 +223,7 @@ export class SubagentManager {
       toolManager,
       status: "initializing",
       messages: [],
+      subagentType: parameters.subagent_type, // Store the subagent type
     };
 
     this.instances.set(subagentId, instance);
@@ -189,8 +232,9 @@ export class SubagentManager {
     this.parentMessageManager.addSubagentBlock(
       subagentId,
       configuration.name,
+      messageManager.getSessionId(),
+      configuration,
       "active",
-      [],
       parameters,
     );
 
@@ -220,15 +264,22 @@ export class SubagentManager {
         status: "active",
       });
 
-      // Set up abort handler
+      // Set up consolidated abort handler to prevent listener accumulation
+      let abortCleanup: (() => void) | undefined;
       if (abortSignal) {
-        abortSignal.addEventListener("abort", () => {
-          this.updateInstanceStatus(instance.subagentId, "aborted");
-          this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
-            status: "aborted",
-            messages: instance.messages,
-          });
-        });
+        abortCleanup = addConsolidatedAbortListener(abortSignal, [
+          () => {
+            // Update status to aborted
+            this.updateInstanceStatus(instance.subagentId, "aborted");
+            this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
+              status: "aborted",
+            });
+          },
+          () => {
+            // Abort the AI execution
+            instance.aiManager.abortAIMessage();
+          },
+        ]);
       }
 
       // Add the user's prompt as a message
@@ -248,32 +299,43 @@ export class SubagentManager {
 
       // Execute the AI request with tool restrictions
       // The AIManager will handle abort signals through its own abort controllers
-      // We need to abort the AI execution if the external abort signal is triggered
+      // Resolve model name for sendAIMessage
+      let resolvedModel: string | undefined;
+      if (
+        instance.configuration.model &&
+        instance.configuration.model !== "inherit"
+      ) {
+        if (instance.configuration.model === "fastModel") {
+          // Use parent's fastModel for special "fastModel" value
+          const parentModelConfig = this.getModelConfig();
+          resolvedModel = parentModelConfig.fastModel;
+        } else {
+          // Use specific model name
+          resolvedModel = instance.configuration.model;
+        }
+      }
+      // For "inherit" or undefined, resolvedModel remains undefined (uses AIManager default)
+
       const executeAI = instance.aiManager.sendAIMessage({
         allowedTools,
-        model:
-          instance.configuration.model !== "inherit"
-            ? instance.configuration.model
-            : undefined,
+        model: resolvedModel,
       });
 
-      // If we have an abort signal, race against it
-      if (abortSignal) {
-        await Promise.race([
-          executeAI,
-          new Promise<never>((_, reject) => {
-            if (abortSignal.aborted) {
-              reject(new Error("Task was aborted"));
-            }
-            abortSignal.addEventListener("abort", () => {
-              // Abort the AI execution
-              instance.aiManager.abortAIMessage();
-              reject(new Error("Task was aborted"));
-            });
-          }),
-        ]);
-      } else {
-        await executeAI;
+      try {
+        // If we have an abort signal, race against it using utilities to prevent listener accumulation
+        if (abortSignal) {
+          await Promise.race([
+            executeAI,
+            createAbortPromise(abortSignal, "Task was aborted"),
+          ]);
+        } else {
+          await executeAI;
+        }
+      } finally {
+        // Clean up abort listeners to prevent memory leaks
+        if (abortCleanup) {
+          abortCleanup();
+        }
       }
 
       // Get the latest messages to extract the response
@@ -292,11 +354,10 @@ export class SubagentManager {
       );
       const response = textBlocks.map((block) => block.content).join("\n");
 
-      // Update status to completed and update parent with final messages
+      // Update status to completed and update parent
       this.updateInstanceStatus(instance.subagentId, "completed");
       this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
         status: "completed",
-        messages: messages,
       });
 
       return response || "Task completed with no text response";
@@ -370,4 +431,162 @@ export class SubagentManager {
   cleanup(): void {
     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
+        });
+
+        // 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,
+          getGatewayConfig: this.getGatewayConfig,
+          getModelConfig: () => ({
+            ...parentModelConfig,
+            agentModel: modelToUse,
+          }),
+          getTokenLimit: this.getTokenLimit,
+          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
+   */
+  private createSubagentCallbacks(subagentId: string) {
+    return {
+      onUserMessageAdded: (params: UserMessageParams) => {
+        // Forward user message events to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentUserMessageAdded) {
+          this.callbacks.onSubagentUserMessageAdded(subagentId, params);
+        }
+      },
+
+      onAssistantMessageAdded: () => {
+        // Forward assistant message events to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentAssistantMessageAdded) {
+          this.callbacks.onSubagentAssistantMessageAdded(subagentId);
+        }
+      },
+
+      onAssistantContentUpdated: (chunk: string, accumulated: string) => {
+        // Forward assistant content updates to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentAssistantContentUpdated) {
+          this.callbacks.onSubagentAssistantContentUpdated(
+            subagentId,
+            chunk,
+            accumulated,
+          );
+        }
+      },
+      onAssistantReasoningUpdated: (chunk: string, accumulated: string) => {
+        // Forward assistant reasoning updates to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentAssistantReasoningUpdated) {
+          this.callbacks.onSubagentAssistantReasoningUpdated(
+            subagentId,
+            chunk,
+            accumulated,
+          );
+        }
+      },
+
+      onToolBlockUpdated: (params: AgentToolBlockUpdateParams) => {
+        // Forward tool block updates to parent via SubagentManager callbacks
+        if (this.callbacks?.onSubagentToolBlockUpdated) {
+          this.callbacks.onSubagentToolBlockUpdated(subagentId, params);
+        }
+      },
+
+      // These callbacks will be handled by the parent agent
+      onMessagesChange: (messages: Message[]) => {
+        const instance = this.instances.get(subagentId);
+        if (instance) {
+          instance.messages = messages;
+          // Forward subagent message changes to parent via callbacks
+          if (this.callbacks?.onSubagentMessagesChange) {
+            this.callbacks.onSubagentMessagesChange(subagentId, messages);
+          }
+        }
+      },
+
+      onSessionIdChange: (newSessionId: string) => {
+        // Update the subagent block with the new session ID
+        this.parentMessageManager.updateSubagentBlock(subagentId, {
+          sessionId: newSessionId,
+        });
+      },
+    };
+  }
 }

package/src/managers/toolManager.ts

@@ -10,30 +10,56 @@ 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 { McpManager } from "./mcpManager.js";
+import { PermissionManager } from "./permissionManager.js";
 import { ChatCompletionFunctionTool } from "openai/resources.js";
-import type { Logger } from "../types/index.js";
+import type {
+  Logger,
+  PermissionMode,
+  PermissionCallback,
+  ILspManager,
+} from "../types/index.js";
 import type { SubagentManager } from "./subagentManager.js";
 import type { SkillManager } from "./skillManager.js";
 
 export interface ToolManagerOptions {
   mcpManager: McpManager;
+  lspManager?: ILspManager;
   logger?: Logger;
+  /** Optional permission manager for handling tool permission checks */
+  permissionManager?: PermissionManager;
+  /** Permission mode for tool execution (defaults to "default") */
+  permissionMode?: PermissionMode;
+  /** Custom permission callback for tool usage */
+  canUseToolCallback?: PermissionCallback;
 }
 
 /**
  * Tool Manager
+ *
+ * Manages tool registration and execution with optional permission system integration.
+ * Supports both built-in tools and MCP (Model Context Protocol) tools.
  */
 class ToolManager {
   private tools = new Map<string, ToolPlugin>();
   private mcpManager: McpManager;
+  private lspManager?: ILspManager;
   private logger?: Logger;
+  private permissionManager?: PermissionManager;
+  private permissionMode?: PermissionMode;
+  private canUseToolCallback?: PermissionCallback;
 
   constructor(options: ToolManagerOptions) {
     this.mcpManager = options.mcpManager;
+    this.lspManager = options.lspManager;
     this.logger = options.logger;
+    this.permissionManager = options.permissionManager;
+    // Store CLI permission mode, let PermissionManager resolve effective mode
+    this.permissionMode = options.permissionMode;
+    this.canUseToolCallback = options.canUseToolCallback;
   }
 
   /**
@@ -83,6 +109,7 @@ class ToolManager {
       lsTool,
       readTool,
       todoWriteTool,
+      lspTool,
     ];
 
     for (const tool of builtInTools) {
@@ -101,22 +128,66 @@ class ToolManager {
     }
   }
 
+  /**
+   * Execute a tool by name with the provided arguments and context
+   *
+   * Enhances the context with permission-related fields before execution:
+   * - permissionMode: The current permission mode (default or bypassPermissions)
+   * - canUseToolCallback: Custom permission callback if provided
+   * - permissionManager: The PermissionManager instance for permission checks
+   *
+   * @param name - Name of the tool to execute
+   * @param args - Arguments to pass to the tool
+   * @param context - Execution context for the tool
+   * @returns Promise resolving to the tool execution result
+   */
   async execute(
     name: string,
     args: Record<string, unknown>,
     context: ToolContext,
   ): Promise<ToolResult> {
+    // Resolve effective permission mode (CLI override > configuration > default)
+    const effectivePermissionMode = this.permissionManager
+      ? this.permissionManager.getCurrentEffectiveMode(this.permissionMode)
+      : this.permissionMode || "default";
+
+    // Enhance context with permission-related fields
+    const enhancedContext: ToolContext = {
+      ...context,
+      permissionMode: effectivePermissionMode,
+      canUseToolCallback: this.canUseToolCallback,
+      permissionManager: this.permissionManager,
+      mcpManager: this.mcpManager,
+      lspManager: this.lspManager,
+    };
+
+    this.logger?.debug("Executing tool with enhanced context", {
+      toolName: name,
+      cliPermissionMode: this.permissionMode,
+      effectivePermissionMode,
+      hasPermissionManager: !!this.permissionManager,
+      hasPermissionCallback: !!this.canUseToolCallback,
+    });
+
     // Check if it's an MCP tool first
     if (this.mcpManager.isMcpTool(name)) {
-      return this.mcpManager.executeMcpToolByRegistry(name, args, context);
+      return this.mcpManager.executeMcpToolByRegistry(
+        name,
+        args,
+        enhancedContext,
+      );
     }
 
     // Check built-in tools
     const plugin = this.tools.get(name);
     if (plugin) {
       try {
-        return await plugin.execute(args, context);
+        return await plugin.execute(args, enhancedContext);
       } catch (error) {
+        this.logger?.error("Tool execution failed", {
+          toolName: name,
+          error: error instanceof Error ? error.message : String(error),
+        });
         return {
           success: false,
           content: "",
@@ -125,6 +196,7 @@ class ToolManager {
       }
     }
 
+    this.logger?.warn("Tool not found", { toolName: name });
     return {
       success: false,
       content: "",
@@ -145,6 +217,26 @@ class ToolManager {
     const mcpToolsConfig = this.mcpManager.getMcpToolsConfig();
     return [...builtInToolsConfig, ...mcpToolsConfig];
   }
+
+  /**
+   * Get the current permission mode
+   */
+  public getPermissionMode(): PermissionMode {
+    if (this.permissionManager) {
+      return this.permissionManager.getCurrentEffectiveMode(
+        this.permissionMode,
+      );
+    }
+    return this.permissionMode || "default";
+  }
+
+  /**
+   * Set the permission mode
+   * @param mode - The new permission mode
+   */
+  public setPermissionMode(mode: PermissionMode): void {
+    this.permissionMode = mode;
+  }
 }
 
 // Export tool registry class and types