@runuai/host 0.9.149 → 0.9.151

package/lib/token-broker-adopt.ts

@@ -0,0 +1,344 @@
+/**
+ * ADR-154, owner answer 3: what Uai Cloud already keeps in Parameter Store
+ * under today's layout (ADR-134), its workspaces' AI accounts
+ * (`<prefix>/orgs/<org>/engines/<kind>/<id>`), is adopted in place. The cloud
+ * lists the accounts through the broker (metadata only: the broker is the
+ * one identity that reads those values), opens a revision for each, and
+ * names it here. People's MCP connections follow separately, once a fence
+ * keeps the v0 task machines that refresh them themselves from doing so.
+ *
+ * One refresher per grant, always (§7):
+ * - before anything is copied, the adoption marker (`<prefix>/adopted/...`) is
+ *   created under the scan's own lock, so no scan run is refreshing the
+ *   account meanwhile and none refreshes it again;
+ * - the full material becomes the revision's broker-only part and what a
+ *   task may hold its lease-readable part, both created (never overwritten)
+ *   with the tags the role policy checks, as a capture's are;
+ * - today's parameter then keeps what a task may hold only (the brokered
+ *   copy): the one refresh token
+ *   left is the revision's, which only the cloud's scheduler names. Until
+ *   the fleet audit each refresh of that revision writes the same there,
+ *   for the runners still reading it (`syncLegacyCopy`).
+ *
+ * A value that can never be adopted as it is answers `unsupported`, before
+ * any marker: it keeps its refresher. A retry after a lost answer finds its
+ * own marker and creates, and carries on where it stopped. Nothing here
+ * returns or logs material, nor an error's message (only its kind).
+ */
+
+import {
+  brokerErrorKind,
+  brokerParameterName,
+  DEFAULT_CREDENTIAL_STORE_PREFIX,
+  leaseParameterName,
+  legacyAdoptedRoot,
+  legacyAdoptionMarkerName,
+  MAX_LEGACY_ACCOUNTS,
+  parseLegacyEngineParameterName,
+  parseLegacyParameterName,
+  type BrokerAdoptAnswer,
+  type BrokerAdoptRequest,
+  type BrokerLegacyAccount,
+  type BrokerLegacyListAnswer,
+  type CredentialLocation,
+  type LegacyEngineParameter,
+} from "../src/cloud-credentials";
+import { buildEngineCapture, captureTags, parseCaptureFlow, type EngineCaptureMaterial } from "./capture-flow";
+import type { RevisionStore, StoredParameter } from "./token-broker-revision";
+import { laneFor } from "./token-lanes/registry";
+import type { CredentialFiles, TokenLane } from "./token-lanes/types";
+
+/** What adoption needs beyond a refresh: create-only writes, today's
+ *  accounts and markers, and the scan's lock. */
+export interface AdoptionStore extends RevisionStore {
+  /** The exact ADR-134 root this store reads and writes. */
+  readonly legacyPrefix: string;
+  /**
+   * Create a parameter with its identity tags, never overwriting:
+   * `matched` when it exists already with exactly these tags (this create,
+   * retried). Throws when it exists as anything else.
+   */
+  create(name: string, value: string, tags: Record<string, string>): Promise<"created" | "matched">;
+  /** Every parameter under today's workspace accounts root, decrypted. */
+  listLegacy(): Promise<Array<{ name: string; value: string }>>;
+  /** Every adoption marker of a workspace account, by the marker's name. */
+  listMarkers(): Promise<Array<{ name: string; value: string }>>;
+  /** The scan's lock (one run at a time), or null while a run holds it. */
+  scanLock(): Promise<(() => Promise<void>) | null>;
+}
+
+/** An AI account's value under today's layout, as the shared host wrote it. */
+export interface LegacyEngineValue {
+  /** The whole envelope: rewritten with only its files changed. */
+  raw: Record<string, unknown>;
+  label: string;
+  authKind: "env" | "config-dir";
+  position: number;
+  env: Record<string, string> | null;
+  files: CredentialFiles | null;
+}
+
+function stringMap(value: unknown): Record<string, string> | null {
+  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
+  const out: Record<string, string> = {};
+  for (const [key, text] of Object.entries(value as Record<string, unknown>)) {
+    if (typeof text !== "string") return null;
+    out[key] = text;
+  }
+  return Object.keys(out).length > 0 ? out : null;
+}
+
+function jsonObject(value: string): Record<string, unknown> | null {
+  try {
+    const parsed: unknown = JSON.parse(value);
+    return parsed && typeof parsed === "object" && !Array.isArray(parsed) ? (parsed as Record<string, unknown>) : null;
+  } catch {
+    return null;
+  }
+}
+
+/** Today's value for exactly this account, or null. */
+export function parseLegacyEngineValue(value: string, legacy: LegacyEngineParameter): LegacyEngineValue | null {
+  const o = jsonObject(value);
+  if (
+    !o ||
+    o.v !== 1 ||
+    o.id !== legacy.accountId ||
+    o.kind !== legacy.engineKind ||
+    o.scope !== `org:${legacy.orgId}` ||
+    typeof o.label !== "string" ||
+    o.label.trim().length === 0 ||
+    (o.authKind !== "env" && o.authKind !== "config-dir")
+  ) {
+    return null;
+  }
+  const env = o.authKind === "env" ? stringMap(o.env) : null;
+  const files = o.authKind === "config-dir" ? stringMap(o.files) : null;
+  if (o.authKind === "env" ? env === null : files === null) return null;
+  return {
+    raw: o,
+    label: o.label.trim(),
+    authKind: o.authKind,
+    position: typeof o.position === "number" && Number.isSafeInteger(o.position) ? o.position : 0,
+    env,
+    files,
+  };
+}
+
+function laneOf(legacy: LegacyEngineParameter, value: LegacyEngineValue): TokenLane | null {
+  if (value.authKind !== "config-dir" || !value.files) return null;
+  return laneFor(
+    { kind: legacy.engineKind, authKind: "config-dir", id: legacy.accountId, scope: "org" },
+    value.files,
+  );
+}
+
+/** Whether an account's material is one a revision can hold, without a
+ *  revision: the capture's own checks, on a stand-in destination. */
+function engineUnsupported(legacy: LegacyEngineParameter, value: LegacyEngineValue): string | null {
+  const lane = laneOf(legacy, value);
+  const location: CredentialLocation = {
+    class: "engine",
+    scope: "org",
+    ownerId: legacy.orgId,
+    engineKind: legacy.engineKind,
+    credentialId: "adoption-check",
+    revision: "adoption-check",
+  };
+  const flow = parseCaptureFlow({
+    credentialId: "adoption-check",
+    revision: "adoption-check",
+    leaseParameter: leaseParameterName(location),
+    brokerParameter: lane ? brokerParameterName(location) : null,
+    orgId: legacy.orgId,
+    userId: null,
+  });
+  if (!flow) return "destination";
+  const built = buildEngineCapture(flow, material(value));
+  return "error" in built ? (value.authKind === "env" ? "env_keys" : "files") : null;
+}
+
+function material(value: LegacyEngineValue): EngineCaptureMaterial {
+  return value.authKind === "env" ? { authKind: "env", env: value.env! } : { authKind: "config-dir", files: value.files! };
+}
+
+function legacyNameOfMarker(markerName: string, prefix: string): string | null {
+  const adoptedRoot = legacyAdoptedRoot(prefix);
+  if (!markerName.startsWith(`${adoptedRoot}/`)) return null;
+  const name = `${prefix}${markerName.slice(adoptedRoot.length)}`;
+  return parseLegacyParameterName(name, prefix) ? name : null;
+}
+
+function parseMarker(value: string): { credentialId: string; revision: string } | null {
+  const o = jsonObject(value);
+  return typeof o?.credentialId === "string" && typeof o.revision === "string"
+    ? { credentialId: o.credentialId, revision: o.revision }
+    : null;
+}
+
+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]);
+}
+
+const unavailable = (error: unknown): { outcome: "unavailable"; error: string } => ({
+  outcome: "unavailable",
+  error: typeof error === "string" ? error : brokerErrorKind(error),
+});
+
+/**
+ * One page of today's workspace AI accounts, metadata only (never a value),
+ * in parameter-name order after `after`: each with its lane, the revision it
+ * was adopted into once it was, and why it cannot be adopted as it is.
+ */
+export async function listLegacyAccounts(store: AdoptionStore, after: string | null = null): Promise<BrokerLegacyListAnswer> {
+  try {
+    const [values, markers] = await Promise.all([store.listLegacy(), store.listMarkers()]);
+    const adopted = new Map<string, { credentialId: string; revision: string }>();
+    for (const marker of markers) {
+      const name = legacyNameOfMarker(marker.name, store.legacyPrefix);
+      const named = parseMarker(marker.value);
+      if (name && named) adopted.set(name, named);
+    }
+    const ordered = values
+      .filter(
+        ({ name }) =>
+          parseLegacyEngineParameterName(name, store.legacyPrefix) !== null && (after === null || name > after),
+      )
+      .sort((a, b) => (a.name < b.name ? -1 : a.name > b.name ? 1 : 0));
+    const accounts: BrokerLegacyAccount[] = [];
+    let next: string | null = null;
+    for (const { name, value } of ordered) {
+      if (accounts.length >= MAX_LEGACY_ACCOUNTS) {
+        next = accounts.at(-1)!.parameter;
+        break;
+      }
+      const legacy = parseLegacyEngineParameterName(name, store.legacyPrefix)!;
+      const parsed = parseLegacyEngineValue(value, legacy);
+      accounts.push({
+        parameter: name,
+        orgId: legacy.orgId,
+        engineKind: legacy.engineKind,
+        accountId: legacy.accountId,
+        label: parsed ? parsed.label.slice(0, 256) : legacy.accountId,
+        authKind: parsed?.authKind ?? "config-dir",
+        position: parsed?.position ?? 0,
+        lane: parsed ? (laneOf(legacy, parsed)?.id ?? null) : null,
+        adopted: adopted.get(name) ?? null,
+        unsupported: parsed ? engineUnsupported(legacy, parsed) : "value",
+      });
+    }
+    return { outcome: "listed", accounts, next };
+  } catch (error) {
+    return unavailable(error);
+  }
+}
+
+/** Adopt one parameter as the revision the cloud named. */
+export async function adoptLegacyAccount(request: BrokerAdoptRequest, store: AdoptionStore): Promise<BrokerAdoptAnswer> {
+  const legacy = parseLegacyParameterName(request.legacyParameter, store.legacyPrefix)!;
+  let releaseCredential: (() => Promise<void>) | null = null;
+  let releaseScan: (() => Promise<void>) | null = null;
+  try {
+    releaseCredential = await store.lock(request.credentialId);
+    if (!releaseCredential) return { outcome: "busy" };
+    // No scan run is refreshing today's parameter while it is copied.
+    releaseScan = await store.scanLock();
+    if (!releaseScan) return { outcome: "busy" };
+
+    const current = await store.get(request.legacyParameter);
+    if (!current) return { outcome: "absent" };
+    const flow = parseCaptureFlow({
+      credentialId: request.credentialId,
+      revision: request.revision,
+      leaseParameter: request.leaseParameter,
+      brokerParameter: request.brokerParameter,
+      orgId: request.orgId,
+      userId: request.userId,
+    });
+    if (!flow) return { outcome: "unsupported", reason: "destination" };
+
+    // Built and checked before the marker: what cannot be adopted keeps its
+    // refresher.
+    const value = parseLegacyEngineValue(current.value, legacy);
+    if (!value) return { outcome: "unsupported", reason: "value" };
+    const parts = buildEngineCapture(flow, material(value));
+    if ("error" in parts) return { outcome: "unsupported", reason: value.authKind === "env" ? "env_keys" : "files" };
+    // Today's parameter keeps the brokered copy only.
+    let rewrite: string | null = null;
+    const lane = laneOf(legacy, value);
+    if (lane && parts.broker !== null && value.files) {
+      const brokered = lane.brokeredCopy(value.files);
+      if (!sameFiles(brokered, value.files)) rewrite = JSON.stringify({ ...value.raw, files: brokered });
+    }
+
+    // 1. The marker: from here nothing refreshes today's parameter.
+    const markerName = legacyAdoptionMarkerName(request.legacyParameter, store.legacyPrefix);
+    const marker = await store.get(markerName);
+    if (marker) {
+      const named = parseMarker(marker.value);
+      if (!named || named.credentialId !== request.credentialId || named.revision !== request.revision) {
+        return unavailable("adopted into another revision");
+      }
+    } else {
+      await store.create(
+        markerName,
+        JSON.stringify({ credentialId: request.credentialId, revision: request.revision }),
+        { uai_credential: request.credentialId, uai_revision: request.revision, uai_part: "adopted" },
+      );
+    }
+
+    // 2. Both parts, broker-only first, as a capture writes them. A retry
+    // matches its own creates and never rebuilds the broker part from what
+    // step 3 left in today's parameter.
+    if (parts.broker !== null) await store.create(request.brokerParameter!, parts.broker, captureTags(flow, "broker"));
+    await store.create(request.leaseParameter, parts.lease, captureTags(flow, "lease"));
+
+    // 3. Today's parameter keeps what a task may hold only.
+    if (rewrite !== null) await store.overwrite(request.legacyParameter, rewrite);
+
+    const lease = await store.get(request.leaseParameter);
+    if (!lease) return unavailable("the lease-readable part is missing");
+    return { outcome: "adopted", version: lease.version, expiresAt: parts.expiresAt };
+  } catch (error) {
+    return unavailable(error);
+  } finally {
+    if (releaseScan) await releaseScan();
+    if (releaseCredential) await releaseCredential();
+  }
+}
+
+/** Whether today's parameter of an adopted revision now agrees with it. */
+export type LegacySync = "synced" | "behind";
+
+/**
+ * Until the fleet audit: keep today's parameter of an adopted AI account
+ * agreeing with its revision, the brokered copy only, for the runners still
+ * reading it. Never creates it, and never fails the refresh it rides on: it
+ * says whether the write landed, and the cloud asks again when it did not.
+ */
+export async function syncLegacyCopy(
+  store: Pick<RevisionStore, "get" | "overwrite">,
+  legacyParameter: string,
+  lane: TokenLane,
+  files: CredentialFiles,
+  legacyPrefix: string = DEFAULT_CREDENTIAL_STORE_PREFIX,
+): Promise<LegacySync> {
+  const legacy = parseLegacyEngineParameterName(legacyParameter, legacyPrefix);
+  if (!legacy) return "synced";
+  try {
+    const current: StoredParameter | null = await store.get(legacyParameter);
+    if (!current) return "synced";
+    const value = parseLegacyEngineValue(current.value, legacy);
+    if (!value || value.authKind !== "config-dir" || !value.files) return "synced";
+    const brokered = lane.brokeredCopy(files);
+    if (sameFiles(brokered, value.files)) return "synced";
+    await store.overwrite(legacyParameter, JSON.stringify({ ...value.raw, files: brokered }));
+    return "synced";
+  } catch (error) {
+    console.warn(`[token-broker] today's copy of ${legacyParameter} was not updated (${brokerErrorKind(error)})`);
+    return "behind";
+  }
+}
+

package/lib/token-broker-lambda.ts

@@ -5,6 +5,14 @@
  * (parameter names to refresh now). Returns the cycle summary; logs one
  * audit line per decision and never a value.
  *
+ * ADR-154 §7: the cloud also invokes it with one revision to refresh
+ * (`{ v: 1, credentialId, revision, brokerParameter, leaseParameter, force }`),
+ * which never scans and never takes the cycle's lock: one lock per credential.
+ * Owner answer 3: `{ v: 1, op: "legacy.list", after }` lists a page of
+ * today's workspace accounts (metadata only) and `{ v: 1, op:
+ * "legacy.adopt", … }` adopts one account into a revision;
+ * an adoption takes its credential's lock and the cycle's.
+ *
  * Lives in lib/ (the package import boundary) and is bundled by
  * host-agent/scripts/token-broker/build.sh (esbuild, AWS SDK left external —
  * the Node runtime provides it). Nothing in the host imports it.
@@ -13,7 +21,16 @@
 import { SSMClient } from "@aws-sdk/client-ssm";
 
 import { runBrokerCycle } from "./token-broker";
-import { acquireBrokerLock, SsmBrokerStorage } from "./token-broker-ssm";
+import {
+  DEFAULT_CREDENTIAL_STORE_PREFIX,
+  parseBrokerAdoptRequest,
+  parseBrokerLegacyListRequest,
+  parseBrokerRevisionRequest,
+  parseCredentialStorePrefix,
+} from "../src/cloud-credentials";
+import { adoptLegacyAccount, listLegacyAccounts } from "./token-broker-adopt";
+import { refreshRevision } from "./token-broker-revision";
+import { acquireBrokerLock, SsmAdoptionStore, SsmBrokerStorage, SsmRevisionStore } from "./token-broker-ssm";
 
 interface BrokerEvent {
   names?: unknown;
@@ -22,8 +39,45 @@ interface BrokerEvent {
 const MAX_FORCED_NAMES = 50;
 
 export async function handler(event: BrokerEvent = {}): Promise<Record<string, unknown>> {
-  const prefix = process.env.UAI_CREDENTIAL_STORE_PREFIX || "/uai";
+  const configuredPrefix = process.env.UAI_CREDENTIAL_STORE_PREFIX || DEFAULT_CREDENTIAL_STORE_PREFIX;
+  const prefix = parseCredentialStorePrefix(configuredPrefix);
+  if (!prefix) throw new Error("invalid UAI_CREDENTIAL_STORE_PREFIX");
   const kmsKeyId = process.env.UAI_CREDENTIAL_STORE_KMS_KEY || "alias/uai-credentials";
+  // ADR-154 §7: the cloud names one revision; the broker refreshes exactly
+  // that one and answers for it. Anything shaped like a request but not one
+  // is refused, never read as a scan.
+  if ((event as { v?: unknown }).v !== undefined) {
+    // Owner answer 3: today's workspace accounts, listed (metadata only)
+    // and adopted one at a time into the revision the cloud names.
+    const op = (event as { op?: unknown }).op;
+    if (op === "legacy.list") {
+      const list = parseBrokerLegacyListRequest(event, prefix);
+      if (!list) return { outcome: "unavailable", error: "malformed request" };
+      const listed = await listLegacyAccounts(new SsmAdoptionStore({ client: new SSMClient({}), kmsKeyId, prefix }), list.after);
+      console.log(
+        `[token-broker] legacy.list ${JSON.stringify({ outcome: listed.outcome, accounts: listed.outcome === "listed" ? listed.accounts.length : 0 })}`,
+      );
+      return listed;
+    }
+    if (op === "legacy.adopt") {
+      const adopt = parseBrokerAdoptRequest(event, prefix);
+      if (!adopt) return { outcome: "unavailable", error: "malformed request" };
+      const adopted = await adoptLegacyAccount(adopt, new SsmAdoptionStore({ client: new SSMClient({}), kmsKeyId, prefix }));
+      console.log(
+        `[token-broker] legacy.adopt ${JSON.stringify({ credentialId: adopt.credentialId, revision: adopt.revision, outcome: adopted.outcome })}`,
+      );
+      return adopted;
+    }
+    const request = parseBrokerRevisionRequest(event, prefix);
+    if (!request) return { outcome: "unavailable", error: "malformed request" };
+    const answer = await refreshRevision(request, new SsmRevisionStore({ client: new SSMClient({}), kmsKeyId }), {
+      legacyPrefix: prefix,
+    });
+    console.log(
+      `[token-broker] revision ${JSON.stringify({ credentialId: request.credentialId, revision: request.revision, force: request.force, outcome: answer.outcome })}`,
+    );
+    return answer;
+  }
   const names = Array.isArray(event.names)
     ? event.names.filter((n): n is string => typeof n === "string" && n.startsWith(`${prefix}/`)).slice(0, MAX_FORCED_NAMES)
     : [];

package/lib/token-broker-local.ts

@@ -11,6 +11,11 @@
  * The loop runs ONLY on owned hosts (`configuredCredentialStore() === null`).
  * A Uai-operated host must not refresh — the broker Lambda does — it picks
  * up refreshed parameters instead (credential-store-pickup.ts).
+ *
+ * ADR-154 §7: and only while the runner is on its own credentials. No runner
+ * refreshes once it leaves them: the loop stops, a cycle under way ends, and
+ * no reactive request refreshes or invokes the broker Lambda (a leased
+ * account's refresh goes through the cloud).
  */
 
 import {
@@ -27,6 +32,7 @@ import { homedir, tmpdir } from "node:os";
 import { join } from "node:path";
 import { spawn } from "node:child_process";
 
+import { onOwnCredentials, whileOnOwnCredentials } from "./cloud-credential-host";
 import { configuredCredentialStore, credentialStorePrefix } from "./credential-store";
 import { getDb, schema } from "./db";
 import {
@@ -212,6 +218,8 @@ function log(line: string): void {
  */
 export async function runLocalBrokerCycle(opts: { force?: ReadonlySet<string>; seams?: LocalBrokerSeams } = {}): Promise<void> {
   const work = async (): Promise<void> => {
+    // A cycle queued behind the one leaving waits for refreshes nothing.
+    if (!onOwnCredentials()) return;
     const storage = new LocalBrokerStorage(opts.seams);
     const refreshedNames = new Set<string>();
     const summary = await runBrokerCycle({
@@ -244,9 +252,10 @@ export async function runLocalBrokerCycle(opts: { force?: ReadonlySet<string>; s
   }
 }
 
-/** Owned hosts only: refresh now, then hourly. Idempotent. */
+/** Owned hosts on their own credentials only: refresh now, then hourly.
+ *  Idempotent. */
 export function startTokenBrokerLoop(): void {
-  if (timer || !ownedHostRefreshes()) return;
+  if (timer || !ownedHostRefreshes() || !onOwnCredentials()) return;
   void runLocalBrokerCycle().catch((err) => log(`cycle failed: ${err instanceof Error ? err.message : String(err)}`));
   timer = setInterval(() => {
     void runLocalBrokerCycle().catch((err) => log(`cycle failed: ${err instanceof Error ? err.message : String(err)}`));
@@ -254,9 +263,12 @@ export function startTokenBrokerLoop(): void {
   timer.unref?.();
 }
 
-export function stopTokenBrokerLoop(): void {
+/** Stop the loop at once (ADR-154 §7: the runner leaves its own
+ *  credentials); resolves once a cycle under way has ended. */
+export async function stopTokenBrokerLoop(): Promise<void> {
   if (timer) clearInterval(timer);
   timer = null;
+  while (running !== null) await running.catch(() => undefined);
 }
 
 /** Whether an engine's error output means its access token is dead now. */
@@ -275,6 +287,18 @@ export async function requestCredentialRefresh(input: {
   kind: string;
   accountId: string | null;
   scope: string;
+}): Promise<boolean> {
+  // ADR-154 §7: off its own credentials the runner refreshes nothing and
+  // invokes no broker; the cloud refreshes, through the task's lease. A
+  // request on them is counted whole (its cycle or its Lambda invocation,
+  // and the pickup after): leaving them waits for it.
+  return (await whileOnOwnCredentials(() => refreshOnOwnCredentials(input))) ?? false;
+}
+
+async function refreshOnOwnCredentials(input: {
+  kind: string;
+  accountId: string | null;
+  scope: string;
 }): Promise<boolean> {
   const key = `${input.kind}:${input.accountId ?? "default"}`;
   const now = Date.now();
@@ -306,7 +330,7 @@ export function setPickupNow(fn: (() => Promise<void>) | null): void {
 }
 
 /** `aws lambda invoke` on the broker function named by UAI_TOKEN_BROKER_FUNCTION. */
-export async function invokeBrokerLambda(names: string[]): Promise<boolean> {
+async function invokeBrokerLambda(names: string[]): Promise<boolean> {
   const fn = env.UAI_TOKEN_BROKER_FUNCTION;
   const region = env.UAI_AWS_REGION;
   if (!fn || !region) {