@runuai/host 0.9.113 → 0.9.114

package/lib/token-broker-local.ts

@@ -0,0 +1,368 @@
+/**
+ * Local storage adapter + the owned-host refresh loop (ADR-135).
+ *
+ * On a host the operator owns, the credentials live on disk: the default
+ * slots (`~/.codex`, `~/.grok`) and every config-dir engine account
+ * (`~/.<kind>-acct-<id>`). This adapter lists them for the broker core,
+ * writes refreshed material back atomically, and keeps the broker's meta in
+ * a sidecar file (`uai-broker.json`) that never travels to a task and never
+ * mirrors to the store.
+ *
+ * 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).
+ */
+
+import {
+  existsSync,
+  mkdtempSync,
+  readdirSync,
+  readFileSync,
+  renameSync,
+  rmSync,
+  statSync,
+  writeFileSync,
+} from "node:fs";
+import { homedir, tmpdir } from "node:os";
+import { join } from "node:path";
+import { spawn } from "node:child_process";
+
+import { configuredCredentialStore, credentialStorePrefix } from "./credential-store";
+import { getDb, schema } from "./db";
+import {
+  BROKER_SIDECAR_FILE,
+  engineAccountParameterName,
+  mirrorEngineAccount,
+} from "./engine-accounts";
+import { runBrokerCycle, type BrokerCredential, type BrokerStorage } from "./token-broker";
+import { lanesForKind } from "./token-lanes/registry";
+import { refreshMetaFrom, type CredentialFiles, type RefreshMeta } from "./token-lanes/types";
+import { env } from "./env";
+
+/** Hourly: Grok tokens live six hours (observed 2026-09-09); the lanes'
+ *  leads decide what is actually refreshed, so most runs do nothing. */
+const CYCLE_INTERVAL_MS = 60 * 60 * 1000;
+/** Kinds whose default slot lives in a home dir the broker may rewrite. */
+const DEFAULT_SLOTS: Record<string, (home: string) => string> = {
+  codex: (home) => join(home, ".codex"),
+  grok: (home) => join(home, ".grok"),
+};
+
+export interface LocalBrokerSeams {
+  ownerHome?: () => string;
+  /** Rows to consider (tests inject); defaults to every stored account. */
+  accounts?: () => Array<{ id: string; kind: string; scope: string; configDir: string | null; authKind: string }>;
+  now?: () => number;
+}
+
+function ownerHome(seams: LocalBrokerSeams): string {
+  return (seams.ownerHome ?? (() => process.env.UAI_OWNER_HOME?.trim() || homedir()))();
+}
+
+function readDirFiles(dir: string): CredentialFiles | null {
+  let names: string[];
+  try {
+    names = readdirSync(dir);
+  } catch {
+    return null;
+  }
+  const files: CredentialFiles = {};
+  for (const name of names) {
+    if (name === BROKER_SIDECAR_FILE || name !== "auth.json") continue;
+    try {
+      const path = join(dir, name);
+      if (!statSync(path).isFile()) continue;
+      files[name] = readFileSync(path, "utf8");
+    } catch {
+      /* unreadable file: the lane will not match */
+    }
+  }
+  return files;
+}
+
+function readSidecar(dir: string): RefreshMeta | null {
+  try {
+    return refreshMetaFrom(JSON.parse(readFileSync(join(dir, BROKER_SIDECAR_FILE), "utf8")));
+  } catch {
+    return null;
+  }
+}
+
+/** temp + rename in the same dir, 0600: a crash never leaves a torn file. */
+export function writeFileAtomic(path: string, content: string): void {
+  const tmp = `${path}.${process.pid}.${Date.now()}.tmp`;
+  writeFileSync(tmp, content, { mode: 0o600 });
+  renameSync(tmp, path);
+}
+
+/** A local credential name the reactive path and tests can address. */
+export function localCredentialName(kind: string, accountId: string | null): string {
+  return accountId ? `account:${kind}:${accountId}` : `default:${kind}`;
+}
+
+export class LocalBrokerStorage implements BrokerStorage {
+  private readonly dirs = new Map<string, { dir: string; accountId: string | null; kind: string }>();
+
+  constructor(private readonly seams: LocalBrokerSeams = {}) {}
+
+  async list(): Promise<BrokerCredential[]> {
+    const out: BrokerCredential[] = [];
+    const home = ownerHome(this.seams);
+    for (const [kind, dirOf] of Object.entries(DEFAULT_SLOTS)) {
+      const dir = dirOf(home);
+      if (!existsSync(join(dir, "auth.json"))) continue;
+      const files = readDirFiles(dir);
+      if (!files) continue;
+      const name = localCredentialName(kind, null);
+      this.dirs.set(name, { dir, accountId: null, kind });
+      out.push({
+        name,
+        header: { kind, authKind: "config-dir", id: "default", scope: "host" },
+        files,
+        refresh: readSidecar(dir),
+      });
+    }
+    const rows =
+      this.seams.accounts?.() ??
+      getDb()
+        .select()
+        .from(schema.engineAccounts)
+        .all()
+        .map((r) => ({
+          id: r.id,
+          kind: r.kind,
+          scope: r.scope ?? "host",
+          configDir: r.configDir,
+          authKind: r.authKind,
+        }));
+    for (const row of rows) {
+      if (row.authKind !== "config-dir" || !row.configDir) continue;
+      const files = readDirFiles(row.configDir);
+      if (!files) continue;
+      const name = localCredentialName(row.kind, row.id);
+      this.dirs.set(name, { dir: row.configDir, accountId: row.id, kind: row.kind });
+      out.push({
+        name,
+        header: { kind: row.kind, authKind: "config-dir", id: row.id, scope: row.scope },
+        files,
+        refresh: readSidecar(row.configDir),
+      });
+    }
+    return out;
+  }
+
+  async save(credential: BrokerCredential, files: CredentialFiles | null, refresh: RefreshMeta): Promise<void> {
+    const target = this.dirs.get(credential.name);
+    if (!target) throw new Error(`unknown local credential ${credential.name}`);
+    if (files) {
+      for (const [name, content] of Object.entries(files)) {
+        if (name.includes("/") || name === BROKER_SIDECAR_FILE) continue;
+        writeFileAtomic(join(target.dir, name), content);
+      }
+    }
+    writeFileAtomic(join(target.dir, BROKER_SIDECAR_FILE), `${JSON.stringify(refresh, null, 2)}\n`);
+    // Mirrors are no-ops without a store; on a Uai-operated host this
+    // adapter does not run at all, so this only matters in tests.
+    if (target.accountId && files) mirrorEngineAccount(target.accountId);
+  }
+
+  kindsRefreshed(names: Iterable<string>): Set<string> {
+    const kinds = new Set<string>();
+    for (const name of names) {
+      const t = this.dirs.get(name);
+      if (t) kinds.add(t.kind);
+    }
+    return kinds;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Loop + reactive entry points.
+// ---------------------------------------------------------------------------
+
+export interface TokenBrokerHooks {
+  /** Re-push brokered copies of `kind` into running tasks. */
+  repush: (kind: string) => Promise<void>;
+  log?: (line: string) => void;
+}
+
+let hooks: TokenBrokerHooks | null = null;
+let timer: ReturnType<typeof setInterval> | null = null;
+let running: Promise<void> | null = null;
+const reactiveLastAt = new Map<string, number>();
+const REACTIVE_MIN_INTERVAL_MS = 10 * 60_000;
+
+export function configureTokenBroker(next: TokenBrokerHooks): void {
+  hooks = next;
+}
+
+/** Whether this host is the refresher for its own credentials. */
+export function ownedHostRefreshes(): boolean {
+  return configuredCredentialStore() === null;
+}
+
+function log(line: string): void {
+  (hooks?.log ?? ((l: string) => console.log(`[token-broker] ${l}`)))(line);
+}
+
+/**
+ * One local cycle: refresh what is due (or forced) and re-push the kinds
+ * that changed into running tasks. Serialized — a reactive request during a
+ * scheduled run waits for it.
+ */
+export async function runLocalBrokerCycle(opts: { force?: ReadonlySet<string>; seams?: LocalBrokerSeams } = {}): Promise<void> {
+  const work = async (): Promise<void> => {
+    const storage = new LocalBrokerStorage(opts.seams);
+    const refreshedNames = new Set<string>();
+    const summary = await runBrokerCycle({
+      storage,
+      force: opts.force,
+      now: opts.seams?.now,
+      log: (line) => {
+        if (line.startsWith("refreshed ")) refreshedNames.add(line.split(" ")[1]!);
+        log(line);
+      },
+    });
+    if (summary.checked > 0) {
+      log(`local cycle: ${summary.checked} checked, ${summary.due} due, ${summary.refreshed} refreshed, ${summary.failed} failed, ${summary.terminal} need sign-in`);
+    }
+    for (const kind of storage.kindsRefreshed(refreshedNames)) {
+      try {
+        await hooks?.repush(kind);
+      } catch (err) {
+        log(`re-push of ${kind} copies failed: ${err instanceof Error ? err.message : String(err)}`);
+      }
+    }
+  };
+  const previous = running ?? Promise.resolve();
+  const current = previous.then(work, work);
+  running = current;
+  try {
+    await current;
+  } finally {
+    if (running === current) running = null;
+  }
+}
+
+/** Owned hosts only: refresh now, then hourly. Idempotent. */
+export function startTokenBrokerLoop(): void {
+  if (timer || !ownedHostRefreshes()) 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)}`));
+  }, CYCLE_INTERVAL_MS);
+  timer.unref?.();
+}
+
+export function stopTokenBrokerLoop(): void {
+  if (timer) clearInterval(timer);
+  timer = null;
+}
+
+/** Whether an engine's error output means its access token is dead now. */
+export function isCredentialInvalidSignal(kind: string, message: string): boolean {
+  return lanesForKind(kind).some((lane) => lane.reactiveSignals.some((re) => re.test(message)));
+}
+
+/**
+ * Reactive path: a task reported a dead token for `kind` (account `accountId`,
+ * or the default slot). Owned host: one forced local refresh + re-push.
+ * Uai-operated host: invoke the broker Lambda for that parameter, then pick
+ * up the new version. Rate-limited per credential; returns whether a refresh
+ * was requested.
+ */
+export async function requestCredentialRefresh(input: {
+  kind: string;
+  accountId: string | null;
+  scope: string;
+}): Promise<boolean> {
+  const key = `${input.kind}:${input.accountId ?? "default"}`;
+  const now = Date.now();
+  const last = reactiveLastAt.get(key) ?? 0;
+  if (now - last < REACTIVE_MIN_INTERVAL_MS) return false;
+  reactiveLastAt.set(key, now);
+  if (ownedHostRefreshes()) {
+    await runLocalBrokerCycle({ force: new Set([localCredentialName(input.kind, input.accountId)]) });
+    return true;
+  }
+  if (!input.accountId) return false;
+  const name = engineAccountParameterName({ id: input.accountId, kind: input.kind, scope: input.scope });
+  if (!name) return false;
+  const invoked = await invokeBrokerLambda([name]);
+  if (invoked) {
+    try {
+      await pickupNow?.();
+    } catch (err) {
+      log(`pickup after reactive refresh failed: ${err instanceof Error ? err.message : String(err)}`);
+    }
+  }
+  return invoked;
+}
+
+let pickupNow: (() => Promise<void>) | null = null;
+/** credential-store-pickup.ts registers its immediate poll here. */
+export function setPickupNow(fn: (() => Promise<void>) | null): void {
+  pickupNow = fn;
+}
+
+/** `aws lambda invoke` on the broker function named by UAI_TOKEN_BROKER_FUNCTION. */
+export async function invokeBrokerLambda(names: string[]): Promise<boolean> {
+  const fn = env.UAI_TOKEN_BROKER_FUNCTION;
+  const region = env.UAI_AWS_REGION;
+  if (!fn || !region) {
+    log("reactive refresh requested but UAI_TOKEN_BROKER_FUNCTION/UAI_AWS_REGION are not set");
+    return false;
+  }
+  const prefix = credentialStorePrefix();
+  const safe = names.filter((n) => n.startsWith(`${prefix}/`));
+  if (safe.length === 0) return false;
+  const outDir = mkdtempSync(join(tmpdir(), "uai-broker-invoke-"));
+  const outFile = join(outDir, "out.json");
+  try {
+    const result = await new Promise<{ status: number | null; stderr: string }>((resolve) => {
+      const child = spawn(
+        "aws",
+        [
+          "--region",
+          region,
+          "--output",
+          "json",
+          "lambda",
+          "invoke",
+          "--function-name",
+          fn,
+          "--cli-binary-format",
+          "raw-in-base64-out",
+          "--payload",
+          JSON.stringify({ names: safe }),
+          outFile,
+        ],
+        { stdio: ["ignore", "ignore", "pipe"] },
+      );
+      let stderr = "";
+      const timer = setTimeout(() => child.kill("SIGKILL"), 150_000);
+      timer.unref?.();
+      child.stderr?.setEncoding("utf8");
+      child.stderr?.on("data", (c: string) => (stderr += c));
+      child.once("error", () => resolve({ status: null, stderr }));
+      child.once("close", (code) => {
+        clearTimeout(timer);
+        resolve({ status: code, stderr });
+      });
+    });
+    if (result.status !== 0) {
+      log(`lambda invoke failed (exit ${result.status ?? "none"}): ${result.stderr.split("\n")[0] ?? ""}`);
+      return false;
+    }
+    let summary = "";
+    try {
+      summary = readFileSync(outFile, "utf8").trim().slice(0, 300);
+    } catch {
+      /* no body */
+    }
+    log(`lambda refreshed on request: ${summary}`);
+    return true;
+  } finally {
+    rmSync(outDir, { recursive: true, force: true });
+  }
+}

package/lib/token-broker-ssm.ts

@@ -0,0 +1,185 @@
+/**
+ * Parameter Store adapter for the broker core (ADR-135) — the storage the
+ * broker Lambda runs against. Reads the ADR-134 engine-account parameters
+ * under `<prefix>/orgs` and `<prefix>/hosts`, hands each to the core as a
+ * credential, and writes refreshed material + the `refresh` block back with
+ * `Overwrite`. Only the Lambda imports this module: it depends on the AWS
+ * SDK the Lambda runtime provides, which the host package does not ship.
+ */
+
+import {
+  DeleteParameterCommand,
+  GetParameterCommand,
+  GetParametersByPathCommand,
+  PutParameterCommand,
+  SSMClient,
+} from "@aws-sdk/client-ssm";
+
+import type { BrokerCredential, BrokerStorage } from "./token-broker";
+import { refreshMetaFrom, type CredentialFiles, type RefreshMeta } from "./token-lanes/types";
+
+export interface SsmBrokerStorageOptions {
+  prefix: string;
+  kmsKeyId: string;
+  client?: SSMClient;
+  region?: string;
+}
+
+interface StoredValue {
+  raw: Record<string, unknown>;
+  version: number | undefined;
+}
+
+/** `/uai/orgs/<o>/engines/<kind>/<id>` and `/uai/hosts/<h>/engines/<kind>/<id>`. */
+export function engineParameterHeader(
+  name: string,
+  prefix: string,
+): { kind: string; id: string; scope: string } | null {
+  if (!name.startsWith(`${prefix}/`)) return null;
+  const parts = name.slice(prefix.length + 1).split("/");
+  if (parts.length !== 5 || parts[2] !== "engines") return null;
+  const [owner, ownerId, , kind, id] = parts as [string, string, string, string, string];
+  if (owner === "orgs") return { kind, id, scope: `org:${ownerId}` };
+  if (owner === "hosts") return { kind, id, scope: "host" };
+  return null;
+}
+
+export class SsmBrokerStorage implements BrokerStorage {
+  private readonly client: SSMClient;
+  private readonly values = new Map<string, StoredValue>();
+
+  constructor(private readonly opts: SsmBrokerStorageOptions) {
+    this.client = opts.client ?? new SSMClient(opts.region ? { region: opts.region } : {});
+  }
+
+  async list(): Promise<BrokerCredential[]> {
+    const out: BrokerCredential[] = [];
+    for (const path of [`${this.opts.prefix}/orgs`, `${this.opts.prefix}/hosts`]) {
+      let nextToken: string | undefined;
+      do {
+        const page = await this.client.send(
+          new GetParametersByPathCommand({
+            Path: path,
+            Recursive: true,
+            WithDecryption: true,
+            NextToken: nextToken,
+          }),
+        );
+        for (const p of page.Parameters ?? []) {
+          if (!p.Name || typeof p.Value !== "string") continue;
+          const header = engineParameterHeader(p.Name, this.opts.prefix);
+          if (!header) continue;
+          let raw: unknown;
+          try {
+            raw = JSON.parse(p.Value);
+          } catch {
+            continue;
+          }
+          if (!raw || typeof raw !== "object" || Array.isArray(raw)) continue;
+          const value = raw as Record<string, unknown>;
+          if (value.v !== 1 || value.authKind !== "config-dir") continue;
+          const files = value.files;
+          if (!files || typeof files !== "object" || Array.isArray(files)) continue;
+          const materials: CredentialFiles = {};
+          for (const [k, v] of Object.entries(files as Record<string, unknown>)) {
+            if (typeof v === "string") materials[k] = v;
+          }
+          this.values.set(p.Name, { raw: value, version: p.Version });
+          out.push({
+            name: p.Name,
+            header: { kind: header.kind, authKind: "config-dir", id: header.id, scope: header.scope },
+            files: materials,
+            refresh: refreshMetaFrom(value.refresh),
+          });
+        }
+        nextToken = page.NextToken;
+      } while (nextToken);
+    }
+    return out;
+  }
+
+  async save(
+    credential: BrokerCredential,
+    files: CredentialFiles | null,
+    refresh: RefreshMeta,
+  ): Promise<void> {
+    const stored = this.values.get(credential.name);
+    const raw = stored?.raw ?? { v: 1 };
+    const next = { ...raw, refresh, ...(files ? { files } : {}) };
+    await this.client.send(
+      new PutParameterCommand({
+        Name: credential.name,
+        Value: JSON.stringify(next),
+        Type: "SecureString",
+        KeyId: this.opts.kmsKeyId,
+        Tier: "Advanced",
+        Overwrite: true,
+      }),
+    );
+    this.values.set(credential.name, { raw: next, version: undefined });
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Mutual exclusion between runs. Reserved concurrency 1 is the first line
+// (deploy-aws.sh sets it when the account's limit allows); this lock is the
+// one that always holds: a `PutParameter` WITHOUT `Overwrite` is atomic, so
+// two runs can never both create `<prefix>/locks/token-broker`. A holder
+// that died keeps the lock at most LOCK_STALE_MS; a later run takes it over.
+// ---------------------------------------------------------------------------
+
+export const LOCK_STALE_MS = 5 * 60_000;
+
+export function brokerLockName(prefix: string): string {
+  return `${prefix}/locks/token-broker`;
+}
+
+export interface BrokerLockOptions {
+  client: SSMClient;
+  prefix: string;
+  now?: () => number;
+  staleMs?: number;
+}
+
+/** Try to take the lock. Returns a release function, or null when another
+ *  run holds a fresh one. */
+export async function acquireBrokerLock(opts: BrokerLockOptions): Promise<(() => Promise<void>) | null> {
+  const name = brokerLockName(opts.prefix);
+  const now = opts.now ?? Date.now;
+  const staleMs = opts.staleMs ?? LOCK_STALE_MS;
+  const token = `${now()}:${Math.random().toString(36).slice(2)}`;
+  const tryCreate = async (overwrite: boolean): Promise<boolean> => {
+    try {
+      await opts.client.send(
+        new PutParameterCommand({ Name: name, Value: token, Type: "String", Overwrite: overwrite }),
+      );
+      return true;
+    } catch (err) {
+      if (!overwrite && (err as { name?: string }).name === "ParameterAlreadyExists") return false;
+      throw err;
+    }
+  };
+  if (!(await tryCreate(false))) {
+    // Someone holds it. Stale?
+    let heldAt = 0;
+    try {
+      const current = await opts.client.send(new GetParameterCommand({ Name: name }));
+      heldAt = Number(current.Parameter?.Value?.split(":")[0] ?? 0) || 0;
+    } catch {
+      heldAt = 0;
+    }
+    if (heldAt && now() - heldAt < staleMs) return null;
+    await tryCreate(true);
+  }
+  return async () => {
+    try {
+      // Only release our own token: a takeover after a stall must not be
+      // undone by the stalled run finishing late.
+      const current = await opts.client.send(new GetParameterCommand({ Name: name }));
+      if (current.Parameter?.Value !== token) return;
+      await opts.client.send(new DeleteParameterCommand({ Name: name }));
+    } catch {
+      /* the stale window covers a failed release */
+    }
+  };
+}