wave-agent-sdk 0.0.7 → 0.0.8

package/src/managers/subagentManager.ts

@@ -7,12 +7,39 @@ 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";
+
+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 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 +49,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;
+  hookManager?: HookManager;
   onUsageAdded?: (usage: Usage) => void;
 }
 
@@ -42,20 +72,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 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.hookManager = options.hookManager;
     this.onUsageAdded = options.onUsageAdded;
   }
 
@@ -126,24 +160,15 @@ 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",
+      parentSessionId: this.parentMessageManager.getSessionId(),
+      subagentType: parameters.subagent_type,
     });
 
     // Use the parent tool manager directly - tool restrictions will be handled by allowedTools parameter
@@ -162,6 +187,8 @@ export class SubagentManager {
       logger: this.logger,
       workdir: this.workdir,
       systemPrompt: configuration.systemPrompt,
+      subagentType: parameters.subagent_type, // Pass subagent type for hook context
+      hookManager: this.hookManager,
       gatewayConfig: this.gatewayConfig,
       modelConfig: {
         ...this.modelConfig,
@@ -181,6 +208,7 @@ export class SubagentManager {
       toolManager,
       status: "initializing",
       messages: [],
+      subagentType: parameters.subagent_type, // Store the subagent type
     };
 
     this.instances.set(subagentId, instance);
@@ -189,8 +217,9 @@ export class SubagentManager {
     this.parentMessageManager.addSubagentBlock(
       subagentId,
       configuration.name,
+      messageManager.getSessionId(),
+      configuration,
       "active",
-      [],
       parameters,
     );
 
@@ -226,7 +255,6 @@ export class SubagentManager {
           this.updateInstanceStatus(instance.subagentId, "aborted");
           this.parentMessageManager.updateSubagentBlock(instance.subagentId, {
             status: "aborted",
-            messages: instance.messages,
           });
         });
       }
@@ -292,11 +320,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 +397,239 @@ 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",
+          parentSessionId: this.parentMessageManager.getSessionId(),
+          subagentType: configuration.name, // Use configuration name for restored sessions
+        });
+
+        // Use the parent tool manager
+        const toolManager = this.parentToolManager;
+
+        // Determine model to use
+        const modelToUse =
+          configuration.model && configuration.model !== "inherit"
+            ? configuration.model
+            : this.modelConfig.agentModel;
+
+        // 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,
+          gatewayConfig: this.gatewayConfig,
+          modelConfig: {
+            ...this.modelConfig,
+            agentModel: modelToUse,
+          },
+          tokenLimit: this.tokenLimit,
+          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);
+
+        // Now restore the session data including sessionId, which will trigger the callback chain
+        const sessionDataObj = {
+          id: sessionData.id,
+          messages: sessionData.messages,
+          version: "1.0.0",
+          metadata: {
+            workdir: this.workdir,
+            startedAt: new Date().toISOString(),
+            lastActiveAt: new Date().toISOString(),
+            latestTotalTokens: 0,
+          },
+        };
+        messageManager.initializeFromSession(sessionDataObj);
+      } 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,
+          );
+        }
+      },
+
+      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,
+        });
+      },
+    };
+  }
+
+  /**
+   * Update configuration for SubagentManager and all active subagents
+   * This method updates configuration for live config reload support
+   * @param newGatewayConfig - New gateway configuration
+   * @param newModelConfig - New model configuration
+   * @param newTokenLimit - New token limit
+   */
+  updateConfiguration(
+    newGatewayConfig: GatewayConfig,
+    newModelConfig: ModelConfig,
+    newTokenLimit: number,
+  ): void {
+    this.logger?.info("Live Config: Updating SubagentManager configuration");
+
+    // Update stored configuration
+    this.gatewayConfig = newGatewayConfig;
+    this.modelConfig = newModelConfig;
+    this.tokenLimit = newTokenLimit;
+
+    // Update all active subagent AIManager instances
+    let updatedCount = 0;
+    for (const [subagentId, instance] of this.instances.entries()) {
+      if (instance.status === "active" || instance.status === "initializing") {
+        try {
+          // For subagents, we need to preserve their model if it was explicitly set
+          const subagentModelConfig = {
+            ...newModelConfig,
+            // If subagent has its own model configured, preserve it
+            agentModel:
+              instance.configuration.model &&
+              instance.configuration.model !== "inherit"
+                ? instance.configuration.model
+                : newModelConfig.agentModel,
+          };
+
+          instance.aiManager.updateConfiguration(
+            newGatewayConfig,
+            subagentModelConfig,
+            newTokenLimit,
+          );
+          updatedCount++;
+
+          this.logger?.debug(
+            `Live Config: Updated configuration for subagent ${subagentId}`,
+          );
+        } catch (error) {
+          this.logger?.error(
+            `Live Config: Failed to update configuration for subagent ${subagentId}: ${(error as Error).message}`,
+          );
+        }
+      }
+    }
+
+    this.logger?.info(
+      `Live Config: SubagentManager configuration updated - ${updatedCount} active subagents updated`,
+    );
+  }
+
+  /**
+   * Get current configuration for debugging
+   */
+  getCurrentConfiguration(): {
+    gatewayConfig: GatewayConfig;
+    modelConfig: ModelConfig;
+    tokenLimit: number;
+    activeSubagents: number;
+  } {
+    const activeSubagents = Array.from(this.instances.values()).filter(
+      (instance) =>
+        instance.status === "active" || instance.status === "initializing",
+    ).length;
+
+    return {
+      gatewayConfig: { ...this.gatewayConfig },
+      modelConfig: { ...this.modelConfig },
+      tokenLimit: this.tokenLimit,
+      activeSubagents,
+    };
+  }
 }