wave-agent-sdk 0.4.0 → 0.6.1

package/src/agent.ts

@@ -1,3 +1,4 @@
+import { ForegroundTaskManager } from "./managers/foregroundTaskManager.js";
 import {
   MessageManager,
   type MessageManagerCallbacks,
@@ -8,14 +9,13 @@ 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";
 import {
-  BackgroundBashManager,
-  type BackgroundBashManagerCallbacks,
-} from "./managers/backgroundBashManager.js";
+  BackgroundTaskManager,
+  type BackgroundTaskManagerCallbacks,
+} from "./managers/backgroundTaskManager.js";
 import { SlashCommandManager } from "./managers/slashCommandManager.js";
 import { PluginManager } from "./managers/pluginManager.js";
 import { HookManager } from "./managers/hookManager.js";
@@ -38,11 +38,14 @@ import type {
   Usage,
   PermissionMode,
   PermissionCallback,
+  BackgroundTask,
+  ForegroundTask,
 } from "./types/index.js";
 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,
@@ -100,10 +103,17 @@ export interface AgentOptions {
 
 export interface AgentCallbacks
   extends MessageManagerCallbacks,
-    BackgroundBashManagerCallbacks,
+    BackgroundTaskManagerCallbacks,
     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;
 }
 
 export class Agent {
@@ -111,7 +121,7 @@ export class Agent {
   private aiManager: AIManager;
 
   private bashManager: BashManager | null = null;
-  private backgroundBashManager: BackgroundBashManager;
+  private backgroundTaskManager: BackgroundTaskManager;
   private logger?: Logger; // Add optional logger property
   private toolManager: ToolManager; // Add tool registry instance
   private mcpManager: McpManager; // Add MCP manager instance
@@ -126,6 +136,8 @@ 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
   private systemPrompt?: string; // Custom system prompt
@@ -200,8 +212,55 @@ export class Agent {
     // Store options for dynamic configuration resolution
     this.options = options;
 
-    this.backgroundBashManager = new BackgroundBashManager({
-      callbacks,
+    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);
+        },
+        onSubagentTaskStopRequested: (subagentId) => {
+          this.backgroundTaskManager.stopTask(subagentId);
+        },
+      },
+      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,
+        onTasksChange: (tasks) => {
+          callbacks.onTasksChange?.(tasks);
+        },
+      },
       workdir: this.workdir,
     });
     this.mcpManager = new McpManager({ callbacks, logger: this.logger }); // Initialize MCP manager
@@ -230,23 +289,11 @@ export class Agent {
       workdir: this.workdir,
     });
 
-    // Initialize MessageManager
-    this.messageManager = new MessageManager({
-      callbacks,
-      workdir: this.workdir,
-      logger: this.logger,
-    });
-
-    // Initialize ReversionManager
+    // ReversionManager depends on MessageManager
     this.reversionManager = new ReversionManager(
       new ReversionService(this.messageManager.getTranscriptPath()),
     );
 
-    // Initialize memory rule manager
-    this.memoryRuleManager = new MemoryRuleManager({
-      workdir: this.workdir,
-    });
-
     // Create a wrapper for canUseTool that triggers notification hooks
     const canUseToolWithNotification: PermissionCallback | undefined =
       options.canUseTool
@@ -298,6 +345,9 @@ 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
     this.liveConfigManager = new LiveConfigManager({
       workdir: this.workdir,
@@ -323,6 +373,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(),
@@ -331,14 +383,18 @@ export class Agent {
       getLanguage: () => this.getLanguage(),
       hookManager: this.hookManager,
       onUsageAdded: (usage) => this.addUsage(usage),
+      backgroundTaskManager: this.backgroundTaskManager,
+      taskManager: this.taskManager,
+      memoryRuleManager: this.memoryRuleManager,
     });
 
     // Initialize AI manager with resolved configuration
     this.aiManager = new AIManager({
       messageManager: this.messageManager,
       toolManager: this.toolManager,
+      taskManager: this.taskManager,
       logger: this.logger,
-      backgroundBashManager: this.backgroundBashManager,
+      backgroundTaskManager: this.backgroundTaskManager,
       hookManager: this.hookManager,
       permissionManager: this.permissionManager,
       callbacks: {
@@ -362,6 +418,8 @@ export class Agent {
     this.slashCommandManager = new SlashCommandManager({
       messageManager: this.messageManager,
       aiManager: this.aiManager,
+      backgroundTaskManager: this.backgroundTaskManager,
+      taskManager: this.taskManager,
       workdir: this.workdir,
       logger: this.logger,
     });
@@ -436,10 +494,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;
@@ -456,34 +510,8 @@ export class Agent {
   }
 
   /** Get combined memory content (project + user + modular rules) */
-  public get combinedMemory(): string {
-    let combined = "";
-    if (this._projectMemoryContent.trim()) {
-      combined += this._projectMemoryContent;
-    }
-    if (this._userMemoryContent.trim()) {
-      if (combined) {
-        combined += "\n\n";
-      }
-      combined += this._userMemoryContent;
-    }
-
-    // Add modular memory rules
-    const filesInContext = this.messageManager.getFilesInContext();
-    const activeRules = this.memoryRuleManager.getActiveRules(filesInContext);
-    if (activeRules.length > 0) {
-      this.logger?.debug(
-        `Active modular rules (${activeRules.length}): ${activeRules.map((r) => r.id).join(", ")}`,
-      );
-      if (combined) {
-        combined += "\n\n";
-      }
-      combined += activeRules.map((r) => r.content).join("\n\n");
-    } else {
-      this.logger?.debug("No active modular rules for current context");
-    }
-
-    return combined;
+  public async getCombinedMemory(): Promise<string> {
+    return this.messageManager.getCombinedMemory();
   }
 
   /** Get AI loading status */
@@ -506,12 +534,25 @@ export class Agent {
     id: string,
     filter?: string,
   ): { stdout: string; stderr: string; status: string } | null {
-    return this.backgroundBashManager.getOutput(id, filter);
+    return this.backgroundTaskManager.getOutput(id, filter);
   }
 
   /** Kill background bash shell */
   public killBackgroundShell(id: string): boolean {
-    return this.backgroundBashManager.killShell(id);
+    return this.backgroundTaskManager.stopTask(id);
+  }
+
+  /** Get background task output */
+  public getBackgroundTaskOutput(
+    id: string,
+    filter?: string,
+  ): { stdout: string; stderr: string; status: string } | null {
+    return this.backgroundTaskManager.getOutput(id, filter);
+  }
+
+  /** Stop background task */
+  public stopBackgroundTask(id: string): boolean {
+    return this.backgroundTaskManager.stopTask(id);
   }
 
   /**
@@ -774,8 +815,18 @@ export class Agent {
       // After main session is restored, restore any associated subagent sessions
       await this.restoreSubagentSessions(sessionToRestore?.messages || []);
 
-      if (sessionToRestore)
+      if (sessionToRestore) {
         this.messageManager.initializeFromSession(sessionToRestore);
+
+        // Update task manager with the root session ID to ensure continuity across compressions
+        this.taskManager.setTaskListId(
+          sessionToRestore.rootSessionId || sessionToRestore.id,
+        );
+
+        // After session is initialized, load tasks for the session
+        const tasks = await this.taskManager.listTasks();
+        this.options.callbacks?.onSessionTasksChange?.(tasks);
+      }
     }
   }
 
@@ -898,6 +949,13 @@ export class Agent {
 
     // 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 {
@@ -906,12 +964,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();
   }
@@ -923,11 +978,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();
@@ -938,14 +988,36 @@ export class Agent {
     this.slashCommandManager.abortCurrentCommand();
   }
 
+  /**
+   * Register a foreground task that can be backgrounded
+   */
+  public registerForegroundTask(task: ForegroundTask): void {
+    this.foregroundTaskManager.registerForegroundTask(task);
+  }
+
+  /**
+   * Unregister a foreground task
+   */
+  public unregisterForegroundTask(id: string): void {
+    this.foregroundTaskManager.unregisterForegroundTask(id);
+  }
+
+  /**
+   * Background the current foreground task
+   */
+  public async backgroundCurrentTask(): Promise<void> {
+    await this.foregroundTaskManager.backgroundCurrentTask();
+    this.options.callbacks?.onBackgroundCurrentTask?.();
+  }
+
   /** Destroy managers, clean up resources */
   public async destroy(): Promise<void> {
     await this.messageManager.saveSession();
     this.abortAIMessage(); // This will abort tools including Task tool (subagents)
     this.abortBashCommand();
     this.abortSlashCommand();
-    // Cleanup background bash manager
-    this.backgroundBashManager.cleanup();
+    // Cleanup background task manager
+    this.backgroundTaskManager.cleanup();
     // Cleanup MCP connections
     await this.mcpManager.cleanup();
     // Cleanup LSP connections
@@ -1027,8 +1099,6 @@ export class Agent {
           // Execute valid slash command
           await this.slashCommandManager.executeCommand(commandId, args);
 
-          // Add slash command to history
-          this.addToInputHistory(command);
           return;
         }
 
@@ -1037,9 +1107,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,
@@ -1096,67 +1163,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 */
@@ -1236,6 +1242,16 @@ export class Agent {
     await this.messageManager.truncateHistory(index, this.reversionManager);
   }
 
+  /**
+   * Get the full message thread including parent sessions
+   */
+  public async getFullMessageThread(): Promise<{
+    messages: Message[];
+    sessionIds: string[];
+  }> {
+    return this.messageManager.getFullMessageThread();
+  }
+
   /**
    * Get the current plan file path (for testing and UI)
    */
@@ -1304,7 +1320,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);
@@ -1316,4 +1332,11 @@ export class Agent {
       this.permissionManager.setPlanFilePath(undefined);
     }
   }
+
+  /**
+   * Get the current task list ID
+   */
+  public get taskListId(): string {
+    return this.taskManager.getTaskListId();
+  }
 }