@runuai/host 0.9.170 → 0.9.172

package/src/main.ts

@@ -1,7 +1,13 @@
 #!/usr/bin/env tsx
+import "./load-env";
+import { discoverCustomConnection, saveCustomConnection, saveCustomConnectionManual } from "../lib/custom-connection-capture";
+import { OpenRouterBudgetPoller } from "../lib/openrouter-budget";
+import { onLeasedEnginesChanged, leasedEnginesFor, isLeasedEngineAccount } from "../lib/credential-consumers/engines";
+import { isOwnedInternalConnectionRuntime } from "../lib/internal-connections";
+import { CUSTOM_CONNECTIONS_FEATURE, ENGINE_CONNECT_CUSTOM_FEATURE, ENGINE_CONNECT_CUSTOM_MANUAL_FEATURE, INTERNAL_CONNECTIONS_FEATURE, connectionOpId, parseCustomConnectionMaterial, parseCustomConnectionEnabled, parseEngineConnectCustomDiscoverFrame, parseEngineConnectCustomSaveFrame, parseEngineConnectCustomManualSaveFrame } from "./custom-connections";
 import type { AssignmentDelivery, DeliveryOrigin } from "./protocol";
+import { parseCaptureFlowWire } from "./cloud-credentials";
 
-import "./load-env";
 import { sampleDiskPressure } from "../lib/disk-pressure";
 import { TASK_SECRETS_FEATURE, isTaskSecretStagingPut } from "./task-secrets";
 import { writeTaskSecretStaging } from "../lib/task-secret-staging";
@@ -181,7 +187,7 @@ import {
   rejectContainerRuntimeSelection,
   reprobeContainerRuntimeMachineIdentity,
 } from "../lib/container-runtime";
-import { persistedTaskEnvironmentsMatchMachine, reconstructPersistedTaskEnvironment } from "../lib/task-environment";
+import { machineTasksEnabled, persistedTaskEnvironmentsMatchMachine, reconstructPersistedTaskEnvironment } from "../lib/task-environment";
 import { reconcileRunningTaskPersonCredentials } from "../lib/task-person-credentials";
 import { logVnodePressure } from "../lib/vnode-pressure";
 import { resolveAppleTunnelRoute } from "./apple-tunnel-route";
@@ -509,6 +515,15 @@ const bridgeUrl =
 
 let ws: WebSocket | null = null;
 let stopping = false;
+const openRouterBudgets = new OpenRouterBudgetPoller({ emit: frame => {
+  if (!stopping && credentialFramesReady && ws?.readyState === WebSocket.OPEN) send(ws, frame);
+} });
+onLeasedEnginesChanged(taskId => {
+  if (stopping) { openRouterBudgets.stop(taskId); return; }
+  const lease = leasedEnginesFor(taskId);
+  openRouterBudgets.sync(taskId, lease?.open ? lease.accounts.flatMap(account => isLeasedEngineAccount(account) && account.authKind === "custom-connection" && account.connectionMaterial?.providerType === "openrouter" ? [{ credentialId: account.credentialId, revision: account.revision, apiKey: account.connectionMaterial.apiKey }] : []) : []);
+});
+
 let fatal = false;
 /** Auth-rejected (4001) retry cadence. A genuinely revoked credential costs
  * one connection attempt per interval; a transient rejection (the bridge's
@@ -523,6 +538,17 @@ let shutdownRequested = false;
 let shutdownExitCode = 0;
 const managedMutationDrain = new ManagedOperationDrain(maybeExit);
 const hostLogManager = createHostLogManager();
+const customConnectionOperations = new Map<AbortController, Promise<void>>();
+
+function runCustomConnectionOperation(operation: (signal: AbortSignal) => Promise<void>): void {
+  void managedMutationDrain.track(async () => {
+    const controller = new AbortController();
+    const work = Promise.resolve().then(() => operation(controller.signal));
+    customConnectionOperations.set(controller, work);
+    try { await work; } finally { customConnectionOperations.delete(controller); }
+  });
+}
+
 const pendingRemoteOperationCleanups = new Set<Promise<void>>();
 let stopActiveRemoteOperations = (): Promise<void> => Promise.resolve();
 
@@ -545,8 +571,10 @@ function trackRemoteOperationCleanup(cleanup: Promise<void>): Promise<void> {
 function stopRemoteOperationManagers(): Promise<void> {
   hostLogManager.stopAll();
   engineConnectManager.stopAll();
+  for (const controller of customConnectionOperations.keys()) controller.abort();
   return trackRemoteOperationCleanup(
     Promise.all([
+      ...customConnectionOperations.values(),
       engineLoginManager.stopAll(),
       // Tunnel upstreams are generation-scoped resources just like login
       // containers; a bridge replacement or shutdown must join their
@@ -784,6 +812,9 @@ const credentialsActivation = new CredentialsActivationAuthority<CredentialsReti
     requestManagedHostRestart();
   },
   log: (line) => console.warn(line),
+  // What stops the switch reaches the cloud too, not only this log.
+  onStatusChange: () => sendCapabilities(),
+  redact: redactHomeDirectory,
 });
 setCredentialsModeSource(() => credentialsActivation.mode());
 /** ADR-154 §13: the steady/operator cadence. Active runners report their
@@ -1030,6 +1061,7 @@ const taskFiles = sidecarMode()
   : null;
 const cloudCredentials = new CloudCredentialRuntime({
   runnerKind: sidecarMode() ? "task-machine" : "owned",
+  internalRuntimeProvider: (assignment) => internalRuntimeProvider(assignment.taskId),
   send: sendCredentialFrame,
   quiesceProcesses: (request) => getOrchestrator().quiesceLeasedEngineAgents(request),
   mcpServedChanged: (taskId) => getOrchestrator().mcpServedChanged(taskId),
@@ -1489,12 +1521,30 @@ void refreshDiskSnapshot();
 const diskTimer = setInterval(() => void refreshDiskSnapshot(), 60_000);
 diskTimer.unref?.();
 
+/** A persisted provider wins over this runner's default selection. Legacy
+ * resources without a proven provider are unknown; only a new/empty task
+ * may use the selected provisioning backend as its authoritative plan. */
+function internalRuntimeProvider(taskId?: string): string | null {
+  if (sidecarMode() || signInBoxMode()) return null;
+  const task = taskId ? getHostTask(taskId) : undefined;
+  if (task?.environmentProvider !== null && task?.environmentProvider !== undefined) return task.environmentProvider;
+  if (task && (task.composeProject !== null || task.worktreePath !== null)) return null;
+  if (machineTasksEnabled()) return null;
+  const runtime = containerRuntimeCapabilityState();
+  if (runtime.status !== "ready") return null;
+  return runtime.provider === "docker" ? "docker-compose" : runtime.provider === "apple-container" ? "apple-container" : null;
+}
+
 /** Build the current host capability advertisement (ADR-021). */
 function buildCapabilities(): HostCapabilities {
   const runtime = containerRuntimeCapabilityState();
   const projectEnvBindings = projectEnvBindingCapability();
   return {
     version: packageVersion(),
+    customConnections: credentialsMode() === "active",
+    engineConnectCustom: signInBoxMode(),
+    engineConnectCustomManual: signInBoxMode(),
+    internalConnections: isOwnedInternalConnectionRuntime(internalRuntimeProvider()),
     ...(diskSnapshot ? { disk: diskSnapshot } : {}),
     protocolFeatures: [
       "host-workspace-file-read-v1",
@@ -1540,6 +1590,9 @@ function buildCapabilities(): HostCapabilities {
       // ADR-161 §1.11: this image's `uai files` reads cloud links. The cloud
       // captures the advertisement at container creation.
       HOST_AGENT_FILES_FEATURE,
+      ...(credentialsMode() === "active" ? [CUSTOM_CONNECTIONS_FEATURE] : []),
+      ...(signInBoxMode() ? [ENGINE_CONNECT_CUSTOM_FEATURE, ENGINE_CONNECT_CUSTOM_MANUAL_FEATURE] : []),
+      ...(isOwnedInternalConnectionRuntime(internalRuntimeProvider()) ? [INTERNAL_CONNECTIONS_FEATURE] : []),
       // ADR-154: only the sign-in box writes cloud credentials.
       ...(signInBoxMode()
         ? [
@@ -1581,9 +1634,18 @@ function buildCapabilities(): HostCapabilities {
     engineAccounts: engineAccountSummaries(),
     engines: engineCatalogCapabilities(),
     githubUsers: connectedUserIds(),
+    // ADR-154 §13: where the switch to cloud credentials stands, and why.
+    credentialsSwitch: credentialsActivation.status(),
   };
 }
 
+/** A path under the home directory, as `~/…`: a status line leaves the
+ *  machine, and the account name in a path is not the cloud's to keep. */
+function redactHomeDirectory(text: string): string {
+  const home = homedir();
+  return home.length > 1 ? text.split(home).join("~") : text;
+}
+
 // engineCliStatuses probes binaries asynchronously; capabilities build
 // synchronously. The snapshot starts empty (cliInstalled: false) and each
 // build kicks a refresh that re-advertises only when the answer changed.
@@ -3046,7 +3108,7 @@ async function connect(): Promise<void> {
     }
 
     const frame = parseCloudFrame(data);
-    if (!frame) return;
+    if (!frame) { replyToInvalidCustomConnectionFrame(socket, data); return; }
     if (stopping && rejectFrameDuringShutdown(socket, frame)) return;
     if (signInBoxMode() && rejectFrameOnSignInBox(socket, frame)) return;
     if (!onOwnCredentials() && rejectCredentialFrameOffOwnCredentials(socket, frame)) return;
@@ -3597,6 +3659,21 @@ async function connect(): Promise<void> {
         );
         break;
       }
+      case "engine.connect.custom.discover": {
+        if (!signInBoxMode()) { send(socket, { kind: "engine.connect.custom.discovered", opId: frame.opId, result: { ok: false, reason: "unsupported" } }); break; }
+        runCustomConnectionOperation(signal => discoverCustomConnection(frame, { signal }).catch(() => ({ ok: false as const, reason: "upstream-error" as const })).then(result => send(socket, { kind: "engine.connect.custom.discovered", opId: frame.opId, result })));
+        break;
+      }
+      case "engine.connect.custom.manual-save": {
+        if (!signInBoxMode()) { send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId, result: { ok: false, reason: "unsupported" } }); break; }
+        runCustomConnectionOperation(signal => saveCustomConnectionManual(frame, configuredCredentialStore(), { signal }).catch(() => ({ ok: false as const, reason: "upstream-error" as const })).then(result => send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId, result })));
+        break;
+      }
+      case "engine.connect.custom.save": {
+        if (!signInBoxMode()) { send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId, result: { ok: false, reason: "unsupported" } }); break; }
+        runCustomConnectionOperation(signal => saveCustomConnection(frame, configuredCredentialStore(), { signal }).catch(() => ({ ok: false as const, reason: "upstream-error" as const })).then(result => send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId, result })));
+        break;
+      }
       case "mcp.capture.probe":
       case "mcp.capture.complete": {
         // ADR-154 §11: only the sign-in box runs an MCP handshake as a
@@ -4819,6 +4896,24 @@ const PICKER_FAILURE_TEXT: Record<GitHubRepositoryAccessErrorCode, string> = {
   invalid_response: "GitHub answered with something unexpected",
 };
 
+function replyToInvalidCustomConnectionFrame(socket: WebSocket, data: RawData): void {
+  let value: unknown;
+  try { value = JSON.parse(rawDataToBuffer(data).toString("utf8")); } catch { return; }
+  if (!value || typeof value !== "object" || Array.isArray(value)) return;
+  const frame = value as Record<string, unknown>;
+  if (!connectionOpId(frame.opId)) return;
+  if (frame.kind === "engine.connect.custom.discover") {
+    send(socket, { kind: "engine.connect.custom.discovered", opId: frame.opId, result: { ok: false, reason: "unsupported" } });
+  } else if (frame.kind === "engine.connect.custom.save" || frame.kind === "engine.connect.custom.manual-save") {
+    const settings = parseCustomConnectionMaterial(frame.settings);
+    const flow = parseCaptureFlowWire(frame.flow);
+    const supported = settings && (frame.kind !== "engine.connect.custom.manual-save" || (settings.providerType === "custom" && settings.internal === true)) && parseCustomConnectionEnabled(frame.enabled) && flow &&
+      flow.location.class === "engine" && flow.location.engineKind === settings.engine && flow.brokerParameter === null;
+    send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId,
+      result: { ok: false, reason: supported ? "invalid-models" : "unsupported" } });
+  }
+}
+
 function parseCloudFrame(data: RawData): CloudToHost | null {
   let parsed: unknown;
   try {
@@ -5080,6 +5175,9 @@ function parseCloudFrame(data: RawData): CloudToHost | null {
   // ADR-154: the v1 lease, credential and task env replica frames, parsed
   // to their exact shapes here; only a runner on cloud credentials acts on
   // them.
+  if (frame.kind === "engine.connect.custom.discover") return parseEngineConnectCustomDiscoverFrame(frame);
+  if (frame.kind === "engine.connect.custom.manual-save") return parseEngineConnectCustomManualSaveFrame(frame);
+  if (frame.kind === "engine.connect.custom.save") return parseEngineConnectCustomSaveFrame(frame);
   if (frame.kind === "lease.grant.v1") return parseLeaseGrantV1Frame(frame);
   if (frame.kind === "lease.unavailable") return parseLeaseUnavailableFrame(frame);
   if (frame.kind === "lease.stale") return parseLeaseStaleFrame(frame);
@@ -5201,6 +5299,13 @@ function rejectFrameDuringShutdown(
 ): boolean {
   const error = "managed host is restarting; retry the operation";
   switch (frame.kind) {
+    case "engine.connect.custom.discover":
+      send(socket, { kind: "engine.connect.custom.discovered", opId: frame.opId, result: { ok: false, reason: "unreachable" } });
+      return true;
+    case "engine.connect.custom.manual-save":
+    case "engine.connect.custom.save":
+      send(socket, { kind: "engine.connect.custom.saved", opId: frame.opId, result: { ok: false, reason: "unreachable" } });
+      return true;
     case "command":
       send(socket, {
         kind: "result",
@@ -5614,6 +5719,12 @@ function expectTaskAgents(value: unknown): TaskAgent[] {
     ) {
       out.accountId = row.accountId;
     }
+    if (row.accountBinding !== undefined) {
+      if (row.accountBinding !== "required" && row.accountBinding !== "preferred") throw new Error("invalid command args: account binding");
+      if (row.accountBinding === "required" && row.kind !== "claude" && row.kind !== "codex") throw new Error("invalid command args: unsupported required account engine");
+      if (row.accountBinding === "required" && !out.accountId) throw new Error("invalid command args: required account missing");
+      out.accountBinding = row.accountBinding;
+    }
     if (row.uaiProvider === "anthropic" || row.uaiProvider === "openai") {
       out.uaiProvider = row.uaiProvider;
     }
@@ -6095,6 +6206,7 @@ function reconnectDelay(): number {
 function requestShutdown(signal: NodeJS.Signals): void {
   console.log(`[host-agent] ${signal} received; shutting down`);
   stopping = true;
+  openRouterBudgets.stop();
   // Cancel human-paced operations without invalidating this socket's exact
   // generation: already-admitted one-shot config writes may still ack while
   // the ordinary in-flight drain completes.
@@ -6107,6 +6219,7 @@ function requestShutdown(signal: NodeJS.Signals): void {
   // it). Exit waits for it.
   void cloudCredentials
     .shutdown()
+    .then(() => openRouterBudgets.drain())
     .catch((error: unknown) =>
       console.warn(`[credentials] shutdown failed: ${error instanceof Error ? error.message : String(error)}`),
     )

package/src/protocol.ts

@@ -1,3 +1,5 @@
+export * from "./custom-connections";
+import type { AccountBinding, EngineConnectCustomDiscoverFrame, EngineConnectCustomDiscoveredFrame, EngineConnectCustomSaveFrame, EngineConnectCustomManualSaveFrame, EngineConnectCustomSavedFrame, OpenRouterBudgetFrame } from "./custom-connections";
 import type { TaskSecretApplied, TaskSecretStagingPut, TaskSecretStagingAck } from "./task-secrets";
 import {
   parseCaptureFlowWire,
@@ -3154,6 +3156,7 @@ export interface TaskAgent {
   effort?: string;
   /** ADR-157 addendum §7: preferred cloud-leased AI account. */
   accountId?: string;
+  accountBinding?: AccountBinding;
   defaultPrompt?: string;
   initialPrompt?: string;
   /** ADR-046 reference material resolved from the agent's persona + team.
@@ -4068,6 +4071,10 @@ export function isHostConfigAckFrame(
  * changes. The cloud caches it per hostId to drive the task-creation pickers.
  */
 export interface HostCapabilities {
+  customConnections?: boolean;
+  engineConnectCustom?: boolean;
+  engineConnectCustomManual?: boolean;
+  internalConnections?: boolean;
   /** Filesystem capacity, refreshed each minute. Null means unmeasurable
    * (for example Docker in a VM); available excludes root-reserved blocks. */
   disk?: {
@@ -4140,6 +4147,99 @@ export interface HostCapabilities {
   // the per-host gh-connected state shown on the host detail page. Re-advertised
   // whenever a token is added/removed. Optional: older hosts omit it.
   githubUsers?: string[];
+  /** ADR-154 §13: where this runner's switch to cloud credentials stands, and
+   *  what stops it. Optional: hosts before 0.9.171 omit it, and the cloud
+   *  then says only that the runner has not switched. */
+  credentialsSwitch?: CredentialsSwitchCapability;
+}
+
+/**
+ * ADR-154 §13: the runner's switch to cloud credentials, as its capability
+ * frame reports it. `state` is the runner's credentials mode. `stuck` is the
+ * latest substantive refusal or blocker while the switch has not finished:
+ * `by: "cloud"` the cloud's answer, `by: "runner"` something on the machine.
+ * A transient one (the container runtime still preparing after a reconnect,
+ * the bridge down) never replaces a substantive one; it shows only when
+ * there is nothing else to say. Bounded and secret-free: the runner's own log
+ * words (task ids, Compose project names, folder paths with the home
+ * directory as `~`), never a credential.
+ */
+export interface CredentialsSwitchCapability {
+  state: "legacy" | "pending" | "active";
+  stuck?: CredentialsSwitchStuck;
+}
+
+export interface CredentialsSwitchStuck {
+  by: "cloud" | "runner";
+  /** The refusal or blocker in one line. */
+  detail: string;
+  /** Each blocker on its own, when the runner found several. */
+  blockers?: string[];
+  /** The `task-*` Compose projects nothing proves Uai's, which the operator
+   *  removes (`docker compose -p <name> down`) or renames: the first
+   *  `MAX_CREDENTIALS_SWITCH_PROJECTS` names, and how many there are. */
+  composeProjects?: { names: string[]; total: number };
+  /** Nothing is wrong yet: the runner could not look (see above). */
+  transient?: boolean;
+  /** Epoch ms this refusal or blocker was first seen, unchanged since. */
+  since: number;
+}
+
+export const MAX_CREDENTIALS_SWITCH_DETAIL_CHARS = 600;
+export const MAX_CREDENTIALS_SWITCH_BLOCKERS = 20;
+export const MAX_CREDENTIALS_SWITCH_BLOCKER_CHARS = 300;
+export const MAX_CREDENTIALS_SWITCH_PROJECTS = 100;
+
+/** The capability, exact and bounded; undefined when it is not one. The
+ *  cloud drops a malformed one alone, never the frame around it. */
+export function parseCredentialsSwitchCapability(value: unknown): CredentialsSwitchCapability | undefined {
+  if (!isPlainRecord(value)) return undefined;
+  const { state, stuck } = value;
+  if (state !== "legacy" && state !== "pending" && state !== "active") return undefined;
+  if (Object.keys(value).some((key) => key !== "state" && key !== "stuck")) return undefined;
+  if (stuck === undefined) return { state };
+  if (!isPlainRecord(stuck)) return undefined;
+  const allowed = new Set(["by", "detail", "blockers", "composeProjects", "transient", "since"]);
+  if (Object.keys(stuck).some((key) => !allowed.has(key))) return undefined;
+  const boundedText = (text: unknown, max: number): text is string =>
+    typeof text === "string" && text.length >= 1 && text.length <= max;
+  if (stuck.by !== "cloud" && stuck.by !== "runner") return undefined;
+  if (!boundedText(stuck.detail, MAX_CREDENTIALS_SWITCH_DETAIL_CHARS)) return undefined;
+  if (typeof stuck.since !== "number" || !Number.isSafeInteger(stuck.since) || stuck.since < 0) return undefined;
+  if (stuck.transient !== undefined && typeof stuck.transient !== "boolean") return undefined;
+  const out: CredentialsSwitchStuck = { by: stuck.by, detail: stuck.detail, since: stuck.since };
+  if (stuck.blockers !== undefined) {
+    if (
+      !Array.isArray(stuck.blockers) ||
+      stuck.blockers.length > MAX_CREDENTIALS_SWITCH_BLOCKERS ||
+      !stuck.blockers.every((blocker) => boundedText(blocker, MAX_CREDENTIALS_SWITCH_BLOCKER_CHARS))
+    ) {
+      return undefined;
+    }
+    out.blockers = [...stuck.blockers];
+  }
+  if (stuck.composeProjects !== undefined) {
+    const projects = stuck.composeProjects;
+    if (!isPlainRecord(projects) || Object.keys(projects).some((key) => key !== "names" && key !== "total")) return undefined;
+    const { names, total } = projects;
+    if (
+      !Array.isArray(names) ||
+      names.length > MAX_CREDENTIALS_SWITCH_PROJECTS ||
+      !names.every((name) => typeof name === "string" && /^task-[a-z0-9][a-z0-9_-]{0,127}$/.test(name)) ||
+      typeof total !== "number" ||
+      !Number.isSafeInteger(total) ||
+      total < names.length
+    ) {
+      return undefined;
+    }
+    out.composeProjects = { names: [...names] as string[], total };
+  }
+  if (stuck.transient === true) out.transient = true;
+  return { state, stuck: out };
+}
+
+function isPlainRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
 }
 
 /** One transcript projection the cloud asks the host to append to (ADR-083). */
@@ -4811,6 +4911,9 @@ export interface TaskMachineRetireAckFrame {
 }
 
 export type CloudToHost =
+  | EngineConnectCustomDiscoverFrame
+  | EngineConnectCustomSaveFrame
+  | EngineConnectCustomManualSaveFrame
   | {
       kind: "command";
       commandId: string;
@@ -4990,6 +5093,9 @@ export interface AssignmentDelivery {
   recipientGeneration: string | null;
 }
 export interface AssignmentTurnEvidence {
+  credentialId?: string | null;
+  credentialRevision?: string | null;
+  connection?: boolean | null;
   leadRunId?: string | null;
   turnId: string;
   assignmentId: string | null;
@@ -5063,6 +5169,9 @@ export type McpOp =
   | { kind: "disconnect"; connectionId: string };
 
 export type HostToCloud =
+  | EngineConnectCustomDiscoveredFrame
+  | EngineConnectCustomSavedFrame
+  | OpenRouterBudgetFrame
   // ADR-103: `bootId` scopes event seqs to one host process lifetime — the
   // cloud's replay watermark only applies within a matching boot. Optional
   // for pre-103 hosts (which also never send seq'd events).
@@ -5265,6 +5374,8 @@ export type HostEvent =
        *  `credentialId` is the leased AI account the turn ran on (host
        *  0.9.167+); older runners omit it and the cloud stores null. */
       usage?: {
+        connection?: boolean;
+        credentialRevision?: string;
         model?: string;
         inputTokens?: number;
         outputTokens?: number;