fix: deliver session broadcasts to focused clients during resume (#1193)

2026-05-31 14:39:15 +01:00
parent af3b26b6d9
commit 0511ef3028
4 changed files with 263 additions and 3 deletions
@@ -269,18 +269,51 @@ func (h *Hub) BroadcastToTask(taskID string, msg *ws.Message) {
 	h.sendToClients(data, clients, msg.Action)
 }

-// BroadcastToSession sends a notification to clients subscribed to a specific session
+// getSessionRecipientsLocked returns the deduped set of clients that should
+// receive a session-scoped broadcast: those subscribed to the session OR
+// focused on it.
+//
+// Focus is the stable "actively viewing this session" signal — it's held for
+// the whole time the task page is open. The ref-counted session.subscribe, by
+// contrast, churns to 0 during task-switch/resume (the sidebar hands the
+// active session off to the task-page hooks, and the resume state transitions
+// re-run the subscription effects). If a client is focused but its subscribe
+// ref-count was transiently dropped, it must still receive session events
+// (e.g. the session.message.updated that marks an agent_boot script_execution
+// completed) — otherwise the UI is stuck until a manual refetch.
+func (h *Hub) getSessionRecipientsLocked(sessionID string) []*Client {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	subs := h.sessionSubscribers[sessionID]
+	focus := h.sessionMode.focusByClient[sessionID]
+	clients := make([]*Client, 0, len(subs)+len(focus))
+	seen := make(map[*Client]struct{}, len(subs)+len(focus))
+	for client := range subs {
+		seen[client] = struct{}{}
+		clients = append(clients, client)
+	}
+	for client := range focus {
+		if _, ok := seen[client]; ok {
+			continue
+		}
+		clients = append(clients, client)
+	}
+	return clients
+}
+
+// BroadcastToSession sends a notification to clients subscribed to OR focused on
+// a specific session. See getSessionRecipientsLocked for why focus is included.
 func (h *Hub) BroadcastToSession(sessionID string, msg *ws.Message) {
 	data, err := json.Marshal(msg)
 	if err != nil {
 		h.logger.Error("Failed to marshal message", zap.Error(err))
 		return
 	}
-	clients := h.getSubscribersLocked(h.sessionSubscribers, sessionID)
+	clients := h.getSessionRecipientsLocked(sessionID)
 	h.logger.Debug("BroadcastToSession",
 		zap.String("session_id", sessionID),
 		zap.String("action", msg.Action),
-		zap.Int("subscriber_count", len(clients)))
+		zap.Int("recipient_count", len(clients)))
 	h.sendToClients(data, clients, msg.Action)
 }

@@ -0,0 +1,119 @@
+package websocket
+
+import (
+	"testing"
+
+	ws "github.com/kandev/kandev/pkg/websocket"
+)
+
+// registerTestClient adds a client to the hub's connected set so broadcasts can
+// reach it (BroadcastToSession only iterates subscriber/focus maps, but the
+// client must be a live connection in real usage).
+func registerTestClient(h *Hub, c *Client) {
+	h.mu.Lock()
+	h.clients[c] = true
+	h.mu.Unlock()
+}
+
+func clientReceived(c *Client) bool {
+	select {
+	case <-c.send:
+		return true
+	default:
+		return false
+	}
+}
+
+// TestBroadcastToSession_ReachesFocusedClient guards the resume regression:
+// a client that is focused on a session but whose ref-counted session.subscribe
+// was dropped (subscriber_count 0) must still receive session-scoped broadcasts
+// such as session.message.updated. Focus is the stable "actively viewing" signal;
+// the subscribe ref-count churns to 0 during task-switch/resume.
+func TestBroadcastToSession_ReachesFocusedClient(t *testing.T) {
+	h := newTestHub(t)
+	c := newTestClient("c1")
+	registerTestClient(h, c)
+
+	// Focused but NOT subscribed — mirrors the dropped-subscription resume race.
+	h.FocusSession(c, "sess-1")
+
+	msg, err := ws.NewNotification(ws.ActionSessionMessageUpdated, map[string]any{"session_id": "sess-1"})
+	if err != nil {
+		t.Fatalf("notification: %v", err)
+	}
+	h.BroadcastToSession("sess-1", msg)
+
+	if !clientReceived(c) {
+		t.Fatal("focused-but-unsubscribed client did not receive session broadcast")
+	}
+}
+
+// TestBroadcastToSession_ReachesSubscribedClient is the baseline: a subscribed
+// (sidebar/background) client still receives broadcasts after the union change.
+func TestBroadcastToSession_ReachesSubscribedClient(t *testing.T) {
+	h := newTestHub(t)
+	c := newTestClient("c1")
+	registerTestClient(h, c)
+
+	h.SubscribeToSession(c, "sess-1")
+
+	msg, err := ws.NewNotification(ws.ActionSessionMessageAdded, map[string]any{"session_id": "sess-1"})
+	if err != nil {
+		t.Fatalf("notification: %v", err)
+	}
+	h.BroadcastToSession("sess-1", msg)
+
+	if !clientReceived(c) {
+		t.Fatal("subscribed client did not receive session broadcast")
+	}
+}
+
+// TestBroadcastToSession_NotDeliveredToUninterestedClient ensures the union does
+// not over-deliver: a client neither subscribed nor focused gets nothing.
+func TestBroadcastToSession_NotDeliveredToUninterestedClient(t *testing.T) {
+	h := newTestHub(t)
+	focused := newTestClient("focused")
+	other := newTestClient("other")
+	registerTestClient(h, focused)
+	registerTestClient(h, other)
+
+	h.FocusSession(focused, "sess-1")
+
+	msg, err := ws.NewNotification(ws.ActionSessionMessageUpdated, map[string]any{"session_id": "sess-1"})
+	if err != nil {
+		t.Fatalf("notification: %v", err)
+	}
+	h.BroadcastToSession("sess-1", msg)
+
+	if !clientReceived(focused) {
+		t.Fatal("focused client did not receive broadcast")
+	}
+	if clientReceived(other) {
+		t.Fatal("uninterested client should not receive broadcast")
+	}
+}
+
+// TestBroadcastToSession_DeliversOnceToSubscribedAndFocusedClient ensures a
+// client that is BOTH subscribed and focused receives exactly one copy (the
+// union must dedupe).
+func TestBroadcastToSession_DeliversOnceToSubscribedAndFocusedClient(t *testing.T) {
+	h := newTestHub(t)
+	c := newTestClient("c1")
+	registerTestClient(h, c)
+
+	h.SubscribeToSession(c, "sess-1")
+	h.FocusSession(c, "sess-1")
+
+	msg, err := ws.NewNotification(ws.ActionSessionMessageUpdated, map[string]any{"session_id": "sess-1"})
+	if err != nil {
+		t.Fatalf("notification: %v", err)
+	}
+	h.BroadcastToSession("sess-1", msg)
+
+	if !clientReceived(c) {
+		t.Fatal("subscribed+focused client did not receive broadcast")
+	}
+	if clientReceived(c) {
+		t.Fatal("subscribed+focused client received a duplicate broadcast")
+	}
+}
@@ -28,10 +28,12 @@ afterEach(() => {
 });
 import {
  hasUserOrAgentMessage,
+  isTurnSettleTransition,
  runBackfillRound,
  autoBackfillUntilUserMessage,
  MAX_AUTO_BACKFILL_PAGES,
 } from "./use-session-messages";
+import type { TaskSessionState } from "@/lib/types/http";

 function makeMessage(overrides: Partial<Message>): Message {
  return {
@@ -112,6 +114,47 @@ describe("hasUserOrAgentMessage", () => {
  });
 });

+describe("isTurnSettleTransition", () => {
+  const settled: TaskSessionState[] = [
+    "IDLE",
+    "WAITING_FOR_INPUT",
+    "COMPLETED",
+    "FAILED",
+    "CANCELLED",
+  ];
+
+  it("is true when leaving RUNNING for a settled state (resume turn ends)", () => {
+    for (const next of settled) {
+      expect(isTurnSettleTransition("RUNNING", next)).toBe(true);
+    }
+  });
+
+  it("is true when leaving STARTING for a settled state (resume boots with no turn)", () => {
+    expect(isTurnSettleTransition("STARTING", "WAITING_FOR_INPUT")).toBe(true);
+  });
+
+  it("is false for the active-phase transition STARTING -> RUNNING", () => {
+    expect(isTurnSettleTransition("STARTING", "RUNNING")).toBe(false);
+  });
+
+  it("is false when staying in a settled state (no churn on WAITING -> WAITING)", () => {
+    expect(isTurnSettleTransition("WAITING_FOR_INPUT", "WAITING_FOR_INPUT")).toBe(false);
+  });
+
+  it("is false when entering an active state", () => {
+    expect(isTurnSettleTransition("WAITING_FOR_INPUT", "RUNNING")).toBe(false);
+    expect(isTurnSettleTransition("IDLE", "STARTING")).toBe(false);
+  });
+
+  it("is false when there is no previous state (initial render)", () => {
+    expect(isTurnSettleTransition(null, "WAITING_FOR_INPUT")).toBe(false);
+  });
+
+  it("is false when the next state is unknown", () => {
+    expect(isTurnSettleTransition("RUNNING", null)).toBe(false);
+  });
+});
+
 describe("runBackfillRound", () => {
  it("returns 'stop' when a user/agent message already exists in the store", async () => {
    const store = makeStore({
@@ -15,6 +15,36 @@ export function hasUserOrAgentMessage(messages: Message[]): boolean {
  );
 }

+// States where a turn (or the agent boot) is actively progressing.
+const ACTIVE_SESSION_STATES: ReadonlySet<TaskSessionState> = new Set(["STARTING", "RUNNING"]);
+// States the session settles into once a turn finishes.
+const SETTLED_SESSION_STATES: ReadonlySet<TaskSessionState> = new Set([
+  "IDLE",
+  "WAITING_FOR_INPUT",
+  "COMPLETED",
+  "FAILED",
+  "CANCELLED",
+]);
+
+/**
+ * True when the session just left an active state for a settled one — i.e. a
+ * turn (or a resume's agent boot) finished. Session-scoped message updates
+ * emitted as the turn winds down (e.g. the `agent_boot` `script_execution`
+ * completion during a resume) can be missed if the live subscription lapsed
+ * during the resume churn, so this is the signal to refetch and reconcile.
+ *
+ * `state_changed` / `turn.completed` are broadcast globally (not session-scoped),
+ * so the client always observes this transition even when its session
+ * subscription was dropped.
+ */
+export function isTurnSettleTransition(
+  prev: TaskSessionState | null,
+  next: TaskSessionState | null,
+): boolean {
+  if (prev === null || next === null) return false;
+  return ACTIVE_SESSION_STATES.has(prev) && SETTLED_SESSION_STATES.has(next);
+}
+
 const debug = createDebugLogger("messages:fetch");

 function summarizeMessages(messages: Message[]): {
@@ -356,6 +386,40 @@ function useSessionSubscription(
  }, [taskSessionId, connectionStatus, store, isSessionStartingOrUnknown]);
 }

+/**
+ * Refetch messages whenever a turn settles (active → settled). During a resume
+ * the agent_boot `script_execution` is created and then marked completed within
+ * ~1s, all server-side; if the live session subscription lapsed in that window
+ * the completion `session.message.updated` is dropped and the entry renders
+ * with a spinner forever (until a manual refresh). The settle transition is
+ * delivered globally, so reconciling messages here recovers any session-scoped
+ * updates missed while the turn was running.
+ */
+function useResyncOnTurnSettle(
+  taskSessionId: string | null,
+  taskSessionState: TaskSessionState | null,
+  connectionStatus: string,
+  store: ReturnType<typeof useAppStoreApi>,
+) {
+  const prevRef = useRef<{ sessionId: string | null; state: TaskSessionState | null }>({
+    sessionId: null,
+    state: null,
+  });
+  useEffect(() => {
+    const prev = prevRef.current;
+    prevRef.current = { sessionId: taskSessionId, state: taskSessionState };
+    if (!taskSessionId || connectionStatus !== "connected") return;
+    const prevState = prev.sessionId === taskSessionId ? prev.state : null;
+    if (!isTurnSettleTransition(prevState, taskSessionState)) return;
+    debug("resync on turn settle", {
+      sessionId: taskSessionId,
+      prev: prevState,
+      next: taskSessionState,
+    });
+    fetchAndStoreMessages(taskSessionId, store).catch(() => {});
+  }, [taskSessionId, taskSessionState, connectionStatus, store]);
+}
+
 export function useSessionMessages(taskSessionId: string | null): UseSessionMessagesReturn {
  const store = useAppStoreApi();
  const messages = useAppStore((state) =>
@@ -444,6 +508,7 @@ export function useSessionMessages(taskSessionId: string | null): UseSessionMess
  const isSessionStartingOrUnknown = taskSessionState === null || taskSessionState === "STARTING";

  useSessionSubscription(taskSessionId, connectionStatus, isSessionStartingOrUnknown, store);
+  useResyncOnTurnSettle(taskSessionId, taskSessionState, connectionStatus, store);
  useVisibilityBackfill(taskSessionId, store);

  const terminalFetchRefs = useMemo(