@runuai/host 0.9.112 → 0.9.114

package/lib/github-tokens.ts

@@ -9,6 +9,7 @@
 
 import { and, eq } from "drizzle-orm";
 
+import { queueStoreWrite, userParameterName } from "./credential-store";
 import { getDb, schema } from "./db";
 import { sealAesGcm, openAesGcm } from "./secrets";
 import { dockerCli } from "./docker-exec";
@@ -446,10 +447,13 @@ export async function onConnectSet(
 }
 
 function deleteStoredTokenRow(userId: string): void {
-  getDb()
+  const result = getDb()
     .delete(schema.githubTokens)
     .where(eq(schema.githubTokens.userId, userId))
     .run();
+  // A connect clears any prior grant first; with none stored there is no
+  // parameter to delete either.
+  if (result.changes > 0) deleteGithubTokenFromStore(userId);
 }
 
 function storeTokenRow(
@@ -484,6 +488,117 @@ function storeTokenRow(
       .onConflictDoUpdate({ target: schema.githubTokens.userId, set: fields })
       .run();
   });
+  mirrorGithubTokenToStore(userId);
+}
+
+// ---------------------------------------------------------------------------
+// Credential store mirror (ADR-134 slice 2): `<prefix>/users/<userId>/github`.
+// The generation rides along so a replacement host imports the row behind the
+// same durable fence (a later disconnect this host never saw still wins).
+// Owned hosts: no-ops.
+// ---------------------------------------------------------------------------
+
+export interface StoredGithubToken {
+  v: 1;
+  kind: "github";
+  userId: string;
+  installationId: number;
+  generation: number;
+  tokenKind: string;
+  token: string;
+  expiresAt: number | null;
+  updatedAt: number;
+}
+
+export function githubTokenParameterName(userId: string): string {
+  return userParameterName(userId, "github");
+}
+
+/** The row as it is mirrored (token decrypted). Never logged. */
+export function githubTokenStoreValue(userId: string): StoredGithubToken | null {
+  const row = getDb()
+    .select()
+    .from(schema.githubTokens)
+    .where(eq(schema.githubTokens.userId, userId))
+    .get();
+  if (!row) return null;
+  return {
+    v: 1,
+    kind: "github",
+    userId,
+    installationId: row.installationId,
+    generation: row.generation,
+    tokenKind: row.kind,
+    token: openAesGcm(
+      Buffer.from(row.refreshTokenCt as Uint8Array),
+      Buffer.from(row.refreshTokenNonce as Uint8Array),
+    ),
+    expiresAt: row.refreshTokenExpiresAt ?? null,
+    updatedAt: row.updatedAt,
+  };
+}
+
+export function setGithubTokenStoreRef(userId: string, storeRef: string): void {
+  getDb()
+    .update(schema.githubTokens)
+    .set({ storeRef })
+    .where(eq(schema.githubTokens.userId, userId))
+    .run();
+}
+
+export function listGithubTokenStoreRows(): Array<{ userId: string; storeRef: string | null }> {
+  return getDb()
+    .select({ userId: schema.githubTokens.userId, storeRef: schema.githubTokens.storeRef })
+    .from(schema.githubTokens)
+    .all();
+}
+
+export function mirrorGithubTokenToStore(userId: string): void {
+  queueStoreWrite(`mirror github token ${userId}`, async (store) => {
+    const value = githubTokenStoreValue(userId);
+    if (!value) return;
+    const name = githubTokenParameterName(userId);
+    await store.put(name, JSON.stringify(value));
+    setGithubTokenStoreRef(userId, name);
+  });
+}
+
+export function deleteGithubTokenFromStore(userId: string): void {
+  queueStoreWrite(`delete github token ${userId}`, (store) =>
+    store.delete(githubTokenParameterName(userId)),
+  );
+}
+
+/**
+ * Reconcile import: a token this host lacks. Claims the stored generation as
+ * a `set` first — refused (false) when this host's fence already moved past
+ * it, so a disconnect that outran the store's delete can never resurrect a
+ * token here.
+ */
+export function importGithubTokenFromStore(
+  value: StoredGithubToken,
+  storeRef: string,
+): boolean {
+  const claim = claimGithubCredentialGeneration(value.userId, value.generation, "set");
+  if (!claim) return false;
+  const sealed = sealAesGcm(value.token);
+  const fields = {
+    installationId: value.installationId,
+    generation: value.generation,
+    refreshTokenCt: sealed.ct,
+    refreshTokenNonce: sealed.nonce,
+    refreshTokenExpiresAt: value.expiresAt,
+    kind: value.tokenKind,
+    updatedAt: value.updatedAt,
+    storeRef,
+  };
+  getDb()
+    .insert(schema.githubTokens)
+    .values({ userId: value.userId, ...fields })
+    .onConflictDoUpdate({ target: schema.githubTokens.userId, set: fields })
+    .run();
+  notifyGithubChange();
+  return true;
 }
 
 async function bestEffortRevoke(
@@ -663,6 +778,7 @@ export function deleteToken(userId: string): void {
     .delete(schema.githubTokens)
     .where(eq(schema.githubTokens.userId, userId))
     .run();
+  deleteGithubTokenFromStore(userId);
   advanceGithubAuthGeneration(userId);
   notifyGithubChange();
 }

package/lib/mcp-connections.ts

@@ -14,6 +14,7 @@
 import { createHash, randomBytes } from "node:crypto";
 import { eq } from "drizzle-orm";
 
+import { queueStoreWrite, userParameterName } from "./credential-store";
 import { getDb, schema } from "./db";
 import { sealAesGcm, openAesGcm } from "./secrets";
 import type { McpOp } from "../src/protocol";
@@ -276,10 +277,7 @@ export async function handleMcpOp(op: McpOp): Promise<McpAck> {
     case "oauth.complete":
       return oauthComplete(op.connectionId, op.code);
     case "disconnect": {
-      getDb()
-        .delete(schema.mcpConnections)
-        .where(eq(schema.mcpConnections.id, op.connectionId))
-        .run();
+      deleteMcpConnectionLocal(op.connectionId);
       return { status: "disconnected" };
     }
   }
@@ -334,6 +332,7 @@ async function probe(op: Extract<McpOp, { kind: "probe" }>): Promise<McpAck> {
         },
       })
       .run();
+    mirrorMcpConnectionToStore(op.connectionId);
     return { status: "connected" };
   }
 
@@ -354,6 +353,7 @@ async function probe(op: Extract<McpOp, { kind: "probe" }>): Promise<McpAck> {
         set: { url: op.url, authKind: "none", status: "connected", updatedAt: now },
       })
       .run();
+    mirrorMcpConnectionToStore(op.connectionId);
     return { status: "connected" };
   }
   if (res.status !== 401 && res.status !== 403) {
@@ -414,6 +414,7 @@ async function probe(op: Extract<McpOp, { kind: "probe" }>): Promise<McpAck> {
     .values({ id: op.connectionId, ...fields })
     .onConflictDoUpdate({ target: schema.mcpConnections.id, set: fields })
     .run();
+  mirrorMcpConnectionToStore(op.connectionId);
   return { status: "auth_required", authorizeUrl: authorize.toString() };
 }
 
@@ -472,6 +473,7 @@ async function oauthComplete(connectionId: string, code: string): Promise<McpAck
     })
     .where(eq(schema.mcpConnections.id, connectionId))
     .run();
+  mirrorMcpConnectionToStore(connectionId);
   return {
     status: "connected",
     scopes: scopes ? scopes.split(" ").filter(Boolean) : undefined,
@@ -611,5 +613,137 @@ async function refreshTokens(
     .set({ secretEnc: pack(JSON.stringify(next)), updatedAt: Date.now() })
     .where(eq(schema.mcpConnections.id, row.id))
     .run();
+  mirrorMcpConnectionToStore(row.id);
   return next;
 }
+
+// ---------------------------------------------------------------------------
+// Credential store mirror (ADR-134 slice 2):
+// `<prefix>/users/<userId>/mcp/<connectionId>`. The connection id is the
+// cloud's, which is org-bound there; on the host the owner is the user.
+// Owned hosts: no-ops.
+// ---------------------------------------------------------------------------
+
+export interface StoredMcpConnection {
+  v: 1;
+  kind: "mcp";
+  id: string;
+  userId: string;
+  url: string;
+  authKind: string;
+  tokenEndpoint: string | null;
+  redirectUri: string | null;
+  clientId: string | null;
+  clientSecret: string | null;
+  pkceVerifier: string | null;
+  /** The sealed JSON, decrypted: header or OAuth tokens. */
+  secret: string | null;
+  status: string;
+  scopes: string | null;
+  updatedAt: number;
+}
+
+export function mcpConnectionParameterName(userId: string, id: string): string {
+  return userParameterName(userId, "mcp", id);
+}
+
+/** The row as it is mirrored (secrets decrypted). Never logged. */
+export function mcpConnectionStoreValue(id: string): StoredMcpConnection | null {
+  const row = getDb()
+    .select()
+    .from(schema.mcpConnections)
+    .where(eq(schema.mcpConnections.id, id))
+    .get();
+  if (!row) return null;
+  return {
+    v: 1,
+    kind: "mcp",
+    id: row.id,
+    userId: row.userId,
+    url: row.url,
+    authKind: row.authKind,
+    tokenEndpoint: row.tokenEndpoint,
+    redirectUri: row.redirectUri,
+    clientId: row.clientId,
+    clientSecret: row.clientSecretEnc ? unpack(row.clientSecretEnc) : null,
+    pkceVerifier: row.pkceVerifierEnc ? unpack(row.pkceVerifierEnc) : null,
+    secret: row.secretEnc ? unpack(row.secretEnc) : null,
+    status: row.status,
+    scopes: row.scopes,
+    updatedAt: row.updatedAt,
+  };
+}
+
+export function setMcpConnectionStoreRef(id: string, storeRef: string): void {
+  getDb()
+    .update(schema.mcpConnections)
+    .set({ storeRef })
+    .where(eq(schema.mcpConnections.id, id))
+    .run();
+}
+
+export function listMcpConnectionStoreRows(): Array<{
+  id: string;
+  userId: string;
+  storeRef: string | null;
+}> {
+  return getDb()
+    .select({
+      id: schema.mcpConnections.id,
+      userId: schema.mcpConnections.userId,
+      storeRef: schema.mcpConnections.storeRef,
+    })
+    .from(schema.mcpConnections)
+    .all();
+}
+
+export function mirrorMcpConnectionToStore(id: string): void {
+  queueStoreWrite(`mirror mcp connection ${id}`, async (store) => {
+    const value = mcpConnectionStoreValue(id);
+    if (!value) return;
+    const name = mcpConnectionParameterName(value.userId, id);
+    await store.put(name, JSON.stringify(value));
+    setMcpConnectionStoreRef(id, name);
+  });
+}
+
+/** Delete the row and its parameter (the disconnect op and reconcile). */
+export function deleteMcpConnectionLocal(id: string): void {
+  const row = getDb()
+    .select({ userId: schema.mcpConnections.userId, storeRef: schema.mcpConnections.storeRef })
+    .from(schema.mcpConnections)
+    .where(eq(schema.mcpConnections.id, id))
+    .get();
+  getDb().delete(schema.mcpConnections).where(eq(schema.mcpConnections.id, id)).run();
+  if (!row) return;
+  const name = row.storeRef ?? mcpConnectionParameterName(row.userId, id);
+  queueStoreWrite(`delete mcp connection ${id}`, (store) => store.delete(name));
+}
+
+/** Reconcile import: a connection this host lacks. Existing rows win. */
+export function importMcpConnectionFromStore(
+  value: StoredMcpConnection,
+  storeRef: string,
+): boolean {
+  const result = getDb()
+    .insert(schema.mcpConnections)
+    .values({
+      id: value.id,
+      userId: value.userId,
+      url: value.url,
+      authKind: value.authKind,
+      tokenEndpoint: value.tokenEndpoint,
+      redirectUri: value.redirectUri,
+      clientId: value.clientId,
+      clientSecretEnc: value.clientSecret === null ? null : pack(value.clientSecret),
+      pkceVerifierEnc: value.pkceVerifier === null ? null : pack(value.pkceVerifier),
+      secretEnc: value.secret === null ? null : pack(value.secret),
+      status: value.status,
+      scopes: value.scopes,
+      updatedAt: value.updatedAt,
+      storeRef,
+    })
+    .onConflictDoNothing({ target: schema.mcpConnections.id })
+    .run();
+  return result.changes > 0;
+}

package/lib/orchestrator.ts

@@ -106,6 +106,7 @@ import {
   type BrowserSetup,
 } from "./browser-testing";
 import { injectCodexIntoContainer } from "./codex-auth";
+import { isCredentialInvalidSignal, requestCredentialRefresh } from "./token-broker-local";
 import {
   codexHomeIsolatable,
   codexHomeSeedCommand,
@@ -2424,6 +2425,29 @@ export class Orchestrator {
         ) {
           break;
         }
+        // ADR-135 reactive path: an engine whose brokered copy just died asks
+        // the one refresher (this host, or the broker Lambda) for a fresh
+        // token, which re-pushes the copy; then the agent respawns on it.
+        // Rate-limited per credential inside requestCredentialRefresh.
+        if (
+          agent &&
+          agent.kind !== "claude" &&
+          isCredentialInvalidSignal(agent.kind, event.message)
+        ) {
+          console.warn(
+            `[orchestrator] ${channel.taskId}/${agentId}: ${agent.kind} credential-invalid pattern matched: ${event.message.slice(0, 300)}`,
+          );
+          const accountId = channel.accountByAgent.get(agentId) ?? null;
+          const requested = await requestCredentialRefresh({
+            kind: agent.kind,
+            accountId: accountId && !accountId.startsWith("default:") ? accountId : null,
+            scope: channel.engineAccountScope,
+          });
+          if (requested) {
+            await this.refreshAgentToken(channel, agent);
+            break;
+          }
+        }
         const isClaude = agent?.kind === "claude";
         if (isClaude && agent && isClaudeAuthRevoked(event.message)) {
           // Token revoked/expired: the env-injected credential only changes on
@@ -2985,6 +3009,34 @@ export class Orchestrator {
    * changes (a re-login). Resets each agent's respawn budget first so an agent
    * that exhausted it during the outage gets a clean chance on the new token.
    */
+  /**
+   * ADR-135: after the broker refreshed `kind`'s material, copy the brokered
+   * accounts into every running task that has an agent of that kind. The
+   * copy is what the agent reads on its next spawn; a live engine process
+   * keeps its access token until it respawns or the reactive path fires.
+   */
+  async refreshEngineAccountCopies(kind: string): Promise<void> {
+    for (const channel of this.channels.values()) {
+      if (!this.isActiveChannel(channel)) continue;
+      if (!channel.roster.some((agent) => agent.kind === kind)) continue;
+      try {
+        await provisionEngineAccounts(
+          channel.containerName,
+          [kind],
+          {},
+          channel.agentHandle ?? undefined,
+          channel.engineAccountScope,
+        );
+      } catch (err) {
+        console.warn(
+          `[orchestrator] ${channel.taskId}: re-push of ${kind} accounts failed: ${
+            err instanceof Error ? err.message : String(err)
+          }`,
+        );
+      }
+    }
+  }
+
   async refreshEngineAccountAgents(kind: string): Promise<void> {
     for (const channel of this.channels.values()) {
       if (!this.isActiveChannel(channel)) continue;

package/lib/ssh.ts

@@ -23,6 +23,7 @@ import { join, resolve } from "node:path";
 
 import { eq } from "drizzle-orm";
 
+import { queueStoreWrite, userParameterName } from "./credential-store";
 import { getDb, schema } from "./db";
 import { env } from "./env";
 import { sealAesGcm, openAesGcm } from "./secrets";
@@ -104,12 +105,102 @@ export function ensureKeyForUser(userId: string): string {
     })
     .onConflictDoNothing({ target: schema.sshKeys.userId })
     .run();
+  mirrorSshKeyToStore(userId);
   // Re-read in case a concurrent task-up won the insert race.
   return getPublicKey(userId) ?? publicKey;
 }
 
 export function deleteKey(userId: string): void {
   getDb().delete(schema.sshKeys).where(eq(schema.sshKeys.userId, userId)).run();
+  queueStoreWrite(`delete ssh key ${userId}`, (store) =>
+    store.delete(sshKeyParameterName(userId)),
+  );
+}
+
+// ---------------------------------------------------------------------------
+// Credential store mirror (ADR-134 slice 2): `<prefix>/users/<userId>/ssh`.
+// A replacement Uai-operated host imports the SAME keypair, so the public key
+// the user pasted into GitHub keeps working. Owned hosts: no-ops.
+// ---------------------------------------------------------------------------
+
+export interface StoredSshKey {
+  v: 1;
+  kind: "ssh";
+  userId: string;
+  publicKey: string;
+  privateKey: string;
+  createdAt: number;
+  updatedAt: number;
+}
+
+export function sshKeyParameterName(userId: string): string {
+  return userParameterName(userId, "ssh");
+}
+
+/** The row as it is mirrored (private key decrypted). Never logged. */
+export function sshKeyStoreValue(userId: string): StoredSshKey | null {
+  const row = getDb()
+    .select()
+    .from(schema.sshKeys)
+    .where(eq(schema.sshKeys.userId, userId))
+    .get();
+  if (!row) return null;
+  return {
+    v: 1,
+    kind: "ssh",
+    userId,
+    publicKey: row.publicKey,
+    privateKey: openAesGcm(
+      Buffer.from(row.privateKeyCt as Uint8Array),
+      Buffer.from(row.privateKeyNonce as Uint8Array),
+    ),
+    createdAt: row.createdAt,
+    updatedAt: row.updatedAt,
+  };
+}
+
+export function setSshKeyStoreRef(userId: string, storeRef: string): void {
+  getDb()
+    .update(schema.sshKeys)
+    .set({ storeRef })
+    .where(eq(schema.sshKeys.userId, userId))
+    .run();
+}
+
+export function listSshKeyStoreRows(): Array<{ userId: string; storeRef: string | null }> {
+  return getDb()
+    .select({ userId: schema.sshKeys.userId, storeRef: schema.sshKeys.storeRef })
+    .from(schema.sshKeys)
+    .all();
+}
+
+export function mirrorSshKeyToStore(userId: string): void {
+  queueStoreWrite(`mirror ssh key ${userId}`, async (store) => {
+    const value = sshKeyStoreValue(userId);
+    if (!value) return;
+    const name = sshKeyParameterName(userId);
+    await store.put(name, JSON.stringify(value));
+    setSshKeyStoreRef(userId, name);
+  });
+}
+
+/** Reconcile import: a keypair this host lacks. Existing rows win. */
+export function importSshKeyFromStore(value: StoredSshKey, storeRef: string): boolean {
+  const sealed = sealAesGcm(value.privateKey);
+  const result = getDb()
+    .insert(schema.sshKeys)
+    .values({
+      userId: value.userId,
+      publicKey: value.publicKey,
+      privateKeyCt: sealed.ct,
+      privateKeyNonce: sealed.nonce,
+      createdAt: value.createdAt,
+      updatedAt: value.updatedAt,
+      storeRef,
+    })
+    .onConflictDoNothing({ target: schema.sshKeys.userId })
+    .run();
+  return result.changes > 0;
 }
 
 /** Decrypt + return the user's private key, or null if they have none. */

package/lib/token-broker-lambda.ts

@@ -0,0 +1,51 @@
+/**
+ * Broker Lambda entrypoint (ADR-135). Runs one broker cycle against the
+ * credential store. Invoked by the six-hourly EventBridge rule, and
+ * synchronously by hosts for the reactive path with `{ "names": [...] }`
+ * (parameter names to refresh now). Returns the cycle summary; logs one
+ * audit line per decision and never a value.
+ *
+ * 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.
+ */
+
+import { SSMClient } from "@aws-sdk/client-ssm";
+
+import { runBrokerCycle } from "./token-broker";
+import { acquireBrokerLock, SsmBrokerStorage } from "./token-broker-ssm";
+
+interface BrokerEvent {
+  names?: unknown;
+}
+
+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 kmsKeyId = process.env.UAI_CREDENTIAL_STORE_KMS_KEY || "alias/uai-credentials";
+  const names = Array.isArray(event.names)
+    ? event.names.filter((n): n is string => typeof n === "string" && n.startsWith(`${prefix}/`)).slice(0, MAX_FORCED_NAMES)
+    : [];
+  const started = Date.now();
+  const client = new SSMClient({});
+  // One run at a time, whatever the account's concurrency settings allow:
+  // two refreshers for the same credential is the bug this exists to end.
+  const release = await acquireBrokerLock({ client, prefix });
+  if (!release) {
+    console.log("[token-broker] another run holds the lock; skipped");
+    return { skipped: "locked", forced: names.length, ms: Date.now() - started };
+  }
+  try {
+    const summary = await runBrokerCycle({
+      storage: new SsmBrokerStorage({ prefix, kmsKeyId, client }),
+      force: new Set(names),
+      log: (line) => console.log(`[token-broker] ${line}`),
+    });
+    const result = { ...summary, forced: names.length, ms: Date.now() - started };
+    console.log(`[token-broker] cycle ${JSON.stringify(result)}`);
+    return result;
+  } finally {
+    await release();
+  }
+}