wave-agent-sdk 0.0.4 → 0.0.6

package/src/agent.ts

@@ -13,15 +13,16 @@ import {
   type BackgroundBashManagerCallbacks,
 } from "./managers/backgroundBashManager.js";
 import { SlashCommandManager } from "./managers/slashCommandManager.js";
-import type { SlashCommand, CustomSlashCommand } from "./types.js";
+import type { SlashCommand, CustomSlashCommand } from "./types/index.js";
 import type {
   Message,
   Logger,
   McpServerStatus,
   GatewayConfig,
   ModelConfig,
-} from "./types.js";
-import { HookManager } from "./hooks/index.js";
+  Usage,
+} from "./types/index.js";
+import { HookManager } from "./managers/hookManager.js";
 import { configResolver } from "./utils/configResolver.js";
 import { configValidator } from "./utils/configValidator.js";
 import { SkillManager } from "./managers/skillManager.js";
@@ -51,6 +52,14 @@ export interface AgentOptions {
   workdir?: string;
   /**Optional custom system prompt - if provided, replaces default system prompt */
   systemPrompt?: string;
+
+  // New: Session directory configuration
+  /**
+   * Optional custom directory for session file storage
+   * @default join(homedir(), ".wave", "sessions")
+   * @example "/path/to/custom/sessions"
+   */
+  sessionDir?: string;
 }
 
 export interface AgentCallbacks
@@ -72,6 +81,7 @@ export class Agent {
   private hookManager: HookManager; // Add hooks manager instance
   private workdir: string; // Working directory
   private systemPrompt?: string; // Custom system prompt
+  private _usages: Usage[] = []; // Usage tracking array
 
   // Configuration properties
   private gatewayConfig: GatewayConfig;
@@ -81,13 +91,21 @@ export class Agent {
   /**
    * Agent constructor - handles configuration resolution and validation
    *
-   * IMPORTANT: Keep this constructor's signature exactly the same as Agent.create()
-   * to maintain API consistency. Both methods should accept the same AgentOptions.
+   * IMPORTANT: This constructor is private. Use Agent.create() instead for proper
+   * async initialization. Keep this constructor's signature exactly the same as
+   * Agent.create() to maintain API consistency.
    *
    * @param options - Configuration options for the Agent instance
+   * @param options.sessionDir - Optional custom directory for session storage
    */
-  constructor(options: AgentOptions) {
-    const { callbacks = {}, logger, workdir, systemPrompt } = options;
+  private constructor(options: AgentOptions) {
+    const {
+      callbacks = {},
+      logger,
+      workdir,
+      systemPrompt,
+      sessionDir,
+    } = options;
 
     // Resolve configuration from constructor args and environment variables
     const gatewayConfig = configResolver.resolveGatewayConfig(
@@ -127,18 +145,14 @@ export class Agent {
       logger: this.logger,
     }); // Initialize tool registry, pass MCP manager
 
-    this.hookManager = new HookManager(
-      this.workdir,
-      undefined,
-      undefined,
-      this.logger,
-    ); // Initialize hooks manager
+    this.hookManager = new HookManager(this.workdir, undefined, this.logger); // Initialize hooks manager
 
     // Initialize MessageManager
     this.messageManager = new MessageManager({
       callbacks,
       workdir: this.workdir,
       logger: this.logger,
+      sessionDir,
     });
 
     // Initialize subagent manager with all dependencies in constructor
@@ -151,6 +165,7 @@ export class Agent {
       gatewayConfig,
       modelConfig,
       tokenLimit,
+      onUsageAdded: (usage) => this.addUsage(usage),
     });
 
     // Initialize AI manager with resolved configuration
@@ -160,7 +175,12 @@ export class Agent {
       logger: this.logger,
       backgroundBashManager: this.backgroundBashManager,
       hookManager: this.hookManager,
-      callbacks,
+      callbacks: {
+        ...callbacks,
+        onUsageAdded: (usage: Usage) => {
+          this.addUsage(usage);
+        },
+      },
       workdir: this.workdir,
       systemPrompt: this.systemPrompt,
       gatewayConfig: this.gatewayConfig,
@@ -192,6 +212,38 @@ export class Agent {
     return this.messageManager.getMessages();
   }
 
+  public get usages(): Usage[] {
+    return [...this._usages]; // Return copy to prevent external modification
+  }
+
+  public get sessionFilePath(): string {
+    return this.messageManager.getTranscriptPath();
+  }
+
+  /**
+   * Rebuild usage array from messages containing usage metadata
+   * Called during session restoration to reconstruct usage tracking
+   */
+  private rebuildUsageFromMessages(): void {
+    this._usages = [];
+    this.messages.forEach((message) => {
+      if (message.role === "assistant" && message.usage) {
+        this._usages.push(message.usage);
+      }
+    });
+    // Trigger callback after rebuilding usage array
+    this.messageManager.triggerUsageChange();
+  }
+
+  /**
+   * Add usage data to the tracking array and trigger callbacks
+   * @param usage Usage data from AI operations
+   */
+  private addUsage(usage: Usage): void {
+    this._usages.push(usage);
+    this.messageManager.triggerUsageChange();
+  }
+
   public get latestTotalTokens(): number {
     return this.messageManager.getlatestTotalTokens();
   }
@@ -242,6 +294,43 @@ export class Agent {
    * @param options - Same AgentOptions interface used by constructor
    * @returns Promise<Agent> - Fully initialized Agent instance
    */
+  /**
+   * Create a new Agent instance with async initialization
+   *
+   * This is the recommended way to create Agent instances. The constructor is private
+   * to ensure proper async initialization of all components.
+   *
+   * @param options - Configuration options for the Agent instance
+   * @param options.apiKey - API key for the AI service (or set WAVE_API_KEY env var)
+   * @param options.baseURL - Base URL for the AI service (or set WAVE_BASE_URL env var)
+   * @param options.sessionDir - Optional custom directory for session file storage.
+   *   If not provided, defaults to ~/.wave/sessions/. Can be relative or absolute path.
+   *   Examples: "./app-sessions", "/var/myapp/sessions", "~/Documents/sessions"
+   * @param options.callbacks - Optional callbacks for various Agent events
+   * @param options.restoreSessionId - Optional session ID to restore from
+   * @param options.continueLastSession - Whether to continue the last session automatically
+   * @param options.logger - Optional custom logger implementation
+   * @param options.messages - Optional initial messages for testing convenience
+   * @param options.workdir - Working directory (defaults to process.cwd())
+   * @param options.systemPrompt - Optional custom system prompt
+   * @returns Promise that resolves to initialized Agent instance
+   *
+   * @example
+   * ```typescript
+   * // Basic usage with default session directory
+   * const agent = await Agent.create({
+   *   apiKey: 'your-api-key',
+   *   baseURL: 'https://api.example.com'
+   * });
+   *
+   * // Custom session directory
+   * const agent = await Agent.create({
+   *   apiKey: 'your-api-key',
+   *   baseURL: 'https://api.example.com',
+   *   sessionDir: './app-sessions'
+   * });
+   * ```
+   */
   static async create(options: AgentOptions): Promise<Agent> {
     // Create Agent instance - configuration resolution and validation now happens in constructor
     const instance = new Agent(options);
@@ -289,7 +378,7 @@ export class Agent {
     // Initialize hooks configuration
     try {
       // Load hooks configuration from user and project settings
-      this.logger?.info("Loading hooks configuration...");
+      this.logger?.debug("Loading hooks configuration...");
       this.hookManager.loadConfigurationFromSettings();
       this.logger?.debug("Hooks system initialized successfully");
     } catch (error) {
@@ -301,12 +390,16 @@ export class Agent {
     if (options?.messages) {
       // If messages are provided, use them directly (useful for testing)
       this.messageManager.setMessages(options.messages);
+      // Rebuild usage array from restored messages
+      this.rebuildUsageFromMessages();
     } else {
       // Otherwise, handle session restoration
       await this.messageManager.handleSessionRestoration(
         options?.restoreSessionId,
         options?.continueLastSession,
       );
+      // Rebuild usage array from restored messages
+      this.rebuildUsageFromMessages();
     }
   }
 
@@ -328,10 +421,9 @@ export class Agent {
 
   /** Unified interrupt method, interrupts both AI messages and command execution */
   public abortMessage(): void {
-    this.abortAIMessage();
+    this.abortAIMessage(); // This will abort tools including Task tool (subagents)
     this.abortBashCommand();
     this.abortSlashCommand();
-    this.abortSubagents();
   }
 
   /** Add to input history */
@@ -349,23 +441,12 @@ export class Agent {
     this.slashCommandManager.abortCurrentCommand();
   }
 
-  /** Interrupt all subagent execution */
-  public abortSubagents(): void {
-    this.subagentManager.abortAllInstances();
-  }
-
-  /** Interrupt specific subagent execution */
-  public abortSubagent(subagentId: string): boolean {
-    return this.subagentManager.abortInstance(subagentId);
-  }
-
   /** Destroy managers, clean up resources */
   public async destroy(): Promise<void> {
-    this.messageManager.saveSession();
-    this.abortAIMessage();
+    await this.messageManager.saveSession();
+    this.abortAIMessage(); // This will abort tools including Task tool (subagents)
     this.abortBashCommand();
     this.abortSlashCommand();
-    this.abortSubagents();
     // Cleanup background bash manager
     this.backgroundBashManager.cleanup();
     // Cleanup MCP connections
@@ -405,26 +486,7 @@ export class Agent {
       // Add user message to history
       this.addToInputHistory(content);
 
-      // Execute UserPromptSubmit hooks before processing the prompt
-      if (this.hookManager) {
-        try {
-          await this.hookManager.executeHooks("UserPromptSubmit", {
-            event: "UserPromptSubmit",
-            projectDir: this.workdir,
-            timestamp: new Date(),
-            // UserPromptSubmit doesn't need toolName
-            sessionId: this.sessionId,
-            transcriptPath: this.messageManager.getTranscriptPath(),
-            cwd: this.workdir,
-            userPrompt: content,
-          });
-        } catch (error) {
-          this.logger?.warn("UserPromptSubmit hooks execution failed:", error);
-          // Continue processing even if hooks fail
-        }
-      }
-
-      // Add user message, will automatically sync to UI
+      // Add user message first, will automatically sync to UI
       this.messageManager.addUserMessage(
         content,
         images?.map((img) => ({
@@ -433,6 +495,44 @@ export class Agent {
         })),
       );
 
+      // Execute UserPromptSubmit hooks after adding the user message
+      if (this.hookManager) {
+        try {
+          const hookResults = await this.hookManager.executeHooks(
+            "UserPromptSubmit",
+            {
+              event: "UserPromptSubmit",
+              projectDir: this.workdir,
+              timestamp: new Date(),
+              // UserPromptSubmit doesn't need toolName
+              sessionId: this.sessionId,
+              transcriptPath: this.messageManager.getTranscriptPath(),
+              cwd: this.workdir,
+              userPrompt: content,
+            },
+          );
+
+          // Process hook results and determine if we should continue
+          const processResult = this.hookManager.processHookResults(
+            "UserPromptSubmit",
+            hookResults,
+            this.messageManager,
+          );
+
+          // If hook processing indicates we should block (exit code 2), stop here
+          if (processResult.shouldBlock) {
+            this.logger?.info(
+              "UserPromptSubmit hook blocked prompt processing with error:",
+              processResult.errorMessage,
+            );
+            return; // Don't send to AI
+          }
+        } catch (error) {
+          this.logger?.warn("UserPromptSubmit hooks execution failed:", error);
+          // Continue processing even if hooks fail
+        }
+      }
+
       // Send AI message
       await this.aiManager.sendAIMessage();
     } catch (error) {

package/src/index.ts

@@ -2,6 +2,7 @@
 export * from "./services/aiService.js";
 export * from "./services/memory.js";
 export * from "./services/session.js";
+export * from "./services/hook.js";
 
 // Export main agent
 export * from "./agent.js";
@@ -15,9 +16,7 @@ export * from "./utils/messageOperations.js";
 export * from "./utils/path.js";
 export * from "./utils/stringUtils.js";
 export * from "./utils/customCommands.js";
-
-// Export hooks system
-export * from "./hooks/index.js";
+export * from "./utils/hookMatcher.js";
 
 // Export types
-export * from "./types.js";
+export * from "./types/index.js";