fix(ui): keep session narration live across chat switches (#112526)

* fix(ui): refcount session message subscriptions * chore(release): keep changelog generation release-owned
2026-08-28 05:16:23 -06:00 · 2026-07-21 22:32:06 -07:00
parent 1a42e005fb
commit 70fd8d23c8
4 changed files with 141 additions and 143 deletions
@@ -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 () => {
@@ -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<string>();
  private subscriptions = new Map<string, NarrationSubscription>();
  private pendingSubscriptions = new Map<string, PendingSubscription>();
-  private deferredSubscriptions = new Map<string, ReturnType<typeof globalThis.setTimeout>>();
  private internalRuntimeBlockDepth = new Map<string, number>();
  private internalRuntimeDelimiterTails = new Map<string, string>();
  // 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<void> {
    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);
    }
  }

@@ -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<void> {
  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<void> {
@@ -536,6 +537,90 @@ export function unsubscribeSessionMessages(
    .then(() => undefined);
 }

+type SessionMessageSubscriptionEntry = {
+  key: string;
+  agentId: string | null;
+  owners: number;
+  result: Promise<SessionMessageSubscription>;
+};
+
+const sessionMessageSubscriptionRegistries = new WeakMap<
+  GatewayBrowserClient,
+  Set<SessionMessageSubscriptionEntry>
+>();
+const sessionMessageSubscriptionOwners = new WeakMap<
+  SessionMessageSubscription,
+  {
+    client: GatewayBrowserClient;
+    entry: SessionMessageSubscriptionEntry;
+    registry: Set<SessionMessageSubscriptionEntry>;
+    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<SessionMessageSubscription> {
+  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<void> {
+  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<SessionMessageSubscription>();

  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();
@@ -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<Omit<AgentsListResult, "agents">> & {
 };

 type ChatSessionMessageSubscriptionState = ChatState & {
-  sessions: Pick<SessionCapability, "subscribeMessages">;
+  sessions: Pick<SessionCapability, "subscribeMessages" | "unsubscribeMessages">;
  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<void> {
-  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);