wave-agent-sdk 0.4.0 → 0.6.1

package/src/managers/backgroundTaskManager.ts

@@ -0,0 +1,305 @@
+import { spawn, type ChildProcess } from "child_process";
+import { BackgroundTask, BackgroundShell } from "../types/processes.js";
+import { stripAnsiColors } from "../utils/stringUtils.js";
+import { logger } from "../utils/globalLogger.js";
+
+export interface BackgroundTaskManagerCallbacks {
+  onTasksChange?: (tasks: BackgroundTask[]) => void;
+}
+
+export interface BackgroundTaskManagerOptions {
+  callbacks?: BackgroundTaskManagerCallbacks;
+  workdir: string;
+}
+
+export class BackgroundTaskManager {
+  private tasks = new Map<string, BackgroundTask>();
+  private nextId = 1;
+  private callbacks: BackgroundTaskManagerCallbacks;
+  private workdir: string;
+
+  constructor(options: BackgroundTaskManagerOptions) {
+    this.callbacks = options.callbacks || {};
+    this.workdir = options.workdir;
+  }
+
+  private notifyTasksChange(): void {
+    this.callbacks.onTasksChange?.(Array.from(this.tasks.values()));
+  }
+
+  public generateId(): string {
+    return `task_${this.nextId++}`;
+  }
+
+  public addTask(task: BackgroundTask): void {
+    this.tasks.set(task.id, task);
+    this.notifyTasksChange();
+  }
+
+  public getTask(id: string): BackgroundTask | undefined {
+    return this.tasks.get(id);
+  }
+
+  public getAllTasks(): BackgroundTask[] {
+    return Array.from(this.tasks.values());
+  }
+
+  public startShell(
+    command: string,
+    timeout?: number,
+  ): { id: string; child: ChildProcess; detach: () => void } {
+    const id = this.generateId();
+    const startTime = Date.now();
+
+    const child = spawn(command, {
+      shell: true,
+      stdio: "pipe",
+      cwd: this.workdir,
+      env: {
+        ...process.env,
+      },
+    });
+
+    const shell: BackgroundShell = {
+      id,
+      type: "shell",
+      process: child,
+      command,
+      startTime,
+      status: "running",
+      stdout: "",
+      stderr: "",
+      onStop: () => {
+        try {
+          if (child.pid) {
+            process.kill(-child.pid, "SIGTERM");
+            setTimeout(() => {
+              if (child.pid && !child.killed) {
+                try {
+                  process.kill(-child.pid, "SIGKILL");
+                } catch (error) {
+                  logger.error("Failed to force kill process:", error);
+                }
+              }
+            }, 1000);
+          } else {
+            child.kill("SIGTERM");
+          }
+        } catch {
+          child.kill("SIGTERM");
+        }
+      },
+    };
+
+    this.tasks.set(id, shell);
+    this.notifyTasksChange();
+
+    // Set up timeout if specified
+    let timeoutHandle: NodeJS.Timeout | undefined;
+    if (timeout && timeout > 0) {
+      timeoutHandle = setTimeout(() => {
+        if (shell.status === "running") {
+          this.stopTask(id);
+        }
+      }, timeout);
+    }
+
+    const onStdout = (data: Buffer | string) => {
+      shell.stdout += stripAnsiColors(data.toString());
+      this.notifyTasksChange();
+    };
+
+    const onStderr = (data: Buffer | string) => {
+      shell.stderr += stripAnsiColors(data.toString());
+      this.notifyTasksChange();
+    };
+
+    const onExit = (code: number | null) => {
+      if (timeoutHandle) {
+        clearTimeout(timeoutHandle);
+      }
+      shell.status = code === 0 ? "completed" : "failed";
+      shell.exitCode = code ?? 0;
+      shell.endTime = Date.now();
+      shell.runtime = shell.endTime - startTime;
+      this.notifyTasksChange();
+    };
+
+    const onError = (error: Error) => {
+      if (timeoutHandle) {
+        clearTimeout(timeoutHandle);
+      }
+      shell.status = "failed";
+      shell.stderr += `\nProcess error: ${stripAnsiColors(error.message)}`;
+      shell.exitCode = 1;
+      shell.endTime = Date.now();
+      shell.runtime = shell.endTime - startTime;
+      this.notifyTasksChange();
+    };
+
+    child.stdout?.on("data", onStdout);
+    child.stderr?.on("data", onStderr);
+    child.on("exit", onExit);
+    child.on("error", onError);
+
+    const detach = () => {
+      child.stdout?.off("data", onStdout);
+      child.stderr?.off("data", onStderr);
+      child.off("exit", onExit);
+      child.off("error", onError);
+      if (timeoutHandle) {
+        clearTimeout(timeoutHandle);
+      }
+      this.tasks.delete(id);
+      this.notifyTasksChange();
+    };
+
+    return { id, child, detach };
+  }
+
+  public adoptProcess(
+    child: ChildProcess,
+    command: string,
+    initialStdout: string = "",
+    initialStderr: string = "",
+  ): string {
+    const id = this.generateId();
+    const startTime = Date.now();
+
+    const shell: BackgroundShell = {
+      id,
+      type: "shell",
+      process: child,
+      command,
+      startTime,
+      status: "running",
+      stdout: initialStdout,
+      stderr: initialStderr,
+      onStop: () => {
+        try {
+          if (child.pid) {
+            process.kill(-child.pid, "SIGTERM");
+            setTimeout(() => {
+              if (child.pid && !child.killed) {
+                try {
+                  process.kill(-child.pid, "SIGKILL");
+                } catch (error) {
+                  logger.error("Failed to force kill process:", error);
+                }
+              }
+            }, 1000);
+          } else {
+            child.kill("SIGTERM");
+          }
+        } catch {
+          child.kill("SIGTERM");
+        }
+      },
+    };
+
+    this.tasks.set(id, shell);
+    this.notifyTasksChange();
+
+    child.stdout?.on("data", (data) => {
+      shell.stdout += stripAnsiColors(data.toString());
+      this.notifyTasksChange();
+    });
+
+    child.stderr?.on("data", (data) => {
+      shell.stderr += stripAnsiColors(data.toString());
+      this.notifyTasksChange();
+    });
+
+    child.on("exit", (code) => {
+      shell.status = code === 0 ? "completed" : "failed";
+      shell.exitCode = code ?? 0;
+      shell.endTime = Date.now();
+      shell.runtime = shell.endTime - startTime;
+      this.notifyTasksChange();
+    });
+
+    child.on("error", (error) => {
+      shell.status = "failed";
+      shell.stderr += `\nProcess error: ${stripAnsiColors(error.message)}`;
+      shell.exitCode = 1;
+      shell.endTime = Date.now();
+      shell.runtime = shell.endTime - startTime;
+      this.notifyTasksChange();
+    });
+
+    return id;
+  }
+
+  public getOutput(
+    id: string,
+    filter?: string,
+  ): { stdout: string; stderr: string; status: string } | null {
+    const task = this.tasks.get(id);
+    if (!task) {
+      return null;
+    }
+
+    let stdout = task.stdout;
+    let stderr = task.stderr;
+
+    // Apply regex filter if provided
+    if (filter) {
+      try {
+        const regex = new RegExp(filter);
+        stdout = stdout
+          .split("\n")
+          .filter((line) => regex.test(line))
+          .join("\n");
+        stderr = stderr
+          .split("\n")
+          .filter((line) => regex.test(line))
+          .join("\n");
+      } catch (error) {
+        logger.warn(`Invalid filter regex: ${filter}`, error);
+      }
+    }
+
+    return {
+      stdout,
+      stderr,
+      status: task.status,
+    };
+  }
+
+  public stopTask(id: string): boolean {
+    const task = this.tasks.get(id);
+    if (!task || task.status !== "running") {
+      return false;
+    }
+
+    if (task.onStop) {
+      try {
+        const result = task.onStop();
+        if (result instanceof Promise) {
+          result.catch((error) => {
+            logger.error("Error in background task onStop callback:", error);
+          });
+        }
+      } catch (error) {
+        logger.error("Error in background task onStop callback:", error);
+      }
+    }
+
+    task.status = "killed";
+    task.endTime = Date.now();
+    task.runtime = task.endTime - task.startTime;
+    this.notifyTasksChange();
+    return true;
+  }
+
+  public cleanup(): void {
+    // Kill all running tasks
+    for (const [id, task] of this.tasks) {
+      if (task.status === "running") {
+        this.stopTask(id);
+      }
+    }
+    this.tasks.clear();
+    this.notifyTasksChange();
+  }
+}

package/src/managers/foregroundTaskManager.ts

@@ -0,0 +1,27 @@
+import { ForegroundTask, IForegroundTaskManager } from "../types/processes.js";
+
+export class ForegroundTaskManager implements IForegroundTaskManager {
+  private activeForegroundTasks: ForegroundTask[] = [];
+
+  public registerForegroundTask(task: ForegroundTask): void {
+    this.activeForegroundTasks.push(task);
+  }
+
+  public unregisterForegroundTask(id: string): void {
+    this.activeForegroundTasks = this.activeForegroundTasks.filter(
+      (t) => t.id !== id,
+    );
+  }
+
+  public async backgroundCurrentTask(): Promise<void> {
+    const tasks = [...this.activeForegroundTasks].reverse();
+    this.activeForegroundTasks = [];
+    for (const task of tasks) {
+      await task.backgroundHandler();
+    }
+  }
+
+  public hasActiveTasks(): boolean {
+    return this.activeForegroundTasks.length > 0;
+  }
+}

package/src/managers/lspManager.ts

@@ -418,7 +418,9 @@ export class LspManager implements ILspManager {
         await this.sendRequest(lspProc, "shutdown", {}, timeout);
         await this.sendNotification(lspProc, "exit", {});
         // Give it a moment to exit
-        await new Promise((resolve) => setTimeout(resolve, 100));
+        if (timeout > 100) {
+          await new Promise((resolve) => setTimeout(resolve, 100));
+        }
       } catch (error) {
         this.logger?.debug(
           `Failed to gracefully shutdown LSP for ${language}: ${error}`,

package/src/managers/mcpManager.ts

@@ -3,7 +3,7 @@ import { join } from "path";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { ChatCompletionFunctionTool } from "openai/resources.js";
-import { createMcpToolPlugin, findToolServer } from "@/utils/mcpUtils.js";
+import { createMcpToolPlugin, findToolServer } from "../utils/mcpUtils.js";
 import type { ToolPlugin, ToolResult, ToolContext } from "../tools/types.js";
 import type {
   Logger,
@@ -476,8 +476,11 @@ export class McpManager {
       const server = findToolServer(tool.name, servers);
 
       if (server) {
-        const plugin = createMcpToolPlugin(tool, server.name, (name, args) =>
-          this.executeMcpTool(name, args),
+        const plugin = createMcpToolPlugin(
+          tool,
+          server.name,
+          (name: string, args: Record<string, unknown>) =>
+            this.executeMcpTool(name, args),
         );
         mcpTools.set(plugin.name, plugin);
       }