@runuai/host 0.9.171 → 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
@@ -1033,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),
@@ -1492,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",
@@ -1543,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()
         ? [
@@ -3058,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;
@@ -3609,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
@@ -4831,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 {
@@ -5092,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);
@@ -5213,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",
@@ -5626,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;
     }
@@ -6107,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.
@@ -6119,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?: {
@@ -4904,6 +4911,9 @@ export interface TaskMachineRetireAckFrame {
 }
 
 export type CloudToHost =
+  | EngineConnectCustomDiscoverFrame
+  | EngineConnectCustomSaveFrame
+  | EngineConnectCustomManualSaveFrame
   | {
       kind: "command";
       commandId: string;
@@ -5083,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;
@@ -5156,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).
@@ -5358,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;