wave-agent-sdk 0.0.4 → 0.0.6

package/src/managers/aiManager.ts

@@ -2,17 +2,23 @@ import { callAgent, compressMessages } from "../services/aiService.js";
 import { getMessagesToCompress } from "../utils/messageOperations.js";
 import { convertMessagesForAPI } from "../utils/convertMessagesForAPI.js";
 import * as memory from "../services/memory.js";
-import type { Logger, GatewayConfig, ModelConfig } from "../types.js";
+import type {
+  Logger,
+  GatewayConfig,
+  ModelConfig,
+  Usage,
+} from "../types/index.js";
 import type { ToolManager } from "./toolManager.js";
 import type { ToolContext, ToolResult } from "../tools/types.js";
 import type { MessageManager } from "./messageManager.js";
 import type { BackgroundBashManager } from "./backgroundBashManager.js";
 import { ChatCompletionMessageFunctionToolCall } from "openai/resources.js";
-import type { HookManager } from "../hooks/index.js";
-import type { ExtendedHookExecutionContext } from "../hooks/types.js";
+import type { HookManager } from "./hookManager.js";
+import type { ExtendedHookExecutionContext } from "../types/hooks.js";
 
 export interface AIManagerCallbacks {
   onCompressionStateChange?: (isCompressing: boolean) => void;
+  onUsageAdded?: (usage: Usage) => void;
 }
 
 export interface AIManagerOptions {
@@ -140,7 +146,7 @@ export class AIManager {
 
     // Check if token limit exceeded - use injected configuration
     if (usage.total_tokens > this.tokenLimit) {
-      this.logger?.info(
+      this.logger?.debug(
         `Token usage exceeded ${this.tokenLimit}, compressing messages...`,
       );
 
@@ -156,7 +162,7 @@ export class AIManager {
 
         this.setIsCompressing(true);
         try {
-          const compressedContent = await compressMessages({
+          const compressionResult = await compressMessages({
             gatewayConfig: this.gatewayConfig,
             modelConfig: this.modelConfig,
             messages: recentChatMessages,
@@ -166,10 +172,26 @@ export class AIManager {
           // Execute message reconstruction and sessionId update after compression
           this.messageManager.compressMessagesAndUpdateSession(
             insertIndex,
-            compressedContent,
+            compressionResult.content,
           );
 
-          this.logger?.info(
+          // Handle usage tracking for compression operations
+          if (compressionResult.usage) {
+            const usage: Usage = {
+              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);
+            }
+          }
+
+          this.logger?.debug(
             `Successfully compressed ${messagesToCompress.length} messages and updated session`,
           );
         } catch (compressError) {
@@ -200,18 +222,29 @@ export class AIManager {
     } = {},
   ): Promise<void> {
     const { recursionDepth = 0, model, allowedTools } = options;
+
     // Only check isLoading for the initial call (recursionDepth === 0)
     if (recursionDepth === 0 && this.isLoading) {
       return;
     }
 
-    // Create new AbortController
-    const abortController = new AbortController();
-    this.abortController = abortController;
+    // Only create new AbortControllers for the initial call (recursionDepth === 0)
+    // For recursive calls, reuse existing controllers to maintain abort signal
+    let abortController: AbortController;
+    let toolAbortController: AbortController;
 
-    // Create separate AbortController for tool execution
-    const toolAbortController = new AbortController();
-    this.toolAbortController = toolAbortController;
+    if (recursionDepth === 0) {
+      // Create new AbortControllers for initial call
+      abortController = new AbortController();
+      this.abortController = abortController;
+
+      toolAbortController = new AbortController();
+      this.toolAbortController = toolAbortController;
+    } else {
+      // Reuse existing controllers for recursive calls
+      abortController = this.abortController!;
+      toolAbortController = this.toolAbortController!;
+    }
 
     // Only set loading state for the initial call
     if (recursionDepth === 0) {
@@ -255,143 +288,171 @@ export class AIManager {
         }
       }
 
-      // Add assistant message at once (including content and tool calls)
-      this.messageManager.addAssistantMessage(content, toolCalls);
-
-      if (toolCalls.length > 0) {
-        for (const functionToolCall of toolCalls) {
-          const toolId = functionToolCall.id || "";
-          // Execute tool
-          try {
-            // Check if already interrupted, skip tool execution if so
-            if (
-              abortController.signal.aborted ||
-              toolAbortController.signal.aborted
-            ) {
-              return;
-            }
-
-            // Safely parse tool parameters, handle tools without parameters
-            let toolArgs: Record<string, unknown> = {};
-            const argsString = functionToolCall.function?.arguments?.trim();
+      // Handle usage tracking for agent operations
+      let usage: Usage | undefined;
+      if (result.usage) {
+        usage = {
+          prompt_tokens: result.usage.prompt_tokens,
+          completion_tokens: result.usage.completion_tokens,
+          total_tokens: result.usage.total_tokens,
+          model: model || this.modelConfig.agentModel,
+          operation_type: "agent",
+        };
+      }
 
-            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);
-              }
-            }
+      // Add assistant message at once (including content, tool calls, and usage)
+      this.messageManager.addAssistantMessage(content, toolCalls, usage);
 
-            // Set tool start execution state
-            const toolName = functionToolCall.function?.name || "";
-            const compactParams = this.generateCompactParams(
-              toolName,
-              toolArgs,
-            );
+      // Notify Agent to add to usage tracking
+      if (usage) {
+        if (this.callbacks?.onUsageAdded) {
+          this.callbacks.onUsageAdded(usage);
+        }
+      }
 
-            this.messageManager.updateToolBlock({
-              toolId,
-              args: JSON.stringify(toolArgs, null, 2),
-              isRunning: true, // isRunning: true
-              name: toolName,
-              compactParams,
-            });
+      if (toolCalls.length > 0) {
+        // Execute all tools in parallel using Promise.all
+        const toolExecutionPromises = toolCalls.map(
+          async (functionToolCall) => {
+            const toolId = functionToolCall.id || "";
 
             try {
-              // Execute PreToolUse hooks before tool execution
-              await this.executePreToolUseHooks(toolName, toolArgs);
-
-              // 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,
-              );
-
-              // 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,
-              });
-
-              // If tool returns diff information, add diff block
+              // Check if already interrupted, skip tool execution if so
               if (
-                toolResult.success &&
-                toolResult.diffResult &&
-                toolResult.filePath
+                abortController.signal.aborted ||
+                toolAbortController.signal.aborted
               ) {
-                this.messageManager.addDiffBlock(
-                  toolResult.filePath,
-                  toolResult.diffResult,
-                );
+                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);
+                }
               }
 
-              // Execute PostToolUse hooks after successful tool completion
-              await this.executePostToolUseHooks(
+              // Set tool start execution state
+              const toolName = functionToolCall.function?.name || "";
+              const compactParams = this.generateCompactParams(
                 toolName,
                 toolArgs,
-                toolResult,
               );
-            } catch (toolError) {
-              const errorMessage =
-                toolError instanceof Error
-                  ? toolError.message
-                  : String(toolError);
 
               this.messageManager.updateToolBlock({
                 toolId,
                 args: JSON.stringify(toolArgs, null, 2),
-                result: `Tool execution failed: ${errorMessage}`,
-                success: false,
-                error: errorMessage,
-                isRunning: false,
+                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
+                }
+
+                // 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,
+                );
+
+                // 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,
+                });
+
+                // 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);
+
+                this.messageManager.updateToolBlock({
+                  toolId,
+                  args: JSON.stringify(toolArgs, null, 2),
+                  result: `Tool execution failed: ${errorMessage}`,
+                  success: false,
+                  error: errorMessage,
+                  isRunning: false,
+                  name: toolName,
+                  compactParams,
+                });
+              }
+            } catch (parseError) {
+              const errorMessage =
+                parseError instanceof Error
+                  ? parseError.message
+                  : String(parseError);
+              this.messageManager.addErrorBlock(
+                `Failed to parse tool arguments for ${functionToolCall.function?.name}: ${errorMessage}`,
+              );
             }
-          } catch (parseError) {
-            // Check if it's a parsing error due to interruption
-            const isAborted =
-              abortController.signal.aborted ||
-              toolAbortController.signal.aborted;
-
-            if (isAborted) {
-              // If interrupted, return directly without showing error
-              return;
-            }
+          },
+        );
 
-            const errorMessage =
-              parseError instanceof Error
-                ? parseError.message
-                : String(parseError);
-            this.messageManager.addErrorBlock(
-              `Failed to parse tool arguments for ${functionToolCall.function?.name}: ${errorMessage}`,
-            );
-          }
-        }
+        // Wait for all tools to complete execution in parallel
+        await Promise.all(toolExecutionPromises);
       }
 
       // Handle token statistics and message compression
@@ -403,12 +464,6 @@ export class AIManager {
         const isCurrentlyAborted =
           abortController.signal.aborted || toolAbortController.signal.aborted;
 
-        // AI service call ends, clear abort controller
-        this.abortController = null;
-
-        // Clear tool AbortController after tool execution completes
-        this.toolAbortController = null;
-
         if (!isCurrentlyAborted) {
           // Recursively call AI service, increment recursion depth, and pass same configuration
           await this.sendAIMessage({
@@ -417,47 +472,57 @@ export class AIManager {
             allowedTools,
           });
         }
-      } else {
-        // Clear abort controller when no tool operations
-        this.abortController = null;
-        this.toolAbortController = null;
       }
     } catch (error) {
-      // Check if error is due to user interrupt operation
-      const isAborted =
-        abortController.signal.aborted ||
-        toolAbortController.signal.aborted ||
-        (error instanceof Error &&
-          (error.name === "AbortError" || error.message.includes("aborted")));
-
-      if (!isAborted) {
-        this.messageManager.addErrorBlock(
-          error instanceof Error ? error.message : "Unknown error occurred",
-        );
-      }
-
-      // Reset abort controller on error
-      this.abortController = null;
-      this.toolAbortController = null;
+      this.messageManager.addErrorBlock(
+        error instanceof Error ? error.message : "Unknown error occurred",
+      );
     } finally {
-      // Only clear loading state for the initial call
+      // Only execute Stop hooks for the initial call
       if (recursionDepth === 0) {
-        this.setIsLoading(false);
+        // Execute Stop 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),
+          // restart the AI conversation cycle
+          if (shouldContinue) {
+            this.logger?.info(
+              "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
+            await this.sendAIMessage({
+              recursionDepth: 1,
+              model,
+              allowedTools,
+            });
+          }
+        }
 
-        // Save session before executing Stop hooks
+        // Save session after all operations (including continuation) are complete
         await this.messageManager.saveSession();
 
-        // Execute Stop hooks when AI response cycle completes
-        await this.executeStopHooks();
+        // 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
+   * @returns Promise<boolean> - true if should continue conversation, false if should stop
    */
-  private async executeStopHooks(): Promise<void> {
-    if (!this.hookManager) return;
+  private async executeStopHooks(): Promise<boolean> {
+    if (!this.hookManager) return false;
 
     try {
       const context: ExtendedHookExecutionContext = {
@@ -472,6 +537,25 @@ export class AIManager {
 
       const results = await this.hookManager.executeHooks("Stop", context);
 
+      // Process hook results to handle exit codes and appropriate responses
+      let shouldContinue = false;
+      if (results.length > 0) {
+        const processResult = this.hookManager.processHookResults(
+          "Stop",
+          results,
+          this.messageManager,
+        );
+
+        // If hook processing indicates we should block (exit code 2), continue conversation
+        if (processResult.shouldBlock) {
+          this.logger?.info(
+            "Stop hook blocked stopping with error:",
+            processResult.errorMessage,
+          );
+          shouldContinue = true;
+        }
+      }
+
       // Log hook execution results for debugging
       if (results.length > 0) {
         this.logger?.debug(
@@ -485,20 +569,25 @@ export class AIManager {
           })),
         );
       }
+
+      return shouldContinue;
     } catch (error) {
       // Hook execution errors should not interrupt the main workflow
       this.logger?.error("Stop hook execution failed:", error);
+      return false;
     }
   }
 
   /**
    * Execute PreToolUse hooks before tool execution
+   * Returns true if hooks allow tool execution, false if blocked
    */
   private async executePreToolUseHooks(
     toolName: string,
     toolInput?: Record<string, unknown>,
-  ): Promise<void> {
-    if (!this.hookManager) return;
+    toolId?: string,
+  ): Promise<boolean> {
+    if (!this.hookManager) return true;
 
     try {
       const context: ExtendedHookExecutionContext = {
@@ -517,6 +606,18 @@ export class AIManager {
         context,
       );
 
+      // Process hook results to handle exit codes and determine if tool should be blocked
+      let shouldContinue = true;
+      if (results.length > 0) {
+        const processResult = this.hookManager.processHookResults(
+          "PreToolUse",
+          results,
+          this.messageManager,
+          toolId, // Pass toolId for proper PreToolUse blocking error handling
+        );
+        shouldContinue = !processResult.shouldBlock;
+      }
+
       // Log hook execution results for debugging
       if (results.length > 0) {
         this.logger?.debug(
@@ -530,9 +631,12 @@ export class AIManager {
           })),
         );
       }
+
+      return shouldContinue;
     } catch (error) {
       // Hook execution errors should not interrupt the main workflow
       this.logger?.error("PreToolUse hook execution failed:", error);
+      return true; // Allow tool execution on hook errors
     }
   }
 
@@ -540,6 +644,7 @@ export class AIManager {
    * Execute PostToolUse hooks after tool completion
    */
   private async executePostToolUseHooks(
+    toolId: string,
     toolName: string,
     toolInput?: Record<string, unknown>,
     toolResponse?: ToolResult,
@@ -564,6 +669,19 @@ export class AIManager {
         context,
       );
 
+      // 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,
+        );
+      }
+
       // Log hook execution results for debugging
       if (results.length > 0) {
         this.logger?.debug(

package/src/managers/backgroundBashManager.ts

@@ -1,5 +1,5 @@
 import { spawn } from "child_process";
-import type { BackgroundShell } from "../types.js";
+import type { BackgroundShell } from "../types/index.js";
 
 export interface BackgroundBashManagerCallbacks {
   onShellsChange?: (shells: BackgroundShell[]) => void;