From 7fdaa65b9710d806f94a4b91e334f7a0d4e9bac4 Mon Sep 17 00:00:00 2001 From: Logan Johnson Date: Fri, 4 Sep 2026 23:03:24 -0400 Subject: [PATCH] fix(desktop): observe receiver closure after subscription readiness Signed-off-by: Logan Johnson --- .../features/agents/desktopLifecycle.test.mjs | 11 +- .../src/features/agents/desktopLifecycle.ts | 34 ++- .../ui/DesktopLifecycleControl.test.mjs | 12 +- desktop/src/shared/api/relayClientSession.ts | 25 +- desktop/src/shared/api/relayClientShared.ts | 15 + desktop/src/shared/api/relayClosedRecovery.ts | 17 +- .../api/relayLiveSubscriptionState.test.mjs | 268 ++++++++++++++++++ 7 files changed, 355 insertions(+), 27 deletions(-) create mode 100644 desktop/src/shared/api/relayLiveSubscriptionState.test.mjs diff --git a/desktop/src/features/agents/desktopLifecycle.test.mjs b/desktop/src/features/agents/desktopLifecycle.test.mjs index 33de28d85b..2a0ff0f314 100644 --- a/desktop/src/features/agents/desktopLifecycle.test.mjs +++ b/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++; diff --git a/desktop/src/features/agents/desktopLifecycle.ts b/desktop/src/features/agents/desktopLifecycle.ts index 1fb321c67a..b81e482605 100644 --- a/desktop/src/features/agents/desktopLifecycle.ts +++ b/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 = () => { diff --git a/desktop/src/features/agents/ui/DesktopLifecycleControl.test.mjs b/desktop/src/features/agents/ui/DesktopLifecycleControl.test.mjs index 3ac174ac53..35e52f7aab 100644 --- a/desktop/src/features/agents/ui/DesktopLifecycleControl.test.mjs +++ b/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++; }; diff --git a/desktop/src/shared/api/relayClientSession.ts b/desktop/src/shared/api/relayClientSession.ts index 2f94a65bf9..a7c0bded02 100644 --- a/desktop/src/shared/api/relayClientSession.ts +++ b/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); diff --git a/desktop/src/shared/api/relayClientShared.ts b/desktop/src/shared/api/relayClientShared.ts index 8bd6379d7a..925c74b573 100644 --- a/desktop/src/shared/api/relayClientShared.ts +++ b/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. diff --git a/desktop/src/shared/api/relayClosedRecovery.ts b/desktop/src/shared/api/relayClosedRecovery.ts index 9d30a233e7..8f1a68287d 100644 --- a/desktop/src/shared/api/relayClosedRecovery.ts +++ b/desktop/src/shared/api/relayClosedRecovery.ts @@ -124,10 +124,24 @@ function recoverLiveSubscriptionFromClosed({ sendReq: (subId: string, filter: RelaySubscriptionFilter) => Promise; }) { 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); diff --git a/desktop/src/shared/api/relayLiveSubscriptionState.test.mjs b/desktop/src/shared/api/relayLiveSubscriptionState.test.mjs new file mode 100644 index 0000000000..587bfae1f2 --- /dev/null +++ b/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(); +});