wave-agent-sdk 0.0.4 → 0.0.5

package/src/agent.ts

@@ -20,6 +20,7 @@ import type {
   McpServerStatus,
   GatewayConfig,
   ModelConfig,
+  Usage,
 } from "./types.js";
 import { HookManager } from "./hooks/index.js";
 import { configResolver } from "./utils/configResolver.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(
@@ -139,6 +157,7 @@ export class Agent {
       callbacks,
       workdir: this.workdir,
       logger: this.logger,
+      sessionDir,
     });
 
     // Initialize subagent manager with all dependencies in constructor
@@ -151,6 +170,7 @@ export class Agent {
       gatewayConfig,
       modelConfig,
       tokenLimit,
+      onUsageAdded: (usage) => this.addUsage(usage),
     });
 
     // Initialize AI manager with resolved configuration
@@ -160,7 +180,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 +217,34 @@ export class Agent {
     return this.messageManager.getMessages();
   }
 
+  public get usages(): Usage[] {
+    return [...this._usages]; // Return copy to prevent external modification
+  }
+
+  /**
+   * 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 +295,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 +379,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 +391,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();
     }
   }
 
@@ -361,7 +455,7 @@ export class Agent {
 
   /** Destroy managers, clean up resources */
   public async destroy(): Promise<void> {
-    this.messageManager.saveSession();
+    await this.messageManager.saveSession();
     this.abortAIMessage();
     this.abortBashCommand();
     this.abortSlashCommand();

package/src/hooks/manager.ts

@@ -174,7 +174,9 @@ export class HookManager implements IHookManager {
 
     const eventConfigs = this.configuration[event];
     if (!eventConfigs || eventConfigs.length === 0) {
-      this.logger?.debug(`[HookManager] No hooks configured for ${event} event`);
+      this.logger?.debug(
+        `[HookManager] No hooks configured for ${event} event`,
+      );
       return [];
     }
 
@@ -261,7 +263,7 @@ export class HookManager implements IHookManager {
       results,
       totalDuration,
     );
-    this.logger?.info(`[HookManager] ${event} execution summary: ${summary}`);
+    this.logger?.debug(`[HookManager] ${event} execution summary: ${summary}`);
 
     return results;
   }

package/src/managers/aiManager.ts

@@ -2,7 +2,7 @@ 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.js";
 import type { ToolManager } from "./toolManager.js";
 import type { ToolContext, ToolResult } from "../tools/types.js";
 import type { MessageManager } from "./messageManager.js";
@@ -13,6 +13,7 @@ import type { ExtendedHookExecutionContext } from "../hooks/types.js";
 
 export interface AIManagerCallbacks {
   onCompressionStateChange?: (isCompressing: boolean) => void;
+  onUsageAdded?: (usage: Usage) => void;
 }
 
 export interface AIManagerOptions {
@@ -140,7 +141,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 +157,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 +167,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) {
@@ -255,8 +272,27 @@ export class AIManager {
         }
       }
 
-      // Add assistant message at once (including content and tool calls)
-      this.messageManager.addAssistantMessage(content, toolCalls);
+      // 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",
+        };
+      }
+
+      // Add assistant message at once (including content, tool calls, and usage)
+      this.messageManager.addAssistantMessage(content, toolCalls, usage);
+
+      // Notify Agent to add to usage tracking
+      if (usage) {
+        if (this.callbacks?.onUsageAdded) {
+          this.callbacks.onUsageAdded(usage);
+        }
+      }
 
       if (toolCalls.length > 0) {
         for (const functionToolCall of toolCalls) {

package/src/managers/mcpManager.ts

@@ -53,7 +53,7 @@ export class McpManager {
     this.workdir = workdir;
 
     if (autoConnect) {
-      this.logger?.info("Initializing MCP servers...");
+      this.logger?.debug("Initializing MCP servers...");
 
       // Ensure MCP configuration is loaded
       const config = await this.ensureConfigLoaded();
@@ -63,10 +63,10 @@ export class McpManager {
         const connectionPromises = Object.keys(config.mcpServers).map(
           async (serverName) => {
             try {
-              this.logger?.info(`Connecting to MCP server: ${serverName}`);
+              this.logger?.debug(`Connecting to MCP server: ${serverName}`);
               const success = await this.connectServer(serverName);
               if (success) {
-                this.logger?.info(
+                this.logger?.debug(
                   `Successfully connected to MCP server: ${serverName}`,
                 );
               } else {
@@ -86,7 +86,7 @@ export class McpManager {
         await Promise.all(connectionPromises);
       }
 
-      this.logger?.info("MCP servers initialization completed");
+      this.logger?.debug("MCP servers initialization completed");
       // Trigger state change callback after initialization
       this.callbacks.onServersChange?.(this.getAllServers());
     }
@@ -255,7 +255,7 @@ export class McpManager {
       };
 
       transport.onclose = () => {
-        this.logger?.info(`MCP Server ${name} transport closed`);
+        this.logger?.debug(`MCP Server ${name} transport closed`);
         this.connections.delete(name);
         this.updateServerStatus(name, {
           status: "disconnected",

package/src/managers/messageManager.ts

@@ -16,7 +16,7 @@ import {
   type UpdateSubagentBlockParams,
   type AgentToolBlockUpdateParams,
 } from "../utils/messageOperations.js";
-import type { Logger, Message } from "../types.js";
+import type { Logger, Message, Usage } from "../types.js";
 import {
   cleanupExpiredSessions,
   getLatestSession,
@@ -32,6 +32,7 @@ export interface MessageManagerCallbacks {
   onSessionIdChange?: (sessionId: string) => void;
   onLatestTotalTokensChange?: (latestTotalTokens: number) => void;
   onUserInputHistoryChange?: (history: string[]) => void;
+  onUsagesChange?: (usages: Usage[]) => void;
   // Incremental callback
   onUserMessageAdded?: (
     content: string,
@@ -71,6 +72,13 @@ export interface MessageManagerOptions {
   callbacks: MessageManagerCallbacks;
   workdir: string;
   logger?: Logger;
+
+  // New: Optional session directory override
+  /**
+   * Custom session directory path
+   * @default join(homedir(), ".wave", "sessions")
+   */
+  sessionDir?: string;
 }
 
 export class MessageManager {
@@ -83,6 +91,7 @@ export class MessageManager {
   private workdir: string;
   private logger?: Logger; // Add optional logger property
   private callbacks: MessageManagerCallbacks;
+  private sessionDir?: string; // Add session directory property
 
   constructor(options: MessageManagerOptions) {
     this.sessionId = randomUUID();
@@ -93,6 +102,7 @@ export class MessageManager {
     this.workdir = options.workdir;
     this.callbacks = options.callbacks;
     this.logger = options.logger;
+    this.sessionDir = options.sessionDir;
   }
 
   // Getter methods
@@ -113,7 +123,7 @@ export class MessageManager {
   }
 
   public getTranscriptPath(): string {
-    return getSessionFilePath(this.sessionId);
+    return getSessionFilePath(this.sessionId, this.sessionDir);
   }
 
   // Setter methods, will trigger callbacks
@@ -140,6 +150,7 @@ export class MessageManager {
         this.workdir,
         this.latestTotalTokens,
         this.sessionStartTime,
+        this.sessionDir,
       );
     } catch (error) {
       this.logger?.error("Failed to save session:", error);
@@ -155,7 +166,7 @@ export class MessageManager {
   ): Promise<void> {
     // Clean up expired sessions first
     try {
-      await cleanupExpiredSessions(this.workdir);
+      await cleanupExpiredSessions(this.workdir, this.sessionDir);
     } catch (error) {
       console.warn("Failed to cleanup expired sessions:", error);
     }
@@ -168,13 +179,16 @@ export class MessageManager {
       let sessionToRestore: SessionData | null = null;
 
       if (restoreSessionId) {
-        sessionToRestore = await loadSession(restoreSessionId);
+        sessionToRestore = await loadSession(restoreSessionId, this.sessionDir);
         if (!sessionToRestore) {
           console.error(`Session not found: ${restoreSessionId}`);
           process.exit(1);
         }
       } else if (continueLastSession) {
-        sessionToRestore = await getLatestSession(this.workdir);
+        sessionToRestore = await getLatestSession(
+          this.workdir,
+          this.sessionDir,
+        );
         if (!sessionToRestore) {
           console.error(
             `No previous session found for workdir: ${this.workdir}`,
@@ -290,11 +304,13 @@ export class MessageManager {
   public addAssistantMessage(
     content?: string,
     toolCalls?: ChatCompletionMessageFunctionToolCall[],
+    usage?: Usage,
   ): void {
     const newMessages = addAssistantMessageToMessages(
       this.messages,
       content,
       toolCalls,
+      usage,
     );
     this.setMessages(newMessages);
     this.callbacks.onAssistantMessageAdded?.(content, toolCalls);
@@ -466,4 +482,17 @@ export class MessageManager {
     };
     this.callbacks.onSubAgentBlockUpdated?.(params.subagentId, params.messages);
   }
+
+  /**
+   * Trigger usage change callback with all usage data from assistant messages
+   */
+  public triggerUsageChange(): void {
+    const usages: Usage[] = [];
+    for (const message of this.messages) {
+      if (message.role === "assistant" && message.usage) {
+        usages.push(message.usage);
+      }
+    }
+    this.callbacks.onUsagesChange?.(usages);
+  }
 }

package/src/managers/skillManager.ts

@@ -36,7 +36,7 @@ export class SkillManager {
    * Initialize the skill manager by discovering available skills
    */
   async initialize(): Promise<void> {
-    this.logger?.info("Initializing SkillManager...");
+    this.logger?.debug("Initializing SkillManager...");
 
     try {
       // Clear existing data before discovery
@@ -66,7 +66,7 @@ export class SkillManager {
       }
 
       this.initialized = true;
-      this.logger?.info(
+      this.logger?.debug(
         `SkillManager initialized with ${this.skillMetadata.size} skills`,
       );
     } catch (error) {
@@ -247,7 +247,7 @@ export class SkillManager {
   ): Promise<{ content: string; context?: SkillInvocationContext }> {
     const { skill_name } = args;
 
-    this.logger?.info(`Invoking skill: ${skill_name}`);
+    this.logger?.debug(`Invoking skill: ${skill_name}`);
 
     try {
       // Load the skill

package/src/managers/slashCommandManager.ts

@@ -92,7 +92,7 @@ export class SlashCommandManager {
         });
       }
 
-      this.logger?.info(`Loaded ${customCommands.length} custom commands`);
+      this.logger?.debug(`Loaded ${customCommands.length} custom commands`);
     } catch (error) {
       this.logger?.warn("Failed to load custom commands:", error);
     }

package/src/managers/subagentManager.ts

@@ -1,6 +1,12 @@
 import { randomUUID } from "crypto";
 import type { SubagentConfiguration } from "../utils/subagentParser.js";
-import type { Message, Logger, GatewayConfig, ModelConfig } from "../types.js";
+import type {
+  Message,
+  Logger,
+  GatewayConfig,
+  ModelConfig,
+  Usage,
+} from "../types.js";
 import { AIManager } from "./aiManager.js";
 import {
   MessageManager,
@@ -27,6 +33,7 @@ export interface SubagentManagerOptions {
   gatewayConfig: GatewayConfig;
   modelConfig: ModelConfig;
   tokenLimit: number;
+  onUsageAdded?: (usage: Usage) => void;
 }
 
 export class SubagentManager {
@@ -40,6 +47,7 @@ export class SubagentManager {
   private gatewayConfig: GatewayConfig;
   private modelConfig: ModelConfig;
   private tokenLimit: number;
+  private onUsageAdded?: (usage: Usage) => void;
 
   constructor(options: SubagentManagerOptions) {
     this.workdir = options.workdir;
@@ -49,6 +57,7 @@ export class SubagentManager {
     this.gatewayConfig = options.gatewayConfig;
     this.modelConfig = options.modelConfig;
     this.tokenLimit = options.tokenLimit;
+    this.onUsageAdded = options.onUsageAdded;
   }
 
   /**
@@ -156,6 +165,9 @@ export class SubagentManager {
         agentModel: modelToUse,
       },
       tokenLimit: this.tokenLimit,
+      callbacks: {
+        onUsageAdded: this.onUsageAdded,
+      },
     });
 
     const instance: SubagentInstance = {
@@ -311,7 +323,7 @@ export class SubagentManager {
         messages: instance.messages,
       });
 
-      this.logger?.info(`Aborted subagent instance: ${subagentId}`);
+      this.logger?.debug(`Aborted subagent instance: ${subagentId}`);
       return true;
     } catch (error) {
       this.logger?.error(

package/src/services/aiService.ts

@@ -214,9 +214,18 @@ export interface CompressMessagesOptions {
   abortSignal?: AbortSignal;
 }
 
+export interface CompressMessagesResult {
+  content: string;
+  usage?: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+}
+
 export async function compressMessages(
   options: CompressMessagesOptions,
-): Promise<string> {
+): Promise<CompressMessagesResult> {
   const { gatewayConfig, modelConfig, messages, abortSignal } = options;
 
   // Create OpenAI client with injected configuration
@@ -294,15 +303,29 @@ For technical conversations, structure as:
       },
     );
 
-    return (
+    const content =
       response.choices[0]?.message?.content?.trim() ||
-      "Failed to compress conversation history"
-    );
+      "Failed to compress conversation history";
+    const usage = response.usage
+      ? {
+          prompt_tokens: response.usage.prompt_tokens,
+          completion_tokens: response.usage.completion_tokens,
+          total_tokens: response.usage.total_tokens,
+        }
+      : undefined;
+
+    return {
+      content,
+      usage,
+    };
   } catch (error) {
     if ((error as Error).name === "AbortError") {
       throw new Error("Compression request was aborted");
     }
     // // logger.error("Failed to compress messages:", error);
-    return "Failed to compress conversation history";
+    return {
+      content: "Failed to compress conversation history",
+      usage: undefined,
+    };
   }
 }

package/src/services/memory.ts

@@ -41,7 +41,7 @@ export const addMemory = async (
     // Write file
     await fs.writeFile(memoryFilePath, updatedContent, "utf-8");
 
-    // logger.info(`Memory added to ${memoryFilePath}:`, message);
+    // logger.debug(`Memory added to ${memoryFilePath}:`, message);
   } catch (error) {
     // logger.error("Failed to add memory:", error);
     throw new Error(`Failed to add memory: ${(error as Error).message}`);
@@ -63,7 +63,7 @@ export const ensureUserMemoryFile = async (): Promise<void> => {
         const initialContent =
           "# User Memory\n\nThis is the user-level memory file, recording important information and context across projects.\n\n";
         await fs.writeFile(USER_MEMORY_FILE, initialContent, "utf-8");
-        // logger.info(`Created user memory file: ${USER_MEMORY_FILE}`);
+        // logger.debug(`Created user memory file: ${USER_MEMORY_FILE}`);
       } else {
         throw error;
       }
@@ -93,7 +93,7 @@ export const addUserMemory = async (message: string): Promise<void> => {
     // Write file
     await fs.writeFile(USER_MEMORY_FILE, updatedContent, "utf-8");
 
-    // logger.info(`User memory added to ${USER_MEMORY_FILE}:`, message);
+    // logger.debug(`User memory added to ${USER_MEMORY_FILE}:`, message);
   } catch (error) {
     // logger.error("Failed to add user memory:", error);
     throw new Error(`Failed to add user memory: ${(error as Error).message}`);