refactor: simplify the stale-conversation gate changes

pauldambra · charlesvien · commit 0a8bdca0eb6d · 2026-07-05T17:18:37.000-07:00
- clamp selectedIndex in useActionSelectorState when options change,
  replacing the consumer-level remount key (fixes the primitive for
  every consumer with live-changing options)
- extract ComposerSlot/ComposerWidth and ConnectingToAgent in
  SessionView, collapsing three copies of the composer-slot wrapper and
  two copies of the connecting spinner
- reuse formatUsd instead of hand-rolled toFixed(2)
- plain functions instead of useCallbacks for the gate handlers (the
  notice isn't memoized and the closed-over usage values change per
  streamed event anyway); single guard on onNewSession
- flatten the gate store's engagement Map to number | null and keep the
  Map reference stable when acknowledge has nothing to release

Generated-By: PostHog Code
Task-Id: 1ec8e82e-b126-48b8-8fd4-54edcaf041d5
diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx
@@ -99,6 +99,54 @@ interface SessionViewProps {
 const DEFAULT_ERROR_MESSAGE =
   "Failed to resume this session. The working directory may have been deleted. Please start a new session.";
 
+function ConnectingToAgent() {
+  return (
+    <>
+      <Spinner size={28} className="animate-spin text-gray-9" />
+      <Text color="gray" className="text-base">
+        Connecting to agent...
+      </Text>
+    </>
+  );
+}
+
+/** Centers composer-slot content at the chat width (or compact padding). */
+function ComposerWidth({
+  compact,
+  children,
+}: {
+  compact: boolean;
+  children: React.ReactNode;
+}) {
+  return (
+    <Box
+      className={compact ? "p-1" : "mx-auto px-2 pb-3"}
+      style={compact ? undefined : { maxWidth: CHAT_CONTENT_MAX_WIDTH }}
+    >
+      {children}
+    </Box>
+  );
+}
+
+/**
+ * Input region replacing the composer: `shrink-0` keeps it from being
+ * compressed by the scroller above, and `min-h-0 overflow-y-auto` lets tall
+ * content scroll inside itself.
+ */
+function ComposerSlot({
+  compact,
+  children,
+}: {
+  compact: boolean;
+  children: React.ReactNode;
+}) {
+  return (
+    <Box className="min-h-0 shrink-0 overflow-y-auto">
+      <ComposerWidth compact={compact}>{children}</ComposerWidth>
+    </Box>
+  );
+}
+
 interface CloudStreamDisconnectedBannerProps {
   errorTitle?: string;
   errorMessage?: string;
@@ -276,35 +324,38 @@ export function SessionView({
   // prompt cache has likely expired (see useStaleConversationGate).
   const staleGate = useStaleConversationGate(sessionId, events);
 
-  const trackStaleGateChoice = useCallback(
-    (choice: StaleConversationGateChoiceProperties["choice"]) =>
-      track(ANALYTICS_EVENTS.STALE_CONVERSATION_GATE_CHOICE, {
-        choice,
-        used_tokens: staleGate.usedTokens,
-        cost_usd: staleGate.costUsd,
-      }),
-    [staleGate.usedTokens, staleGate.costUsd],
-  );
+  // Plain functions, not useCallbacks: the notice isn't memoized, and the
+  // values they close over (usage, cost) change with every streamed event.
+  const trackStaleGateChoice = (
+    choice: StaleConversationGateChoiceProperties["choice"],
+  ) =>
+    track(ANALYTICS_EVENTS.STALE_CONVERSATION_GATE_CHOICE, {
+      choice,
+      used_tokens: staleGate.usedTokens,
+      cost_usd: staleGate.costUsd,
+    });
 
-  const handleStaleCompact = useCallback(() => {
+  const handleStaleCompact = () => {
     if (!isOnline) {
       showOfflineToast();
       return;
     }
     trackStaleGateChoice("compact");
     staleGate.onContinue();
     onSendPrompt("/compact");
-  }, [isOnline, trackStaleGateChoice, staleGate.onContinue, onSendPrompt]);
+  };
 
-  const handleStaleContinue = useCallback(() => {
+  const handleStaleContinue = () => {
     trackStaleGateChoice("continue");
     staleGate.onContinue();
-  }, [trackStaleGateChoice, staleGate.onContinue]);
+  };
 
-  const handleStaleNewSession = useCallback(() => {
-    trackStaleGateChoice("new_session");
-    onNewSession?.();
-  }, [trackStaleGateChoice, onNewSession]);
+  const handleStaleNewSession = onNewSession
+    ? () => {
+        trackStaleGateChoice("new_session");
+        onNewSession();
+      }
+    : undefined;
 
   const [isDraggingFile, setIsDraggingFile] = useState(false);
   const editorRef = useRef<PromptInputHandle>(null);
@@ -611,42 +662,20 @@ export function SessionView({
                   // Replaces the composer (and any pending permission — answering
                   // one also resumes the costly turn) until the user chooses.
                   isRunning ? (
-                    <Box className="min-h-0 shrink-0 overflow-y-auto">
-                      <Box
-                        className={compact ? "p-1" : "mx-auto px-2 pb-3"}
-                        style={
-                          compact
+                    <ComposerSlot compact={compact}>
+                      <StaleConversationCostNotice
+                        usedTokens={staleGate.usedTokens}
+                        lastActivityAt={staleGate.lastActivityAt}
+                        costUsd={staleGate.costUsd}
+                        onContinue={handleStaleContinue}
+                        onCompact={
+                          firstPendingPermission
                             ? undefined
-                            : { maxWidth: CHAT_CONTENT_MAX_WIDTH }
+                            : handleStaleCompact
                         }
-                      >
-                        <StaleConversationCostNotice
-                          // Remount when the option set changes so the
-                          // selector's highlighted index can't dereference or
-                          // fire a different option than the one shown.
-                          key={
-                            firstPendingPermission
-                              ? "no-compact"
-                              : "with-compact"
-                          }
-                          usedTokens={staleGate.usedTokens}
-                          lastActivityAt={staleGate.lastActivityAt}
-                          costUsd={staleGate.costUsd}
-                          onContinue={handleStaleContinue}
-                          // A queued /compact would land after the pending
-                          // permission resumes the costly turn — paying the
-                          // reload twice — so the option hides until then.
-                          onCompact={
-                            firstPendingPermission
-                              ? undefined
-                              : handleStaleCompact
-                          }
-                          onNewSession={
-                            onNewSession ? handleStaleNewSession : undefined
-                          }
-                        />
-                      </Box>
-                    </Box>
+                        onNewSession={handleStaleNewSession}
+                      />
+                    </ComposerSlot>
                   ) : (
                     // While reconnecting the gate still covers the composer
                     // slot: handoff can leave pendingPermissions set, and the
@@ -657,33 +686,18 @@ export function SessionView({
                       gap="2"
                       className="min-h-[66px]"
                     >
-                      <Spinner size={28} className="animate-spin text-gray-9" />
-                      <Text color="gray" className="text-base">
-                        Connecting to agent...
-                      </Text>
+                      <ConnectingToAgent />
                     </Flex>
                   )
                 ) : firstPendingPermission ? (
-                  // This box replaces the composer while a permission is pending, so it's an input
-                  // region: `shrink-0` keeps it from being compressed by the scroller above, and
-                  // `min-h-0 overflow-y-auto` lets a tall permission prompt scroll inside itself.
-                  <Box className="min-h-0 shrink-0 overflow-y-auto">
-                    <Box
-                      className={compact ? "p-1" : "mx-auto px-2 pb-3"}
-                      style={
-                        compact
-                          ? undefined
-                          : { maxWidth: CHAT_CONTENT_MAX_WIDTH }
-                      }
-                    >
-                      <PermissionSelector
-                        toolCall={firstPendingPermission.toolCall}
-                        options={firstPendingPermission.options}
-                        onSelect={handlePermissionSelect}
-                        onCancel={handlePermissionCancel}
-                      />
-                    </Box>
-                  </Box>
+                  <ComposerSlot compact={compact}>
+                    <PermissionSelector
+                      toolCall={firstPendingPermission.toolCall}
+                      options={firstPendingPermission.options}
+                      onSelect={handlePermissionSelect}
+                      onCancel={handlePermissionCancel}
+                    />
+                  </ComposerSlot>
                 ) : (
                   <Box className="relative">
                     <Box
@@ -693,10 +707,7 @@ export function SessionView({
                           : "opacity-100"
                       }`}
                     >
-                      <Spinner size={28} className="animate-spin text-gray-9" />
-                      <Text color="gray" className="text-base">
-                        Connecting to agent...
-                      </Text>
+                      <ConnectingToAgent />
                     </Box>
                     <Box
                       className={`transition-all duration-300 ease-out ${
@@ -705,14 +716,7 @@ export function SessionView({
                           : "pointer-events-none translate-y-4 opacity-0"
                       }`}
                     >
-                      <Box
-                        className={compact ? "p-1" : "mx-auto px-2 pb-3"}
-                        style={
-                          compact
-                            ? undefined
-                            : { maxWidth: CHAT_CONTENT_MAX_WIDTH }
-                        }
-                      >
+                      <ComposerWidth compact={compact}>
                         {taskId && <QueuedMessagesDock taskId={taskId} />}
                         <PromptInput
                           ref={editorRef}
@@ -764,7 +768,7 @@ export function SessionView({
                           onBashCommand={onBashCommand}
                           onCancel={onCancelPrompt}
                         />
-                      </Box>
+                      </ComposerWidth>
                     </Box>
                   </Box>
                 )}
diff --git a/packages/ui/src/features/sessions/components/StaleConversationCostNotice.tsx b/packages/ui/src/features/sessions/components/StaleConversationCostNotice.tsx
@@ -1,3 +1,4 @@
+import { formatUsd } from "@posthog/core/billing/spendAnalysisFormat";
 import { formatRelativeTimeLong } from "@posthog/shared";
 import { formatTokensCompact } from "@posthog/ui/features/sessions/contextColors";
 import { ActionSelector } from "@posthog/ui/primitives/ActionSelector";
@@ -40,7 +41,7 @@ export function StaleConversationCostNotice({
       ? `was last active ${formatRelativeTimeLong(lastActivityAt)}`
       : "has been idle";
   const spent =
-    costUsd !== null ? ` (≈$${costUsd.toFixed(2)} spent so far)` : "";
+    costUsd !== null ? ` (≈${formatUsd(costUsd)} spent so far)` : "";
   return (
     <ActionSelector
       title="Continue this large, idle conversation?"
diff --git a/packages/ui/src/features/sessions/staleConversationGateStore.test.ts b/packages/ui/src/features/sessions/staleConversationGateStore.test.ts
@@ -20,14 +20,14 @@ describe("useStaleConversationGateStore", () => {
 
   it("engages a single session with its last-activity snapshot", () => {
     state().engage("s1", 1000);
-    expect(state().engagedSessions.get("s1")).toEqual({ lastActivityAt: 1000 });
+    expect(state().engagedSessions.get("s1")).toBe(1000);
     expect(engaged("s2")).toBe(false);
   });
 
   it("keeps the first engagement snapshot when engaged again", () => {
     state().engage("s1", 1000);
     state().engage("s1", 2000);
-    expect(state().engagedSessions.get("s1")).toEqual({ lastActivityAt: 1000 });
+    expect(state().engagedSessions.get("s1")).toBe(1000);
   });
 
   it("does not re-engage an acknowledged session", () => {
@@ -62,4 +62,11 @@ describe("useStaleConversationGateStore", () => {
     const second = state().acknowledgedSessions;
     expect(second).toBe(first);
   });
+
+  it("acknowledging a never-engaged session keeps the engaged Map reference", () => {
+    state().engage("s2", 1000);
+    const before = state().engagedSessions;
+    state().acknowledge("s1");
+    expect(state().engagedSessions).toBe(before);
+  });
 });
diff --git a/packages/ui/src/features/sessions/staleConversationGateStore.ts b/packages/ui/src/features/sessions/staleConversationGateStore.ts
@@ -3,12 +3,13 @@ import { create } from "zustand";
 interface StaleConversationGateState {
   /**
    * Sessions where the gate engaged, keyed to the last-activity time observed
-   * at that moment. Latched because reconnecting to a stale session
-   * immediately appends freshly-stamped events (usage updates, handshakes)
-   * that would otherwise make the conversation look active again and dismiss
-   * the warning before the user has chosen.
+   * at that moment (null when no activity was observed). Latched because
+   * reconnecting to a stale session immediately appends freshly-stamped
+   * events (usage updates, handshakes) that would otherwise make the
+   * conversation look active again and dismiss the warning before the user
+   * has chosen.
    */
-  engagedSessions: Map<string, { lastActivityAt: number | null }>;
+  engagedSessions: Map<string, number | null>;
   /** Sessions where the user accepted the "large + idle = costly" warning. */
   acknowledgedSessions: Set<string>;
 }
@@ -40,7 +41,7 @@ export const useStaleConversationGateStore =
           return state;
         }
         const next = new Map(state.engagedSessions);
-        next.set(sessionId, { lastActivityAt });
+        next.set(sessionId, lastActivityAt);
         return { engagedSessions: next };
       }),
 
@@ -49,6 +50,9 @@ export const useStaleConversationGateStore =
         if (state.acknowledgedSessions.has(sessionId)) return state;
         const nextAcknowledged = new Set(state.acknowledgedSessions);
         nextAcknowledged.add(sessionId);
+        if (!state.engagedSessions.has(sessionId)) {
+          return { acknowledgedSessions: nextAcknowledged };
+        }
         const nextEngaged = new Map(state.engagedSessions);
         nextEngaged.delete(sessionId);
         return {
diff --git a/packages/ui/src/features/sessions/useStaleConversationGate.ts b/packages/ui/src/features/sessions/useStaleConversationGate.ts
@@ -36,7 +36,10 @@ export function useStaleConversationGate(
   const contextUsage = useContextUsage(events);
   const { data: currentUser } = useMeQuery();
   const isStaff = currentUser?.is_staff === true;
-  const engagement = useStaleConversationGateStore((s) =>
+  const engaged = useStaleConversationGateStore((s) =>
+    s.engagedSessions.has(sessionId),
+  );
+  const engagedLastActivityAt = useStaleConversationGateStore((s) =>
     s.engagedSessions.get(sessionId),
   );
   const acknowledged = useStaleConversationGateStore((s) =>
@@ -67,11 +70,13 @@ export function useStaleConversationGate(
 
   return {
     // shouldEngage covers the first paint before the effect latches;
-    // engagement covers every render after (reconnect events flip
+    // engaged covers every render after (reconnect events flip
     // shouldEngage back off — see the hook doc).
-    active: shouldEngage || engagement !== undefined,
+    active: shouldEngage || engaged,
     usedTokens,
-    lastActivityAt: engagement ? engagement.lastActivityAt : liveLastActivityAt,
+    lastActivityAt: engaged
+      ? (engagedLastActivityAt ?? null)
+      : liveLastActivityAt,
     costUsd: contextUsage?.cost?.amount ?? null,
     onContinue,
   };
diff --git a/packages/ui/src/primitives/action-selector/useActionSelectorState.ts b/packages/ui/src/primitives/action-selector/useActionSelectorState.ts
@@ -103,6 +103,12 @@ export function useActionSelectorState({
     isEditing && selectedOption && needsCustomInput(selectedOption);
   const canSubmitOrAdvance = checkedOptions.size > 0;
 
+  // Options can change while mounted (a consumer adding/removing a choice);
+  // clamp so the highlight and Enter never reference past the end of the list.
+  useEffect(() => {
+    setSelectedIndex((i) => Math.min(i, Math.max(0, numOptions - 1)));
+  }, [numOptions]);
+
   useEffect(() => {
     if (!isInteractiveElementInDifferentCell(containerRef)) {
       containerRef.current?.focus();
