feat(cloud-worker): mount subscription logins for the run and write refreshed tokens back

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased - Major]
 
+### Added
+
+- `agent-relay cloud worker` can run Codex and Grok jobs with the workspace's subscription logins when the control plane provides them; the login files exist only while the run does.
+- `agent-relay cloud worker` saves refreshed Codex and Grok login tokens for later runs.
+
 ### Changed
 
 - The core (`agent-relay-probe`) talks to one Agent Relay site, fixed when it is built: `AGENT_RELAY_SITE_URL` at build time (`https://agentrelay.com` when unset; an https origin, or `http://localhost[:port]` / `http://127.0.0.1[:port]`), and a build for a site the core could not trust fails to build. Nothing changes the site at run time, in any build.

packages/cli/src/cli/commands/cloud-worker.ts

@@ -20,6 +20,14 @@ import {
 } from '@agent-relay/cloud';
 
 import { defaultExit } from '../lib/exit.js';
+import {
+  isCliCredentialsPayload,
+  cliLoginHome,
+  materializeCliCredentials,
+  removeLoginHome,
+  writeBackCliCredentials,
+  type MaterializedCliCredentials,
+} from '../lib/cloud-worker-cli-credentials.js';
 
 type ExitFn = (code: number) => never;
 
@@ -80,11 +88,15 @@ function safePathSegment(value: string): string {
 
 function buildWorkerRuntimeEnv(
   payload: WorkerWorkflowPayload,
-  deps: CloudWorkerDependencies
+  deps: CloudWorkerDependencies,
+  cliCredentialEnv: Record<string, string> = {}
 ): NodeJS.ProcessEnv {
   return {
     ...deps.env,
     ...payload.envSecrets,
+    // Materialised subscription logins (CODEX_HOME/GROK_HOME pointing into the
+    // run dir) win over anything inherited or delivered under those names.
+    ...cliCredentialEnv,
     AGENT_RELAY_CLOUD_WORKER_RUN_ID: payload.runId,
     RELAY_WORKSPACE_ID: payload.relayWorkspaceId,
     RELAY_API_KEY: payload.relaycastApiKey,
@@ -217,6 +229,9 @@ export function createDefaultAssignmentRunner(deps: CloudWorkerDependencies): Ex
     // directory/env and then invokes relayflows with the equivalent run flags.
     const runDir = path.join(cloudWorkerStateDir(deps.env), 'runs', payload.runId);
     const keepRunDir = deps.env.AGENT_RELAY_WORKER_KEEP_RUN_DIR === '1';
+    const cliCredentials = isCliCredentialsPayload(payload.cliCredentials) ? payload.cliCredentials : null;
+    let materialized: MaterializedCliCredentials | null = null;
+    let writeBackDone = false;
     try {
       await fsp.mkdir(runDir, { recursive: true, mode: 0o700 });
       await fsp.chmod(runDir, 0o700).catch(() => undefined);
@@ -225,12 +240,23 @@ export function createDefaultAssignmentRunner(deps: CloudWorkerDependencies): Ex
       const workflowPath = path.join(runDir, safeFileName(payload.workflowFileName));
       await writeSecretFile(workflowPath, payload.workflow);
 
+      if (cliCredentials) {
+        materialized = await materializeCliCredentials({
+          worker,
+          credentials: cliCredentials,
+          runDir,
+          fetchImpl: deps.fetchImpl,
+          signal,
+          log: deps.log,
+        });
+      }
+
       const relayflowsCli = deps.resolveRelayflowsCliEntrypoint();
       const result = await runChild({
         command: process.execPath,
         args: relayflowsArgs(relayflowsCli, workflowPath, payload),
         cwd: runDir,
-        env: buildWorkerRuntimeEnv(payload, deps),
+        env: buildWorkerRuntimeEnv(payload, deps, materialized?.env),
         deps,
         signal,
       });
@@ -239,16 +265,56 @@ export function createDefaultAssignmentRunner(deps: CloudWorkerDependencies): Ex
         throw new Error(`Workflow runner exited with code ${result.exitCode}`);
       }
 
+      const failedWriteBacks = await storeRefreshedLogins();
       return {
         exitCode: result.exitCode,
         durationMs: result.durationMs,
-        summary: `Workflow ${payload.runId} completed.`,
+        summary: failedWriteBacks.length
+          ? `Workflow ${payload.runId} completed. The refreshed ${failedWriteBacks.join(', ')} login could not be stored; sign in again if the next run fails.`
+          : `Workflow ${payload.runId} completed.`,
       };
+    } catch (error) {
+      // A failed or cancelled run still stores a refreshed login (Grok rotates
+      // its single-use refresh token), and says so when it could not.
+      const failedWriteBacks = await storeRefreshedLogins();
+      if (failedWriteBacks.length && error instanceof Error) {
+        const warning = `The refreshed ${failedWriteBacks.join(', ')} login could not be stored; sign in again if the next run fails.`;
+        // Assignment reports keep the first 500 characters: lead with the
+        // warning when the original error would push it out.
+        error.message =
+          error.message.length + warning.length + 1 > 500
+            ? `${warning} ${error.message}`
+            : `${error.message} ${warning}`;
+      }
+      throw error;
     } finally {
+      await storeRefreshedLogins();
+      // The login files are removed even when the run directory is kept for
+      // debugging, and even when materialisation failed partway.
+      if (cliCredentials) {
+        const loginHome = materialized?.home ?? cliLoginHome(runDir);
+        await removeLoginHome(loginHome).catch(() => {
+          deps.error("[cloud-worker] could not remove the run's CLI login directory; delete it manually");
+        });
+      }
       if (!keepRunDir) {
         await fsp.rm(runDir, { recursive: true, force: true }).catch(() => undefined);
       }
     }
+
+    /** Writes back refreshed logins once; returns the CLIs whose login could not be stored. */
+    async function storeRefreshedLogins(): Promise<string[]> {
+      if (!cliCredentials || !materialized?.files.length || writeBackDone) return [];
+      writeBackDone = true;
+      const outcome = await writeBackCliCredentials({
+        worker,
+        credentials: cliCredentials,
+        materialized,
+        fetchImpl: deps.fetchImpl,
+        log: deps.log,
+      }).catch(() => ({ written: 0, failed: materialized?.files ?? [] }));
+      return outcome.failed.map((file) => file.cli);
+    }
   };
 }