wave-agent-sdk 0.0.8 → 0.0.11

package/src/managers/messageManager.ts

@@ -2,7 +2,6 @@ import {
   addAssistantMessageToMessages,
   updateToolBlockInMessage,
   addErrorBlockToMessage,
-  addDiffBlockToMessage,
   addUserMessageToMessages,
   extractUserInputHistory,
   addMemoryBlockToMessage,
@@ -21,9 +20,6 @@ import type { SubagentConfiguration } from "../utils/subagentParser.js";
 import type { Logger, Message, Usage } from "../types/index.js";
 import { join } from "path";
 import {
-  cleanupExpiredSessionsFromJsonl,
-  getLatestSessionFromJsonl,
-  loadSessionFromJsonl,
   appendMessages,
   createSession,
   generateSessionId,
@@ -45,8 +41,9 @@ export interface MessageManagerCallbacks {
   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;
@@ -80,7 +77,6 @@ export interface MessageManagerOptions {
   workdir: string;
   logger?: Logger;
   sessionType?: "main" | "subagent";
-  parentSessionId?: string;
   subagentType?: string;
 }
 
@@ -90,7 +86,6 @@ 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
@@ -98,7 +93,6 @@ export class MessageManager {
   private transcriptPath: string; // Cached transcript path
   private savedMessageCount: number; // Track how many messages have been saved to prevent duplication
   private sessionType: "main" | "subagent";
-  private parentSessionId?: string;
   private subagentType?: string;
 
   constructor(options: MessageManagerOptions) {
@@ -106,14 +100,12 @@ export class MessageManager {
     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.savedMessageCount = 0; // Initialize saved message count tracker
     this.sessionType = options.sessionType || "main";
-    this.parentSessionId = options.parentSessionId;
     this.subagentType = options.subagentType;
 
     // Compute and cache the transcript path
@@ -137,10 +129,6 @@ export class MessageManager {
     return [...this.userInputHistory];
   }
 
-  public getSessionStartTime(): string {
-    return this.sessionStartTime;
-  }
-
   public getWorkdir(): string {
     return this.workdir;
   }
@@ -173,6 +161,7 @@ export class MessageManager {
       this.savedMessageCount = 0;
       // Recompute transcript path since session ID changed
       this.transcriptPath = this.computeTranscriptPath();
+
       this.callbacks.onSessionIdChange?.(sessionId);
     }
   }
@@ -182,13 +171,7 @@ export class MessageManager {
    */
   private async createSessionIfNeeded(): Promise<void> {
     try {
-      await createSession(
-        this.sessionId,
-        this.workdir,
-        this.sessionType,
-        this.parentSessionId,
-        this.subagentType,
-      );
+      await createSession(this.sessionId, this.workdir, this.sessionType);
     } catch (error) {
       this.logger?.error("Failed to create session:", error);
     }
@@ -223,6 +206,7 @@ export class MessageManager {
         this.sessionId,
         unsavedMessages, // Only append new messages
         this.workdir,
+        this.sessionType,
       );
 
       // Update the saved message count
@@ -232,61 +216,10 @@ export class MessageManager {
     }
   }
 
-  /**
-   * Handle session restoration logic
-   */
-  public async handleSessionRestoration(
-    restoreSessionId?: string,
-    continueLastSession?: boolean,
-  ): Promise<void> {
-    // Clean up expired sessions first
-    cleanupExpiredSessionsFromJsonl(this.workdir).catch((error) => {
-      this.logger?.warn("Failed to cleanup expired sessions:", error);
-    });
-
-    if (!restoreSessionId && !continueLastSession) {
-      return;
-    }
-
-    try {
-      let sessionToRestore: SessionData | null = null;
-
-      if (restoreSessionId) {
-        // Use only JSONL format - no legacy support
-        sessionToRestore = await loadSessionFromJsonl(
-          restoreSessionId,
-          this.workdir,
-        );
-        if (!sessionToRestore) {
-          console.error(`Session not found: ${restoreSessionId}`);
-          process.exit(1);
-        }
-      } else if (continueLastSession) {
-        // Use only JSONL format - no legacy support
-        sessionToRestore = await getLatestSessionFromJsonl(this.workdir);
-        if (!sessionToRestore) {
-          console.error(
-            `No previous session found for workdir: ${this.workdir}`,
-          );
-          process.exit(1);
-        }
-      }
-
-      if (sessionToRestore) {
-        console.log(`Restoring session: ${sessionToRestore.id}`);
-
-        // Initialize from session data
-        this.initializeFromSession(sessionToRestore);
-      }
-    } catch (error) {
-      console.error("Failed to restore session:", error);
-      process.exit(1);
-    }
-  }
-
   public setlatestTotalTokens(latestTotalTokens: number): void {
     if (this.latestTotalTokens !== latestTotalTokens) {
       this.latestTotalTokens = latestTotalTokens;
+
       this.callbacks.onLatestTotalTokensChange?.(latestTotalTokens);
     }
   }
@@ -304,7 +237,6 @@ export class MessageManager {
     this.setUserInputHistory([]);
     this.setSessionId(generateSessionId());
     this.setlatestTotalTokens(0);
-    this.sessionStartTime = new Date().toISOString();
     this.savedMessageCount = 0; // Reset saved message count
   }
 
@@ -314,9 +246,6 @@ export class MessageManager {
     this.setMessages([...sessionData.messages]);
     this.setlatestTotalTokens(sessionData.metadata.latestTotalTokens);
 
-    // Restore the original session start time
-    this.sessionStartTime = sessionData.metadata.startedAt;
-
     // Extract user input history from session messages
     this.setUserInputHistory(extractUserInputHistory(sessionData.messages));
 
@@ -359,11 +288,13 @@ export class MessageManager {
     content?: string,
     toolCalls?: ChatCompletionMessageFunctionToolCall[],
     usage?: Usage,
-    metadata?: Record<string, unknown>,
+    additionalFields?: Record<string, unknown>,
   ): void {
-    const metadataRecord = metadata
+    const additionalFieldsRecord = additionalFields
       ? Object.fromEntries(
-          Object.entries(metadata).filter(([, value]) => value !== undefined),
+          Object.entries(additionalFields).filter(
+            ([, value]) => value !== undefined,
+          ),
         )
       : undefined;
 
@@ -372,7 +303,7 @@ export class MessageManager {
       content,
       toolCalls,
       usage,
-      metadataRecord,
+      additionalFieldsRecord,
     );
     this.setMessages(newMessages);
     this.callbacks.onAssistantMessageAdded?.();
@@ -380,8 +311,10 @@ export class MessageManager {
     // Note: Subagent-specific callbacks are now handled by SubagentManager
   }
 
-  public mergeAssistantMetadata(metadata: Record<string, unknown>): void {
-    if (!metadata || Object.keys(metadata).length === 0) {
+  public mergeAssistantAdditionalFields(
+    additionalFields: Record<string, unknown>,
+  ): void {
+    if (!additionalFields || Object.keys(additionalFields).length === 0) {
       return;
     }
 
@@ -389,22 +322,22 @@ export class MessageManager {
     for (let i = newMessages.length - 1; i >= 0; i--) {
       const message = newMessages[i];
       if (message.role === "assistant") {
-        const mergedMetadata = {
-          ...(message.metadata || {}),
+        const mergedAdditionalFields = {
+          ...(message.additionalFields || {}),
         } as Record<string, unknown>;
 
-        for (const [key, value] of Object.entries(metadata)) {
+        for (const [key, value] of Object.entries(additionalFields)) {
           if (value === undefined) {
             continue;
           }
-          mergedMetadata[key] = value;
+          mergedAdditionalFields[key] = value;
         }
 
-        if (Object.keys(mergedMetadata).length === 0) {
+        if (Object.keys(mergedAdditionalFields).length === 0) {
           return;
         }
 
-        message.metadata = mergedMetadata;
+        message.additionalFields = mergedAdditionalFields;
         this.setMessages(newMessages);
         return;
       }
@@ -432,19 +365,6 @@ export class MessageManager {
     // Note: Subagent-specific callbacks are now handled by SubagentManager
   }
 
-  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));
-  }
-
   public addErrorBlock(error: string): void {
     const newMessages = addErrorBlockToMessage({
       messages: this.messages,
@@ -640,7 +560,7 @@ export class MessageManager {
       };
     } else {
       // Add new text block if none exists
-      lastMessage.blocks.unshift({
+      lastMessage.blocks.push({
         type: "text",
         content: newAccumulatedContent,
       });
@@ -654,6 +574,56 @@ export class MessageManager {
     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