feat: wire real AgentOS execution

2026-07-28 16:52:49 +05:30
parent 4d9f6da41b
commit 4edd456b5b
17 changed files with 1090 additions and 308 deletions
--- a/packages/agents/src/app.ts
+++ b/packages/agents/src/app.ts
@@ -1,4 +1,5 @@
 import { parseAgentEnv } from "@code/env/agent";
+import { WorkAttemptExecutionError } from "@code/primitives/execution-runtime";
 import { registerProvider } from "@flue/runtime";
 import type { Fetchable } from "@flue/runtime/routing";
 import { flue } from "@flue/runtime/routing";
@@ -42,8 +43,23 @@ app.post("/internal/work-attempts/execute", async (context) => {
  try {
    return context.json(await executeAgentOsAttempt(await context.req.json()));
  } catch (error) {
+    const failure =
+      error instanceof WorkAttemptExecutionError
+        ? error
+        : new WorkAttemptExecutionError({
+            message:
+              error instanceof Error ? error.message : "Execution failed",
+            reason: "HarnessFailed",
+            retryable: false,
+          });
    return context.json(
-      { error: error instanceof Error ? error.message : "Execution failed" },
+      {
+        error: {
+          message: failure.message,
+          reason: failure.reason,
+          retryable: failure.retryable,
+        },
+      },
      500
    );
  }
@@ -55,7 +71,11 @@ app.post("/internal/work-attempts/:workspaceKey/cancel", async (context) => {
  ) {
    return context.json({ error: "Unauthorized" }, 401);
  }
-  await cancelAgentOsAttempt(context.req.param("workspaceKey"));
+  const body = (await context.req.json()) as { attemptId?: unknown };
+  if (typeof body.attemptId !== "string" || body.attemptId.length === 0) {
+    return context.json({ error: "attemptId is required" }, 400);
+  }
+  await cancelAgentOsAttempt(context.req.param("workspaceKey"), body.attemptId);
  return context.json({ cancelled: true });
 });

--- a/packages/agents/src/runtime/agent-os.ts
+++ b/packages/agents/src/runtime/agent-os.ts
@@ -7,13 +7,24 @@ import {
  codexSessionEnv,
  makeCodexAgentOsConfig,
 } from "../../../primitives/src/agent-os";
-import { decodeWorkAttemptExecutionInput } from "../../../primitives/src/execution-runtime";
+import {
+  decodeWorkAttemptExecutionInput,
+  WorkAttemptExecutionError,
+} from "../../../primitives/src/execution-runtime";
 import type {
  ExecutionEvent,
  WorkAttemptExecutionResult,
 } from "../../../primitives/src/execution-runtime";

-const workspace = agentOS(makeCodexAgentOsConfig() as never);
+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");
+    }
+  },
+  software: codexConfig.software,
+});

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

@@ -43,6 +54,32 @@ const requireSuccess = (
  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",
@@ -52,162 +89,189 @@ const gitAuthEnv = (username: string, credential: string) => ({
 export const executeAgentOsAttempt = async (
  rawInput: unknown
 ): Promise<WorkAttemptExecutionResult> => {
-  const input = await Effect.runPromise(
-    decodeWorkAttemptExecutionInput(rawInput)
-  );
-  const env = parseAgentEnv(process.env);
-  const client = createClient<typeof runtimeRegistry>({
-    endpoint: env.RIVET_PUBLIC_ENDPOINT ?? env.RIVET_ENDPOINT,
-  });
-  const vm = client.workspace.getOrCreate([input.workspaceKey]);
-  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,
-      }
+  try {
+    const input = await Effect.runPromise(
+      decodeWorkAttemptExecutionInput(rawInput)
    );
-    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,
-  });
-  events.push(
-    event(4, "harness.progress", "Codex implementation turn completed", {
-      stopReason: String(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 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
    );
-    const tree = requireSuccess(
-      await vm.execArgv("git", ["write-tree"], {
+
+    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",
      }),
-      "Candidate tree creation"
+      "Base revision lookup"
    );
-    candidateRevision = requireSuccess(
-      await vm.exec(
-        `printf %s ${shellQuote(`Zopu candidate for ${input.attemptId}`)} | git commit-tree ${shellQuote(tree)} -p ${shellQuote(baseRevision)}`,
-        {
+    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",
-          env: {
-            GIT_AUTHOR_EMAIL: "agent@zopu.dev",
-            GIT_AUTHOR_NAME: "Zopu Agent",
-            GIT_COMMITTER_EMAIL: "agent@zopu.dev",
-            GIT_COMMITTER_NAME: "Zopu Agent",
-          },
+        }),
+        "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,
        }
      ),
-      "Candidate revision creation"
+      event(6, "runtime.completed", "AgentOS execution completed")
    );
-    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",
-  };
+    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) => {
+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]).cancelPrompt({});
+  await client.workspace
+    .getOrCreate([workspaceKey], {
+      params: { token: env.RIVET_WORKSPACE_TOKEN },
+    })
+    .cancelPrompt({ sessionId: `codex-${attemptId}` });
 };