refactor: remove Zustand nodeStates mirror, subscribe to Jotai atoms directly#417
Conversation
…directly After the Jotai atom refactor (#392), useNodeStatesSubscription was half-vestigial — it only patched edge alerts (slow-input-timeout) but still subscribed to all of Zustand nodeStates with a 100ms throttle and forced a dual-write in websocket.ts. - Rename useNodeStatesSubscription → useEdgeAlertSubscription - Subscribe directly to per-node Jotai state atoms via sessionStore.sub() instead of the Zustand store, using queueMicrotask to coalesce rapid atom changes from a single batchWriteNodeStates flush - Drop the batchUpdateNodeStatesMulti call from websocket.ts (the Jotai write is the only one needed) - Remove the 100ms throttle — edge alerts are infrequent and the throttle added unnecessary latency - Migrate SessionInfoChip and SessionItem to read node states from Jotai atoms (new useSessionNodeStates hook) instead of Zustand - Remove batchUpdateNodeStatesMulti and batchUpdateNodeStates from sessionStore (dead code after the above changes) Closes #397 Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
![staging-devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/870844?s=60&v=4){kind=small}
There was a problem hiding this comment.
🚩 Zustand nodeStates are now effectively legacy initial-state storage
After websocket.ts stopped calling batchUpdateNodeStatesMulti, live nodestatechanged events no longer update useSessionStore().sessions[*].nodeStates; they only update Jotai atoms. Current monitor runtime consumers changed by this PR (SessionItem.tsx via useSessionNodeStates, nodes via useNodeStateFromAtom, and edge alerts via useEdgeAlertSubscription) read from atoms, so this did not look like an immediate UI bug. However, sessionStore.ts still exposes nodeStates and updateNodeState, and its tests still validate that API, so future code using getSession(...).nodeStates will observe only pipeline-seeded/legacy state rather than live WebSocket state.
(Refers to lines 9-12)
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
There was a problem hiding this comment.
📝 Info: Edge alert updates intentionally only toggle alert presence
The edge-alert patcher returns the existing edge when shouldWarn === isCurrentlyWarned, so an edge that remains in slow_input_timeout will not get refreshed tooltip details if newly_slow_pins, sync_timeout_ms, or the slow-pin set changes while the edge remains warned. This behavior existed in the old logic as well, so I did not flag it as a PR-introduced bug, but it means tooltip metadata can lag until the warning clears and reappears.
(Refers to line 149)
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
| queueMicrotask(() => { | ||
| pending = false; | ||
| readAll(); | ||
| }); |
There was a problem hiding this comment.
🟡 Queued atom flush can update state after the hook has unsubscribed
onAtomChange schedules readAll() in a microtask, but the cleanup only unsubscribes atom listeners and does not cancel or guard an already-queued microtask. If the session item/chip unmounts or switches to a different sessionId after an atom change schedules the microtask, the stale callback still calls setNodeStates for the old session. This can briefly render stale status/issue information on recycled components and performs a React state update after the subscription has been torn down.
Prompt for agents
In ui/src/hooks/useSessionNodeStates.ts, guard the queued microtask so it becomes a no-op after the effect cleanup runs. The hook currently unsubscribes from Jotai atoms but cannot cancel queueMicrotask callbacks already scheduled by onAtomChange. Add an effect-local disposed/cancelled flag that is set in the cleanup and checked before readAll() is invoked from the microtask, while preserving the existing unsubscribe behavior.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
There was a problem hiding this comment.
Fixed in 79216db — added a disposed flag that's set in the cleanup function and checked inside the microtask callback before calling readAll().
| queueMicrotask(() => { | ||
| pendingFlush = false; | ||
| applyPatch(); | ||
| }); |
There was a problem hiding this comment.
🟡 Stale edge-alert microtasks can patch the wrong session graph
onAtomChange queues applyPatch() without any cancellation guard, while cleanup only unsubscribes atom listeners. If an atom change schedules the microtask and the selected session/topology changes before it runs, the stale callback still reads pipelineRef.current (which now points at the new graph) using the old selectedSessionId captured by this effect and then calls setEdges. That can remove or add slow-input alert metadata on the newly selected/topology graph using node states from the previous subscription.
Prompt for agents
In ui/src/hooks/useEdgeAlertSubscription.ts, make queued atom-change flushes no-op after the effect has been cleaned up. The current cleanup unsubscribes listeners but cannot cancel queueMicrotask callbacks already scheduled by the old selectedSessionId/topoKey closure. Add an effect-local disposed/cancelled flag checked inside the queued callback before resetting pendingFlush and before calling applyPatch(), and set that flag in the cleanup along with unsubscribing.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
There was a problem hiding this comment.
Fixed in 79216db — added a disposed flag that's set in the cleanup function and checked inside the microtask callback before calling applyPatch().
| const unsubs = nodeIds.map((id) => | ||
| sessionStore.sub(nodeStateAtom(nodeKey(sessionId, id)), onAtomChange) | ||
| ); | ||
|
|
||
| return () => unsubs.forEach((u) => u()); | ||
| }, [sessionId, pipeline]); |
There was a problem hiding this comment.
📝 Info: New hooks resubscribe to all node atoms whenever the pipeline object changes
Both useSessionNodeStates and useEdgeAlertSubscription derive nodeIds from the current pipeline and subscribe to each per-node atom. Because the dependency is the entire pipeline object in useSessionNodeStates and topoKey in useEdgeAlertSubscription, low-frequency topology changes are handled, and live node-state changes no longer wake the large Zustand store. The tradeoff is that a full unsubscribe/resubscribe occurs when the pipeline object identity changes even if node IDs are unchanged; that is probably acceptable for the intended low-frequency pipeline updates, but reviewers may want to keep an eye on this if setPipeline starts being called more often.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
Prevent stale microtask callbacks from running after effect cleanup when the session or topology changes between scheduling and execution. Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
Port from #409: - Fully remove nodeStates field from SessionData and related actions - Add renderHook tests for useEdgeAlertSubscription and useSessionNodeStates - Add shallow-equality guard in useSessionNodeStates for reference stability - Reset topoEffectRanRef on effect re-run in useEdgeAlertSubscription - Remove stale nodeStates comment in useSession.ts - Clean up nodeState-related tests in sessionStore test files Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
2 participants
Summary
After the Jotai atom refactor in #392,
useNodeStatesSubscriptionwas half-vestigial — it only patched edge alerts but still subscribed to all of ZustandnodeStateswith a 100ms throttle and forced a dual-write inwebsocket.ts. This PR removes the ZustandnodeStatesmirror entirely and subscribes directly to Jotai atoms.Changes
useNodeStatesSubscription→useEdgeAlertSubscription(reflects actual responsibility)sessionStore.sub(nodeStateAtom(...))instead of ZustandnodeStatesbatchUpdateNodeStatesMulticall fromwebsocket.ts— the Jotai write is the only one needednodeStatesfrom ZustandSessionDataand all related actions (updateNodeState,batchUpdateNodeStates,batchUpdateNodeStatesMulti)SessionItem/SessionInfoChipfrom ZustandnodeStatesto a newuseSessionNodeStateshook that aggregates Jotai atoms keyed by pipeline node IDsuseSessionNodeStatesprevents unnecessary re-renders when atom values haven't changed (e.g. duplicate writes within a single RAF flush)topoEffectRanRefon effect re-run to prevent stale patches during session transitionsuseEdgeAlertSubscriptionanduseSessionNodeStates— covering edge patching gating, recovery, reference stability, and session status integrationSubsumes #409 (closed). Closes #397.
Review & Testing Checklist for Human
SessionItem/SessionInfoChipdon't re-render on every RAF flush when node states haven't changed (shallow-equality guard)Notes
useSessionNodeStateshook reads node IDs from Zustand pipeline (low-frequency structural data) and states from per-node Jotai atoms (high-frequency updates), following the state management split documented inagent_docs/ui-development.mdawait act(async () => ...)for atom changes because the hooks usequeueMicrotaskcoalescing — microtasks need an async act boundary to flush within the testLink to Devin session: https://staging.itsdev.in/sessions/f56b91a30f8240d69b062e71d013d755
Requested by: @streamer45
Devin Review
7983eac(HEAD isa376ab5)