runtimeuse 0.2.0

package/src/session.ts

@@ -0,0 +1,319 @@
+import fs from "fs";
+import path from "path";
+import { WebSocket } from "ws";
+
+import type { AgentHandler, MessageSender } from "./agent-handler.js";
+import { ArtifactManager } from "./artifact-manager.js";
+import type { UploadTracker } from "./upload-tracker.js";
+import type {
+  InvocationMessage,
+  IncomingMessage,
+  OutgoingMessage,
+  ResultMessage,
+} from "./types.js";
+import { sleep } from "./utils.js";
+import { createLogger, defaultLogger, type Logger } from "./logger.js";
+import { DEFAULT_ARTIFACTS_DIR } from "./constants.js";
+import CommandHandler from "./command-handler.js";
+import DownloadHandler from "./download-handler.js";
+
+export interface SessionConfig {
+  handler: AgentHandler;
+  uploadTracker: UploadTracker;
+  defaultModel?: string;
+  uploadTimeoutMs?: number;
+  artifactWaitMs?: number;
+  postInvocationDelayMs?: number;
+  logger?: Logger;
+}
+
+export class WebSocketSession {
+  private readonly ws: WebSocket;
+  private readonly config: SessionConfig;
+  private readonly abortController = new AbortController();
+  private artifactManager: ArtifactManager | null = null;
+  private invocationReceived = false;
+  private logger: Logger;
+  private readonly downloadHandler: DownloadHandler;
+
+  constructor(ws: WebSocket, config: SessionConfig) {
+    this.ws = ws;
+    this.config = config;
+    this.logger = config.logger ?? defaultLogger;
+    this.downloadHandler = new DownloadHandler(this.logger);
+  }
+
+  run(): Promise<void> {
+    return new Promise<void>((resolve) => {
+      this.ws.on("message", async (rawData) => {
+        this.logger.log("Received new WS message");
+        try {
+          const message: IncomingMessage = JSON.parse(rawData.toString());
+          await this.handleMessage(message, resolve);
+        } catch (error) {
+          this.logger.error("Error processing message:", error);
+          this.send({
+            message_type: "error_message",
+            error: String(error),
+            metadata: {},
+          });
+        }
+      });
+
+      this.ws.on("close", async () => {
+        this.logger.log("WebSocket connection closed");
+        this.abortController.abort();
+        await this.artifactManager?.stopWatching();
+        await this.config.uploadTracker.waitForAll(
+          this.config.uploadTimeoutMs ?? 30_000,
+        );
+        resolve();
+      });
+
+      this.ws.on("error", (error) => {
+        this.logger.error("WebSocket error:", error);
+      });
+    });
+  }
+
+  private async handleMessage(
+    message: IncomingMessage,
+    resolve: () => void,
+  ): Promise<void> {
+    if (
+      !this.invocationReceived &&
+      message.message_type !== "invocation_message"
+    ) {
+      throw new Error(
+        "Received non-invocation message before invocation message! Received: " +
+          JSON.stringify(message),
+      );
+    }
+
+    switch (message.message_type) {
+      case "artifact_upload_response_message":
+        try {
+          await this.artifactManager?.handleUploadResponse(message);
+        } catch (error) {
+          this.logger.error("Error uploading artifact:", error);
+          this.send({
+            message_type: "error_message",
+            error: String(error),
+            metadata: {},
+          });
+        }
+        break;
+
+      case "cancel_message":
+        this.logger.log("Received cancel message. Aborting agent execution...");
+        this.abortController.abort();
+        this.ws.close();
+        break;
+
+      case "invocation_message":
+        if (this.invocationReceived) {
+          throw new Error("Received multiple invocation messages!");
+        }
+        this.invocationReceived = true;
+        await this.executeInvocation(message);
+        if (process.env.NODE_ENV !== "test") {
+          await sleep(this.config.postInvocationDelayMs ?? 3_000);
+        }
+        await this.finalize();
+        resolve();
+        break;
+    }
+  }
+
+  private async executeInvocation(message: InvocationMessage): Promise<void> {
+    this.logger = createLogger(message.source_id);
+    this.config.uploadTracker.setLogger(this.logger);
+
+    const artifactsDir = message.artifacts_dir ?? DEFAULT_ARTIFACTS_DIR;
+    this.artifactManager = new ArtifactManager({
+      artifactsDir,
+      uploadTracker: this.config.uploadTracker,
+      send: (msg) => this.send(msg),
+    });
+    this.artifactManager.setLogger(this.logger);
+
+    const outputFormat = JSON.parse(message.output_format_json_schema_str) as {
+      type: "json_schema";
+      schema: Record<string, unknown>;
+    };
+    const model =
+      message.preferred_model ?? this.config.defaultModel ?? "default";
+
+    this.logger.log(`Received invocation: model=${model}`);
+
+    try {
+      if (message.runtime_environment_downloadables) {
+        this.logger.log("Downloading runtime environment downloadables...");
+        for (const downloadable of message.runtime_environment_downloadables) {
+          await this.downloadHandler.download(
+            downloadable.download_url,
+            downloadable.working_dir,
+          );
+        }
+      }
+      if (message.pre_agent_invocation_commands) {
+        const stdoutFile = path.join(artifactsDir, `cmd-stdout.out`);
+        const stderrFile = path.join(artifactsDir, `cmd-stderr.out`);
+        const stdoutStream = fs.createWriteStream(stdoutFile);
+        const stderrStream = fs.createWriteStream(stderrFile);
+
+        for (const command of message.pre_agent_invocation_commands) {
+          try {
+            this.logger.log(
+              `Executing command: ${command.command} in directory: ${command.cwd}`,
+            );
+            const commandHandler = new CommandHandler(
+              command,
+              this.logger,
+              stdoutStream,
+              stderrStream,
+              this.abortController,
+              (stdout) =>
+                this.send({
+                  message_type: "assistant_message",
+                  text_blocks: [stdout],
+                }),
+              (stderr) =>
+                this.send({
+                  message_type: "assistant_message",
+                  text_blocks: [stderr],
+                }),
+            );
+            try {
+              const result = await commandHandler.execute();
+              if (result.exitCode === 0) {
+                this.send({
+                  message_type: "result_message",
+                  structured_output: {
+                    success: true,
+                    steps: [],
+                    summary: "Workflow completed successfully",
+                  },
+                  metadata: {},
+                });
+                return;
+              } else {
+                this.logger.error(
+                  "Command failed with exit code:",
+                  result.exitCode,
+                );
+                this.send({
+                  message_type: "result_message",
+                  structured_output: {
+                    success: false,
+                    steps: [],
+                    summary: "Workflow completed with failure",
+                  },
+                  metadata: {},
+                });
+                return;
+              }
+            } catch (error) {
+              this.logger.error("Error executing command:", error);
+              this.send({
+                message_type: "error_message",
+                error: String(error),
+                metadata: {},
+              });
+              this.send({
+                message_type: "result_message",
+                structured_output: {
+                  success: false,
+                  steps: [],
+                  summary: "Workflow completed with failure",
+                },
+                metadata: {},
+              });
+              throw error;
+            }
+          } catch (error) {
+            this.logger.error("Error executing command:", error);
+            throw error;
+          }
+        }
+        stderrStream.close();
+        stdoutStream.close();
+      }
+
+      const sender: MessageSender = {
+        sendAssistantMessage: (textBlocks: string[]) => {
+          this.send({
+            message_type: "assistant_message",
+            text_blocks: textBlocks,
+          });
+        },
+        sendErrorMessage: (
+          error: string,
+          metadata?: Record<string, unknown>,
+        ) => {
+          this.send({
+            message_type: "error_message",
+            error,
+            metadata: metadata ?? {},
+          });
+        },
+      };
+
+      const agentResult = await this.config.handler.run(
+        {
+          systemPrompt: message.system_prompt,
+          userPrompt: message.user_prompt,
+          outputFormat,
+          model,
+          secrets: message.secrets_to_redact,
+          env: message.agent_env ?? {},
+          signal: this.abortController.signal,
+          logger: this.logger,
+        },
+        sender,
+      );
+
+      const resultMessage: ResultMessage = {
+        message_type: "result_message",
+        metadata: agentResult.metadata ?? {},
+        structured_output: agentResult.structuredOutput,
+      };
+      this.logger.log("Sending result message:", JSON.stringify(resultMessage));
+      this.send(resultMessage);
+    } catch (error) {
+      if (this.abortController.signal.aborted) {
+        this.ws.close();
+        this.logger.log("Agent execution aborted.");
+        return;
+      }
+      this.logger.error("Error in agent execution:", error);
+      this.send({
+        message_type: "error_message",
+        error: String(error),
+        metadata: {},
+      });
+    }
+  }
+
+  private async finalize(): Promise<void> {
+    await this.artifactManager?.stopWatching();
+
+    if (!this.abortController.signal.aborted) {
+      await this.artifactManager?.waitForPendingRequests(
+        this.config.artifactWaitMs ?? 60_000,
+      );
+    }
+
+    await this.config.uploadTracker.waitForAll(
+      this.config.uploadTimeoutMs ?? 30_000,
+    );
+    this.logger.log("All artifacts uploaded.");
+    this.ws.close();
+  }
+
+  private send(data: OutgoingMessage): void {
+    if (this.ws.readyState === WebSocket.OPEN) {
+      this.ws.send(JSON.stringify(data));
+    }
+  }
+}

package/src/storage.ts

@@ -0,0 +1,28 @@
+import fs from "fs";
+import { defaultLogger, type Logger } from "./logger.js";
+
+export async function uploadFile(
+  filePath: string,
+  presignedUrl: string,
+  contentType: string,
+  logger: Logger = defaultLogger,
+): Promise<boolean> {
+  try {
+    const content = fs.readFileSync(filePath);
+
+    const response = await fetch(presignedUrl, {
+      method: "PUT",
+      headers: {
+        "Content-Type": contentType,
+      },
+      body: content,
+    });
+
+    logger.log(`File ${filePath} uploaded with status: ${response.status}`);
+
+    return response.status === 200;
+  } catch (error) {
+    logger.error("Upload error:", error);
+    return false;
+  }
+}

package/src/types.ts

@@ -0,0 +1,101 @@
+interface Command {
+  command: string;
+  cwd?: string;
+  env?: Record<string, string>;
+}
+
+interface RuntimeEnvironmentDownloadable {
+  download_url: string;
+  working_dir: string;
+}
+
+interface InvocationMessage {
+  message_type: "invocation_message";
+  source_id: string;
+  system_prompt: string;
+  user_prompt: string;
+  secrets_to_redact: string[];
+  agent_env?: Record<string, string>;
+  output_format_json_schema_str: string;
+  preferred_model?: string;
+  artifacts_dir?: string;
+  pre_agent_invocation_commands?: Command[];
+  post_agent_invocation_commands?: Command[];
+  runtime_environment_downloadables?: RuntimeEnvironmentDownloadable[];
+}
+
+interface CancelMessage {
+  message_type: "cancel_message";
+}
+
+interface ResultMessage {
+  message_type: "result_message";
+  metadata?: Record<string, unknown>;
+  structured_output: Record<string, unknown>;
+  [key: string]: unknown;
+}
+
+interface AssistantMessage {
+  message_type: "assistant_message";
+  text_blocks: string[];
+}
+
+interface ArtifactUploadRequestMessage {
+  message_type: "artifact_upload_request_message";
+  artifact_type:
+    | "screenshot"
+    | "video"
+    | "tool_calls"
+    | "javascript"
+    | "python"
+    | "shellscript"
+    | "other";
+  filename: string;
+  filepath: string;
+}
+
+interface ArtifactUploadResponseMessage {
+  message_type: "artifact_upload_response_message";
+  filename: string;
+  filepath: string;
+  presigned_url: string;
+  content_type: string;
+}
+
+interface ErrorMessage {
+  message_type: "error_message";
+  error: string;
+  metadata: Record<string, unknown>;
+}
+
+interface StreamEndMessage {
+  message_type: "stream_end";
+  content: string;
+}
+
+type OutgoingMessage =
+  | ResultMessage
+  | AssistantMessage
+  | ArtifactUploadRequestMessage
+  | ErrorMessage
+  | StreamEndMessage;
+
+type IncomingMessage =
+  | InvocationMessage
+  | ArtifactUploadResponseMessage
+  | CancelMessage;
+
+export type {
+  IncomingMessage,
+  OutgoingMessage,
+  InvocationMessage,
+  CancelMessage,
+  ResultMessage,
+  AssistantMessage,
+  ArtifactUploadRequestMessage,
+  ArtifactUploadResponseMessage,
+  ErrorMessage,
+  StreamEndMessage,
+  Command,
+  RuntimeEnvironmentDownloadable,
+};

package/src/upload-tracker.test.ts

@@ -0,0 +1,112 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { UploadTracker } from "./upload-tracker.js";
+
+describe("UploadTracker", () => {
+  let tracker: UploadTracker;
+
+  beforeEach(() => {
+    tracker = new UploadTracker();
+  });
+
+  it("starts with size 0", () => {
+    expect(tracker.size).toBe(0);
+  });
+
+  it("tracks a pending upload", () => {
+    tracker.track(new Promise(() => {}));
+    expect(tracker.size).toBe(1);
+  });
+
+  it("tracks multiple pending uploads", () => {
+    tracker.track(new Promise(() => {}));
+    tracker.track(new Promise(() => {}));
+    tracker.track(new Promise(() => {}));
+    expect(tracker.size).toBe(3);
+  });
+
+  it("removes upload after it resolves", async () => {
+    let resolve!: () => void;
+    const promise = new Promise<void>((r) => {
+      resolve = r;
+    });
+    tracker.track(promise);
+    expect(tracker.size).toBe(1);
+
+    resolve();
+    await promise;
+    await flushMicrotasks();
+    expect(tracker.size).toBe(0);
+  });
+
+  it("removes upload after it rejects", async () => {
+    let reject!: (e: Error) => void;
+    const promise = new Promise<void>((_, r) => {
+      reject = r;
+    });
+    vi.spyOn(console, "error").mockImplementation(() => {});
+
+    tracker.track(promise);
+    expect(tracker.size).toBe(1);
+
+    reject(new Error("fail"));
+    await flushMicrotasks();
+    expect(tracker.size).toBe(0);
+  });
+
+  it("logs error when a tracked upload rejects", async () => {
+    const consoleSpy = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    let reject!: (e: Error) => void;
+    const promise = new Promise<void>((_, r) => {
+      reject = r;
+    });
+    tracker.track(promise);
+
+    const error = new Error("upload failed");
+    reject(error);
+    await flushMicrotasks();
+
+    expect(consoleSpy).toHaveBeenCalledWith(
+      "Error uploading artifact:",
+      error,
+    );
+  });
+
+  describe("waitForAll", () => {
+    it("resolves immediately when no pending uploads", async () => {
+      await tracker.waitForAll(1000);
+    });
+
+    it("resolves when all uploads complete", async () => {
+      let resolve1!: () => void;
+      let resolve2!: () => void;
+      tracker.track(new Promise<void>((r) => { resolve1 = r; }));
+      tracker.track(new Promise<void>((r) => { resolve2 = r; }));
+
+      let waited = false;
+      const waitPromise = tracker.waitForAll(5000).then(() => {
+        waited = true;
+      });
+
+      resolve1();
+      resolve2();
+      await waitPromise;
+      expect(waited).toBe(true);
+    });
+
+    it("times out if uploads do not complete", async () => {
+      tracker.track(new Promise(() => {}));
+
+      const start = Date.now();
+      await tracker.waitForAll(50);
+      expect(Date.now() - start).toBeGreaterThanOrEqual(40);
+      expect(tracker.size).toBe(1);
+    });
+  });
+});
+
+async function flushMicrotasks() {
+  await new Promise((r) => setTimeout(r, 0));
+}

package/src/upload-tracker.ts

@@ -0,0 +1,30 @@
+import { defaultLogger, type Logger } from "./logger.js";
+
+export class UploadTracker {
+  private readonly pending = new Set<Promise<any>>();
+  private logger: Logger = defaultLogger;
+
+  setLogger(logger: Logger): void {
+    this.logger = logger;
+  }
+
+  track(promise: Promise<any>): void {
+    this.pending.add(promise);
+    promise
+      .catch((error) => this.logger.error("Error uploading artifact:", error))
+      .finally(() => this.pending.delete(promise));
+  }
+
+  async waitForAll(timeoutMs: number): Promise<void> {
+    if (this.pending.size === 0) return;
+    this.logger.log(`Waiting for ${this.pending.size} uploads...`);
+    await Promise.race([
+      Promise.allSettled(this.pending),
+      new Promise<void>((r) => setTimeout(r, timeoutMs)),
+    ]);
+  }
+
+  get size(): number {
+    return this.pending.size;
+  }
+}

package/src/utils.test.ts

@@ -0,0 +1,120 @@
+import { describe, it, expect } from "vitest";
+import { redactSecrets } from "./utils.js";
+
+describe("redactSecrets", () => {
+  describe("strings", () => {
+    it("replaces a secret within a string", () => {
+      expect(redactSecrets("token is abc123", ["abc123"])).toBe(
+        "token is [REDACTED]",
+      );
+    });
+
+    it("replaces multiple occurrences of the same secret", () => {
+      expect(redactSecrets("abc123 and abc123", ["abc123"])).toBe(
+        "[REDACTED] and [REDACTED]",
+      );
+    });
+
+    it("replaces multiple different secrets", () => {
+      expect(redactSecrets("key=SECRET1 pass=SECRET2", ["SECRET1", "SECRET2"])).toBe(
+        "key=[REDACTED] pass=[REDACTED]",
+      );
+    });
+
+    it("returns string unchanged when no secrets match", () => {
+      expect(redactSecrets("nothing here", ["xyz"])).toBe("nothing here");
+    });
+
+    it("skips empty-string secrets", () => {
+      expect(redactSecrets("hello", [""])).toBe("hello");
+    });
+  });
+
+  describe("arrays", () => {
+    it("redacts secrets inside array elements", () => {
+      expect(redactSecrets(["key=SECRET", "ok"], ["SECRET"])).toEqual([
+        "key=[REDACTED]",
+        "ok",
+      ]);
+    });
+
+    it("handles nested arrays", () => {
+      expect(redactSecrets([["SECRET"]], ["SECRET"])).toEqual([
+        ["[REDACTED]"],
+      ]);
+    });
+  });
+
+  describe("objects", () => {
+    it("redacts secrets in object values", () => {
+      expect(
+        redactSecrets({ token: "my-SECRET-value", count: 5 }, ["SECRET"]),
+      ).toEqual({ token: "my-[REDACTED]-value", count: 5 });
+    });
+
+    it("redacts secrets in deeply nested objects", () => {
+      const input = {
+        level1: {
+          level2: {
+            value: "contains SECRET here",
+          },
+        },
+      };
+      expect(redactSecrets(input, ["SECRET"])).toEqual({
+        level1: {
+          level2: {
+            value: "contains [REDACTED] here",
+          },
+        },
+      });
+    });
+
+    it("handles mixed objects with arrays", () => {
+      const input = {
+        args: ["--token=SECRET", "--verbose"],
+        env: { API_KEY: "SECRET" },
+      };
+      expect(redactSecrets(input, ["SECRET"])).toEqual({
+        args: ["--token=[REDACTED]", "--verbose"],
+        env: { API_KEY: "[REDACTED]" },
+      });
+    });
+  });
+
+  describe("non-string primitives", () => {
+    it("returns numbers unchanged", () => {
+      expect(redactSecrets(42, ["42"])).toBe(42);
+    });
+
+    it("returns booleans unchanged", () => {
+      expect(redactSecrets(true, ["true"])).toBe(true);
+    });
+
+    it("returns null unchanged", () => {
+      expect(redactSecrets(null, ["null"])).toBeNull();
+    });
+
+    it("returns undefined unchanged", () => {
+      expect(redactSecrets(undefined, ["undefined"])).toBeUndefined();
+    });
+  });
+
+  describe("edge cases", () => {
+    it("returns value unchanged when secrets list is empty", () => {
+      const input = { key: "value" };
+      expect(redactSecrets(input, [])).toBe(input);
+    });
+
+    it("handles overlapping secrets (longer secret first)", () => {
+      expect(
+        redactSecrets("my-secret-key", ["my-secret-key", "secret"]),
+      ).toBe("[REDACTED]");
+    });
+
+    it("handles overlapping secrets (shorter secret first)", () => {
+      expect(
+        redactSecrets("my-secret-key", ["secret", "my-secret-key"]),
+      ).toBe("my-[REDACTED]-key");
+    });
+  });
+});