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

import { parseAgentEnv } from "@code/env/agent";
import { agentOS, setup } from "@rivet-dev/agentos";
import { createClient } from "@rivet-dev/agentos/client";
import { Effect } from "effect";

import {
  codexSessionEnv,
  makeCodexAgentOsConfig,
} from "../../../primitives/src/agent-os";
import {
  decodeWorkAttemptExecutionInput,
  WorkAttemptExecutionError,
} from "../../../primitives/src/execution-runtime";
import type {
  ExecutionEvent,
  WorkAttemptExecutionResult,
} from "../../../primitives/src/execution-runtime";

const codexConfig = makeCodexAgentOsConfig();
const workspace = agentOS<undefined, { token: string }>({
  onBeforeConnect: (_context, params) => {
    if (params.token !== process.env.RIVET_WORKSPACE_TOKEN) {
      throw new Error("Unauthorized workspace connection");
    }
  },
  options: {
    actionTimeout: 10 * 60 * 1000,
  },
  permissions: codexConfig.permissions,
  software: codexConfig.software,
});

export const runtimeRegistry = setup({ use: { workspace } });

const shellQuote = (value: string): string =>
  `'${value.replaceAll("'", `'"'"'`)}'`;

const event = (
  sequence: number,
  kind: ExecutionEvent["kind"],
  message: string,
  metadata: Record<string, string> = {}
): ExecutionEvent => ({
  kind,
  message,
  metadata,
  occurredAt: Date.now(),
  sequence,
});

const requireSuccess = (
  result: { exitCode: number; stderr: string; stdout: string },
  operation: string
) => {
  if (result.exitCode !== 0) {
    throw new Error(`${operation} failed: ${result.stderr || result.stdout}`);
  }
  return result.stdout.trim();
};

const executionError = (
  message: string,
  reason: WorkAttemptExecutionError["reason"],
  retryable: boolean
) => new WorkAttemptExecutionError({ message, reason, retryable });

const classifyRuntimeFailure = (cause: unknown): WorkAttemptExecutionError => {
  if (cause instanceof WorkAttemptExecutionError) {
    return cause;
  }
  const message = cause instanceof Error ? cause.message : "Execution failed";
  if (/auth|credential|401|403/iu.test(message)) {
    return executionError(message, "Authentication", false);
  }
  if (/clone|checkout|repository|git /iu.test(message)) {
    return executionError(message, "RepositoryFailed", false);
  }
  if (/timeout|timed out/iu.test(message)) {
    return executionError(message, "Timeout", true);
  }
  if (/connect|network|unavailable|502|503|504/iu.test(message)) {
    return executionError(message, "ProviderUnavailable", true);
  }
  return executionError(message, "HarnessFailed", false);
};

const gitAuthEnv = (username: string, credential: string) => ({
  GIT_CONFIG_COUNT: "1",
  GIT_CONFIG_KEY_0: "http.extraHeader",
  GIT_CONFIG_VALUE_0: `Authorization: Basic ${Buffer.from(`${username}:${credential}`).toString("base64")}`,
});

export const executeAgentOsAttempt = async (
  rawInput: unknown
): Promise<WorkAttemptExecutionResult> => {
  try {
    const input = await Effect.runPromise(
      decodeWorkAttemptExecutionInput(rawInput)
    );
    const env = parseAgentEnv(process.env);
    const client = createClient<typeof runtimeRegistry>({
      disableMetadataLookup: true,
      endpoint: env.RIVET_PUBLIC_ENDPOINT ?? env.RIVET_ENDPOINT,
    });
    const vm = client.workspace.getOrCreate([input.workspaceKey], {
      params: { token: env.RIVET_WORKSPACE_TOKEN },
    });
    const events: ExecutionEvent[] = [
      event(0, "runtime.preparing", "AgentOS workspace selected", {
        workspaceKey: input.workspaceKey,
      }),
    ];
    const authEnv = gitAuthEnv(
      input.auth.username ??
        (input.auth.provider === "github" ? "x-access-token" : "git"),
      input.auth.credential
    );

    await vm.mkdir("/workspace", { recursive: true });
    if (!(await vm.exists("/workspace/repository/.git"))) {
      events.push(event(1, "repository.cloning", "Cloning project repository"));
      const clone = await vm.execArgv(
        "git",
        ["clone", input.repositoryUrl, "/workspace/repository"],
        {
          cwd: "/workspace",
          env: authEnv,
        }
      );
      requireSuccess(clone, "Repository clone");
    }

    const checkout = await vm.exec(
      `git fetch origin ${shellQuote(input.baseBranch)} && git checkout -B ${shellQuote(`zopu/${input.runId}`)} ${shellQuote(`origin/${input.baseBranch}`)}`,
      { cwd: "/workspace/repository", env: authEnv }
    );
    requireSuccess(checkout, "Repository checkout");
    const baseRevision = requireSuccess(
      await vm.execArgv("git", ["rev-parse", "HEAD"], {
        cwd: "/workspace/repository",
      }),
      "Base revision lookup"
    );
    events.push(
      event(2, "repository.ready", "Repository checkout is ready", {
        baseRevision,
      })
    );

    const sessionId = `codex-${input.attemptId}`;
    await vm.openSession({
      additionalInstructions:
        "Work only inside /workspace/repository. Do not reveal credentials. Make the requested change and run focused verification. Do not push or open a pull request.",
      agent: "codex",
      cwd: "/workspace/repository",
      env: codexSessionEnv({
        apiKey: env.AGENT_MODEL_API_KEY,
        baseUrl: env.AGENT_MODEL_BASE_URL,
        model: env.AGENT_MODEL_NAME,
      }),
      permissionPolicy: "allow_all",
      sessionId,
    });
    events.push(
      event(3, "harness.started", "Codex implementation session started")
    );
    const promptResult = await vm.prompt({
      content: [{ text: input.prompt, type: "text" }],
      idempotencyKey: input.attemptId,
      sessionId,
    });
    if (promptResult.stopReason !== "end_turn") {
      const reason =
        promptResult.stopReason === "cancelled" ? "Cancelled" : "HarnessFailed";
      throw executionError(
        `Codex stopped with ${promptResult.stopReason}`,
        reason,
        promptResult.stopReason === "max_tokens" ||
          promptResult.stopReason === "max_turn_requests"
      );
    }
    events.push(
      event(4, "harness.progress", "Codex implementation turn completed", {
        stopReason: promptResult.stopReason,
      })
    );

    const status = requireSuccess(
      await vm.execArgv("git", ["status", "--porcelain"], {
        cwd: "/workspace/repository",
      }),
      "Changed file lookup"
    );
    const diff = requireSuccess(
      await vm.execArgv("git", ["diff", "--binary", "HEAD"], {
        cwd: "/workspace/repository",
      }),
      "Diff collection"
    );
    const changedFiles = status
      .split("\n")
      .filter(Boolean)
      .map((line) => line.slice(3).trim());
    let candidateRevision = baseRevision;
    if (changedFiles.length > 0) {
      requireSuccess(
        await vm.execArgv("git", ["add", "-A"], {
          cwd: "/workspace/repository",
        }),
        "Candidate staging"
      );
      const tree = requireSuccess(
        await vm.execArgv("git", ["write-tree"], {
          cwd: "/workspace/repository",
        }),
        "Candidate tree creation"
      );
      candidateRevision = requireSuccess(
        await vm.exec(
          `printf %s ${shellQuote(`Zopu candidate for ${input.attemptId}`)} | git commit-tree ${shellQuote(tree)} -p ${shellQuote(baseRevision)}`,
          {
            cwd: "/workspace/repository",
            env: {
              GIT_AUTHOR_EMAIL: "agent@zopu.dev",
              GIT_AUTHOR_NAME: "Zopu Agent",
              GIT_COMMITTER_EMAIL: "agent@zopu.dev",
              GIT_COMMITTER_NAME: "Zopu Agent",
            },
          }
        ),
        "Candidate revision creation"
      );
      requireSuccess(
        await vm.execArgv("git", ["reset"], { cwd: "/workspace/repository" }),
        "Candidate index reset"
      );
    }
    events.push(
      event(
        5,
        "repository.changed",
        `${changedFiles.length} changed file(s) collected`,
        {
          candidateRevision,
        }
      ),
      event(6, "runtime.completed", "AgentOS execution completed")
    );

    return {
      baseRevision,
      candidateRevision,
      changedFiles,
      diff,
      environmentId: input.workspaceKey,
      events,
      summary:
        changedFiles.length > 0
          ? `Codex changed ${changedFiles.length} file(s)`
          : "Codex completed without repository changes",
    };
  } catch (error) {
    throw classifyRuntimeFailure(error);
  }
};

export const cancelAgentOsAttempt = async (
  workspaceKey: string,
  attemptId: string
) => {
  const env = parseAgentEnv(process.env);
  const client = createClient<typeof runtimeRegistry>({
    disableMetadataLookup: true,
    endpoint: env.RIVET_PUBLIC_ENDPOINT ?? env.RIVET_ENDPOINT,
  });
  await client.workspace
    .getOrCreate([workspaceKey], {
      params: { token: env.RIVET_WORKSPACE_TOKEN },
    })
    .cancelPrompt({ sessionId: `codex-${attemptId}` });
};