agent-relay-orchestrator 0.77.1 → 0.78.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-relay-orchestrator",
-  "version": "0.77.1",
+  "version": "0.78.1",
   "description": "Agent Relay orchestrator — manages agent lifecycle across hosts",
   "type": "module",
   "bin": {
@@ -16,7 +16,7 @@
     "test": "bun test"
   },
   "dependencies": {
-    "agent-relay-sdk": "0.2.52"
+    "agent-relay-sdk": "0.2.54"
   },
   "devDependencies": {
     "@types/bun": "latest",

package/src/config.ts

@@ -42,6 +42,14 @@ export function bunBinFromEnv(): string | undefined {
   return process.env.AGENT_RELAY_BUN_BIN;
 }
 
+export function codexCommandFromEnv(): string | undefined {
+  return process.env.AGENT_RELAY_CODEX_COMMAND;
+}
+
+export function providerHomeRootFromEnv(): string {
+  return process.env.AGENT_RELAY_PROVIDER_HOME_ROOT || join(homedir(), ".agent-relay", "provider-homes");
+}
+
 export function disableSystemdSupervisor(): boolean {
   return process.env.AGENT_RELAY_DISABLE_SYSTEMD_SUPERVISOR === "1";
 }

package/src/index.ts

@@ -9,6 +9,7 @@ import { recoverManagedAgents } from "./recovery";
 import { ProviderProbeCache } from "./provider-probe";
 import { sweepEmptyWorkspaceContainers, workspacesRoot } from "./workspace-probe";
 import { startOrchestratorMaintenanceScheduler } from "./maintenance";
+import { OrchestratorQuotaPoller } from "./quota-poller";
 
 const args = process.argv.slice(2);
 
@@ -50,6 +51,7 @@ const config = loadConfig();
 const probeCache = new ProviderProbeCache(config);
 const relay = createRelayClient(config, probeCache);
 const control = createControlHandler(config, relay);
+const quotaPoller = new OrchestratorQuotaPoller(config, relay);
 
 const POLL_INTERVAL_MS = 3_000;
 const REGISTER_RETRY_MS = 5_000;
@@ -81,6 +83,7 @@ async function startup(): Promise<void> {
 
   // Host-local maintenance must run where the agent processes and tmux sockets live.
   startOrchestratorMaintenanceScheduler();
+  quotaPoller.start();
 
   // Sweep empty workspace container dirs left behind by prior cleanups (#280).
   const swept = sweepEmptyWorkspaceContainers(workspacesRoot(config.baseDir));
@@ -207,6 +210,7 @@ async function shutdown(): Promise<void> {
   if (healthCheckTimer) clearInterval(healthCheckTimer);
   if (guestReaperTimer) clearInterval(guestReaperTimer);
   if (apiServer) apiServer.stop();
+  quotaPoller.stop();
   relay.stopHeartbeatLoop();
   process.exit(0);
 }

package/src/quota-poller.ts

@@ -0,0 +1,446 @@
+import { homedir } from "node:os";
+import { existsSync, readdirSync } from "node:fs";
+import { createServer } from "node:net";
+import { join } from "node:path";
+import {
+  QUOTA_FAILURE_LOG_INTERVAL_MS,
+  QUOTA_FAST_RETRY_MS,
+  QUOTA_POLL_INTERVAL_MS,
+  QuotaCollectionError,
+  collectClaudeQuotaSample,
+  collectCodexQuotaSample,
+  providerQuotaErrorFromCollectorError,
+  quotaRetryAfterMs,
+  resolveStableClaudeQuotaIdentity,
+  resolveStableCodexQuotaIdentityFromHome,
+  type ProviderQuotaIdentity,
+  type ProviderQuotaSample,
+} from "agent-relay-sdk/provider-quota";
+import type { ProviderQuotaLeaseAcquireInput, ProviderQuotaUpdateInput } from "agent-relay-sdk";
+import { errMessage } from "agent-relay-sdk";
+import { codexCommandFromEnv, providerHomeRootFromEnv, type OrchestratorConfig } from "./config";
+
+const QUOTA_LEASE_TTL_MS = 90_000;
+const QUOTA_LEASE_RENEW_MS = 30_000;
+const CODEX_APP_SERVER_CONNECT_ATTEMPTS = 40;
+const CODEX_APP_SERVER_CONNECT_RETRY_MS = 250;
+
+type QuotaRelay = {
+  acquireProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput): Promise<{
+    acquired: boolean;
+    lease?: { leaseToken: string; expiresAt: number };
+    retryAfterMs?: number;
+  }>;
+  releaseProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput & { leaseToken: string }): Promise<unknown>;
+  reportProviderQuota(input: ProviderQuotaUpdateInput): Promise<unknown>;
+  connected: boolean;
+};
+
+type QuotaCandidate = ProviderQuotaIdentity & {
+  appServerUrl?: string;
+  codexHome?: string;
+};
+
+type QuotaPollState = {
+  leaseToken?: string;
+  leaseExpiresAt?: number;
+  nextPollAt?: number;
+  lastAttemptAt?: number;
+  lastLog?: { key: string; at: number };
+};
+
+export class OrchestratorQuotaPoller {
+  private timer?: Timer;
+  private active = false;
+  private inFlight = false;
+  private readonly states = new Map<string, QuotaPollState>();
+  private readonly logStates = new Map<string, { key: string; at: number }>();
+
+  constructor(
+    private readonly config: OrchestratorConfig,
+    private readonly relay: QuotaRelay,
+    private readonly options: {
+      intervalMs?: number;
+      fetchImpl?: typeof fetch;
+      now?: () => number;
+      log?: (message: string) => void;
+      codexRateLimitsRead?: (appServerUrl: string) => Promise<unknown>;
+    } = {},
+  ) {}
+
+  start(): void {
+    if (this.active) return;
+    this.active = true;
+    this.schedule(1_000);
+  }
+
+  stop(): void {
+    this.active = false;
+    if (this.timer) clearTimeout(this.timer);
+    this.timer = undefined;
+    for (const [key, state] of this.states) {
+      if (!state.leaseToken) continue;
+      const [provider, accountKey] = splitStateKey(key);
+      void this.relay.releaseProviderQuotaLease(this.config.id, { provider, accountKey, leaseToken: state.leaseToken }).catch(() => {});
+    }
+    this.states.clear();
+    this.logStates.clear();
+  }
+
+  async tick(): Promise<void> {
+    if (this.inFlight || !this.active || !this.relay.connected) return;
+    this.inFlight = true;
+    try {
+      const candidates = await this.discoverCandidates();
+      await this.releaseRemovedCandidates(candidates);
+      for (const candidate of candidates) {
+        await this.processCandidate(candidate);
+      }
+    } finally {
+      this.inFlight = false;
+      this.schedule(this.options.intervalMs ?? QUOTA_LEASE_RENEW_MS);
+    }
+  }
+
+  private schedule(delayMs: number): void {
+    if (this.timer) clearTimeout(this.timer);
+    this.timer = undefined;
+    if (!this.active) return;
+    this.timer = setTimeout(() => {
+      this.timer = undefined;
+      void this.tick();
+    }, Math.max(1_000, delayMs));
+  }
+
+  private async discoverCandidates(): Promise<QuotaCandidate[]> {
+    const candidates: QuotaCandidate[] = [];
+    if (this.config.providers.includes("claude")) {
+      candidates.push(...await this.discoverClaudeCandidates());
+    }
+    if (this.config.providers.includes("codex")) {
+      candidates.push(...await this.discoverCodexCandidates());
+    }
+    const deduped = new Map<string, QuotaCandidate>();
+    for (const candidate of candidates) {
+      deduped.set(candidateStateKey(candidate), candidate);
+    }
+    return [...deduped.values()];
+  }
+
+  private async discoverClaudeCandidates(): Promise<QuotaCandidate[]> {
+    const paths = [
+      join(this.config.env.CLAUDE_CONFIG_DIR || process.env.CLAUDE_CONFIG_DIR || join(homedir(), ".claude"), ".credentials.json"),
+      ...providerHomeCredentialPaths("claude", ".credentials.json"),
+    ];
+    const candidates: QuotaCandidate[] = [];
+    for (const credentialsPath of paths) {
+      const identity = await resolveStableClaudeQuotaIdentity({ credentialsPath });
+      if (identity) candidates.push(identity);
+    }
+    if (candidates.length === 0) {
+      this.logOnce("claude:no-stable-credentials", "quota refresh skipped for claude: no materialized Claude OAuth credentials found");
+    }
+    return candidates;
+  }
+
+  private async discoverCodexCandidates(): Promise<QuotaCandidate[]> {
+    const homes = [
+      this.config.env.CODEX_HOME || process.env.CODEX_HOME || join(homedir(), ".codex"),
+      ...providerHomeConfigDirs("codex", "auth.json"),
+    ];
+    const appServerUrl = this.config.env.CODEX_APP_SERVER_URL || process.env.CODEX_APP_SERVER_URL;
+    const candidates: QuotaCandidate[] = [];
+    for (const codexHome of homes) {
+      const identity = await resolveStableCodexQuotaIdentityFromHome({ codexHome });
+      if (identity) candidates.push({ ...identity, codexHome, ...(appServerUrl ? { appServerUrl } : {}) });
+    }
+    if (candidates.length === 0) {
+      this.logOnce("codex:no-stable-auth", "quota refresh skipped for codex: no account id found in Codex auth.json");
+    }
+    return candidates;
+  }
+
+  private async releaseRemovedCandidates(candidates: QuotaCandidate[]): Promise<void> {
+    const live = new Set(candidates.map(candidateStateKey));
+    for (const [key, state] of this.states) {
+      if (live.has(key)) continue;
+      this.states.delete(key);
+      if (!state.leaseToken) continue;
+      const [provider, accountKey] = splitStateKey(key);
+      await this.relay.releaseProviderQuotaLease(this.config.id, { provider, accountKey, leaseToken: state.leaseToken }).catch(() => {});
+    }
+  }
+
+  private async processCandidate(candidate: QuotaCandidate): Promise<void> {
+    const state = this.stateFor(candidate);
+    const now = this.now();
+    if (!await this.ensureLease(candidate, state, now)) return;
+    if (state.nextPollAt !== undefined && state.nextPollAt > now) return;
+
+    const lastAttemptAt = now;
+    try {
+      const sample = await this.collect(candidate);
+      const quota = sample.quota;
+      if (!quota) throw new QuotaCollectionError("creds_not_ready", `${candidate.provider} quota source is not ready`);
+      const update: ProviderQuotaUpdateInput = {
+        provider: candidate.provider,
+        accountKey: sample.accountKey ?? candidate.accountKey,
+        quota,
+        lastAttemptAt: quota.updatedAt,
+        sourceAgentId: this.sourceAgentId(),
+      };
+      await this.relay.reportProviderQuota(update);
+      state.lastAttemptAt = update.lastAttemptAt;
+      state.nextPollAt = now + QUOTA_POLL_INTERVAL_MS;
+    } catch (error) {
+      const retryAfterMs = quotaRetryAfterMs(error);
+      const lastError = providerQuotaErrorFromCollectorError(error, retryAfterMs);
+      const retryDelayMs = retryAfterMs ?? (state.lastAttemptAt ? QUOTA_POLL_INTERVAL_MS : QUOTA_FAST_RETRY_MS);
+      state.lastAttemptAt = lastAttemptAt;
+      state.nextPollAt = now + retryDelayMs;
+      await this.relay.reportProviderQuota({
+        provider: candidate.provider,
+        accountKey: candidate.accountKey,
+        lastAttemptAt,
+        lastError,
+        sourceAgentId: this.sourceAgentId(),
+      }).catch((publishError) => this.log(`quota status publish failed: ${errMessage(publishError)}`));
+      this.logFailure(candidate, error, retryAfterMs);
+    }
+  }
+
+  private async ensureLease(candidate: QuotaCandidate, state: QuotaPollState, now: number): Promise<boolean> {
+    if (state.leaseToken && state.leaseExpiresAt && state.leaseExpiresAt - now > QUOTA_LEASE_RENEW_MS) return true;
+    const result = await this.relay.acquireProviderQuotaLease(this.config.id, {
+      provider: candidate.provider,
+      accountKey: candidate.accountKey,
+      ...(state.leaseToken ? { leaseToken: state.leaseToken } : {}),
+      ttlMs: QUOTA_LEASE_TTL_MS,
+    });
+    if (!result.acquired || !result.lease) {
+      state.leaseToken = undefined;
+      state.leaseExpiresAt = undefined;
+      state.nextPollAt = now + Math.min(result.retryAfterMs ?? QUOTA_LEASE_RENEW_MS, QUOTA_LEASE_RENEW_MS);
+      return false;
+    }
+    state.leaseToken = result.lease.leaseToken;
+    state.leaseExpiresAt = result.lease.expiresAt;
+    return true;
+  }
+
+  private async collect(candidate: QuotaCandidate): Promise<ProviderQuotaSample> {
+    if (candidate.provider === "claude") {
+      return collectClaudeQuotaSample({
+        agentId: this.sourceAgentId(),
+        credentialsPath: candidate.credentialsPath,
+        fetchImpl: this.options.fetchImpl,
+      });
+    }
+    if (candidate.provider === "codex") {
+      return collectCodexQuotaSample({
+        agentId: this.sourceAgentId(),
+        rateLimitsRead: () => this.options.codexRateLimitsRead
+          ? this.options.codexRateLimitsRead(candidate.appServerUrl ?? "")
+          : candidate.appServerUrl
+            ? codexRateLimitsRead(candidate.appServerUrl)
+            : codexRateLimitsReadFromHome(candidate.codexHome),
+      });
+    }
+    return {};
+  }
+
+  private stateFor(candidate: QuotaCandidate): QuotaPollState {
+    const key = candidateStateKey(candidate);
+    const existing = this.states.get(key);
+    if (existing) return existing;
+    const created: QuotaPollState = {};
+    this.states.set(key, created);
+    return created;
+  }
+
+  private sourceAgentId(): string {
+    return `orchestrator-${this.config.id}`;
+  }
+
+  private now(): number {
+    return this.options.now?.() ?? Date.now();
+  }
+
+  private log(message: string): void {
+    (this.options.log ?? ((line) => console.error(`[orchestrator] ${line}`)))(message);
+  }
+
+  private logFailure(candidate: QuotaCandidate, error: unknown, retryAfterMs: number | undefined): void {
+    const state = this.stateFor(candidate);
+    const key = retryAfterMs !== undefined ? `retry-after:${retryAfterMs}` : errMessage(error);
+    const now = this.now();
+    if (state.lastLog?.key === key && now - state.lastLog.at < QUOTA_FAILURE_LOG_INTERVAL_MS) return;
+    state.lastLog = { key, at: now };
+    const suffix = retryAfterMs !== undefined ? `; retrying in ${Math.round(retryAfterMs / 1000)}s` : "";
+    this.log(`quota refresh failed for ${candidate.provider}/${candidate.accountKey}${suffix}: ${errMessage(error)}`);
+  }
+
+  private logOnce(key: string, message: string): void {
+    const lastLog = this.logStates.get(key);
+    const now = this.now();
+    if (lastLog && now - lastLog.at < QUOTA_FAILURE_LOG_INTERVAL_MS) return;
+    this.logStates.set(key, { key, at: now });
+    this.log(message);
+  }
+}
+
+async function codexRateLimitsRead(appServerUrl: string, attempts = 1): Promise<unknown> {
+  const client = new JsonRpcWebSocket(appServerUrl);
+  try {
+    await connectWithRetry(client, attempts);
+    await client.request("initialize", {
+      clientInfo: { name: "agent-relay-orchestrator", title: "Agent Relay Orchestrator" },
+      capabilities: { experimentalApi: true },
+    }).catch(() => undefined);
+    return await client.request("account/rateLimits/read");
+  } finally {
+    client.close();
+  }
+}
+
+async function codexRateLimitsReadFromHome(codexHome: string | undefined): Promise<unknown> {
+  if (!codexHome) throw new QuotaCollectionError("creds_not_ready", "Codex home is not configured");
+  const appServerUrl = await freeLoopbackWsUrl();
+  const command = codexCommandFromEnv() || "codex";
+  const proc = Bun.spawn([command, "app-server", "--listen", appServerUrl], {
+    env: {
+      ...process.env,
+      CODEX_HOME: codexHome,
+      CODEX_APP_SERVER_URL: appServerUrl,
+    },
+    stdin: "ignore",
+    stdout: "ignore",
+    stderr: "ignore",
+  });
+  try {
+    return await codexRateLimitsRead(appServerUrl, CODEX_APP_SERVER_CONNECT_ATTEMPTS);
+  } finally {
+    proc.kill();
+    await Promise.race([proc.exited.catch(() => undefined), Bun.sleep(2_000)]).catch(() => undefined);
+  }
+}
+
+async function connectWithRetry(client: JsonRpcWebSocket, attempts: number): Promise<void> {
+  let lastError: unknown;
+  for (let attempt = 0; attempt < attempts; attempt += 1) {
+    try {
+      await client.connect();
+      return;
+    } catch (error) {
+      lastError = error;
+      await Bun.sleep(CODEX_APP_SERVER_CONNECT_RETRY_MS);
+    }
+  }
+  throw lastError ?? new Error("websocket connect failed");
+}
+
+async function freeLoopbackWsUrl(): Promise<string> {
+  const port = await new Promise<number>((resolve, reject) => {
+    const server = createServer();
+    server.on("error", reject);
+    server.listen(0, "127.0.0.1", () => {
+      const address = server.address();
+      const port = typeof address === "object" && address ? address.port : undefined;
+      server.close(() => port ? resolve(port) : reject(new Error("failed to reserve loopback port")));
+    });
+  });
+  return `ws://127.0.0.1:${port}`;
+}
+
+function providerHomeRoot(): string {
+  return providerHomeRootFromEnv();
+}
+
+function providerHomeConfigDirs(provider: "claude" | "codex", markerFile: string): string[] {
+  const root = join(providerHomeRoot(), provider);
+  const dirs: string[] = [];
+  for (const profile of safeReadDir(root)) {
+    const profileDir = join(root, profile);
+    for (const instance of safeReadDir(profileDir)) {
+      const dir = join(profileDir, instance);
+      if (existsSync(join(dir, markerFile))) dirs.push(dir);
+    }
+  }
+  return dirs;
+}
+
+function providerHomeCredentialPaths(provider: "claude", markerFile: string): string[] {
+  return providerHomeConfigDirs(provider, markerFile).map((dir) => join(dir, markerFile));
+}
+
+function safeReadDir(path: string): string[] {
+  try {
+    return readdirSync(path, { withFileTypes: true })
+      .filter((entry) => entry.isDirectory())
+      .map((entry) => entry.name);
+  } catch {
+    return [];
+  }
+}
+
+class JsonRpcWebSocket {
+  private ws!: WebSocket;
+  private nextId = 1;
+  private pending = new Map<number, { resolve: (value: unknown) => void; reject: (error: unknown) => void }>();
+
+  constructor(private readonly url: string) {}
+
+  async connect(): Promise<void> {
+    await new Promise<void>((resolve, reject) => {
+      const ws = new WebSocket(this.url);
+      this.ws = ws;
+      ws.onopen = () => resolve();
+      ws.onerror = () => reject(new Error("websocket error"));
+      ws.onclose = (event) => {
+        const error = new Error(`websocket closed code=${event.code} reason=${event.reason || "(none)"}`);
+        for (const pending of this.pending.values()) pending.reject(error);
+        this.pending.clear();
+      };
+      ws.onmessage = (event) => this.handleMessage(String(event.data));
+    });
+  }
+
+  request(method: string, params?: unknown): Promise<unknown> {
+    const id = this.nextId++;
+    const promise = new Promise<unknown>((resolve, reject) => {
+      this.pending.set(id, { resolve, reject });
+    });
+    this.ws.send(JSON.stringify({ id, method, params }));
+    return promise;
+  }
+
+  close(): void {
+    this.ws?.close();
+  }
+
+  private handleMessage(raw: string): void {
+    let message: unknown;
+    try {
+      message = JSON.parse(raw) as unknown;
+    } catch {
+      return;
+    }
+    if (!message || typeof message !== "object" || !("id" in message)) return;
+    const id = Number((message as { id: unknown }).id);
+    const pending = this.pending.get(id);
+    if (!pending) return;
+    this.pending.delete(id);
+    const response = message as { result?: unknown; error?: { message?: string; code?: number } };
+    if (response.error) pending.reject(new Error(`${response.error.message ?? "JSON-RPC error"} (${response.error.code ?? "unknown"})`));
+    else pending.resolve(response.result);
+  }
+}
+
+function candidateStateKey(candidate: Pick<QuotaCandidate, "provider" | "accountKey">): string {
+  return `${candidate.provider}\0${candidate.accountKey}`;
+}
+
+function splitStateKey(key: string): [string, string] {
+  const [provider, accountKey] = key.split("\0");
+  return [provider ?? "", accountKey ?? ""];
+}

package/src/relay.ts

@@ -2,7 +2,7 @@ import type { OrchestratorConfig } from "./config";
 import type { ProviderProbeCache } from "./provider-probe";
 import { detectSelfSupervision } from "./self-supervision";
 import { GIT_SHA, ORCHESTRATOR_PROTOCOL_VERSION, VERSION, runtimeMetadata } from "./version";
-import type { AgentLifecycle, WorkspaceMetadata, WorkspaceMode, ManagedSessionExitDiagnostics as SdkManagedSessionExitDiagnostics } from "agent-relay-sdk";
+import type { AgentLifecycle, ProviderQuotaLeaseAcquireInput, ProviderQuotaLeaseAcquireResult, ProviderQuotaUpdateInput, WorkspaceMetadata, WorkspaceMode, ManagedSessionExitDiagnostics as SdkManagedSessionExitDiagnostics } from "agent-relay-sdk";
 import { ReconnectionManager, RelayHttpClient } from "agent-relay-sdk";
 
 export interface RelayClient {
@@ -11,6 +11,9 @@ export interface RelayClient {
   updateManagedAgents(agents: ManagedAgentReport[], exitedAgents?: ManagedSessionExitDiagnostics[]): Promise<void>;
   pollCommands(): Promise<RelayCommand[]>;
   updateCommand(commandId: string, status: string, result?: Record<string, unknown>, error?: string): Promise<boolean>;
+  acquireProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput): Promise<ProviderQuotaLeaseAcquireResult>;
+  releaseProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput & { leaseToken: string }): Promise<{ released: boolean }>;
+  reportProviderQuota(input: ProviderQuotaUpdateInput): Promise<unknown>;
   setApiUrl(url: string): void;
   startHeartbeatLoop(): void;
   stopHeartbeatLoop(): void;
@@ -239,6 +242,15 @@ export function createRelayClient(config: OrchestratorConfig, probeCache: Provid
         return null;
       }
     },
+    acquireProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput): Promise<ProviderQuotaLeaseAcquireResult> {
+      return http.acquireProviderQuotaLease(orchestratorId, input);
+    },
+    releaseProviderQuotaLease(orchestratorId: string, input: ProviderQuotaLeaseAcquireInput & { leaseToken: string }): Promise<{ released: boolean }> {
+      return http.releaseProviderQuotaLease(orchestratorId, input);
+    },
+    reportProviderQuota(input: ProviderQuotaUpdateInput): Promise<unknown> {
+      return http.upsertProviderQuota(input);
+    },
     get connected() { return connected; },
   };
 }