@runuai/host 0.9.146 → 0.9.148

package/lib/agents/transport.ts

@@ -33,9 +33,9 @@ import {
 } from "node:fs";
 import { basename, dirname, join } from "node:path";
 
-import { and, eq } from "drizzle-orm";
+import { and, eq, isNull } from "drizzle-orm";
 
-import { getDb, schema } from "../db";
+import { getDb, schema, type Db } from "../db";
 import { taskWorkspaceDir } from "../env";
 import { machineSessionAuthorityEnv } from "../runtime-authority";
 import { requireContainerRuntimeOperational } from "../runtime-guard";
@@ -56,9 +56,12 @@ import {
   type LineHandler,
 } from "./proc";
 import { isEnvironmentTransportProvider } from "../task-environment/types";
+import { attachCompatibilityKeysMatch } from "./attach-compatibility";
 
 /** The shared surface adapters program against (LineProcess's shape). */
 export interface LineTransport {
+  /** Exact durable row generation owned by this transport, when durable. */
+  readonly durableSessionDir?: string;
   onLine(handler: LineHandler): void;
   onExit(handler: ExitHandler): void;
   writeLine(value: unknown): void;
@@ -84,7 +87,7 @@ export interface AgentTransportOptions {
    * protocol state can't outlive the host process (codex handshake).
    */
   allowAttach: boolean;
-  /** Host-owned launch-policy generation. Both absent remains compatible. */
+  /** Host-owned launch-settings generation. Both absent remains compatible. */
   attachCompatibilityKey?: string;
   kind: string;
   debugLabel?: string;
@@ -120,6 +123,73 @@ function durableEnabled(): boolean {
   return process.env.UAI_DURABLE_SESSIONS !== "0";
 }
 
+/** Durable rows are keyed by task + agent for lookup, but callbacks belong to
+ * one exact generation. A delayed offset/exit from the predecessor must never
+ * mutate the replacement row after its upsert. */
+function durableRowCallbacks(
+  db: Db,
+  taskId: string,
+  agentId: string,
+  sessionDir: string,
+): {
+  persistOffset(offset: number): void;
+  markClosed(): void;
+} {
+  const ownedRow = () =>
+    and(
+      eq(schema.hostAgentSessions.taskId, taskId),
+      eq(schema.hostAgentSessions.agentId, agentId),
+      eq(schema.hostAgentSessions.sessionDir, sessionDir),
+    );
+  return {
+    persistOffset(offset: number): void {
+      db.update(schema.hostAgentSessions)
+        .set({ outboxOffset: offset, updatedAt: Date.now() })
+        .where(ownedRow())
+        .run();
+    },
+    markClosed(): void {
+      db.update(schema.hostAgentSessions)
+        .set({ status: "closed", updatedAt: Date.now() })
+        .where(ownedRow())
+        .run();
+    },
+  };
+}
+
+/** Claim a compatible durable row and promote a one-release legacy key to the
+ * rich launch fingerprint. The exact generation/key fence prevents a stale
+ * host process from rewriting a replacement row after reading its snapshot. */
+function claimCompatibleAttachRow(
+  db: Db,
+  taskId: string,
+  agentId: string,
+  sessionDir: string,
+  stored: string | null | undefined,
+  desired: string | null | undefined,
+): boolean {
+  if (!attachCompatibilityKeysMatch(stored, desired)) return false;
+  const storedKey = stored ?? null;
+  const desiredKey = desired ?? null;
+  if (storedKey === desiredKey) return true;
+  if (desiredKey === null) return false;
+  const result = db
+    .update(schema.hostAgentSessions)
+    .set({ attachCompatibilityKey: desiredKey, updatedAt: Date.now() })
+    .where(
+      and(
+        eq(schema.hostAgentSessions.taskId, taskId),
+        eq(schema.hostAgentSessions.agentId, agentId),
+        eq(schema.hostAgentSessions.sessionDir, sessionDir),
+        storedKey === null
+          ? isNull(schema.hostAgentSessions.attachCompatibilityKey)
+          : eq(schema.hostAgentSessions.attachCompatibilityKey, storedKey),
+      ),
+    )
+    .run();
+  return result.changes === 1;
+}
+
 export function createAgentTransport(opts: AgentTransportOptions): LineTransport {
   // ADR-121: machine-backed sessions ride the environment's own transport —
   // the container-runtime preflight is a docker/apple concern a machine task
@@ -161,29 +231,6 @@ export function createAgentTransport(opts: AgentTransportOptions): LineTransport
     )
     .get();
 
-  const persistOffset = (offset: number): void => {
-    db.update(schema.hostAgentSessions)
-      .set({ outboxOffset: offset, updatedAt: Date.now() })
-      .where(
-        and(
-          eq(schema.hostAgentSessions.taskId, opts.taskId),
-          eq(schema.hostAgentSessions.agentId, opts.agentId),
-        ),
-      )
-      .run();
-  };
-  const markClosed = (): void => {
-    db.update(schema.hostAgentSessions)
-      .set({ status: "closed", updatedAt: Date.now() })
-      .where(
-        and(
-          eq(schema.hostAgentSessions.taskId, opts.taskId),
-          eq(schema.hostAgentSessions.agentId, opts.agentId),
-        ),
-      )
-      .run();
-  };
-
   const tailKey = `${opts.taskId}:${opts.agentId}`;
 
   // ---- Attach: a previous host process left this agent's runner alive. ----
@@ -192,10 +239,23 @@ export function createAgentTransport(opts: AgentTransportOptions): LineTransport
     row &&
     row.status === "running" &&
     row.containerName === opts.environment.durableIdentity &&
-    (row.attachCompatibilityKey ?? null) ===
-      (opts.attachCompatibilityKey ?? null) &&
-    heartbeatFresh(row.sessionDir)
+    row.kind === opts.kind &&
+    heartbeatFresh(row.sessionDir) &&
+    claimCompatibleAttachRow(
+      db,
+      opts.taskId,
+      opts.agentId,
+      row.sessionDir,
+      row.attachCompatibilityKey,
+      opts.attachCompatibilityKey,
+    )
   ) {
+    const { persistOffset, markClosed } = durableRowCallbacks(
+      db,
+      opts.taskId,
+      opts.agentId,
+      row.sessionDir,
+    );
     const proc = new DurableProcess({
       hostSessionDir: row.sessionDir,
       initialOutboxOffset: row.outboxOffset,
@@ -240,6 +300,12 @@ export function createAgentTransport(opts: AgentTransportOptions): LineTransport
     opts.agentId,
     basename(sessionDir),
   );
+  const { persistOffset, markClosed } = durableRowCallbacks(
+    db,
+    opts.taskId,
+    opts.agentId,
+    sessionDir,
+  );
   const runnerInEnvironment = join(environmentSessionDir, "runner.mjs");
   try {
     copyFileSync(runnerScriptPath(), runnerInSession);
@@ -350,28 +416,6 @@ function machineDurableTransport(opts: AgentTransportOptions): LineTransport {
       ),
     )
     .get();
-  const persistOffset = (offset: number): void => {
-    db.update(schema.hostAgentSessions)
-      .set({ outboxOffset: offset, updatedAt: Date.now() })
-      .where(
-        and(
-          eq(schema.hostAgentSessions.taskId, opts.taskId),
-          eq(schema.hostAgentSessions.agentId, opts.agentId),
-        ),
-      )
-      .run();
-  };
-  const markClosed = (): void => {
-    db.update(schema.hostAgentSessions)
-      .set({ status: "closed", updatedAt: Date.now() })
-      .where(
-        and(
-          eq(schema.hostAgentSessions.taskId, opts.taskId),
-          eq(schema.hostAgentSessions.agentId, opts.agentId),
-        ),
-      )
-      .run();
-  };
   const tailKey = `${opts.taskId}:${opts.agentId}`;
 
   // ---- Attach: a previous host process left this machine's runner alive. --
@@ -382,9 +426,22 @@ function machineDurableTransport(opts: AgentTransportOptions): LineTransport {
     row &&
     row.status === "running" &&
     row.containerName === environment.durableIdentity &&
-    (row.attachCompatibilityKey ?? null) ===
-      (opts.attachCompatibilityKey ?? null)
+    row.kind === opts.kind &&
+    claimCompatibleAttachRow(
+      db,
+      opts.taskId,
+      opts.agentId,
+      row.sessionDir,
+      row.attachCompatibilityKey,
+      opts.attachCompatibilityKey,
+    )
   ) {
+    const { persistOffset, markClosed } = durableRowCallbacks(
+      db,
+      opts.taskId,
+      opts.agentId,
+      row.sessionDir,
+    );
     const proc = new MachineDurableProcess({
       environment,
       sessionDir: row.sessionDir,
@@ -406,6 +463,12 @@ function machineDurableTransport(opts: AgentTransportOptions): LineTransport {
     environment.descriptor.workspacePath,
     opts.agentId,
   );
+  const { persistOffset, markClosed } = durableRowCallbacks(
+    db,
+    opts.taskId,
+    opts.agentId,
+    sessionDir,
+  );
   const proc = new MachineDurableProcess({
     environment,
     sessionDir,

package/lib/agents/types.ts

@@ -180,9 +180,9 @@ export interface AgentSessionFactory {
     /** Host-derived restricted execution policy (ADR-083). Never supplied by
      *  the browser or stored on the roster entry. */
     executionProfile?: "communicator";
-    /** Durable runners are keyed by task+agent. A host-owned policy change
-     *  supplies a new key so restart recovery replaces, rather than attaches,
-     *  a process launched with incompatible argv/preamble/tooling. */
+    /** Durable runners are keyed by task+agent. Host-owned launch settings
+     *  supply a compatibility key so restart recovery replaces, rather than
+     *  attaches, a process launched with incompatible model/effort/policy. */
     attachCompatibilityKey?: string;
     /** ADR-048: extra per-agent process environment (e.g. this agent's own
      *  UAI_TASK_TOKEN so its `uai` CLI carries only its own permissions). */