diff --git a/apps/backend/internal/gateway/websocket/hub.go b/apps/backend/internal/gateway/websocket/hub.go index 4c2c803ff3..aa1b183475 100644 --- a/apps/backend/internal/gateway/websocket/hub.go +++ b/apps/backend/internal/gateway/websocket/hub.go @@ -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) } diff --git a/apps/backend/internal/gateway/websocket/hub_broadcast_test.go b/apps/backend/internal/gateway/websocket/hub_broadcast_test.go new file mode 100644 index 0000000000..50da21d6f3 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/hub_broadcast_test.go @@ -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") + } +} diff --git a/apps/web/hooks/domains/session/use-session-messages.test.ts b/apps/web/hooks/domains/session/use-session-messages.test.ts index 4e72736631..a6be60e997 100644 --- a/apps/web/hooks/domains/session/use-session-messages.test.ts +++ b/apps/web/hooks/domains/session/use-session-messages.test.ts @@ -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 { 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({ diff --git a/apps/web/hooks/domains/session/use-session-messages.ts b/apps/web/hooks/domains/session/use-session-messages.ts index eb99bee07d..736003083b 100644 --- a/apps/web/hooks/domains/session/use-session-messages.ts +++ b/apps/web/hooks/domains/session/use-session-messages.ts @@ -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 = new Set(["STARTING", "RUNNING"]); +// States the session settles into once a turn finishes. +const SETTLED_SESSION_STATES: ReadonlySet = 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, +) { + 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(