fix(desktop): observe receiver closure after subscription readiness

Signed-off-by: Logan Johnson <loganj@squareup.com>

desktop/src/features/agents/desktopLifecycle.test.mjs

@@ -342,10 +342,17 @@ test("readiness timeout is distinct, late EOSE recovers, CLOSED retires old call
   let notify, deliver;
   let closed = 0,
     ready = 0;
-  f.relay.subscribeLive = async (_filter, event, onReady, timeout) => {
+  f.relay.subscribeLive = async (
+    _filter,
+    event,
+    _onReady,
+    timeout,
+    options,
+  ) => {
     assert.equal(timeout, 5000);
     deliver = event;
-    notify = onReady;
+    assert.equal(options.closedRecovery, "explicit");
+    notify = options.onState;
     notify("timeout");
     return () => {
       closed++;

desktop/src/features/agents/desktopLifecycle.ts

@@ -335,22 +335,26 @@ export async function receiveLifecycle(
             pending--;
           });
       },
-      (readiness) => {
-        if (!valid()) return;
-        subscriptionReady = readiness === "eose";
-        if (readiness === "closed") {
-          stopped = true;
-          stopSubscription();
-          onError(
-            "Desktop lifecycle receiver subscription closed. Retry the receiver to accept new requests.",
-          );
-        } else if (readiness === "timeout") {
-          onError(
-            "Desktop lifecycle subscription readiness timed out. Delivery is unconfirmed.",
-          );
-        } else if (synced) onReady();
-      },
+      undefined,
       5000,
+      {
+        closedRecovery: "explicit",
+        onState: (readiness) => {
+          if (!valid()) return;
+          subscriptionReady = readiness === "eose";
+          if (readiness === "closed") {
+            stopped = true;
+            stopSubscription();
+            onError(
+              "Desktop lifecycle receiver subscription closed. Retry the receiver to accept new requests.",
+            );
+          } else if (readiness === "timeout") {
+            onError(
+              "Desktop lifecycle subscription readiness timed out. Delivery is unconfirmed.",
+            );
+          } else if (synced) onReady();
+        },
+      },
     ),
   );
   const close = () => {

desktop/src/features/agents/ui/DesktopLifecycleControl.test.mjs

@@ -166,10 +166,16 @@ test("receiver failure is a scope-owned notification, not pre-shell layout", asy
       });
     return [];
   };
-  relayClient.subscribeLive = async (_filter, _event, onReadiness) => {
+  relayClient.subscribeLive = async (
+    _filter,
+    _event,
+    _onReadiness,
+    _timeout,
+    options,
+  ) => {
     subscribed++;
-    readiness = onReadiness;
-    onReadiness("eose");
+    readiness = options.onState;
+    readiness("eose");
     return () => {
       closed++;
     };

desktop/src/shared/api/relayClientSession.ts

@@ -16,6 +16,7 @@ import {
   getTextPayload,
   toRelayFrames,
   type ConnectionState,
+  type LiveSubscriptionOptions,
   type LiveSubscriptionReadiness,
   type PendingEvent,
   type RelaySubscription,
@@ -419,8 +420,15 @@ export class RelayClient {
     onEvent: (event: RelayEvent) => void,
     onReady?: (readiness: LiveSubscriptionReadiness) => void,
     readinessTimeoutMs?: number,
+    options?: LiveSubscriptionOptions,
   ) {
-    return this.subscribe(filter, onEvent, onReady, readinessTimeoutMs);
+    return this.subscribe(
+      filter,
+      onEvent,
+      onReady,
+      readinessTimeoutMs,
+      options,
+    );
   }
   async subscribeToChannelMentionEvents(
     channelId: string,
@@ -609,6 +617,7 @@ export class RelayClient {
     onEvent: (event: RelayEvent) => void,
     onReady?: (readiness: LiveSubscriptionReadiness) => void,
     readinessTimeoutMs = 250,
+    options: LiveSubscriptionOptions = {},
   ) {
     await this.ensureConnected();
 
@@ -621,16 +630,18 @@ export class RelayClient {
         resolve();
       };
     });
-    const fallbackTimeout = window.setTimeout(
-      () => resolveReady("timeout"),
-      readinessTimeoutMs,
-    );
+    const fallbackTimeout = window.setTimeout(() => {
+      options.onState?.("timeout");
+      resolveReady("timeout");
+    }, readinessTimeoutMs);
 
     this.subscriptions.set(subId, {
       mode: "live",
       filter,
       onEvent,
       resolveReady,
+      onState: options.onState,
+      closedRecovery: options.closedRecovery ?? "shared",
     });
 
     try {
@@ -1066,8 +1077,12 @@ export class RelayClient {
         continue;
       }
       subscription.resolveReady?.("closed");
+      subscription.onState?.("closed");
       subscription.resolveReady = undefined;
       clearClosedRetry(subscription);
+      if (subscription.closedRecovery === "explicit") {
+        this.subscriptions.delete(subId);
+      }
     }
     for (const [eventId, pendingEvent] of this.pendingEvents) {
       window.clearTimeout(pendingEvent.timeout);

desktop/src/shared/api/relayClientShared.ts

@@ -70,11 +70,26 @@ type FirstEventSubscription = {
 
 export type LiveSubscriptionReadiness = "eose" | "closed" | "timeout";
 
+/**
+ * Optional lifecycle policy for a live subscription.
+ *
+ * Most subscriptions keep the shared reconnect/CLOSED recovery behavior. A
+ * command receiver can instead request explicit recovery: every CLOSED retires
+ * that subscription and `onState` remains observable after initial EOSE so the
+ * owning UI can offer a deliberate fresh subscription.
+ */
+export type LiveSubscriptionOptions = {
+  onState?: (state: LiveSubscriptionReadiness) => void;
+  closedRecovery?: "shared" | "explicit";
+};
+
 type LiveSubscription = {
   mode: "live";
   filter: RelaySubscriptionFilter;
   onEvent: (event: RelayEvent) => void;
   resolveReady?: (readiness: LiveSubscriptionReadiness) => void;
+  onState?: (state: LiveSubscriptionReadiness) => void;
+  closedRecovery: "shared" | "explicit";
   lastSeenCreatedAt?: number;
   /**
    * Lower bound of a reconnect backfill window that has not yet completed.

desktop/src/shared/api/relayClosedRecovery.ts

@@ -124,10 +124,24 @@ function recoverLiveSubscriptionFromClosed({
   sendReq: (subId: string, filter: RelaySubscriptionFilter) => Promise<void>;
 }) {
   subscription.resolveReady?.("closed");
+  subscription.onState?.("closed");
   subscription.resolveReady = undefined;
 
   const closedClass = classifyRelayClosed(message);
 
+  if (closedClass === "rate-limited") {
+    const hintSeconds = parseRateLimitHint(message);
+    activateRateLimit(hintSeconds);
+  }
+
+  if (subscription.closedRecovery === "explicit") {
+    // Command receivers must not survive CLOSED into shared re-subscription.
+    // Their owner presents an explicit fresh-receiver action instead.
+    clearClosedRetry(subscription);
+    subscriptions.delete(subId);
+    return;
+  }
+
   if (closedClass === "terminal") {
     // Auth/access/filter failure — permanently remove the subscription so it
     // doesn't silently loop.
@@ -146,9 +160,7 @@ function recoverLiveSubscriptionFromClosed({
   let delayMs = backoffMs;
 
   if (closedClass === "rate-limited") {
-    // Activate the gate so concurrent operations back off too.
     const hintSeconds = parseRateLimitHint(message);
-    activateRateLimit(hintSeconds);
     // Use the gate's actual remaining time so a shorter hint arriving under a
     // longer active gate does not schedule a premature retry that just gets
     // another CLOSED. The fallback covers the gate-inactive edge case
@@ -259,6 +271,7 @@ export function handleSubscriptionEose({
     if (generation !== undefined)
       markReconnectLiveEose(subscription, generation);
     subscription.resolveReady?.("eose");
+    subscription.onState?.("eose");
     subscription.resolveReady = undefined;
     subscription.closedRetryAttempt = 0;
     clearClosedRetry(subscription);

desktop/src/shared/api/relayLiveSubscriptionState.test.mjs

@@ -0,0 +1,268 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+let fakeNow = 0;
+let nextTimerId = 1;
+const pendingTimers = new Map();
+const sentFrames = [];
+
+globalThis.window = {
+  setTimeout: (fn, ms) => {
+    const id = nextTimerId++;
+    pendingTimers.set(id, { fn, fireAt: fakeNow + ms });
+    return id;
+  },
+  clearTimeout: (id) => pendingTimers.delete(id),
+  __TAURI_INTERNALS__: {
+    invoke: async (command, args) => {
+      if (command === "plugin:websocket|send") sentFrames.push(args);
+    },
+  },
+};
+Date.now = () => fakeNow;
+
+const { RelayClient } = await import("./relayClientSession.ts");
+const { receiveLifecycle } = await import("@/features/agents/desktopLifecycle");
+const { resetRateLimitGate } = await import("./relayRateLimitGate.ts");
+
+function resetHarness() {
+  fakeNow = 0;
+  nextTimerId = 1;
+  pendingTimers.clear();
+  sentFrames.length = 0;
+  resetRateLimitGate();
+}
+
+function connectedClient() {
+  const client = new RelayClient();
+  client.wsId = 7;
+  return client;
+}
+
+function sentProtocolFrames(type) {
+  return sentFrames
+    .map(({ message }) => JSON.parse(message.data))
+    .filter((frame) => frame[0] === type);
+}
+
+async function flushUntil(predicate, attempts = 40) {
+  for (let attempt = 0; attempt < attempts; attempt++) {
+    if (predicate()) return;
+    await Promise.resolve();
+  }
+  assert.fail("condition did not become true before the microtask limit");
+}
+
+async function flushMicrotasks(attempts = 10) {
+  for (let attempt = 0; attempt < attempts; attempt++) await Promise.resolve();
+}
+
+function tickTo(time) {
+  fakeNow = time;
+  for (;;) {
+    const due = [...pendingTimers.entries()].filter(
+      ([, timer]) => timer.fireAt <= fakeNow,
+    );
+    if (!due.length) return;
+    for (const [id, timer] of due) {
+      if (!pendingTimers.delete(id)) continue;
+      timer.fn();
+    }
+  }
+}
+
+function deliver(client, frame) {
+  return client.handleWsMessage(
+    { type: "Text", data: JSON.stringify(frame) },
+    client.connectionGeneration,
+  );
+}
+
+async function openLive(client, options, onEvent = () => {}, onReady) {
+  const opened = client.subscribeLive(
+    { kinds: [50182, 50180], authors: ["owner"], limit: 0 },
+    onEvent,
+    onReady,
+    5000,
+    options,
+  );
+  await flushUntil(() => sentProtocolFrames("REQ").length > 0);
+  const subId = sentProtocolFrames("REQ").at(-1)[1];
+  return { opened, subId };
+}
+
+test("persistent state reports timeout, late EOSE, then CLOSED through RelayClient", async () => {
+  resetHarness();
+  const client = connectedClient();
+  const states = [];
+  const readiness = [];
+  const { opened, subId } = await openLive(
+    client,
+    { closedRecovery: "explicit", onState: (state) => states.push(state) },
+    () => {},
+    (state) => readiness.push(state),
+  );
+
+  tickTo(5000);
+  const close = await opened;
+  assert.deepEqual(states, ["timeout"]);
+  assert.deepEqual(readiness, ["timeout"]);
+
+  await deliver(client, ["EOSE", subId]);
+  await deliver(client, ["CLOSED", subId, "restricted: access revoked"]);
+
+  assert.deepEqual(states, ["timeout", "eose", "closed"]);
+  assert.deepEqual(readiness, ["timeout", "eose"]);
+  assert.equal(client.subscriptions.has(subId), false);
+  await close();
+});
+
+for (const [label, message] of [
+  ["terminal", "restricted: access revoked"],
+  ["retryable", "error: storage temporarily unavailable"],
+]) {
+  test(`explicit recovery retires an EOSE-ready ${label} CLOSED without re-REQ`, async () => {
+    resetHarness();
+    const client = connectedClient();
+    const states = [];
+    const { opened, subId } = await openLive(client, {
+      closedRecovery: "explicit",
+      onState: (state) => states.push(state),
+    });
+    await deliver(client, ["EOSE", subId]);
+    const close = await opened;
+
+    await deliver(client, ["CLOSED", subId, message]);
+    tickTo(60_000);
+    await Promise.resolve();
+
+    assert.deepEqual(states, ["eose", "closed"]);
+    assert.equal(client.subscriptions.has(subId), false);
+    assert.equal(sentProtocolFrames("REQ").length, 1);
+    await close();
+  });
+}
+
+test("connection reset reports CLOSED and retires only explicit-recovery subscriptions", async () => {
+  resetHarness();
+  const client = connectedClient();
+  const explicitStates = [];
+  const explicit = await openLive(client, {
+    closedRecovery: "explicit",
+    onState: (state) => explicitStates.push(state),
+  });
+  await deliver(client, ["EOSE", explicit.subId]);
+  const closeExplicit = await explicit.opened;
+
+  const ordinary = await openLive(client);
+  await deliver(client, ["EOSE", ordinary.subId]);
+  const closeOrdinary = await ordinary.opened;
+
+  client.resetConnection(new Error("fixture connection reset"));
+
+  assert.deepEqual(explicitStates, ["eose", "closed"]);
+  assert.equal(client.subscriptions.has(explicit.subId), false);
+  assert.equal(
+    client.subscriptions.has(ordinary.subId),
+    true,
+    "default live subscribers retain shared reconnect recovery",
+  );
+
+  await closeExplicit();
+  await closeOrdinary();
+  client.disconnect();
+});
+
+test("default live subscribers retain shared retry after retryable CLOSED", async () => {
+  resetHarness();
+  const client = connectedClient();
+  const ordinary = await openLive(client);
+  await deliver(client, ["EOSE", ordinary.subId]);
+  const close = await ordinary.opened;
+
+  await deliver(client, [
+    "CLOSED",
+    ordinary.subId,
+    "error: storage temporarily unavailable",
+  ]);
+  assert.equal(client.subscriptions.has(ordinary.subId), true);
+
+  tickTo(1000);
+  await flushUntil(() => sentProtocolFrames("REQ").length === 2);
+  assert.equal(client.subscriptions.has(ordinary.subId), true);
+
+  await close();
+});
+
+test("real explicit CLOSED fences queued lifecycle work and a deliberate retry is fresh", async () => {
+  resetHarness();
+  const client = connectedClient();
+  let resolveHistory;
+  client.fetchEvents = () =>
+    new Promise((resolve) => {
+      resolveHistory = resolve;
+    });
+  const ipcCalls = [];
+  const ipc = async (command) => {
+    ipcCalls.push(command);
+    return null;
+  };
+  const errors = [];
+  const receiving = receiveLifecycle(
+    { owner: "owner", community: "wss://one.example" },
+    () => true,
+    (error) => errors.push(error),
+    ipc,
+    client,
+  );
+
+  await flushUntil(() => sentProtocolFrames("REQ").length === 1);
+  const retiredSubId = sentProtocolFrames("REQ")[0][1];
+  await deliver(client, ["EOSE", retiredSubId]);
+  await flushUntil(() => resolveHistory !== undefined);
+  await deliver(client, [
+    "EVENT",
+    retiredSubId,
+    { id: "queued-old", kind: 50182, created_at: 1 },
+  ]);
+  tickTo(20);
+  await deliver(client, ["CLOSED", retiredSubId, "restricted: access revoked"]);
+  resolveHistory([]);
+
+  await assert.rejects(receiving, /receiver is unavailable/);
+  await flushMicrotasks();
+  assert.equal(
+    ipcCalls.includes("receive_desktop_lifecycle"),
+    false,
+    "queued work from the retired receiver must not execute",
+  );
+  assert.match(errors.at(-1), /subscription closed/);
+
+  client.fetchEvents = async () => [];
+  const retried = receiveLifecycle(
+    { owner: "owner", community: "wss://one.example" },
+    () => true,
+    (error) => errors.push(error),
+    ipc,
+    client,
+  );
+  await flushUntil(() => sentProtocolFrames("REQ").length === 2);
+  const freshSubId = sentProtocolFrames("REQ")[1][1];
+  assert.notEqual(freshSubId, retiredSubId);
+  await deliver(client, ["EOSE", freshSubId]);
+  const close = await retried;
+
+  await deliver(client, [
+    "EVENT",
+    freshSubId,
+    { id: "new-live", kind: 50182, created_at: 2 },
+  ]);
+  tickTo(40);
+  await flushUntil(() => ipcCalls.includes("receive_desktop_lifecycle"));
+  assert.equal(
+    ipcCalls.filter((command) => command === "receive_desktop_lifecycle")
+      .length,
+    1,
+  );
+  await close();
+});