wave-agent-sdk 0.0.6 → 0.0.8

package/src/managers/aiManager.ts

@@ -1,6 +1,7 @@
 import { callAgent, compressMessages } from "../services/aiService.js";
 import { getMessagesToCompress } from "../utils/messageOperations.js";
 import { convertMessagesForAPI } from "../utils/convertMessagesForAPI.js";
+import { calculateComprehensiveTotalTokens } from "../utils/tokenCalculation.js";
 import * as memory from "../services/memory.js";
 import type {
   Logger,
@@ -30,6 +31,7 @@ export interface AIManagerOptions {
   callbacks?: AIManagerCallbacks;
   workdir: string;
   systemPrompt?: string;
+  subagentType?: string; // Optional subagent type for hook context
   // Resolved configuration
   gatewayConfig: GatewayConfig;
   modelConfig: ModelConfig;
@@ -47,6 +49,7 @@ export class AIManager {
   private hookManager?: HookManager;
   private workdir: string;
   private systemPrompt?: string;
+  private subagentType?: string; // Store subagent type for hook context
 
   // Configuration properties
   private gatewayConfig: GatewayConfig;
@@ -61,6 +64,7 @@ export class AIManager {
     this.logger = options.logger;
     this.workdir = options.workdir;
     this.systemPrompt = options.systemPrompt;
+    this.subagentType = options.subagentType; // Store subagent type
     this.callbacks = options.callbacks ?? {};
 
     // Store resolved configuration
@@ -72,6 +76,74 @@ export class AIManager {
   private isCompressing: boolean = false;
   private callbacks: AIManagerCallbacks;
 
+  /**
+   * Update gateway configuration at runtime for live config reload
+   * @param newConfig - New gateway configuration
+   */
+  updateGatewayConfig(newConfig: GatewayConfig): void {
+    this.logger?.info(
+      `Live Config: Updating AIManager gateway config - baseURL: ${newConfig.baseURL}`,
+    );
+    this.gatewayConfig = newConfig;
+  }
+
+  /**
+   * Update model configuration at runtime for live config reload
+   * @param newConfig - New model configuration
+   */
+  updateModelConfig(newConfig: ModelConfig): void {
+    this.logger?.info(
+      `Live Config: Updating AIManager model config - agent: ${newConfig.agentModel}, fast: ${newConfig.fastModel}`,
+    );
+    this.modelConfig = newConfig;
+  }
+
+  /**
+   * Update token limit at runtime for live config reload
+   * @param newLimit - New token limit
+   */
+  updateTokenLimit(newLimit: number): void {
+    this.logger?.info(
+      `Live Config: Updating AIManager token limit: ${newLimit}`,
+    );
+    this.tokenLimit = newLimit;
+  }
+
+  /**
+   * Update all configurations at once for live config reload
+   * @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 all AIManager configuration");
+    this.gatewayConfig = newGatewayConfig;
+    this.modelConfig = newModelConfig;
+    this.tokenLimit = newTokenLimit;
+    this.logger?.info(
+      `Live Config: Configuration updated - model: ${newModelConfig.agentModel}, tokenLimit: ${newTokenLimit}`,
+    );
+  }
+
+  /**
+   * Get current configuration for debugging
+   */
+  getCurrentConfiguration(): {
+    gatewayConfig: GatewayConfig;
+    modelConfig: ModelConfig;
+    tokenLimit: number;
+  } {
+    return {
+      gatewayConfig: { ...this.gatewayConfig },
+      modelConfig: { ...this.modelConfig },
+      tokenLimit: this.tokenLimit,
+    };
+  }
+
   /**
    * Get filtered tool configuration
    */
@@ -117,7 +189,7 @@ export class AIManager {
   private generateCompactParams(
     toolName: string,
     toolArgs: Record<string, unknown>,
-  ): string | undefined {
+  ): string {
     try {
       const toolPlugin = this.toolManager
         .list()
@@ -131,21 +203,27 @@ export class AIManager {
     } catch (error) {
       this.logger?.warn("Failed to generate compactParams", error);
     }
-    return undefined;
+    return "";
   }
 
   // Private method to handle token statistics and message compression
   private async handleTokenUsageAndCompression(
-    usage: { total_tokens: number } | undefined,
+    usage: Usage | undefined,
     abortController: AbortController,
   ): Promise<void> {
     if (!usage) return;
 
-    // Update token statistics - display latest token usage
-    this.messageManager.setlatestTotalTokens(usage.total_tokens);
+    // Update token statistics - display comprehensive token usage including cache tokens
+    const comprehensiveTotalTokens = calculateComprehensiveTotalTokens(usage);
+    this.messageManager.setlatestTotalTokens(comprehensiveTotalTokens);
 
     // Check if token limit exceeded - use injected configuration
-    if (usage.total_tokens > this.tokenLimit) {
+    if (
+      usage.total_tokens +
+        (usage.cache_read_input_tokens || 0) +
+        (usage.cache_creation_input_tokens || 0) >
+      this.tokenLimit
+    ) {
       this.logger?.debug(
         `Token usage exceeded ${this.tokenLimit}, compressing messages...`,
       );
@@ -160,6 +238,9 @@ export class AIManager {
       if (messagesToCompress.length > 0) {
         const recentChatMessages = convertMessagesForAPI(messagesToCompress);
 
+        // Save session before compression to preserve original messages
+        await this.messageManager.saveSession();
+
         this.setIsCompressing(true);
         try {
           const compressionResult = await compressMessages({
@@ -169,26 +250,28 @@ export class AIManager {
             abortSignal: abortController.signal,
           });
 
-          // Execute message reconstruction and sessionId update after compression
-          this.messageManager.compressMessagesAndUpdateSession(
-            insertIndex,
-            compressionResult.content,
-          );
-
           // Handle usage tracking for compression operations
+          let compressionUsage: Usage | undefined;
           if (compressionResult.usage) {
-            const usage: Usage = {
+            compressionUsage = {
               prompt_tokens: compressionResult.usage.prompt_tokens,
               completion_tokens: compressionResult.usage.completion_tokens,
               total_tokens: compressionResult.usage.total_tokens,
               model: this.modelConfig.fastModel,
               operation_type: "compress",
             };
+          }
 
-            // Notify Agent to add to usage tracking
-            if (this.callbacks?.onUsageAdded) {
-              this.callbacks.onUsageAdded(usage);
-            }
+          // Execute message reconstruction and sessionId update after compression
+          this.messageManager.compressMessagesAndUpdateSession(
+            insertIndex,
+            compressionResult.content,
+            compressionUsage,
+          );
+
+          // Notify Agent to add to usage tracking
+          if (compressionUsage && this.callbacks?.onUsageAdded) {
+            this.callbacks.onUsageAdded(compressionUsage);
           }
 
           this.logger?.debug(
@@ -228,6 +311,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;
@@ -262,7 +348,10 @@ export class AIManager {
         this.workdir,
       );
 
-      // Call AI service (non-streaming)
+      // Add assistant message first (for streaming updates)
+      this.messageManager.addAssistantMessage();
+
+      // Call AI service with streaming callbacks
       const result = await callAgent({
         gatewayConfig: this.gatewayConfig,
         modelConfig: this.modelConfig,
@@ -274,18 +363,47 @@ export class AIManager {
         tools: this.getFilteredToolsConfig(allowedTools), // Pass filtered tool configuration
         model: model, // Use passed model
         systemPrompt: this.systemPrompt, // Pass custom system prompt
+        // Streaming callbacks
+        onContentUpdate: (content: string) => {
+          this.messageManager.updateCurrentMessageContent(content);
+        },
+        onToolUpdate: (toolCall) => {
+          // Use parametersChunk as compact param for better performance
+          // No need to extract params or generate compact params during streaming
+          this.logger?.debug("Tool streaming update:", toolCall);
+
+          // Update tool block with streaming parameters using parametersChunk as compact param
+          this.messageManager.updateToolBlock({
+            id: toolCall.id,
+            name: toolCall.name,
+            parameters: toolCall.parameters,
+            parametersChunk: toolCall.parametersChunk,
+            compactParams: toolCall.parameters?.split("\n").pop()?.slice(-30),
+            stage: toolCall.stage || "streaming", // Default to streaming if stage not provided
+          });
+        },
       });
 
-      // Collect content and tool calls
-      const content = result.content || "";
-      const toolCalls: ChatCompletionMessageFunctionToolCall[] = [];
+      // Log finish reason and response headers if available
+      if (result.finish_reason) {
+        this.logger?.debug(
+          `AI response finished with reason: ${result.finish_reason}`,
+        );
+      }
+      if (
+        result.response_headers &&
+        Object.keys(result.response_headers).length > 0
+      ) {
+        this.logger?.debug("AI response headers:", result.response_headers);
+      }
 
-      if (result.tool_calls) {
-        for (const toolCall of result.tool_calls) {
-          if (toolCall.type === "function") {
-            toolCalls.push(toolCall);
-          }
-        }
+      if (result.metadata && Object.keys(result.metadata).length > 0) {
+        this.messageManager.mergeAssistantMetadata(result.metadata);
+      }
+
+      // Handle result content from non-streaming mode
+      if (result.content) {
+        this.messageManager.updateCurrentMessageContent(result.content);
       }
 
       // Handle usage tracking for agent operations
@@ -297,156 +415,183 @@ export class AIManager {
           total_tokens: result.usage.total_tokens,
           model: model || this.modelConfig.agentModel,
           operation_type: "agent",
+          // Preserve cache fields if present
+          ...(result.usage.cache_read_input_tokens !== undefined && {
+            cache_read_input_tokens: result.usage.cache_read_input_tokens,
+          }),
+          ...(result.usage.cache_creation_input_tokens !== undefined && {
+            cache_creation_input_tokens:
+              result.usage.cache_creation_input_tokens,
+          }),
+          ...(result.usage.cache_creation && {
+            cache_creation: result.usage.cache_creation,
+          }),
         };
       }
 
-      // Add assistant message at once (including content, tool calls, and usage)
-      this.messageManager.addAssistantMessage(content, toolCalls, usage);
-
-      // Notify Agent to add to usage tracking
+      // Set usage on the assistant message if available
       if (usage) {
+        const messages = this.messageManager.getMessages();
+        const lastMessage = messages[messages.length - 1];
+        if (lastMessage && lastMessage.role === "assistant") {
+          lastMessage.usage = usage;
+          this.messageManager.setMessages(messages);
+        }
+
+        // Notify Agent to add to usage tracking
         if (this.callbacks?.onUsageAdded) {
           this.callbacks.onUsageAdded(usage);
         }
       }
 
+      // Collect tool calls for processing
+      const toolCalls: ChatCompletionMessageFunctionToolCall[] = [];
+      if (result.tool_calls) {
+        for (const toolCall of result.tool_calls) {
+          if (toolCall.type === "function") {
+            toolCalls.push(toolCall);
+          }
+        }
+      }
+
       if (toolCalls.length > 0) {
         // Execute all tools in parallel using Promise.all
         const toolExecutionPromises = toolCalls.map(
           async (functionToolCall) => {
             const toolId = functionToolCall.id || "";
 
-            try {
-              // Check if already interrupted, skip tool execution if so
-              if (
-                abortController.signal.aborted ||
-                toolAbortController.signal.aborted
-              ) {
+            // Check if already interrupted, skip tool execution if so
+            if (
+              abortController.signal.aborted ||
+              toolAbortController.signal.aborted
+            ) {
+              return;
+            }
+
+            const toolName = functionToolCall.function?.name || "";
+            // Safely parse tool parameters, handle tools without parameters
+            let toolArgs: Record<string, unknown> = {};
+            const argsString = functionToolCall.function?.arguments?.trim();
+
+            if (!argsString || argsString === "") {
+              // Tool without parameters, use empty object
+              toolArgs = {};
+            } else {
+              try {
+                toolArgs = JSON.parse(argsString);
+              } catch (parseError) {
+                // For non-empty but malformed JSON, still throw exception
+                const errorMessage = `Failed to parse tool arguments, finish_reason: ${result.finish_reason}`;
+                const fullErrorMessage = `${errorMessage}\nAI response headers:, ${JSON.stringify(result.response_headers)}`;
+                this.logger?.error(fullErrorMessage, parseError);
+                this.messageManager.updateToolBlock({
+                  id: toolId,
+                  parameters: argsString,
+                  result: errorMessage,
+                  success: false,
+                  error: fullErrorMessage,
+                  stage: "end",
+                  name: toolName,
+                  compactParams: "",
+                });
                 return;
               }
+            }
 
-              // Safely parse tool parameters, handle tools without parameters
-              let toolArgs: Record<string, unknown> = {};
-              const argsString = functionToolCall.function?.arguments?.trim();
-
-              if (!argsString || argsString === "") {
-                // Tool without parameters, use empty object
-                toolArgs = {};
-              } else {
-                try {
-                  toolArgs = JSON.parse(argsString);
-                } catch (parseError) {
-                  // For non-empty but malformed JSON, still throw exception
-                  const errorMessage = `Failed to parse tool arguments: ${argsString}`;
-                  this.logger?.error(errorMessage, parseError);
-                  throw new Error(errorMessage);
-                }
-              }
+            const compactParams = this.generateCompactParams(
+              toolName,
+              toolArgs,
+            );
 
-              // Set tool start execution state
-              const toolName = functionToolCall.function?.name || "";
-              const compactParams = this.generateCompactParams(
+            // Emit running stage for non-streaming tool calls (tool execution about to start)
+            this.messageManager.updateToolBlock({
+              id: toolId,
+              stage: "running",
+              name: toolName,
+              compactParams,
+              parameters: argsString,
+              parametersChunk: "",
+            });
+
+            try {
+              // Execute PreToolUse hooks before tool execution
+              const shouldExecuteTool = await this.executePreToolUseHooks(
                 toolName,
                 toolArgs,
-              );
-
-              this.messageManager.updateToolBlock({
                 toolId,
-                args: JSON.stringify(toolArgs, null, 2),
-                isRunning: true, // isRunning: true
-                name: toolName,
-                compactParams,
-              });
+              );
 
-              try {
-                // Execute PreToolUse hooks before tool execution
-                const shouldExecuteTool = await this.executePreToolUseHooks(
-                  toolName,
-                  toolArgs,
-                  toolId,
+              // If PreToolUse hooks blocked execution, skip tool execution
+              if (!shouldExecuteTool) {
+                this.logger?.info(
+                  `Tool ${toolName} execution blocked by PreToolUse hooks`,
                 );
+                return; // Skip this tool and return from this map function
+              }
 
-                // If PreToolUse hooks blocked execution, skip tool execution
-                if (!shouldExecuteTool) {
-                  this.logger?.info(
-                    `Tool ${toolName} execution blocked by PreToolUse hooks`,
-                  );
-                  return; // Skip this tool and return from this map function
-                }
-
-                // Create tool execution context
-                const context: ToolContext = {
-                  abortSignal: toolAbortController.signal,
-                  backgroundBashManager: this.backgroundBashManager,
-                  workdir: this.workdir,
-                };
-
-                // Execute tool
-                const toolResult = await this.toolManager.execute(
-                  functionToolCall.function?.name || "",
-                  toolArgs,
-                  context,
-                );
+              // Create tool execution context
+              const context: ToolContext = {
+                abortSignal: toolAbortController.signal,
+                backgroundBashManager: this.backgroundBashManager,
+                workdir: this.workdir,
+              };
 
-                // Update message state - tool execution completed
-                this.messageManager.updateToolBlock({
-                  toolId,
-                  args: JSON.stringify(toolArgs, null, 2),
-                  result:
-                    toolResult.content ||
-                    (toolResult.error ? `Error: ${toolResult.error}` : ""),
-                  success: toolResult.success,
-                  error: toolResult.error,
-                  isRunning: false, // isRunning: false
-                  name: toolName,
-                  shortResult: toolResult.shortResult,
-                  compactParams,
-                });
+              // Execute tool
+              const toolResult = await this.toolManager.execute(
+                functionToolCall.function?.name || "",
+                toolArgs,
+                context,
+              );
 
-                // If tool returns diff information, add diff block
-                if (
-                  toolResult.success &&
-                  toolResult.diffResult &&
-                  toolResult.filePath
-                ) {
-                  this.messageManager.addDiffBlock(
-                    toolResult.filePath,
-                    toolResult.diffResult,
-                  );
-                }
-
-                // Execute PostToolUse hooks after successful tool completion
-                await this.executePostToolUseHooks(
-                  toolId,
-                  toolName,
-                  toolArgs,
-                  toolResult,
-                );
-              } catch (toolError) {
-                const errorMessage =
-                  toolError instanceof Error
-                    ? toolError.message
-                    : String(toolError);
+              // Update message state - tool execution completed
+              this.messageManager.updateToolBlock({
+                id: toolId,
+                parameters: argsString,
+                result:
+                  toolResult.content ||
+                  (toolResult.error ? `Error: ${toolResult.error}` : ""),
+                success: toolResult.success,
+                error: toolResult.error,
+                stage: "end",
+                name: toolName,
+                shortResult: toolResult.shortResult,
+              });
 
-                this.messageManager.updateToolBlock({
-                  toolId,
-                  args: JSON.stringify(toolArgs, null, 2),
-                  result: `Tool execution failed: ${errorMessage}`,
-                  success: false,
-                  error: errorMessage,
-                  isRunning: false,
-                  name: toolName,
-                  compactParams,
-                });
+              // If tool returns diff information, add diff block
+              if (
+                toolResult.success &&
+                toolResult.diffResult &&
+                toolResult.filePath
+              ) {
+                this.messageManager.addDiffBlock(
+                  toolResult.filePath,
+                  toolResult.diffResult,
+                );
               }
-            } catch (parseError) {
-              const errorMessage =
-                parseError instanceof Error
-                  ? parseError.message
-                  : String(parseError);
-              this.messageManager.addErrorBlock(
-                `Failed to parse tool arguments for ${functionToolCall.function?.name}: ${errorMessage}`,
+
+              // Execute PostToolUse hooks after successful tool completion
+              await this.executePostToolUseHooks(
+                toolId,
+                toolName,
+                toolArgs,
+                toolResult,
               );
+            } catch (toolError) {
+              const errorMessage =
+                toolError instanceof Error
+                  ? toolError.message
+                  : String(toolError);
+
+              this.messageManager.updateToolBlock({
+                id: toolId,
+                parameters: JSON.stringify(toolArgs, null, 2),
+                result: `Tool execution failed: ${errorMessage}`,
+                success: false,
+                error: errorMessage,
+                stage: "end",
+                name: toolName,
+                compactParams,
+              });
             }
           },
         );
@@ -478,70 +623,74 @@ export class AIManager {
         error instanceof Error ? error.message : "Unknown error occurred",
       );
     } finally {
-      // Only execute Stop hooks for the initial call
+      // Only execute cleanup and hooks for the initial call
       if (recursionDepth === 0) {
-        // Execute Stop hooks only if the operation was not aborted
+        // Save session in each recursion to ensure message persistence
+        await this.messageManager.saveSession();
+        // Set loading to false first
+        this.setIsLoading(false);
+
+        // Clear abort controllers
+        this.abortController = null;
+        this.toolAbortController = null;
+
+        // Execute Stop/SubagentStop hooks only if the operation was not aborted
         const isCurrentlyAborted =
           abortController.signal.aborted || toolAbortController.signal.aborted;
 
         if (!isCurrentlyAborted) {
           const shouldContinue = await this.executeStopHooks();
 
-          // If Stop hooks indicate we should continue (due to blocking errors),
+          // If Stop/SubagentStop hooks indicate we should continue (due to blocking errors),
           // restart the AI conversation cycle
           if (shouldContinue) {
             this.logger?.info(
-              "Stop hooks indicate issues need fixing, continuing conversation...",
+              `${this.subagentType ? "SubagentStop" : "Stop"} hooks indicate issues need fixing, continuing conversation...`,
             );
 
             // Restart the conversation to let AI fix the issues
-            // Use recursionDepth = 1 to prevent Stop hooks from running again in continuation
+            // Use recursionDepth = 0 to set loading false again for continuation
             await this.sendAIMessage({
-              recursionDepth: 1,
+              recursionDepth: 0,
               model,
               allowedTools,
             });
           }
         }
-
-        // Save session after all operations (including continuation) are complete
-        await this.messageManager.saveSession();
-
-        // Clear abort controllers and loading state after all operations are complete
-        this.abortController = null;
-        this.toolAbortController = null;
-
-        // Set loading to false at the very end, after all operations including continuation
-        this.setIsLoading(false);
       }
     }
   }
 
   /**
-   * Execute Stop hooks when AI response cycle completes
+   * Execute Stop or SubagentStop hooks when AI response cycle completes
+   * Uses "SubagentStop" hook name when triggered by a subagent, otherwise uses "Stop"
    * @returns Promise<boolean> - true if should continue conversation, false if should stop
    */
   private async executeStopHooks(): Promise<boolean> {
     if (!this.hookManager) return false;
 
     try {
+      // Use "SubagentStop" hook name when triggered by a subagent, otherwise use "Stop"
+      const hookName = this.subagentType ? "SubagentStop" : "Stop";
+
       const context: ExtendedHookExecutionContext = {
-        event: "Stop",
+        event: hookName,
         projectDir: this.workdir,
         timestamp: new Date(),
         sessionId: this.messageManager.getSessionId(),
         transcriptPath: this.messageManager.getTranscriptPath(),
         cwd: this.workdir,
+        subagentType: this.subagentType, // Include subagent type in hook context
         // Stop hooks don't need toolName, toolInput, toolResponse, or userPrompt
       };
 
-      const results = await this.hookManager.executeHooks("Stop", context);
+      const results = await this.hookManager.executeHooks(hookName, context);
 
       // Process hook results to handle exit codes and appropriate responses
       let shouldContinue = false;
       if (results.length > 0) {
         const processResult = this.hookManager.processHookResults(
-          "Stop",
+          hookName,
           results,
           this.messageManager,
         );
@@ -549,7 +698,7 @@ export class AIManager {
         // If hook processing indicates we should block (exit code 2), continue conversation
         if (processResult.shouldBlock) {
           this.logger?.info(
-            "Stop hook blocked stopping with error:",
+            `${hookName} hook blocked stopping with error:`,
             processResult.errorMessage,
           );
           shouldContinue = true;
@@ -559,7 +708,7 @@ export class AIManager {
       // Log hook execution results for debugging
       if (results.length > 0) {
         this.logger?.debug(
-          `Executed ${results.length} Stop hook(s):`,
+          `Executed ${results.length} ${hookName} hook(s):`,
           results.map((r) => ({
             success: r.success,
             duration: r.duration,
@@ -573,7 +722,10 @@ export class AIManager {
       return shouldContinue;
     } catch (error) {
       // Hook execution errors should not interrupt the main workflow
-      this.logger?.error("Stop hook execution failed:", error);
+      this.logger?.error(
+        `${this.subagentType ? "SubagentStop" : "Stop"} hook execution failed:`,
+        error,
+      );
       return false;
     }
   }
@@ -599,6 +751,7 @@ export class AIManager {
         transcriptPath: this.messageManager.getTranscriptPath(),
         cwd: this.workdir,
         toolInput,
+        subagentType: this.subagentType, // Include subagent type in hook context
       };
 
       const results = await this.hookManager.executeHooks(
@@ -614,6 +767,7 @@ export class AIManager {
           results,
           this.messageManager,
           toolId, // Pass toolId for proper PreToolUse blocking error handling
+          JSON.stringify(toolInput || {}, null, 2), // Pass serialized tool parameters
         );
         shouldContinue = !processResult.shouldBlock;
       }
@@ -662,6 +816,7 @@ export class AIManager {
         cwd: this.workdir,
         toolInput,
         toolResponse,
+        subagentType: this.subagentType, // Include subagent type in hook context
       };
 
       const results = await this.hookManager.executeHooks(
@@ -671,14 +826,11 @@ export class AIManager {
 
       // Process hook results to handle exit codes and update tool results
       if (results.length > 0) {
-        const originalToolResult = toolResponse?.content || "";
-
         this.hookManager.processHookResults(
           "PostToolUse",
           results,
           this.messageManager,
           toolId,
-          originalToolResult,
         );
       }