zopu-code/packages/agents/src/sandboxes/agent-os.ts

import { api } from "@code/backend/convex/_generated/api";
import type { Id } from "@code/backend/convex/_generated/dataModel";
import { parseAgentEnv } from "@code/env/agent";
import {
  decodeIssueWorkspaceResult,
  makeIssueWorkspacePlan,
} from "@code/primitives/project-workspace";
import { createSandboxSessionEnv } from "@flue/runtime";
import type { FileStat, SandboxApi, SandboxFactory } from "@flue/runtime";
import { ConvexHttpClient } from "convex/browser";
import { Effect } from "effect";
import * as v from "valibot";

const execResultSchema = v.object({
  exitCode: v.number(),
  stderr: v.string(),
  stdout: v.string(),
});

const statResultSchema = v.object({
  isDirectory: v.boolean(),
  isSymbolicLink: v.boolean(),
  mtimeMs: v.number(),
  size: v.number(),
});

interface ExecuteOptions {
  readonly signal?: AbortSignal;
  readonly timeoutMs?: number;
}

class AgentOsSandboxApi implements SandboxApi {
  readonly #actorKey: string[];
  readonly #client: ConvexHttpClient;
  readonly #daemonId: string;

  constructor(
    client: ConvexHttpClient,
    daemonId: string,
    actorKey: readonly string[]
  ) {
    this.#client = client;
    this.#daemonId = daemonId;
    this.#actorKey = [...actorKey];
  }

  async #execute(
    method: string,
    args: unknown[],
    options: ExecuteOptions = {}
  ): Promise<unknown> {
    const commandId = await this.#client.mutation(api.daemonCommands.enqueue, {
      actorKey: this.#actorKey,
      args,
      daemonId: this.#daemonId,
      method,
    });
    const timeoutMs = options.timeoutMs ?? 120_000;
    const deadline = Date.now() + timeoutMs;

    while (Date.now() < deadline) {
      if (options.signal?.aborted) {
        // oxlint-disable-next-line no-await-in-loop -- cancellation must settle before the poll exits.
        await this.#client.mutation(api.daemonCommands.cancel, { commandId });
        throw options.signal.reason ?? new Error("AgentOS command aborted");
      }
      // oxlint-disable-next-line no-await-in-loop -- each query observes the next durable command state.
      const command = await this.#client.query(api.daemonCommands.get, {
        commandId,
      });
      if (!command) {
        throw new Error(`AgentOS command ${commandId} disappeared`);
      }
      if (command.status === "succeeded") {
        return command.result;
      }
      if (command.status === "failed" || command.status === "cancelled") {
        throw new Error(
          command.error ?? `AgentOS command ${method} ${command.status}`
        );
      }
      const { promise, resolve } = Promise.withResolvers<null>();
      setTimeout(() => resolve(null), 200);
      // oxlint-disable-next-line no-await-in-loop -- polling intentionally waits between sequential reads.
      await promise;
    }

    await this.#client.mutation(api.daemonCommands.cancel, { commandId });
    throw new Error(`AgentOS command ${method} timed out after ${timeoutMs}ms`);
  }

  async readFile(path: string): Promise<string> {
    return new TextDecoder().decode(await this.readFileBuffer(path));
  }

  async readFileBuffer(path: string): Promise<Uint8Array> {
    const result = await this.#execute("readFile", [path]);
    if (result instanceof ArrayBuffer) {
      return new Uint8Array(result);
    }
    if (ArrayBuffer.isView(result)) {
      return new Uint8Array(
        result.buffer.slice(
          result.byteOffset,
          result.byteOffset + result.byteLength
        )
      );
    }
    throw new TypeError(
      `AgentOS readFile returned non-binary data for ${path}`
    );
  }

  async writeFile(path: string, content: string | Uint8Array): Promise<void> {
    const serialized =
      typeof content === "string" ? content : Uint8Array.from(content).buffer;
    await this.#execute("writeFile", [path, serialized]);
  }

  async stat(path: string): Promise<FileStat> {
    const result = v.parse(
      statResultSchema,
      await this.#execute("stat", [path])
    );
    return {
      isDirectory: result.isDirectory,
      isFile: !result.isDirectory && !result.isSymbolicLink,
      isSymbolicLink: result.isSymbolicLink,
      mtime: new Date(result.mtimeMs),
      size: result.size,
    };
  }

  async readdir(path: string): Promise<string[]> {
    return v.parse(v.array(v.string()), await this.#execute("readdir", [path]));
  }

  async exists(path: string): Promise<boolean> {
    return v.parse(v.boolean(), await this.#execute("exists", [path]));
  }

  async mkdir(
    path: string,
    options?: { readonly recursive?: boolean }
  ): Promise<void> {
    if (!options?.recursive) {
      await this.#execute("mkdir", [path]);
      return;
    }
    const quotedPath = `'${path.replaceAll("'", `'"'"'`)}'`;
    const result = v.parse(
      execResultSchema,
      await this.#execute("exec", [`mkdir -p ${quotedPath}`])
    );
    if (result.exitCode !== 0) {
      throw new Error(result.stderr || `Could not create ${path}`);
    }
  }

  async rm(
    path: string,
    options?: { readonly force?: boolean; readonly recursive?: boolean }
  ): Promise<void> {
    if (options?.force && !(await this.exists(path))) {
      return;
    }
    await this.#execute("deleteFile", [
      path,
      { recursive: options?.recursive },
    ]);
  }

  async exec(
    command: string,
    options?: {
      readonly cwd?: string;
      readonly env?: Record<string, string>;
      readonly signal?: AbortSignal;
      readonly timeoutMs?: number;
    }
  ): Promise<{ exitCode: number; stderr: string; stdout: string }> {
    return v.parse(
      execResultSchema,
      await this.#execute(
        "exec",
        [
          command,
          {
            ...(options?.cwd === undefined ? {} : { cwd: options.cwd }),
            ...(options?.env === undefined ? {} : { env: options.env }),
          },
        ],
        { signal: options?.signal, timeoutMs: options?.timeoutMs }
      )
    );
  }
}

export const agentOs = (
  runtimeEnv: Record<string, string | undefined>
): SandboxFactory => {
  const env = parseAgentEnv(runtimeEnv);
  const client = new ConvexHttpClient(env.CONVEX_URL);

  return {
    async createSessionEnv({ id }) {
      // Flue controls this id from the issue document selected by the web app.
      const issueId = id as Id<"projectIssues">;
      const run = await client.mutation(api.agentWorkspace.ensureRun, {
        daemonId: env.DAEMON_ID,
        issueId,
        token: env.FLUE_DB_TOKEN,
      });
      if (!run) {
        throw new Error(`Could not initialize AgentOS work run for ${id}`);
      }
      const sandbox = new AgentOsSandboxApi(
        client,
        env.DAEMON_ID,
        run.actorKey
      );
      const context = await client.query(api.agentWorkspace.get, {
        issueId,
        token: env.FLUE_DB_TOKEN,
      });

      const plan = await Effect.runPromise(
        makeIssueWorkspacePlan({
          artifacts: context.artifacts.map((artifact) => ({
            content: artifact.content,
            path: artifact.path,
          })),
          branchName: context.run?.branchName,
          checkoutPath: context.run?.checkoutPath,
          contextFiles: context.contextDocuments.map((document) => ({
            content: document.content,
            kind: document.kind,
            path: document.path,
          })),
          defaultBranch:
            context.run?.baseBranch ?? context.source?.defaultBranch,
          issueBody: context.issue.body,
          issueId: String(context.issue._id),
          issueNumber: context.issue.number,
          issueTitle: context.issue.title,
          sourceUrl: context.run?.sourceUrl ?? context.source?.url,
        })
      );

      const [controlDirectory, checkoutOperation, ...stagingOperations] =
        plan.operations;
      if (!controlDirectory || controlDirectory._tag !== "Mkdir") {
        throw new Error(
          "Issue workspace plan must start with a control directory"
        );
      }
      await sandbox.mkdir(controlDirectory.path, { recursive: true });
      if (!checkoutOperation || checkoutOperation._tag !== "Exec") {
        throw new Error("Issue workspace plan must include a checkout command");
      }
      const checkoutCommandResult = await sandbox.exec(
        checkoutOperation.command,
        {
          cwd: checkoutOperation.cwd,
          timeoutMs: checkoutOperation.timeoutMs,
        }
      );

      await Promise.all(
        stagingOperations
          .filter((operation) => operation._tag === "Mkdir")
          .map((operation) =>
            sandbox.mkdir(operation.path, { recursive: true })
          )
      );
      await Promise.all(
        stagingOperations
          .filter((operation) => operation._tag === "WriteFile")
          .map((operation) =>
            sandbox.writeFile(operation.path, operation.content)
          )
      );

      await Effect.runPromise(
        decodeIssueWorkspaceResult({
          command: checkoutCommandResult,
          plan,
        })
      );
      return createSandboxSessionEnv(sandbox, "/workspace");
    },
  };
};