@runuai/host 0.9.150 → 0.9.152

package/lib/token-broker-revision.ts

@@ -0,0 +1,310 @@
+/**
+ * ADR-154 §7: the broker's exact-name entry point.
+ *
+ * The cloud is the only scheduler: it knows which revisions are active, so it
+ * names one revision — its broker-only part under `/uai-broker` and its
+ * lease-readable part under `/uai/lease` — and the broker refreshes exactly
+ * that one. It never lists `/uai-broker`, so a retired revision is never
+ * refreshed.
+ *
+ * Every invoke, holding the credential's lock:
+ * 1. makes the lease-readable copy agree with the broker-only material,
+ *    rewriting it if a crash left it behind;
+ * 2. refreshes when due (or when the cloud forces it, for a runner's
+ *    credential-invalid report), writing the broker-only part first and the
+ *    lease-readable copy second, so a crash between them keeps the new
+ *    refresh token and leaves the previous access token valid;
+ * 3. answers for that one credential, never with material: `refreshed` or
+ *    `fresh` with the lease-readable parameter's version and the expiry,
+ *    `needs_sign_in`, `busy`, or `unavailable`.
+ *
+ * An adopted revision (`legacyParameter`, owner answer 3) also keeps today's
+ * parameter it came from agreeing with it, the brokered copy only, until the
+ * fleet audit (token-broker-adopt.ts).
+ *
+ * Storage is an interface; the Lambda binds it to SSM (token-broker-ssm.ts).
+ */
+
+import {
+  brokerErrorKind,
+  CREDENTIAL_LANE_SCHEDULES,
+  parseBrokerParameterName,
+  parseLeasedValue,
+  serializeLeasedValue,
+  type BrokerRevisionAnswer,
+  type BrokerRevisionRequest,
+  type CredentialLocation,
+} from "../src/cloud-credentials";
+import { parseBrokerEngineValue, type BrokerEngineValue } from "./capture-flow";
+import { syncLegacyCopy, type LegacySync } from "./token-broker-adopt";
+import { MCP_OAUTH_LANE, parseBrokerMcpValue, postTokenEndpoint, type BrokerMcpValue, oauthFailure } from "./mcp-capture";
+import type { SafeHttpDeps } from "./safe-http";
+import { backoffMs, createLaneHttp, isRefreshDue } from "./token-broker";
+import { laneById } from "./token-lanes/registry";
+import { isTerminalOAuthError, type CredentialFiles, type LaneHttp, type TokenLane } from "./token-lanes/types";
+
+export interface StoredParameter {
+  value: string;
+  version: number;
+}
+
+/** The storage one invoke needs: read and overwrite two named parameters,
+ *  and a lock per credential. */
+export interface RevisionStore {
+  get(name: string): Promise<StoredParameter | null>;
+  /** Overwrite a parameter this invoke has just read; the new version. */
+  overwrite(name: string, value: string): Promise<number>;
+  /** The credential's lock, or null while another invoke holds a fresh one. */
+  lock(credentialId: string): Promise<(() => Promise<void>) | null>;
+}
+
+const MAX_ERROR_CHARS = 512;
+
+function leaseCopy(location: CredentialLocation, lane: TokenLane, files: CredentialFiles): string {
+  return serializeLeasedValue({
+    v: 1,
+    part: "lease",
+    location,
+    class: "engine",
+    authKind: "config-dir",
+    files: lane.brokeredCopy(files),
+    expiresAt: lane.expiresAt(files),
+  });
+}
+
+function sameFiles(a: CredentialFiles, b: CredentialFiles): boolean {
+  const keys = Object.keys(a).sort();
+  const other = Object.keys(b).sort();
+  return keys.length === other.length && keys.every((k, i) => k === other[i] && a[k] === b[k]);
+}
+
+/** Refresh one named revision. Never throws for a provider or storage
+ *  failure: the answer says what happened. */
+export async function refreshRevision(
+  request: BrokerRevisionRequest,
+  store: RevisionStore,
+  options: {
+    now?: () => number;
+    http?: (lane: TokenLane) => LaneHttp;
+    /** The SSRF rule's seams, for an MCP server's token endpoint. */
+    safeHttp?: SafeHttpDeps;
+    /** ADR-134 root of an adopted compatibility copy. */
+    legacyPrefix?: string;
+  } = {},
+): Promise<BrokerRevisionAnswer> {
+  const now = options.now ?? Date.now;
+  const location = parseBrokerParameterName(request.brokerParameter)!;
+  if (location.class === "mcp") return refreshMcpRevision(request, store, location, options);
+  let release: (() => Promise<void>) | null = null;
+  try {
+    release = await store.lock(request.credentialId);
+    if (!release) return { outcome: "busy" };
+
+    const brokerPart = await store.get(request.brokerParameter);
+    const broker = brokerPart ? parseBrokerEngineValue(brokerPart.value, location) : null;
+    if (!broker) return { outcome: "unavailable", error: "the revision's broker part is missing or unreadable" };
+    const lane = laneById(broker.lane);
+    if (!lane) return { outcome: "unavailable", error: `no lane ${broker.lane}` };
+
+    // 1. The lease-readable copy agrees with the broker-only material. It is
+    // rewritten, never created: a missing one was deleted with its retired
+    // revision, which must not come back.
+    const leasePart = await store.get(request.leaseParameter);
+    if (!leasePart) return { outcome: "unavailable", error: "the revision's lease-readable part is missing" };
+    const current = parseLeasedValue(leasePart.value, location);
+    const expected = lane.brokeredCopy(broker.files);
+    let version = leasePart.version;
+    if (
+      !current ||
+      current.class !== "engine" ||
+      current.authKind !== "config-dir" ||
+      !sameFiles(current.files, expected)
+    ) {
+      version = await store.overwrite(request.leaseParameter, leaseCopy(location, lane, broker.files));
+    }
+    // An adopted revision: today's parameter agrees with it too.
+    let legacy = request.legacyParameter
+      ? await syncLegacyCopy(store, request.legacyParameter, lane, broker.files, options.legacyPrefix)
+      : null;
+    const withLegacy = <T extends object>(answer: T): T & { legacy?: LegacySync } =>
+      legacy === null ? answer : { ...answer, legacy };
+
+    if (broker.refresh.needsSignIn) return { outcome: "needs_sign_in" };
+    const at = now();
+    if (!isRefreshDue(lane, broker.files, broker.refresh, at, request.force)) {
+      return withLegacy({
+        outcome: "fresh" as const,
+        version,
+        expiresAt: lane.expiresAt(broker.files) ?? broker.refresh.expiresAt,
+      });
+    }
+
+    // 2. Refresh: the broker-only part first.
+    const http = options.http?.(lane) ?? createLaneHttp(lane.endpoints);
+    const result = await lane.refresh(broker.files, http, at);
+    if (!result.ok) {
+      const next: BrokerEngineValue = {
+        ...broker,
+        refresh: {
+          ...broker.refresh,
+          attempts: broker.refresh.attempts + 1,
+          lastAttemptAt: at,
+          lastError: result.reason.slice(0, MAX_ERROR_CHARS),
+          needsSignIn: result.terminal,
+        },
+      };
+      await store.overwrite(request.brokerParameter, JSON.stringify(next));
+      return result.terminal
+        ? { outcome: "needs_sign_in" }
+        : { outcome: "unavailable", error: result.reason.slice(0, MAX_ERROR_CHARS) };
+    }
+    const refreshed: BrokerEngineValue = {
+      ...broker,
+      files: { ...result.files },
+      refresh: {
+        ...broker.refresh,
+        expiresAt: result.expiresAt,
+        refreshedAt: at,
+        attempts: 0,
+        lastAttemptAt: at,
+        lastError: null,
+        needsSignIn: false,
+      },
+    };
+    await store.overwrite(request.brokerParameter, JSON.stringify(refreshed));
+    version = await store.overwrite(request.leaseParameter, leaseCopy(location, lane, result.files));
+    if (request.legacyParameter) {
+      legacy = await syncLegacyCopy(store, request.legacyParameter, lane, result.files, options.legacyPrefix);
+    }
+    return withLegacy({ outcome: "refreshed" as const, version, expiresAt: result.expiresAt });
+  } catch (error) {
+    // The error's kind only: a message can quote what a request carried.
+    return { outcome: "unavailable", error: brokerErrorKind(error) };
+  } finally {
+    if (release) await release();
+  }
+}
+
+function mcpLeaseCopy(location: CredentialLocation, broker: BrokerMcpValue): string {
+  return serializeLeasedValue({
+    v: 1,
+    part: "lease",
+    location,
+    class: "mcp",
+    authKind: "oauth",
+    accessToken: broker.accessToken,
+    tokenType: "Bearer",
+    expiresAt: broker.refresh.expiresAt,
+  });
+}
+
+function mcpRefreshDue(broker: BrokerMcpValue, now: number, forced: boolean): boolean {
+  if (broker.refresh.needsSignIn) return false;
+  if (forced) return true;
+  const { attempts, lastAttemptAt } = broker.refresh;
+  if (attempts > 0 && lastAttemptAt !== null && now - lastAttemptAt < backoffMs(attempts)) return false;
+  const schedule = CREDENTIAL_LANE_SCHEDULES[MCP_OAUTH_LANE]!;
+  if (broker.refresh.expiresAt !== null) return broker.refresh.expiresAt - now <= schedule.leadMs;
+  return now - (broker.refresh.refreshedAt ?? 0) >= (schedule.everyMs ?? schedule.leadMs);
+}
+
+/**
+ * ADR-154 §11: an MCP connection's sign-in, refreshed at its token endpoint
+ * through the SSRF rule. The same order as an engine's: the lease-readable
+ * copy agrees with the broker-only material first, then a due refresh writes
+ * the broker-only part before the copy. A sign-in that gave no refresh token,
+ * or that the provider calls dead, needs a new sign-in.
+ */
+async function refreshMcpRevision(
+  request: BrokerRevisionRequest,
+  store: RevisionStore,
+  location: CredentialLocation,
+  options: { now?: () => number; safeHttp?: SafeHttpDeps },
+): Promise<BrokerRevisionAnswer> {
+  const now = options.now ?? Date.now;
+  let release: (() => Promise<void>) | null = null;
+  try {
+    release = await store.lock(request.credentialId);
+    if (!release) return { outcome: "busy" };
+    const brokerPart = await store.get(request.brokerParameter);
+    const broker = brokerPart ? parseBrokerMcpValue(brokerPart.value, location) : null;
+    if (!broker) return { outcome: "unavailable", error: "the revision's broker part is missing or unreadable" };
+    const leasePart = await store.get(request.leaseParameter);
+    if (!leasePart) return { outcome: "unavailable", error: "the revision's lease-readable part is missing" };
+    const current = parseLeasedValue(leasePart.value, location);
+    let version = leasePart.version;
+    if (
+      !current ||
+      current.class !== "mcp" ||
+      current.authKind !== "oauth" ||
+      current.accessToken !== broker.accessToken ||
+      current.expiresAt !== broker.refresh.expiresAt
+    ) {
+      version = await store.overwrite(request.leaseParameter, mcpLeaseCopy(location, broker));
+    }
+    if (broker.refresh.needsSignIn) return { outcome: "needs_sign_in" };
+    const at = now();
+    if (!mcpRefreshDue(broker, at, request.force)) {
+      return { outcome: "fresh", version, expiresAt: broker.refresh.expiresAt };
+    }
+    const markFailed = async (reason: string, terminal: boolean): Promise<void> => {
+      const next: BrokerMcpValue = {
+        ...broker,
+        refresh: {
+          ...broker.refresh,
+          attempts: broker.refresh.attempts + 1,
+          lastAttemptAt: at,
+          lastError: reason.slice(0, MAX_ERROR_CHARS),
+          needsSignIn: terminal,
+        },
+      };
+      await store.overwrite(request.brokerParameter, JSON.stringify(next));
+    };
+    if (broker.refreshToken === null) {
+      await markFailed("the server gave no refresh token", true);
+      return { outcome: "needs_sign_in" };
+    }
+    const { status, json } = await postTokenEndpoint(
+      broker.tokenEndpoint,
+      {
+        grant_type: "refresh_token",
+        refresh_token: broker.refreshToken,
+        client_id: broker.clientId,
+        resource: broker.resource,
+      },
+      { clientId: broker.clientId, clientSecret: broker.clientSecret },
+      { http: options.safeHttp },
+    );
+    if (status < 200 || status >= 300 || typeof json.access_token !== "string") {
+      const terminal = status >= 400 && isTerminalOAuthError(status, json);
+      // A standard OAuth error code or the status: never the endpoint's own text.
+      const reason = oauthFailure(json, status);
+      await markFailed(reason, terminal);
+      return terminal ? { outcome: "needs_sign_in" } : { outcome: "unavailable", error: reason.slice(0, MAX_ERROR_CHARS) };
+    }
+    const expiresAt = typeof json.expires_in === "number" ? at + json.expires_in * 1000 : null;
+    const refreshed: BrokerMcpValue = {
+      ...broker,
+      accessToken: json.access_token,
+      // Rotating servers send a new one; the rest keep the old.
+      refreshToken: typeof json.refresh_token === "string" ? json.refresh_token : broker.refreshToken,
+      refresh: {
+        ...broker.refresh,
+        expiresAt,
+        refreshedAt: at,
+        attempts: 0,
+        lastAttemptAt: at,
+        lastError: null,
+        needsSignIn: false,
+      },
+    };
+    await store.overwrite(request.brokerParameter, JSON.stringify(refreshed));
+    version = await store.overwrite(request.leaseParameter, mcpLeaseCopy(location, refreshed));
+    return { outcome: "refreshed", version, expiresAt };
+  } catch (error) {
+    // The error's kind only: a message can quote what a request carried.
+    return { outcome: "unavailable", error: brokerErrorKind(error) };
+  } finally {
+    if (release) await release();
+  }
+}