From 70fd8d23c8705df30d88fd463ef56b40894a33b1 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 21 Jul 2026 22:32:06 -0700 Subject: [PATCH] fix(ui): keep session narration live across chat switches (#112526) * fix(ui): refcount session message subscriptions * chore(release): keep changelog generation release-owned --- .../app-sidebar-session-narration.test.ts | 6 +- .../app-sidebar-session-narration.ts | 101 +++------------ ui/src/lib/sessions/index.ts | 117 ++++++++++++++++-- ui/src/pages/chat/chat-history.ts | 60 +++------ 4 files changed, 141 insertions(+), 143 deletions(-) diff --git a/ui/src/components/app-sidebar-session-narration.test.ts b/ui/src/components/app-sidebar-session-narration.test.ts index 8fbb3ac7f2a7..f54aac99a351 100644 --- a/ui/src/components/app-sidebar-session-narration.test.ts +++ b/ui/src/components/app-sidebar-session-narration.test.ts @@ -668,7 +668,7 @@ describe("SidebarSessionNarrationController", () => { expect(updates.at(-1)?.has("agent:main:run")).toBe(false); }); - it("does not let a stale subscribe completion remove replacement ownership", async () => { + it("releases a stale subscribe completion independently of replacement ownership", async () => { const completions: Array<{ resolve: (subscription: { key: string; agentId: null }) => void; promise: Promise<{ key: string; agentId: null }>; @@ -703,9 +703,9 @@ describe("SidebarSessionNarrationController", () => { completions[0]?.resolve({ key: "agent:main:run", agentId: null }); await Promise.resolve(); - expect(unsubscribeMessages).not.toHaveBeenCalled(); - controller.disconnect(); expect(unsubscribeMessages).toHaveBeenCalledTimes(1); + controller.disconnect(); + expect(unsubscribeMessages).toHaveBeenCalledTimes(2); }); it("rebinds an active global session when the selected agent changes", async () => { diff --git a/ui/src/components/app-sidebar-session-narration.ts b/ui/src/components/app-sidebar-session-narration.ts index 1cd8a0f86ebf..f64d40956d36 100644 --- a/ui/src/components/app-sidebar-session-narration.ts +++ b/ui/src/components/app-sidebar-session-narration.ts @@ -37,8 +37,6 @@ type NarrationSubscription = { type PendingSubscription = { agentId: string | null; - connectionIdentity: object; - source: SessionCapability; operationId: symbol; }; @@ -113,12 +111,10 @@ export class SidebarSessionNarrationController { private connectionIdentity: object | null = null; private connected = false; private enabled = false; - private openSessionKey = ""; private agentId = "main"; private desiredKeys = new Set(); private subscriptions = new Map(); private pendingSubscriptions = new Map(); - private deferredSubscriptions = new Map>(); private internalRuntimeBlockDepth = new Map(); private internalRuntimeDelimiterTails = new Map(); // Chars of the FULL cumulative assistant stream consumed so far, per session. @@ -140,26 +136,22 @@ export class SidebarSessionNarrationController { ) {} sync(input: SidebarNarrationSyncInput): void { - const previousOpenSessionKey = this.openSessionKey; const connectionChanged = this.connectionIdentity !== input.connectionIdentity; const sourceChanged = this.source !== input.source; const disconnected = !input.connected || !input.connectionIdentity || !input.source; if (connectionChanged || sourceChanged || disconnected) { - // A replaced or closed socket already discarded its server-side set. - // Never send cleanup through a new connection for ownership from the old one. - this.resetSubscriptions({ unsubscribe: !connectionChanged && this.connected }); + this.resetSubscriptions(); } this.source = input.source; this.connectionIdentity = input.connectionIdentity; this.connected = input.connected; this.enabled = input.enabled; - this.openSessionKey = input.openSessionKey.trim(); this.agentId = normalizeAgentId(input.agentId); if (disconnected || !input.enabled) { this.desiredKeys = new Set(); - this.resetSubscriptions({ unsubscribe: !disconnected }); + this.resetSubscriptions(); this.clearAllLines(); return; } @@ -167,7 +159,8 @@ export class SidebarSessionNarrationController { const candidates = input.rows .map((row, index) => ({ row, index })) .filter( - ({ row }) => rowIsRunning(row) && !areUiSessionKeysEquivalent(row.key, this.openSessionKey), + ({ row }) => + rowIsRunning(row) && !areUiSessionKeysEquivalent(row.key, input.openSessionKey.trim()), ) .toSorted( (left, right) => rowRecency(right.row) - rowRecency(left.row) || left.index - right.index, @@ -179,14 +172,7 @@ export class SidebarSessionNarrationController { if (nextDesired.has(key)) { continue; } - // The chat pane and sidebar share a per-connection Set, not a refcount. - // Hand an opening row to chat without deleting the subscription it now owns. - const ownedAgentId = - this.subscriptions.get(key)?.subscription.agentId ?? - this.pendingSubscriptions.get(key)?.agentId ?? - null; - const handedToChat = this.subscriptionScopeMatchesOpenChat(key, ownedAgentId); - this.releaseKey(key, { unsubscribe: !handedToChat }); + this.releaseKey(key); } this.desiredKeys = nextDesired; @@ -198,13 +184,12 @@ export class SidebarSessionNarrationController { (this.subscriptions.has(key) && ownedAgentId !== targetAgentId) || (this.pendingSubscriptions.has(key) && pendingAgentId !== targetAgentId) ) { - this.releaseKey(key, { unsubscribe: true }); + this.releaseKey(key); } if (this.subscriptions.has(key) || this.pendingSubscriptions.has(key)) { continue; } - const chatJustReleased = areUiSessionKeysEquivalent(key, previousOpenSessionKey); - this.scheduleSubscription(key, chatJustReleased); + void this.subscribeKey(key); } } @@ -227,34 +212,16 @@ export class SidebarSessionNarrationController { disconnect(): void { this.desiredKeys = new Set(); - this.resetSubscriptions({ unsubscribe: this.connected }); + this.resetSubscriptions(); this.clearAllLines(); this.connected = false; } - private scheduleSubscription(key: string, defer: boolean): void { - if (!defer) { - void this.subscribeKey(key); - return; - } - if (this.deferredSubscriptions.has(key)) { - return; - } - // Chat unsubscribes during the same Lit update. Queueing this request makes - // the wire order unsubscribe -> subscribe when a running row leaves chat. - const timer = globalThis.setTimeout(() => { - this.deferredSubscriptions.delete(key); - void this.subscribeKey(key); - }, 0); - this.deferredSubscriptions.set(key, timer); - } - private async subscribeKey(key: string): Promise { const source = this.source; - const connectionIdentity = this.connectionIdentity; if ( !source || - !connectionIdentity || + !this.connectionIdentity || !this.connected || !this.enabled || !this.desiredKeys.has(key) @@ -263,40 +230,24 @@ export class SidebarSessionNarrationController { } const operationId = Symbol(key); const agentId = this.subscriptionAgentId(key); - this.pendingSubscriptions.set(key, { agentId, connectionIdentity, operationId, source }); + this.pendingSubscriptions.set(key, { agentId, operationId }); try { const subscription = await source.subscribeMessages(key, { agentId: agentId ?? undefined, }); const pending = this.pendingSubscriptions.get(key); - if (pending?.operationId !== operationId || pending.source !== source) { - const completedAgentId = subscription.agentId ?? null; - const replacementOwnsSameScope = - this.desiredKeys.has(key) && - (this.subscriptions.get(key)?.subscription.agentId === completedAgentId || - this.pendingSubscriptions.get(key)?.agentId === completedAgentId || - (this.deferredSubscriptions.has(key) && - this.subscriptionAgentId(key) === completedAgentId)); - if ( - !replacementOwnsSameScope && - connectionIdentity === this.connectionIdentity && - !this.subscriptionScopeMatchesOpenChat(key, completedAgentId) - ) { - await source.unsubscribeMessages(subscription).catch(() => undefined); - } + if (pending?.operationId !== operationId) { + await source.unsubscribeMessages(subscription).catch(() => undefined); return; } this.pendingSubscriptions.delete(key); if ( source !== this.source || - connectionIdentity !== this.connectionIdentity || !this.connected || !this.enabled || !this.desiredKeys.has(key) ) { - if (!this.subscriptionScopeMatchesOpenChat(key, subscription.agentId ?? null)) { - await source.unsubscribeMessages(subscription).catch(() => undefined); - } + await source.unsubscribeMessages(subscription).catch(() => undefined); return; } this.subscriptions.set(key, { source, subscription }); @@ -312,36 +263,20 @@ export class SidebarSessionNarrationController { return isUiGlobalSessionKey(key) ? this.agentId : null; } - private subscriptionScopeMatchesOpenChat(key: string, agentId: string | null): boolean { - if (!areUiSessionKeysEquivalent(key, this.openSessionKey)) { - return false; - } - return !isUiGlobalSessionKey(key) || agentId === this.subscriptionAgentId(key); - } - - private releaseKey(key: string, options: { unsubscribe: boolean }): void { - const deferred = this.deferredSubscriptions.get(key); - if (deferred) { - globalThis.clearTimeout(deferred); - this.deferredSubscriptions.delete(key); - } + private releaseKey(key: string): void { this.pendingSubscriptions.delete(key); const owned = this.subscriptions.get(key); this.subscriptions.delete(key); - if (owned && options.unsubscribe) { + if (owned) { void owned.source.unsubscribeMessages(owned.subscription).catch(() => undefined); } this.clearLine(key); } - private resetSubscriptions(options: { unsubscribe: boolean }): void { - const keys = new Set([ - ...this.subscriptions.keys(), - ...this.pendingSubscriptions.keys(), - ...this.deferredSubscriptions.keys(), - ]); + private resetSubscriptions(): void { + const keys = new Set([...this.subscriptions.keys(), ...this.pendingSubscriptions.keys()]); for (const key of keys) { - this.releaseKey(key, options); + this.releaseKey(key); } } diff --git a/ui/src/lib/sessions/index.ts b/ui/src/lib/sessions/index.ts index 39d459e20c66..eecdef5f11a5 100644 --- a/ui/src/lib/sessions/index.ts +++ b/ui/src/lib/sessions/index.ts @@ -44,6 +44,7 @@ import { } from "./reconcile.ts"; import { areUiSessionKeysEquivalent, + isUiGlobalSessionKey, normalizeAgentId, parseAgentSessionKey, resolveUiSelectedGlobalAgentId, @@ -178,7 +179,7 @@ type SessionConnectionScope = { type SessionCreateReconciliation = "blocking" | "background"; -type SessionMessageSubscription = { +export type SessionMessageSubscription = { key: string; agentId?: string | null; }; @@ -506,7 +507,7 @@ function subscribeSessionGateway(client: SessionRequestClient): Promise { return client.request("sessions.subscribe", {}).then(() => undefined); } -async function subscribeSessionMessages( +async function requestSessionMessageSubscription( client: SessionRequestClient, key: string, options: { agentId?: string | null } = {}, @@ -524,7 +525,7 @@ async function subscribeSessionMessages( }; } -export function unsubscribeSessionMessages( +function requestSessionMessageUnsubscribe( client: SessionRequestClient, subscription: SessionMessageSubscription, ): Promise { @@ -536,6 +537,90 @@ export function unsubscribeSessionMessages( .then(() => undefined); } +type SessionMessageSubscriptionEntry = { + key: string; + agentId: string | null; + owners: number; + result: Promise; +}; + +const sessionMessageSubscriptionRegistries = new WeakMap< + GatewayBrowserClient, + Set +>(); +const sessionMessageSubscriptionOwners = new WeakMap< + SessionMessageSubscription, + { + client: GatewayBrowserClient; + entry: SessionMessageSubscriptionEntry; + registry: Set; + onRelease: (subscription: SessionMessageSubscription) => void; + } +>(); + +function resetSessionMessageSubscriptionRegistry(client: GatewayBrowserClient): void { + sessionMessageSubscriptionRegistries.get(client)?.clear(); + sessionMessageSubscriptionRegistries.delete(client); +} + +async function acquireSessionMessageSubscription( + client: GatewayBrowserClient, + key: string, + options: { agentId?: string | null } = {}, + onRelease: (subscription: SessionMessageSubscription) => void = () => undefined, +): Promise { + const normalizedKey = key.trim(); + const agentId = + isUiGlobalSessionKey(normalizedKey) && options.agentId?.trim() + ? normalizeAgentId(options.agentId) + : null; + const registry = sessionMessageSubscriptionRegistries.get(client) ?? new Set(); + sessionMessageSubscriptionRegistries.set(client, registry); + let entry = [...registry].find( + (candidate) => + candidate.agentId === agentId && areUiSessionKeysEquivalent(candidate.key, normalizedKey), + ); + if (!entry) { + const result = requestSessionMessageSubscription(client, normalizedKey, { agentId }); + entry = { key: normalizedKey, agentId, owners: 0, result }; + registry.add(entry); + void result.catch(() => registry.delete(entry!)); + } + entry.owners += 1; + try { + const resolved = await entry.result; + const subscription: SessionMessageSubscription = { + key: resolved.key, + agentId: resolved.agentId ?? null, + }; + sessionMessageSubscriptionOwners.set(subscription, { client, entry, registry, onRelease }); + return subscription; + } catch (error) { + entry.owners -= 1; + throw error; + } +} + +async function releaseSessionMessageSubscription( + subscription: SessionMessageSubscription, +): Promise { + const owner = sessionMessageSubscriptionOwners.get(subscription); + if (!owner) { + return; + } + sessionMessageSubscriptionOwners.delete(subscription); + owner.onRelease(subscription); + owner.entry.owners -= 1; + if ( + owner.entry.owners > 0 || + sessionMessageSubscriptionRegistries.get(owner.client) !== owner.registry || + !owner.registry.delete(owner.entry) + ) { + return; + } + await requestSessionMessageUnsubscribe(owner.client, subscription); +} + async function listSessionCheckpoints( client: SessionRequestClient, key: string, @@ -746,6 +831,7 @@ export function createSessionCapability(gateway: SessionGateway): SessionCapabil let hasSeededListOptions = false; const listeners = new Set<(next: SessionState) => void>(); const createdListeners = new Set<(key: string) => void>(); + const ownedMessageSubscriptions = new Set(); const captureConnection = (): SessionConnectionScope | null => { const snapshot = gateway.snapshot; @@ -1537,20 +1623,21 @@ export function createSessionCapability(gateway: SessionGateway): SessionCapabil if (!scope) { throw new Error("Session message subscription requires an active Gateway connection"); } - const subscription = await subscribeSessionMessages(scope.client, key, options); + const subscription = await acquireSessionMessageSubscription( + scope.client, + key, + options, + (released) => ownedMessageSubscriptions.delete(released), + ); + ownedMessageSubscriptions.add(subscription); if (!isCurrentConnection(scope)) { + await releaseSessionMessageSubscription(subscription).catch(() => undefined); throw new Error("Session message subscription completed on a replaced Gateway connection"); } return subscription; }; - const unsubscribeMessages = async (subscription: SessionMessageSubscription) => { - const scope = captureConnection(); - if (!scope) { - return; - } - await unsubscribeSessionMessages(scope.client, subscription); - }; + const unsubscribeMessages = releaseSessionMessageSubscription; const listCheckpoints = async ( key: string, @@ -1665,6 +1752,7 @@ export function createSessionCapability(gateway: SessionGateway): SessionCapabil }; const stopGateway = gateway.subscribe((next) => { + const previousClient = connectionClient; const connectionChanged = next.client !== connectionClient || next.connected !== connectionConnected; connectionClient = next.client; @@ -1672,6 +1760,10 @@ export function createSessionCapability(gateway: SessionGateway): SessionCapabil if (connectionChanged) { const hadPullRequestSummaries = pullRequestSummaries.size > 0; connectionEpoch += 1; + if (previousClient) { + resetSessionMessageSubscriptionRegistry(previousClient); + } + ownedMessageSubscriptions.clear(); invalidateGroupsLoad(); swarmActivity.clear(); inFlight = null; @@ -1823,6 +1915,9 @@ export function createSessionCapability(gateway: SessionGateway): SessionCapabil return () => listeners.delete(listener); }, dispose() { + for (const subscription of ownedMessageSubscriptions) { + void releaseSessionMessageSubscription(subscription).catch(() => undefined); + } disposed = true; connectionEpoch += 1; invalidateGroupsLoad(); diff --git a/ui/src/pages/chat/chat-history.ts b/ui/src/pages/chat/chat-history.ts index 1eb247ccd8e7..7df3c012aa4f 100644 --- a/ui/src/pages/chat/chat-history.ts +++ b/ui/src/pages/chat/chat-history.ts @@ -29,9 +29,9 @@ import { isGatewayMethodAdvertised } from "../../lib/gateway-methods.ts"; import { isSessionRunActive } from "../../lib/session-run-state.ts"; import { scopedAgentParamsForSession, - unsubscribeSessionMessages, visibleSessionMatches, type SessionCapability, + type SessionMessageSubscription, } from "../../lib/sessions/index.ts"; import { areUiSessionKeysEquivalent, @@ -377,12 +377,11 @@ type ChatAgentsListSnapshot = Partial> & { }; type ChatSessionMessageSubscriptionState = ChatState & { - sessions: Pick; + sessions: Pick; sessionsResult?: SessionsListResult | null; sessionsError?: string | null; chatSessionMessageSubscriptionRequestedKey?: string | null; - chatSessionMessageSubscriptionKey?: string | null; - chatSessionMessageSubscriptionAgentId?: string | null; + chatSessionMessageSubscription?: SessionMessageSubscription | null; }; export type ChatHistoryResult = { @@ -651,21 +650,6 @@ function isCurrentSelectedSessionMessageSubscriptionSync( ); } -async function unsubscribeSelectedSessionMessageBestEffort( - client: GatewayBrowserClient, - key: string, - agentId?: string | null, -): Promise { - try { - await unsubscribeSessionMessages(client, { - key, - agentId: isUiGlobalSessionKey(key) ? agentId : null, - }); - } catch { - // Cleanup is best effort when a stale subscription completion loses ownership. - } -} - export async function syncSelectedSessionMessageSubscription( state: ChatSessionMessageSubscriptionState, opts?: { force?: boolean }, @@ -682,16 +666,18 @@ export async function syncSelectedSessionMessageSubscription( const previousRequestedKey = normalizeSubscriptionKey( state.chatSessionMessageSubscriptionRequestedKey, ); - const previousCanonicalKey = normalizeSubscriptionKey(state.chatSessionMessageSubscriptionKey); + const previousSubscription = state.chatSessionMessageSubscription ?? null; + const previousCanonicalKey = normalizeSubscriptionKey(previousSubscription?.key); const previousSelectedKey = previousRequestedKey ?? previousCanonicalKey; const nextSubscriptionAgentId = resolveSelectedSessionMessageSubscriptionAgentId(state, nextKey); const selectedAgentChanged = nextSubscriptionAgentId !== null && previousSelectedKey === nextKey && - (state.chatSessionMessageSubscriptionAgentId ?? null) !== nextSubscriptionAgentId; + (previousSubscription?.agentId ?? null) !== nextSubscriptionAgentId; const selectedKeyChanged = previousSelectedKey !== null && previousSelectedKey !== nextKey; const shouldUnsubscribePrevious = - previousCanonicalKey !== null && (selectedKeyChanged || selectedAgentChanged); + previousSubscription !== null && + (opts?.force === true || selectedKeyChanged || selectedAgentChanged); const shouldSubscribe = opts?.force === true || selectedKeyChanged || @@ -709,19 +695,12 @@ export async function syncSelectedSessionMessageSubscription( requestedAgentId: nextSubscriptionAgentId, }); try { - if (shouldUnsubscribePrevious && previousCanonicalKey) { - await unsubscribeSessionMessages(client, { - key: previousCanonicalKey, - agentId: - isUiGlobalSessionKey(previousCanonicalKey) && state.chatSessionMessageSubscriptionAgentId - ? state.chatSessionMessageSubscriptionAgentId - : null, - }); + if (shouldUnsubscribePrevious && previousSubscription) { if (isCurrent()) { - state.chatSessionMessageSubscriptionKey = null; state.chatSessionMessageSubscriptionRequestedKey = null; - state.chatSessionMessageSubscriptionAgentId = null; + state.chatSessionMessageSubscription = null; } + await state.sessions.unsubscribeMessages(previousSubscription); } if (!shouldSubscribe || !isCurrent()) { return; @@ -730,23 +709,12 @@ export async function syncSelectedSessionMessageSubscription( agentId: nextSubscriptionAgentId ?? undefined, }); if (!isCurrent()) { - const staleKeyChanged = - normalizeSubscriptionKey(state.chatSessionMessageSubscriptionKey) !== subscribed.key; - const staleAgentChanged = - isUiGlobalSessionKey(subscribed.key) && - (state.chatSessionMessageSubscriptionAgentId ?? null) !== subscribed.agentId; - if (staleKeyChanged || staleAgentChanged) { - await unsubscribeSelectedSessionMessageBestEffort( - client, - subscribed.key, - subscribed.agentId, - ); - } + // Generation advances before awaiting, so only the newest lease can reach assignment below. + await state.sessions.unsubscribeMessages(subscribed).catch(() => undefined); return; } state.chatSessionMessageSubscriptionRequestedKey = nextKey; - state.chatSessionMessageSubscriptionKey = subscribed.key; - state.chatSessionMessageSubscriptionAgentId = subscribed.agentId; + state.chatSessionMessageSubscription = subscribed; } catch (err) { if (isCurrent()) { state.sessionsError = String(err);