wave-agent-sdk 0.5.0 → 0.6.2

package/src/agent.ts

@@ -9,7 +9,6 @@ import {
   SubagentManager,
   type SubagentManagerCallbacks,
 } from "./managers/subagentManager.js";
-import * as memory from "./services/memory.js";
 import { McpManager, type McpManagerCallbacks } from "./managers/mcpManager.js";
 import { LspManager } from "./managers/lspManager.js";
 import { BashManager } from "./managers/bashManager.js";
@@ -46,11 +45,11 @@ import { MemoryRuleManager } from "./managers/MemoryRuleManager.js";
 import { LiveConfigManager } from "./managers/liveConfigManager.js";
 import { configValidator } from "./utils/configValidator.js";
 import { SkillManager } from "./managers/skillManager.js";
+import { TaskManager } from "./services/taskManager.js";
 import {
   loadSessionFromJsonl,
   handleSessionRestoration,
 } from "./services/session.js";
-import type { SubagentConfiguration } from "./utils/subagentParser.js";
 import { setGlobalLogger } from "./utils/globalLogger.js";
 import { ConfigurationService } from "./services/configurationService.js";
 import * as fs from "fs/promises";
@@ -107,7 +106,12 @@ export interface AgentCallbacks
     McpManagerCallbacks,
     SubagentManagerCallbacks {
   onTasksChange?: (tasks: BackgroundTask[]) => void;
+  onSessionTasksChange?: (tasks: import("./types/tasks.js").Task[]) => void;
   onPermissionModeChange?: (mode: PermissionMode) => void;
+  onSubagentLatestTotalTokensChange?: (
+    subagentId: string,
+    tokens: number,
+  ) => void;
   onBackgroundCurrentTask?: () => void;
 }
 
@@ -131,6 +135,7 @@ export class Agent {
   private reversionManager: ReversionManager;
   private memoryRuleManager: MemoryRuleManager; // Add memory rule manager instance
   private liveConfigManager: LiveConfigManager; // Add live configuration manager
+  private taskManager: TaskManager;
   private foregroundTaskManager: ForegroundTaskManager;
   private configurationService: ConfigurationService; // Add configuration service
   private workdir: string; // Working directory
@@ -208,6 +213,43 @@ export class Agent {
 
     this.foregroundTaskManager = new ForegroundTaskManager();
 
+    // Initialize memory rule manager
+    this.memoryRuleManager = new MemoryRuleManager({
+      workdir: this.workdir,
+    });
+
+    // Initialize MessageManager
+    this.messageManager = new MessageManager({
+      callbacks: {
+        ...callbacks,
+        onSessionIdChange: (sessionId) => {
+          // When session ID changes (e.g. due to compression),
+          // we update the task manager to use the root session ID
+          // to ensure the task list remains consistent.
+          this.taskManager.setTaskListId(
+            this.messageManager.getRootSessionId(),
+          );
+          callbacks.onSessionIdChange?.(sessionId);
+        },
+      },
+      workdir: this.workdir,
+      logger: this.logger,
+      memoryRuleManager: this.memoryRuleManager,
+    });
+
+    // Resolve taskListId once during construction to ensure stability
+    const resolvedTaskListId =
+      this.configurationService.getEnvironmentVars().WAVE_TASK_LIST_ID ||
+      process.env.WAVE_TASK_LIST_ID ||
+      this.messageManager.getRootSessionId();
+
+    this.taskManager = new TaskManager(resolvedTaskListId);
+    this.taskManager.on("tasksChange", async () => {
+      const tasks = await this.taskManager.listTasks();
+      this.options.callbacks?.onSessionTasksChange?.(tasks);
+    });
+
+    // Initialize BackgroundTaskManager
     this.backgroundTaskManager = new BackgroundTaskManager({
       callbacks: {
         ...callbacks,
@@ -243,20 +285,7 @@ export class Agent {
       workdir: this.workdir,
     });
 
-    // Initialize memory rule manager
-    this.memoryRuleManager = new MemoryRuleManager({
-      workdir: this.workdir,
-    });
-
-    // Initialize MessageManager
-    this.messageManager = new MessageManager({
-      callbacks,
-      workdir: this.workdir,
-      logger: this.logger,
-      memoryRuleManager: this.memoryRuleManager,
-    });
-
-    // Initialize ReversionManager
+    // ReversionManager depends on MessageManager
     this.reversionManager = new ReversionManager(
       new ReversionService(this.messageManager.getTranscriptPath()),
     );
@@ -312,6 +341,7 @@ export class Agent {
       reversionManager: this.reversionManager,
       permissionMode: options.permissionMode, // Let PermissionManager handle defaultMode resolution
       canUseToolCallback: canUseToolWithNotification,
+      taskManager: this.taskManager,
       backgroundTaskManager: this.backgroundTaskManager,
       foregroundTaskManager: this.foregroundTaskManager,
     }); // Initialize tool registry with permission support
@@ -328,7 +358,6 @@ export class Agent {
     this.subagentManager = new SubagentManager({
       workdir: this.workdir,
       parentToolManager: this.toolManager,
-      parentMessageManager: this.messageManager,
       callbacks: {
         onSubagentUserMessageAdded: callbacks.onSubagentUserMessageAdded,
         onSubagentAssistantMessageAdded:
@@ -339,6 +368,8 @@ export class Agent {
           callbacks.onSubagentAssistantReasoningUpdated,
         onSubagentToolBlockUpdated: callbacks.onSubagentToolBlockUpdated,
         onSubagentMessagesChange: callbacks.onSubagentMessagesChange,
+        onSubagentLatestTotalTokensChange:
+          callbacks.onSubagentLatestTotalTokensChange,
       }, // Pass subagent callbacks for forwarding
       logger: this.logger,
       getGatewayConfig: () => this.getGatewayConfig(),
@@ -348,6 +379,7 @@ export class Agent {
       hookManager: this.hookManager,
       onUsageAdded: (usage) => this.addUsage(usage),
       backgroundTaskManager: this.backgroundTaskManager,
+      taskManager: this.taskManager,
       memoryRuleManager: this.memoryRuleManager,
     });
 
@@ -355,6 +387,7 @@ export class Agent {
     this.aiManager = new AIManager({
       messageManager: this.messageManager,
       toolManager: this.toolManager,
+      taskManager: this.taskManager,
       logger: this.logger,
       backgroundTaskManager: this.backgroundTaskManager,
       hookManager: this.hookManager,
@@ -381,6 +414,7 @@ export class Agent {
       messageManager: this.messageManager,
       aiManager: this.aiManager,
       backgroundTaskManager: this.backgroundTaskManager,
+      taskManager: this.taskManager,
       workdir: this.workdir,
       logger: this.logger,
     });
@@ -455,10 +489,6 @@ export class Agent {
     return this.messageManager.getlatestTotalTokens();
   }
 
-  public get userInputHistory(): string[] {
-    return this.messageManager.getUserInputHistory();
-  }
-
   /** Get working directory */
   public get workingDirectory(): string {
     return this.workdir;
@@ -777,90 +807,18 @@ export class Agent {
       // Rebuild usage array from restored messages
       this.rebuildUsageFromMessages(sessionToRestore?.messages || []);
 
-      // After main session is restored, restore any associated subagent sessions
-      await this.restoreSubagentSessions(sessionToRestore?.messages || []);
-
-      if (sessionToRestore)
+      if (sessionToRestore) {
         this.messageManager.initializeFromSession(sessionToRestore);
-    }
-  }
-
-  /**
-   * Restore subagent sessions associated with the current main session
-   * This method is called after the main session is restored to load any subagent sessions
-   */
-  private async restoreSubagentSessions(messages: Message[]): Promise<void> {
-    try {
-      // Only attempt to restore subagent sessions if we have messages (session was restored)
-      if (messages.length === 0) {
-        return;
-      }
-
-      // Extract sessionId -> subagentId mapping from SubagentBlocks
-      const subagentBlockMap = new Map<
-        string,
-        { subagentId: string; configuration: SubagentConfiguration }
-      >(); // sessionId -> { subagentId, configuration }
-
-      for (const message of messages) {
-        if (message.role === "assistant" && message.blocks) {
-          for (const block of message.blocks) {
-            if (
-              block.type === "subagent" &&
-              block.sessionId &&
-              block.subagentId &&
-              block.configuration
-            ) {
-              subagentBlockMap.set(block.sessionId, {
-                subagentId: block.subagentId,
-                configuration: block.configuration,
-              });
-            }
-          }
-        }
-      }
-
-      if (subagentBlockMap.size === 0) {
-        return; // No subagent blocks found
-      }
 
-      // Load subagent sessions using sessionIds
-      const subagentSessions = [];
-      for (const [sessionId, blockData] of subagentBlockMap) {
-        try {
-          const sessionData = await loadSessionFromJsonl(
-            sessionId,
-            this.messageManager.getWorkdir(),
-            "subagent",
-          );
-          if (sessionData) {
-            subagentSessions.push({
-              sessionData,
-              subagentId: blockData.subagentId, // Use the subagentId from SubagentBlock
-              configuration: blockData.configuration, // Include configuration
-            });
-          }
-        } catch (error) {
-          this.logger?.warn(
-            `Failed to load subagent session ${sessionId}:`,
-            error,
-          );
-        }
-      }
-
-      if (subagentSessions.length > 0) {
-        this.logger?.debug(
-          `Found ${subagentSessions.length} subagent sessions to restore`,
+        // Update task manager with the root session ID to ensure continuity across compressions
+        this.taskManager.setTaskListId(
+          sessionToRestore.rootSessionId || sessionToRestore.id,
         );
 
-        // Restore subagent sessions through the SubagentManager
-        await this.subagentManager.restoreSubagentSessions(subagentSessions);
-
-        this.logger?.debug("Subagent sessions restored successfully");
+        // After session is initialized, load tasks for the session
+        const tasks = await this.taskManager.listTasks();
+        this.options.callbacks?.onSessionTasksChange?.(tasks);
       }
-    } catch (error) {
-      this.logger?.warn("Failed to restore subagent sessions:", error);
-      // Don't throw error to prevent app startup failure
     }
   }
 
@@ -898,12 +856,18 @@ export class Agent {
     this.abortMessage(); // Abort any running operations
     this.subagentManager.cleanup(); // Clean up active subagents
 
-    // 5. Rebuild usage and restore subagents (in correct order)
+    // 5. Rebuild usage (in correct order)
     this.rebuildUsageFromMessages(sessionData.messages);
-    await this.restoreSubagentSessions(sessionData.messages);
 
     // 6. Initialize session state last
     this.messageManager.initializeFromSession(sessionData);
+
+    // Update task manager with the root session ID to ensure continuity across compressions
+    this.taskManager.setTaskListId(sessionData.rootSessionId || sessionData.id);
+
+    // 7. Load tasks for the restored session
+    const tasks = await this.taskManager.listTasks();
+    this.options.callbacks?.onSessionTasksChange?.(tasks);
   }
 
   public abortAIMessage(): void {
@@ -912,12 +876,9 @@ export class Agent {
 
   /** Execute bash command */
   public async executeBashCommand(command: string): Promise<void> {
-    // Add user message to history (but not displayed in UI)
-    this.addToInputHistory(`!${command}`);
     await this.bashManager?.executeCommand(command);
   }
 
-  /** Clear messages and input history */
   public clearMessages(): void {
     this.messageManager.clearMessages();
   }
@@ -929,11 +890,6 @@ export class Agent {
     this.abortSlashCommand();
   }
 
-  /** Add to input history */
-  private addToInputHistory(input: string): void {
-    this.messageManager.addToInputHistory(input);
-  }
-
   /** Interrupt bash command execution */
   public abortBashCommand(): void {
     this.bashManager?.abortCommand();
@@ -964,12 +920,6 @@ export class Agent {
   public async backgroundCurrentTask(): Promise<void> {
     await this.foregroundTaskManager.backgroundCurrentTask();
     this.options.callbacks?.onBackgroundCurrentTask?.();
-
-    // If there was no foreground task (e.g. a tool that doesn't register one),
-    // or even if there was, we should stop the AI recursion loop.
-    if (!this.foregroundTaskManager.hasActiveTasks()) {
-      this.aiManager.abortRecursion();
-    }
   }
 
   /** Destroy managers, clean up resources */
@@ -1061,8 +1011,6 @@ export class Agent {
           // Execute valid slash command
           await this.slashCommandManager.executeCommand(commandId, args);
 
-          // Add slash command to history
-          this.addToInputHistory(command);
           return;
         }
 
@@ -1071,9 +1019,6 @@ export class Agent {
       }
 
       // Handle normal AI message
-      // Add user message to history
-      this.addToInputHistory(content);
-
       // Add user message first, will automatically sync to UI
       this.messageManager.addUserMessage({
         content,
@@ -1130,67 +1075,6 @@ export class Agent {
     }
   }
 
-  /** Save memory to project or user memory file */
-  public async saveMemory(
-    message: string,
-    type: "project" | "user",
-  ): Promise<void> {
-    try {
-      // Ensure the message starts with # for memory functions
-      const formattedMessage = message.startsWith("#")
-        ? message
-        : `#${message}`;
-
-      if (type === "project") {
-        await memory.addMemory(formattedMessage, this.workdir);
-        // Update internal state after successful save
-        this._projectMemoryContent = await memory.readMemoryFile(this.workdir);
-      } else {
-        await memory.addUserMemory(formattedMessage);
-        // Update internal state after successful save
-        this._userMemoryContent = await memory.getUserMemoryContent();
-      }
-
-      // Add successful MemoryBlock to the last assistant message
-      const memoryText = message.substring(1).trim();
-      const typeLabel = type === "project" ? "Project Memory" : "User Memory";
-      const storagePath = "AGENTS.md";
-
-      const messages = this.messageManager.getMessages();
-      const lastMessage = messages[messages.length - 1];
-      if (lastMessage && lastMessage.role === "assistant") {
-        lastMessage.blocks.push({
-          type: "memory",
-          content: `${typeLabel}: ${memoryText}`,
-          isSuccess: true,
-          memoryType: type,
-          storagePath,
-        });
-        this.messageManager.setMessages(messages);
-      }
-    } catch (error) {
-      // Add failed MemoryBlock to the last assistant message
-      const memoryText = message.substring(1).trim();
-      const typeLabel = type === "project" ? "Project Memory" : "User Memory";
-      const storagePath = "AGENTS.md";
-      const errorMessage =
-        error instanceof Error ? error.message : String(error);
-
-      const messages = this.messageManager.getMessages();
-      const lastMessage = messages[messages.length - 1];
-      if (lastMessage && lastMessage.role === "assistant") {
-        lastMessage.blocks.push({
-          type: "memory",
-          content: `${typeLabel}: ${memoryText} - Error: ${errorMessage}`,
-          isSuccess: false,
-          memoryType: type,
-          storagePath,
-        });
-        this.messageManager.setMessages(messages);
-      }
-    }
-  }
-
   // ========== MCP Management Methods ==========
 
   /** Get all MCP server status */
@@ -1268,6 +1152,18 @@ export class Agent {
    */
   public async truncateHistory(index: number): Promise<void> {
     await this.messageManager.truncateHistory(index, this.reversionManager);
+    // After truncating history, the task list might have changed, so refresh it.
+    await this.taskManager.refreshTasks();
+  }
+
+  /**
+   * Get the full message thread including parent sessions
+   */
+  public async getFullMessageThread(): Promise<{
+    messages: Message[];
+    sessionIds: string[];
+  }> {
+    return this.messageManager.getFullMessageThread();
   }
 
   /**
@@ -1338,7 +1234,7 @@ export class Agent {
   private handlePlanModeTransition(mode: PermissionMode): void {
     if (mode === "plan") {
       this.planManager
-        .getOrGeneratePlanFilePath()
+        .getOrGeneratePlanFilePath(this.messageManager.getRootSessionId())
         .then(({ path }) => {
           this.logger?.debug("Plan file path generated", { path });
           this.permissionManager.setPlanFilePath(path);
@@ -1350,4 +1246,11 @@ export class Agent {
       this.permissionManager.setPlanFilePath(undefined);
     }
   }
+
+  /**
+   * Get the current task list ID
+   */
+  public get taskListId(): string {
+    return this.taskManager.getTaskListId();
+  }
 }

package/src/constants/subagents.ts

@@ -0,0 +1,4 @@
+export const BASH_SUBAGENT_TYPE = "Bash";
+export const EXPLORE_SUBAGENT_TYPE = "Explore";
+export const PLAN_SUBAGENT_TYPE = "Plan";
+export const GENERAL_PURPOSE_SUBAGENT_TYPE = "general-purpose";

package/src/constants/tools.ts

@@ -13,5 +13,8 @@ export const MULTI_EDIT_TOOL_NAME = "MultiEdit";
 export const READ_TOOL_NAME = "Read";
 export const SKILL_TOOL_NAME = "Skill";
 export const TASK_TOOL_NAME = "Task";
-export const TODO_WRITE_TOOL_NAME = "TodoWrite";
+export const TASK_CREATE_TOOL_NAME = "TaskCreate";
+export const TASK_GET_TOOL_NAME = "TaskGet";
+export const TASK_UPDATE_TOOL_NAME = "TaskUpdate";
+export const TASK_LIST_TOOL_NAME = "TaskList";
 export const WRITE_TOOL_NAME = "Write";

package/src/index.ts

@@ -30,6 +30,7 @@ export * from "./utils/promptHistory.js";
 export * from "./utils/stringUtils.js";
 export * from "./utils/customCommands.js";
 export * from "./utils/hookMatcher.js";
+export * from "./utils/tokenCalculation.js";
 
 // Export types
 export * from "./types/index.js";