Add token refresh endpoints to agent server

packages/agent/scripts/dev-server.ts

@@ -0,0 +1,121 @@
+#!/usr/bin/env tsx
+// Boots AgentServer locally for poking with curl. Mocks the PostHog API via
+// msw and stamps a fake session so /command works without a real backend or
+// agent runtime spawn. Use only for endpoint smoke tests.
+
+import { execFileSync } from "node:child_process";
+import { generateKeyPairSync } from "node:crypto";
+import { mkdtempSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import jwt from "jsonwebtoken";
+import { setupServer } from "msw/node";
+import { AgentServer } from "../src/server/agent-server";
+import { SANDBOX_CONNECTION_AUDIENCE } from "../src/server/jwt";
+import { createPostHogHandlers } from "../src/test/mocks/msw-handlers";
+
+const PORT = Number(process.env.PORT ?? 3001);
+const TASK_ID = process.env.TASK_ID ?? "test-task-id";
+const RUN_ID = process.env.RUN_ID ?? "test-run-id";
+const PROJECT_ID = Number(process.env.PROJECT_ID ?? 1);
+const API_URL = "http://localhost:8000";
+
+const { privateKey, publicKey } = generateKeyPairSync("rsa", {
+  modulusLength: 2048,
+  publicKeyEncoding: { type: "spki", format: "pem" },
+  privateKeyEncoding: { type: "pkcs8", format: "pem" },
+});
+
+const mswServer = setupServer(...createPostHogHandlers({ baseUrl: API_URL }));
+mswServer.listen({ onUnhandledRequest: "bypass" });
+
+const repoPath = mkdtempSync(join(tmpdir(), "agent-server-dev-"));
+execFileSync("git", ["init", "-q"], { cwd: repoPath });
+execFileSync("git", ["config", "user.email", "dev@local"], { cwd: repoPath });
+execFileSync("git", ["config", "user.name", "Dev"], { cwd: repoPath });
+execFileSync("git", ["config", "commit.gpgsign", "false"], { cwd: repoPath });
+execFileSync("git", ["commit", "-q", "--allow-empty", "-m", "init"], {
+  cwd: repoPath,
+});
+
+const server = new AgentServer({
+  port: PORT,
+  jwtPublicKey: publicKey,
+  repositoryPath: repoPath,
+  apiUrl: API_URL,
+  apiKey: "dev-api-key",
+  projectId: PROJECT_ID,
+  mode: "interactive",
+  taskId: TASK_ID,
+  runId: RUN_ID,
+});
+
+// start() runs autoInitializeSession() which would spawn a real Claude/Codex
+// runtime. We don't need that to exercise the HTTP routes — no-op it.
+(
+  server as unknown as { autoInitializeSession: () => Promise<void> }
+).autoInitializeSession = async () => {
+  console.log("[dev-server] skipped autoInitializeSession");
+};
+
+await server.start();
+
+// /command's route checks session.payload.run_id; set_token's handler only
+// needs `this.session` to be truthy. Stamp the minimum shape that satisfies
+// both. Other commands (user_message, cancel, etc.) will still fail because
+// they touch session.clientConnection — that's fine.
+(server as unknown as { session: unknown }).session = {
+  payload: {
+    run_id: RUN_ID,
+    task_id: TASK_ID,
+    team_id: PROJECT_ID,
+    user_id: 0,
+    distinct_id: "dev",
+    mode: "interactive",
+  },
+};
+
+const token = jwt.sign(
+  {
+    run_id: RUN_ID,
+    task_id: TASK_ID,
+    team_id: PROJECT_ID,
+    user_id: 1,
+    distinct_id: "dev",
+    mode: "interactive",
+    aud: SANDBOX_CONNECTION_AUDIENCE,
+  },
+  privateKey,
+  { algorithm: "RS256", expiresIn: 3600 },
+);
+
+const base = `http://localhost:${PORT}`;
+console.log("");
+console.log(`agent-server listening on ${base}`);
+console.log(`repo: ${repoPath}`);
+console.log("");
+console.log(`export TOKEN='${token}'`);
+console.log("");
+console.log("# health (no auth)");
+console.log(`curl ${base}/health`);
+console.log("");
+console.log("# /gh (loopback only, no auth)");
+console.log(
+  `curl -X POST ${base}/gh -H 'Content-Type: application/json' -d '{"args":["--version"]}'`,
+);
+console.log("");
+console.log("# set_token via /command (JWT required)");
+console.log(
+  `curl -X POST ${base}/command -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/json' -d '{"jsonrpc":"2.0","id":1,"method":"set_token","params":{"token":"ghp_dev_test"}}'`,
+);
+console.log("");
+console.log("ctrl-c to stop.");
+
+const shutdown = async () => {
+  console.log("\n[dev-server] shutting down");
+  mswServer.close();
+  await server.stop();
+  process.exit(0);
+};
+process.on("SIGINT", shutdown);
+process.on("SIGTERM", shutdown);

packages/agent/src/server/agent-server.test.ts

@@ -382,6 +382,88 @@ describe("AgentServer HTTP Mode", () => {
     }, 20000);
   });
 
+  describe("set_token command", () => {
+    // Build an AgentServer that is NOT stored in the outer `server` ref so the
+    // afterEach hook doesn't try to clean up its (fake) session and cascade
+    // cleanup failures into the test output.
+    const buildBareServer = () => {
+      const bare = new AgentServer({
+        port,
+        jwtPublicKey: TEST_PUBLIC_KEY,
+        repositoryPath: repo.path,
+        apiUrl: "http://localhost:8000",
+        apiKey: "test-api-key",
+        projectId: 1,
+        mode: "interactive",
+        taskId: "test-task-id",
+        runId: "test-run-id",
+      }) as unknown as {
+        session: unknown;
+        executeCommand(
+          method: string,
+          params: Record<string, unknown>,
+        ): Promise<unknown>;
+      };
+      // Pass executeCommand's "no active session" guard.
+      bare.session = {};
+      return bare;
+    };
+
+    it.each([
+      { method: "set_token", token: "ghp_refreshed" },
+      { method: "posthog/set_token", token: "ghp_aliased" },
+    ])(
+      "updates GH_TOKEN and GITHUB_TOKEN via $method",
+      async ({ method, token }) => {
+        const previousGh = process.env.GH_TOKEN;
+        const previousGithub = process.env.GITHUB_TOKEN;
+        delete process.env.GH_TOKEN;
+        delete process.env.GITHUB_TOKEN;
+
+        try {
+          const result = await buildBareServer().executeCommand(method, {
+            token,
+          });
+
+          expect(result).toEqual({ updated: true });
+          expect(process.env.GH_TOKEN).toBe(token);
+          expect(process.env.GITHUB_TOKEN).toBe(token);
+        } finally {
+          if (previousGh === undefined) delete process.env.GH_TOKEN;
+          else process.env.GH_TOKEN = previousGh;
+          if (previousGithub === undefined) delete process.env.GITHUB_TOKEN;
+          else process.env.GITHUB_TOKEN = previousGithub;
+        }
+      },
+    );
+  });
+
+  describe("POST /gh", () => {
+    // Each case rejects with 400 (validation), proving 1) the body schema is
+    // enforced and 2) the route doesn't require a JWT — a 401 would mean the
+    // route is gated, which would defeat the wrapper-script use case.
+    it.each([
+      { name: "malformed JSON", body: "not-json" },
+      { name: "missing args", body: JSON.stringify({ cwd: "/tmp" }) },
+      { name: "empty args array", body: JSON.stringify({ args: [] }) },
+      { name: "empty body (no JWT required)", body: JSON.stringify({}) },
+    ])(
+      "rejects $name with 400",
+      async ({ body }) => {
+        await createServer().start();
+
+        const response = await fetch(`http://localhost:${port}/gh`, {
+          method: "POST",
+          headers: { "Content-Type": "application/json" },
+          body,
+        });
+
+        expect(response.status).toBe(400);
+      },
+      20000,
+    );
+  });
+
   describe("session lifecycle", () => {
     it("emits _posthog/run_started after session initialization", async () => {
       await createServer().start();

packages/agent/src/server/agent-server.ts

@@ -12,6 +12,7 @@ import {
   PROTOCOL_VERSION,
 } from "@agentclientprotocol/sdk";
 import { type ServerType, serve } from "@hono/node-server";
+import { getConnInfo } from "@hono/node-server/conninfo";
 import { getCurrentBranch } from "@posthog/git/queries";
 import { Hono } from "hono";
 import { z } from "zod";
@@ -53,8 +54,10 @@ import {
   normalizeCloudPromptContent,
   promptBlocksToText,
 } from "./cloud-prompt";
+import { isLoopbackAddress, runGh } from "./gh-exec";
 import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt";
 import {
+  ghRequestSchema,
   handoffLocalGitStateSchema,
   jsonRpcRequestSchema,
   validateCommandParams,
@@ -450,6 +453,43 @@ export class AgentServer {
       }
     });
 
+    // Sandbox-local exec for `gh`. Codex agents run in an isolated child process
+    // whose environment is captured at spawn time, so refreshing GH_TOKEN in the
+    // agent server doesn't reach them. Their gh-wrapper script calls this route
+    // to run gh against the agent server's freshly-refreshed env. Loopback-only;
+    // intentionally not JWT-protected so the wrapper has nothing to forward.
+    app.post("/gh", async (c) => {
+      const remote = getConnInfo(c).remote.address;
+      if (!isLoopbackAddress(remote)) {
+        this.logger.warn("Rejected non-loopback /gh request", { remote });
+        return c.json({ error: "Forbidden" }, 403);
+      }
+
+      const rawBody = await c.req.json().catch(() => null);
+      const parsed = ghRequestSchema.safeParse(rawBody);
+      if (!parsed.success) {
+        return c.json({ error: parsed.error.message }, 400);
+      }
+
+      const { args, cwd, timeoutMs } = parsed.data;
+      try {
+        const result = await runGh(args, {
+          cwd: cwd ?? this.config.repositoryPath ?? process.cwd(),
+          timeoutMs: timeoutMs ?? 60_000,
+          logger: this.logger,
+        });
+        return c.json(result);
+      } catch (error) {
+        this.logger.error("Failed to run gh", error);
+        return c.json(
+          {
+            error: error instanceof Error ? error.message : "Unknown error",
+          },
+          500,
+        );
+      }
+    });
+
     app.notFound((c) => {
       return c.json({ error: "Not found" }, 404);
     });
@@ -694,6 +734,15 @@ export class AgentServer {
         );
       }
 
+      case "posthog/set_token":
+      case "set_token": {
+        const token = params.token as string;
+        process.env.GH_TOKEN = token;
+        process.env.GITHUB_TOKEN = token;
+        this.logger.info("GH token refreshed");
+        return { updated: true };
+      }
+
       case POSTHOG_NOTIFICATIONS.PERMISSION_RESPONSE:
       case "permission_response": {
         const requestId = params.requestId as string;