wave-agent-sdk 0.0.7 → 0.0.10

package/src/managers/messageManager.ts

@@ -1,9 +1,7 @@
-import { randomUUID } from "crypto";
 import {
   addAssistantMessageToMessages,
   updateToolBlockInMessage,
   addErrorBlockToMessage,
-  addDiffBlockToMessage,
   addUserMessageToMessages,
   extractUserInputHistory,
   addMemoryBlockToMessage,
@@ -18,16 +16,18 @@ import {
   type UpdateSubagentBlockParams,
   type AgentToolBlockUpdateParams,
 } from "../utils/messageOperations.js";
+import type { SubagentConfiguration } from "../utils/subagentParser.js";
 import type { Logger, Message, Usage } from "../types/index.js";
+import { join } from "path";
 import {
-  cleanupExpiredSessions,
-  getLatestSession,
-  loadSession,
-  saveSession,
+  appendMessages,
+  createSession,
+  generateSessionId,
   SessionData,
-  getSessionFilePath,
+  SESSION_DIR,
 } from "../services/session.js";
 import { ChatCompletionMessageFunctionToolCall } from "openai/resources.js";
+import { pathEncoder } from "../utils/pathEncoder.js";
 
 export interface MessageManagerCallbacks {
   onMessagesChange?: (messages: Message[]) => void;
@@ -37,12 +37,13 @@ export interface MessageManagerCallbacks {
   onUsagesChange?: (usages: Usage[]) => void;
   // Incremental callback
   onUserMessageAdded?: (params: UserMessageParams) => void;
-  onAssistantMessageAdded?: (
-    content?: string,
-    toolCalls?: ChatCompletionMessageFunctionToolCall[],
-  ) => void;
+  // MODIFIED: Remove arguments for separation of concerns
+  onAssistantMessageAdded?: () => void;
+  // NEW: Streaming content callback - FR-001: receives chunk and accumulated content
+  onAssistantContentUpdated?: (chunk: string, accumulated: string) => void;
+  // NEW: Streaming reasoning callback
+  onAssistantReasoningUpdated?: (chunk: string, accumulated: string) => void;
   onToolBlockUpdated?: (params: AgentToolBlockUpdateParams) => void;
-  onDiffBlockAdded?: (filePath: string, diffResult: string) => void;
   onErrorBlockAdded?: (error: string) => void;
   onCompressBlockAdded?: (insertIndex: number, content: string) => void;
   onCompressionStateChange?: (isCompressing: boolean) => void;
@@ -67,7 +68,6 @@ export interface MessageManagerCallbacks {
   ) => void;
   onSubAgentBlockUpdated?: (
     subagentId: string,
-    messages: Message[],
     status: "active" | "completed" | "error" | "aborted",
   ) => void;
 }
@@ -76,13 +76,8 @@ 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;
+  sessionType?: "main" | "subagent";
+  subagentType?: string;
 }
 
 export class MessageManager {
@@ -91,22 +86,30 @@ export class MessageManager {
   private messages: Message[];
   private latestTotalTokens: number;
   private userInputHistory: string[];
-  private sessionStartTime: string;
   private workdir: string;
+  private encodedWorkdir: string; // Cached encoded workdir
   private logger?: Logger; // Add optional logger property
   private callbacks: MessageManagerCallbacks;
-  private sessionDir?: string; // Add session directory property
+  private transcriptPath: string; // Cached transcript path
+  private savedMessageCount: number; // Track how many messages have been saved to prevent duplication
+  private sessionType: "main" | "subagent";
+  private subagentType?: string;
 
   constructor(options: MessageManagerOptions) {
-    this.sessionId = randomUUID();
+    this.sessionId = generateSessionId();
     this.messages = [];
     this.latestTotalTokens = 0;
     this.userInputHistory = [];
-    this.sessionStartTime = new Date().toISOString();
     this.workdir = options.workdir;
+    this.encodedWorkdir = pathEncoder.encodeSync(this.workdir); // Cache encoded workdir
     this.callbacks = options.callbacks;
     this.logger = options.logger;
-    this.sessionDir = options.sessionDir;
+    this.savedMessageCount = 0; // Initialize saved message count tracker
+    this.sessionType = options.sessionType || "main";
+    this.subagentType = options.subagentType;
+
+    // Compute and cache the transcript path
+    this.transcriptPath = this.computeTranscriptPath();
   }
 
   // Getter methods
@@ -126,100 +129,97 @@ export class MessageManager {
     return [...this.userInputHistory];
   }
 
+  public getWorkdir(): string {
+    return this.workdir;
+  }
+
+  public getSessionDir(): string {
+    return SESSION_DIR;
+  }
+
   public getTranscriptPath(): string {
-    return getSessionFilePath(this.sessionId, this.sessionDir);
+    return this.transcriptPath;
+  }
+
+  /**
+   * Compute the transcript path using cached encoded workdir
+   * Called during construction and when sessionId changes
+   */
+  private computeTranscriptPath(): string {
+    const baseDir = join(SESSION_DIR, this.encodedWorkdir);
+
+    // All sessions now go in the same directory
+    // Session type is determined by metadata, not file path
+    return join(baseDir, `${this.sessionId}.jsonl`);
   }
 
   // Setter methods, will trigger callbacks
   public setSessionId(sessionId: string): void {
     if (this.sessionId !== sessionId) {
       this.sessionId = sessionId;
+      // Reset saved message count for new session
+      this.savedMessageCount = 0;
+      // Recompute transcript path since session ID changed
+      this.transcriptPath = this.computeTranscriptPath();
+
       this.callbacks.onSessionIdChange?.(sessionId);
     }
   }
 
-  public setMessages(messages: Message[]): void {
-    this.messages = [...messages];
-    this.callbacks.onMessagesChange?.([...messages]);
-  }
-
   /**
-   * Save current session
+   * Create session if needed (async helper)
    */
-  public async saveSession(): Promise<void> {
+  private async createSessionIfNeeded(): Promise<void> {
     try {
-      await saveSession(
-        this.sessionId,
-        this.messages,
-        this.workdir,
-        this.latestTotalTokens,
-        this.sessionStartTime,
-        this.sessionDir,
-      );
+      await createSession(this.sessionId, this.workdir, this.sessionType);
     } catch (error) {
-      this.logger?.error("Failed to save session:", error);
+      this.logger?.error("Failed to create session:", error);
     }
   }
 
+  public setMessages(messages: Message[]): void {
+    this.messages = [...messages];
+    this.callbacks.onMessagesChange?.([...messages]);
+  }
+
   /**
-   * Handle session restoration logic
+   * Save current session
    */
-  public async handleSessionRestoration(
-    restoreSessionId?: string,
-    continueLastSession?: boolean,
-  ): Promise<void> {
-    // Clean up expired sessions first
+  public async saveSession(): Promise<void> {
     try {
-      await cleanupExpiredSessions(this.workdir, this.sessionDir);
-    } catch (error) {
-      this.logger?.warn("Failed to cleanup expired sessions:", error);
-    }
+      // Only save messages that haven't been saved yet
+      const unsavedMessages = this.messages.slice(this.savedMessageCount);
 
-    if (!restoreSessionId && !continueLastSession) {
-      return;
-    }
-
-    try {
-      let sessionToRestore: SessionData | null = null;
+      if (unsavedMessages.length === 0) {
+        // No new messages to save
+        return;
+      }
 
-      if (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,
-          this.sessionDir,
-        );
-        if (!sessionToRestore) {
-          console.error(
-            `No previous session found for workdir: ${this.workdir}`,
-          );
-          process.exit(1);
-        }
+      // Create session if needed (only when we have messages to save)
+      if (this.savedMessageCount === 0) {
+        // This is the first time saving messages, so create the session
+        await this.createSessionIfNeeded();
       }
 
-      if (sessionToRestore) {
-        console.log(`Restoring session: ${sessionToRestore.id}`);
+      // Use JSONL format for new sessions
+      await appendMessages(
+        this.sessionId,
+        unsavedMessages, // Only append new messages
+        this.workdir,
+        this.sessionType,
+      );
 
-        // Initialize from session data
-        this.initializeFromSession(
-          sessionToRestore.id,
-          sessionToRestore.messages,
-          sessionToRestore.metadata.latestTotalTokens,
-        );
-      }
+      // Update the saved message count
+      this.savedMessageCount = this.messages.length;
     } catch (error) {
-      console.error("Failed to restore session:", error);
-      process.exit(1);
+      this.logger?.error("Failed to save session:", error);
     }
   }
 
   public setlatestTotalTokens(latestTotalTokens: number): void {
     if (this.latestTotalTokens !== latestTotalTokens) {
       this.latestTotalTokens = latestTotalTokens;
+
       this.callbacks.onLatestTotalTokensChange?.(latestTotalTokens);
     }
   }
@@ -235,23 +235,23 @@ export class MessageManager {
   public clearMessages(): void {
     this.setMessages([]);
     this.setUserInputHistory([]);
-    this.setSessionId(randomUUID());
+    this.setSessionId(generateSessionId());
     this.setlatestTotalTokens(0);
-    this.sessionStartTime = new Date().toISOString();
+    this.savedMessageCount = 0; // Reset saved message count
   }
 
   // Initialize state from session data
-  public initializeFromSession(
-    sessionId: string,
-    messages: Message[],
-    latestTotalTokens: number,
-  ): void {
-    this.setSessionId(sessionId);
-    this.setMessages([...messages]);
-    this.setlatestTotalTokens(latestTotalTokens);
+  public initializeFromSession(sessionData: SessionData): void {
+    this.setSessionId(sessionData.id);
+    this.setMessages([...sessionData.messages]);
+    this.setlatestTotalTokens(sessionData.metadata.latestTotalTokens);
 
     // Extract user input history from session messages
-    this.setUserInputHistory(extractUserInputHistory(messages));
+    this.setUserInputHistory(extractUserInputHistory(sessionData.messages));
+
+    // Set saved message count to the number of loaded messages since they're already saved
+    // This must be done after setSessionId which resets it to 0
+    this.savedMessageCount = sessionData.messages.length;
   }
 
   // Add to input history
@@ -280,21 +280,68 @@ export class MessageManager {
     });
     this.setMessages(newMessages);
     this.callbacks.onUserMessageAdded?.(params);
+
+    // Note: Subagent-specific callbacks are now handled by SubagentManager
   }
 
   public addAssistantMessage(
     content?: string,
     toolCalls?: ChatCompletionMessageFunctionToolCall[],
     usage?: Usage,
+    additionalFields?: Record<string, unknown>,
   ): void {
+    const additionalFieldsRecord = additionalFields
+      ? Object.fromEntries(
+          Object.entries(additionalFields).filter(
+            ([, value]) => value !== undefined,
+          ),
+        )
+      : undefined;
+
     const newMessages = addAssistantMessageToMessages(
       this.messages,
       content,
       toolCalls,
       usage,
+      additionalFieldsRecord,
     );
     this.setMessages(newMessages);
-    this.callbacks.onAssistantMessageAdded?.(content, toolCalls);
+    this.callbacks.onAssistantMessageAdded?.();
+
+    // Note: Subagent-specific callbacks are now handled by SubagentManager
+  }
+
+  public mergeAssistantAdditionalFields(
+    additionalFields: Record<string, unknown>,
+  ): void {
+    if (!additionalFields || Object.keys(additionalFields).length === 0) {
+      return;
+    }
+
+    const newMessages = [...this.messages];
+    for (let i = newMessages.length - 1; i >= 0; i--) {
+      const message = newMessages[i];
+      if (message.role === "assistant") {
+        const mergedAdditionalFields = {
+          ...(message.additionalFields || {}),
+        } as Record<string, unknown>;
+
+        for (const [key, value] of Object.entries(additionalFields)) {
+          if (value === undefined) {
+            continue;
+          }
+          mergedAdditionalFields[key] = value;
+        }
+
+        if (Object.keys(mergedAdditionalFields).length === 0) {
+          return;
+        }
+
+        message.additionalFields = mergedAdditionalFields;
+        this.setMessages(newMessages);
+        return;
+      }
+    }
   }
 
   public updateToolBlock(params: AgentToolBlockUpdateParams): void {
@@ -305,27 +352,17 @@ export class MessageManager {
       result: params.result,
       success: params.success,
       error: params.error,
-      isRunning: params.isRunning,
+      stage: params.stage,
       name: params.name,
       shortResult: params.shortResult,
       images: params.images,
       compactParams: params.compactParams,
+      parametersChunk: params.parametersChunk,
     });
     this.setMessages(newMessages);
     this.callbacks.onToolBlockUpdated?.(params);
-  }
 
-  public addDiffBlock(
-    filePath: string,
-    diffResult: Array<{ value: string; added?: boolean; removed?: boolean }>,
-  ): void {
-    const newMessages = addDiffBlockToMessage({
-      messages: this.messages,
-      path: filePath,
-      diffResult,
-    });
-    this.setMessages(newMessages);
-    this.callbacks.onDiffBlockAdded?.(filePath, JSON.stringify(diffResult));
+    // Note: Subagent-specific callbacks are now handled by SubagentManager
   }
 
   public addErrorBlock(error: string): void {
@@ -343,6 +380,7 @@ export class MessageManager {
   public compressMessagesAndUpdateSession(
     insertIndex: number,
     compressedContent: string,
+    usage?: Usage,
   ): void {
     const currentMessages = this.messages;
 
@@ -353,8 +391,10 @@ export class MessageManager {
         {
           type: "compress",
           content: compressedContent,
+          sessionId: this.sessionId,
         },
       ],
+      ...(usage && { usage }),
     };
 
     // Convert negative index to positive index
@@ -368,7 +408,7 @@ export class MessageManager {
     ];
 
     // Update sessionId
-    this.setSessionId(randomUUID());
+    this.setSessionId(generateSessionId());
 
     // Set new message list
     this.setMessages(newMessages);
@@ -428,8 +468,9 @@ export class MessageManager {
   public addSubagentBlock(
     subagentId: string,
     subagentName: string,
+    sessionId: string,
+    configuration: SubagentConfiguration,
     status: "active" | "completed" | "error" = "active",
-    subagentMessages: Message[] = [],
     parameters: {
       description: string;
       prompt: string;
@@ -440,8 +481,9 @@ export class MessageManager {
       messages: this.messages,
       subagentId,
       subagentName,
+      sessionId,
       status,
-      subagentMessages,
+      configuration,
     };
     const updatedMessages = addSubagentBlockToMessage(params);
     this.setMessages(updatedMessages);
@@ -452,7 +494,7 @@ export class MessageManager {
     subagentId: string,
     updates: Partial<{
       status: "active" | "completed" | "error" | "aborted";
-      messages: Message[];
+      sessionId: string;
     }>,
   ): void {
     const updatedMessages = updateSubagentBlockInMessage(
@@ -465,13 +507,8 @@ export class MessageManager {
       messages: this.messages,
       subagentId,
       status: updates.status || "active",
-      subagentMessages: updates.messages || [],
     };
-    this.callbacks.onSubAgentBlockUpdated?.(
-      params.subagentId,
-      params.messages,
-      params.status,
-    );
+    this.callbacks.onSubAgentBlockUpdated?.(params.subagentId, params.status);
   }
 
   /**
@@ -487,6 +524,106 @@ export class MessageManager {
     this.callbacks.onUsagesChange?.(usages);
   }
 
+  /**
+   * Update the current assistant message content during streaming
+   * This method updates the last assistant message's content without creating a new message
+   * FR-001: Tracks and provides both chunk (new content) and accumulated (total content)
+   */
+  public updateCurrentMessageContent(newAccumulatedContent: string): void {
+    if (this.messages.length === 0) return;
+
+    const lastMessage = this.messages[this.messages.length - 1];
+    if (lastMessage.role !== "assistant") return;
+
+    // Get the current content to calculate the chunk
+    const textBlockIndex = lastMessage.blocks.findIndex(
+      (block) => block.type === "text",
+    );
+    const currentContent =
+      textBlockIndex >= 0
+        ? (
+            lastMessage.blocks[textBlockIndex] as {
+              type: "text";
+              content: string;
+            }
+          ).content || ""
+        : "";
+
+    // Calculate the chunk (new content since last update)
+    const chunk = newAccumulatedContent.slice(currentContent.length);
+
+    if (textBlockIndex >= 0) {
+      // Update existing text block
+      lastMessage.blocks[textBlockIndex] = {
+        type: "text",
+        content: newAccumulatedContent,
+      };
+    } else {
+      // Add new text block if none exists
+      lastMessage.blocks.push({
+        type: "text",
+        content: newAccumulatedContent,
+      });
+    }
+
+    // FR-001: Trigger callbacks with chunk and accumulated content
+    this.callbacks.onAssistantContentUpdated?.(chunk, newAccumulatedContent);
+
+    // Note: Subagent-specific callbacks are now handled by SubagentManager
+
+    this.callbacks.onMessagesChange?.([...this.messages]); // Still need to notify of changes
+  }
+
+  /**
+   * Update the current assistant message reasoning during streaming
+   * This method updates the last assistant message's reasoning content without creating a new message
+   */
+  public updateCurrentMessageReasoning(newAccumulatedReasoning: string): void {
+    if (this.messages.length === 0) return;
+
+    const lastMessage = this.messages[this.messages.length - 1];
+    if (lastMessage.role !== "assistant") return;
+
+    // Get the current reasoning content to calculate the chunk
+    const reasoningBlockIndex = lastMessage.blocks.findIndex(
+      (block) => block.type === "reasoning",
+    );
+    const currentReasoning =
+      reasoningBlockIndex >= 0
+        ? (
+            lastMessage.blocks[reasoningBlockIndex] as {
+              type: "reasoning";
+              content: string;
+            }
+          ).content || ""
+        : "";
+
+    // Calculate the chunk (new content since last update)
+    const chunk = newAccumulatedReasoning.slice(currentReasoning.length);
+
+    if (reasoningBlockIndex >= 0) {
+      // Update existing reasoning block
+      lastMessage.blocks[reasoningBlockIndex] = {
+        type: "reasoning",
+        content: newAccumulatedReasoning,
+      };
+    } else {
+      // Add new reasoning block if none exists
+      lastMessage.blocks.push({
+        type: "reasoning",
+        content: newAccumulatedReasoning,
+      });
+    }
+
+    // Trigger callbacks with chunk and accumulated reasoning content
+    this.callbacks.onAssistantReasoningUpdated?.(
+      chunk,
+      newAccumulatedReasoning,
+    );
+
+    this.callbacks.onMessagesChange?.([...this.messages]); // Still need to notify of changes
+  }
+
   /**
    * Remove the last user message from the conversation
    * Used for hook error handling when the user prompt needs to be erased