Preserve background tasks across daemon reconnects

apps/server/src/internal/session-owner-side-effects.ts

@@ -87,24 +87,32 @@ export async function handleHostSessionOpened(
     args.previousSession &&
     args.previousSession.id !== args.openedSession.id
   ) {
+    const sameDaemonInstance =
+      args.previousSession.instanceId === args.openedSession.instanceId;
     deps.hub.cancelPendingDaemonDisconnect(args.previousSession.id);
 
     if (args.previousSession.status === "active") {
-      deps.hub.closeDaemonSession(args.previousSession.id, "replaced");
+      if (sameDaemonInstance) {
+        // A reconnect opens the replacement session before its new WebSocket.
+        // Close only the superseded socket: sending session-close would make
+        // this same daemon shut down every resident provider runtime.
+        deps.hub.closeDaemonSessionSocket(args.previousSession.id, "replaced");
+      } else {
+        deps.hub.closeDaemonSession(args.previousSession.id, "replaced");
+      }
       deps.terminalSessions.handleDaemonSessionClosed({
         sessionId: args.previousSession.id,
       });
     }
 
     interruptPendingInteractionsForHostThreads(deps, {
       hostId: args.hostId,
-      reason:
-        args.previousSession.instanceId === args.openedSession.instanceId
-          ? DAEMON_DISCONNECTED_PENDING_INTERACTION_REASON
-          : DAEMON_RESTARTED_PENDING_INTERACTION_REASON,
+      reason: sameDaemonInstance
+        ? DAEMON_DISCONNECTED_PENDING_INTERACTION_REASON
+        : DAEMON_RESTARTED_PENDING_INTERACTION_REASON,
     });
 
-    if (args.previousSession.instanceId !== args.openedSession.instanceId) {
+    if (!sameDaemonInstance) {
       interruptActiveThreadsForHost(deps, {
         hostId: args.hostId,
         reason: "host-daemon-restarted",

apps/server/src/ws/hub.ts

@@ -536,13 +536,23 @@ export class NotificationHub implements DbNotifier {
   closeDaemonSession(
     sessionId: string,
     reason: HostDaemonSessionCloseReason,
+  ): void {
+    const entry = this.daemonSessions.get(sessionId);
+    if (entry) {
+      entry.socket.send(JSON.stringify({ type: "session-close", reason }));
+    }
+    this.closeDaemonSessionSocket(sessionId, reason);
+  }
+
+  closeDaemonSessionSocket(
+    sessionId: string,
+    reason: HostDaemonSessionCloseReason,
   ): void {
     this.cancelPendingDaemonDisconnect(sessionId);
     const entry = this.daemonSessions.get(sessionId);
     if (!entry) {
       return;
     }
-    entry.socket.send(JSON.stringify({ type: "session-close", reason }));
     entry.socket.close(1000, reason);
     this.unregisterDaemon(sessionId);
   }

apps/server/test/internal/background-task-reconciliation.test.ts

@@ -18,6 +18,7 @@ import {
   seedThreadFixture,
   seedTurnStarted,
 } from "../helpers/seed.js";
+import { createMockHubSocket } from "../helpers/mock-hub-socket.js";
 import { withTestHarness, type TestAppHarness } from "../helpers/test-app.js";
 
 function backgroundTaskItemData(args: {
@@ -367,6 +368,78 @@ describe("background-task lifecycle reconciliation triggers", () => {
     });
   });
 
+  it("does not tell a live daemon to shut down when that same instance reconnects", async () => {
+    await withTestHarness(async (harness) => {
+      const { host, session, thread } = seedOpenBackgroundTaskThread(harness);
+      const previousSocket = createMockHubSocket();
+      harness.deps.hub.registerDaemon(session.id, host.id, previousSocket);
+
+      const response = await harness.app.request("/internal/session/open", {
+        method: "POST",
+        headers: internalAuthHeaders(harness, {
+          hostId: host.id,
+          hostType: host.type,
+        }),
+        body: JSON.stringify({
+          hostId: host.id,
+          instanceId: session.instanceId,
+          hostName: host.name,
+          hostType: host.type,
+          hasMachineCredential: false,
+          platform: "darwin",
+          dataDir: "/tmp/host-daemon-task-live-same-instance",
+          protocolVersion: HOST_DAEMON_PROTOCOL_VERSION,
+          activeThreads: [],
+        }),
+      });
+
+      expect(response.status).toBe(201);
+      expect(previousSocket.messages).toEqual([]);
+      expect(previousSocket.closed).toEqual([
+        { code: 1000, reason: "replaced" },
+      ]);
+      expect(listSettledBackgroundTaskItems(harness, thread.id)).toEqual([]);
+    });
+  });
+
+  it("still tells a superseded daemon instance to shut down", async () => {
+    await withTestHarness(async (harness) => {
+      const { host, session, thread } = seedOpenBackgroundTaskThread(harness);
+      const previousSocket = createMockHubSocket();
+      harness.deps.hub.registerDaemon(session.id, host.id, previousSocket);
+
+      const response = await harness.app.request("/internal/session/open", {
+        method: "POST",
+        headers: internalAuthHeaders(harness, {
+          hostId: host.id,
+          hostType: host.type,
+        }),
+        body: JSON.stringify({
+          hostId: host.id,
+          instanceId: "instance-restarted",
+          hostName: host.name,
+          hostType: host.type,
+          hasMachineCredential: false,
+          platform: "darwin",
+          dataDir: "/tmp/host-daemon-task-live-restarted",
+          protocolVersion: HOST_DAEMON_PROTOCOL_VERSION,
+          activeThreads: [],
+        }),
+      });
+
+      expect(response.status).toBe(201);
+      expect(previousSocket.messages).toEqual([
+        JSON.stringify({ type: "session-close", reason: "replaced" }),
+      ]);
+      expect(previousSocket.closed).toEqual([
+        { code: 1000, reason: "replaced" },
+      ]);
+      expect(listSettledBackgroundTaskItems(harness, thread.id)).toEqual([
+        { status: "interrupted", taskStatus: "stopped" },
+      ]);
+    });
+  });
+
   it("settles open tasks after the disconnect grace elapses without a reconnect", async () => {
     await withTestHarness(async (harness) => {
       const { session, thread } = seedOpenBackgroundTaskThread(harness);

packages/host-daemon-contract/src/commands.ts

@@ -35,7 +35,7 @@ import {
   providerCliStatusResponseSchema,
 } from "./local.js";
 
-export const HOST_DAEMON_PROTOCOL_VERSION = 80 as const;
+export const HOST_DAEMON_PROTOCOL_VERSION = 81 as const;
 
 export {
   BRANCH_LIST_LIMIT_MAX,

packages/host-daemon-contract/test/contract.test.ts

@@ -1038,12 +1038,10 @@ describe("host-daemon local schemas", () => {
 });
 
 describe("host-daemon command schemas", () => {
-  // Version 80 adds structured provider account rate-limit events in runtime
-  // session messages on top of version 79's canonical Codex repository skills.
-  // The bump updates enrolled daemons before they receive the combined wire
-  // contract.
-  it("uses protocol version 80 for provider rate-limit events", () => {
-    expect(HOST_DAEMON_PROTOCOL_VERSION).toBe(80);
+  // Version 81 makes same-instance daemon session replacement reconnect-safe
+  // so enrolled daemons update before receiving the revised server behavior.
+  it("uses protocol version 81 for reconnect-safe session replacement", () => {
+    expect(HOST_DAEMON_PROTOCOL_VERSION).toBe(81);
   });
 
   it("binds Plan cancellation to a required turn id and typed result", () => {

tests/integration/real/provider-background-survival.test.ts

@@ -0,0 +1,112 @@
+import fs from "node:fs/promises";
+import path from "node:path";
+import { describe, expect, it } from "vitest";
+import { scaleTimeoutMs } from "../helpers/time.js";
+import {
+  createRealThread,
+  pathExists,
+  REAL_POLL_INTERVAL_MS,
+  sendAndWaitForIdle,
+} from "./provider-smoke-harness.js";
+
+const PROVIDER_ID = "claude-code";
+const BACKGROUND_SURVIVAL_TEST_TIMEOUT_MS = scaleTimeoutMs(240_000);
+
+async function waitForCondition(
+  predicate: () => boolean | Promise<boolean>,
+  timeoutMs: number,
+  message: string,
+): Promise<void> {
+  const startedAt = Date.now();
+  while (Date.now() - startedAt <= timeoutMs) {
+    if (await predicate()) {
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, REAL_POLL_INTERVAL_MS));
+  }
+  throw new Error(message);
+}
+
+describe("real Claude background process integration", () => {
+  it(
+    "survives a same-daemon session reconnect and follow-up",
+    async () => {
+      const { environment, harness, thread } = await createRealThread({
+        providerId: PROVIDER_ID,
+        workspace: { path: null, type: "unmanaged" },
+      });
+
+      try {
+        if (!environment.path) {
+          throw new Error("Expected an unmanaged workspace path");
+        }
+        const markerPath = path.join(
+          environment.path,
+          "background-survived-reconnect.txt",
+        );
+
+        await sendAndWaitForIdle({
+          harness,
+          providerId: PROVIDER_ID,
+          threadId: thread.id,
+          text: `Use the Bash tool exactly once to execute this exact command: sleep 40; printf survived > ${JSON.stringify(markerPath)}. Set Bash's run_in_background parameter to true. As soon as Bash accepts it, reply exactly STARTED. Do not poll it, wait for it, or use another tool.`,
+        });
+        expect(await pathExists(markerPath)).toBe(false);
+
+        // Cross the host daemon's 10-second shell-environment refresh TTL so
+        // the follow-up exercises the same idle-gap maintenance path as the
+        // report, while the background command remains live.
+        await new Promise((resolve) => setTimeout(resolve, 12_000));
+        expect(await pathExists(markerPath)).toBe(false);
+
+        const runtimeBeforeReconnect = harness.daemonApp.runtimeManager.get(
+          environment.id,
+        )?.runtime;
+        const previousSessionId = harness.daemonApp.connection.sessionId;
+        if (!runtimeBeforeReconnect || !previousSessionId) {
+          throw new Error("Expected a live runtime and daemon session");
+        }
+
+        harness.daemonApp.connection.handleSessionInvalidated({
+          code: "inactive_session",
+          observedSessionId: previousSessionId,
+          source: "postEvents",
+        });
+        await waitForCondition(
+          () => {
+            const sessionId = harness.daemonApp.connection.sessionId;
+            return (
+              sessionId !== null &&
+              sessionId !== previousSessionId &&
+              harness.hub.getDaemonSessionIdForHost(harness.hostId) ===
+                sessionId
+            );
+          },
+          10_000,
+          "Daemon did not register the replacement session socket",
+        );
+
+        expect(
+          harness.daemonApp.runtimeManager.get(environment.id)?.runtime,
+        ).toBe(runtimeBeforeReconnect);
+
+        await sendAndWaitForIdle({
+          harness,
+          providerId: PROVIDER_ID,
+          threadId: thread.id,
+          text: "Reply exactly FOLLOWUP. Do not use any tools.",
+        });
+        await waitForCondition(
+          () => pathExists(markerPath),
+          45_000,
+          "Background command did not write its marker after reconnect",
+        );
+
+        expect(await fs.readFile(markerPath, "utf8")).toBe("survived");
+      } finally {
+        await harness.cleanup();
+      }
+    },
+    BACKGROUND_SURVIVAL_TEST_TIMEOUT_MS,
+  );
+});