agent-relay-orchestrator 0.77.0 → 0.78.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-relay-orchestrator",
-  "version": "0.77.0",
+  "version": "0.78.0",
   "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.53"
   },
   "devDependencies": {
     "@types/bun": "latest",

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,318 @@
+import { homedir } from "node:os";
+import { join } from "node:path";
+import {
+  QUOTA_FAILURE_LOG_INTERVAL_MS,
+  QUOTA_FAST_RETRY_MS,
+  QUOTA_POLL_INTERVAL_MS,
+  QuotaCollectionError,
+  collectClaudeQuotaSample,
+  collectCodexQuotaSample,
+  providerQuotaErrorFromCollectorError,
+  quotaRetryAfterMs,
+  resolveClaudeQuotaIdentity,
+  resolveCodexQuotaIdentityFromHome,
+  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 type { OrchestratorConfig } from "./config";
+
+const QUOTA_LEASE_TTL_MS = 90_000;
+const QUOTA_LEASE_RENEW_MS = 30_000;
+
+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;
+};
+
+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>();
+
+  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();
+  }
+
+  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")) {
+      const configDir = this.config.env.CLAUDE_CONFIG_DIR || process.env.CLAUDE_CONFIG_DIR || join(homedir(), ".claude");
+      candidates.push(await resolveClaudeQuotaIdentity({ configDir, host: this.config.hostname }));
+    }
+    if (this.config.providers.includes("codex")) {
+      const codexHome = this.config.env.CODEX_HOME || process.env.CODEX_HOME || join(homedir(), ".codex");
+      candidates.push({
+        ...(await resolveCodexQuotaIdentityFromHome({ codexHome, host: this.config.hostname })),
+        appServerUrl: this.config.env.CODEX_APP_SERVER_URL || process.env.CODEX_APP_SERVER_URL,
+      });
+    }
+    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") {
+      if (!candidate.appServerUrl) throw new QuotaCollectionError("creds_not_ready", "Codex app-server URL is not configured");
+      return collectCodexQuotaSample({
+        agentId: this.sourceAgentId(),
+        rateLimitsRead: () => this.options.codexRateLimitsRead
+          ? this.options.codexRateLimitsRead(candidate.appServerUrl!)
+          : codexRateLimitsRead(candidate.appServerUrl!),
+      });
+    }
+    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)}`);
+  }
+}
+
+async function codexRateLimitsRead(appServerUrl: string): Promise<unknown> {
+  const client = new JsonRpcWebSocket(appServerUrl);
+  try {
+    await client.connect();
+    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();
+  }
+}
+
+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; },
   };
 }