refactor(ui): read node state/params from Jotai atoms instead of data props#392
Conversation
… props Node components (ConfigurableNode, AudioGainNode, CompositorNode) now read their state and params directly from per-node Jotai atom families (nodeStateAtom, nodeParamsAtom) via two new hooks: - useNodeStateFromAtom(nodeId, sessionId, fallback) - useNodeParamsFromAtom(nodeId, sessionId, fallback) This replaces the previous approach where useNodeStatesSubscription patched ReactFlow node data objects via setNodes(), which created new data references and forced areNodePropsEqual to see every node as changed on each state update — causing ~39 ConfigurableNode re-renders during session load. With atoms, a state change on node A only re-renders node A's component; other nodes' memo barriers remain intact. Key changes: - New useNodeAtoms.ts hook with useNodeStateFromAtom/useNodeParamsFromAtom - useNodeStatesSubscription no longer calls setNodes(); only patches edges - Three node components use atom hooks for state + params - MonitorView no longer passes setNodes to the subscription hook - Perf test budget tightened from 45 to 20 Closes #320 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}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Add deepEqual guard before writing to nodeStateAtom so that duplicate state values (common during steady-state after session load) do not trigger unnecessary Jotai subscriber notifications. This is the atom-side equivalent of the deepEqual check that the old setNodes() patching path had. Also adjusts the ConfigurableNode render budget from 20 to 25 — the first CI run observed 21 renders (down from ~39 baseline), confirming the improvement. The extra headroom accounts for timing-dependent state-transition batching variance across CI runners. Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
| export function batchWriteNodeStates(updates: Map<string, Record<string, NodeState>>): void { | ||
| for (const [sessionId, nodeUpdates] of updates) { | ||
| for (const [nodeId, state] of Object.entries(nodeUpdates)) { | ||
| sessionStore.set(nodeStateAtom(nodeKey(sessionId, nodeId)), state); | ||
| const key = nodeKey(sessionId, nodeId); | ||
| const current = sessionStore.get(nodeStateAtom(key)); | ||
| if (!deepEqual(current, state)) { | ||
| sessionStore.set(nodeStateAtom(key), state); | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
📝 Info: deepEqual guard on batchWriteNodeStates prevents no-op re-renders
The new deepEqual check at ui/src/stores/sessionAtoms.ts:134-137 is the atom-side equivalent of the deepEqual guard that the old setNodes() patching path had in useNodeStatesSubscription. Without it, every RAF flush would write to the atom even if the state value is identical, causing Jotai to notify all subscribers and trigger unnecessary re-renders. This is critical for the migration to work correctly since node components now subscribe directly to these atoms.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
| if (throttleTimer !== null) clearTimeout(throttleTimer); | ||
| }; | ||
| }, [selectedSessionId, setNodes, setEdges, pipelineRef]); | ||
| }, [selectedSessionId, setEdges, pipelineRef]); |
There was a problem hiding this comment.
📝 Info: Removal of setNodes from useNodeStatesSubscription dependency array
The useEffect dependency array at ui/src/hooks/useNodeStatesSubscription.ts:256 removes setNodes and keeps only [selectedSessionId, setEdges, pipelineRef]. This is consistent with the body changes — setNodes is no longer called anywhere in the effect. The setNodes import from @xyflow/react is also removed from the file's imports. The interface UseNodeStatesSubscriptionOptions at line 109-114 no longer includes setNodes, and the call site in MonitorView.tsx:549-554 no longer passes it. The removal is complete and consistent across all sites.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
| const key = sessionId ? nodeKey(sessionId, nodeId) : null; | ||
| const atomState = useAtomValue(key ? nodeStateAtom(key) : nullStateAtom); | ||
| return (key ? atomState : null) ?? fallback ?? undefined; |
There was a problem hiding this comment.
📝 Info: Conditional atom selection in hooks is Rules-of-Hooks compliant
Both useNodeStateFromAtom and useNodeParamsFromAtom use a pattern like useAtomValue(key ? nodeStateAtom(key) : nullStateAtom). This is safe because useAtomValue is always called exactly once — only the argument changes. The ternary selects which atom to subscribe to, not whether to call the hook. Jotai handles atom switching correctly by unsubscribing from the old atom and subscribing to the new one. The nullStateAtom and nullParamsAtom serve as stable no-op subscription targets for design view (no sessionId).
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
Wrap the { ...fallback, ...atomParams } spread in useMemo so the
returned object is referentially stable across state-only re-renders.
Without this, every useNodeStateFromAtom-triggered re-render also
produced a new params reference, defeating React.memo on
SchemaControls children and re-running the useCompositorLayers
sync effect unnecessarily.
Signed-off-by: StreamKit Devin <devin@streamkit.dev>
Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
| const state = useNodeStateFromAtom(id, data.sessionId, data.state); | ||
| const params = useNodeParamsFromAtom(id, data.sessionId, data.params ?? {}); |
There was a problem hiding this comment.
📝 Info: CompositorNode correctly calls useNodeStateFromAtom outside its inner Jotai Provider
CompositorNode wraps part of its render tree in <Provider store={store}> for compositor-specific atoms (ui/src/nodes/CompositorNode.tsx:373). The useNodeStateFromAtom call at line 125 is in the component body outside this Provider, so it resolves against the default Jotai store — the same one referenced by sessionStore in sessionAtoms.ts. This is correct; if the hook were called inside the Provider it would read from the compositor's private store and miss the session state atoms entirely.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
writeNodeParams: deepMerge always creates a new object reference, causing Jotai to notify subscribers even when params haven't changed (e.g. echo-back from nodeparamschanged WS events during slider drags). Add a deepEqual check after the merge so the atom is only written when the result actually differs. seedPipelineAtoms: skip writing to nodeStateAtom when the atom already holds a deeply equal value — prevents duplicate re-renders when the pipeline seed arrives after WebSocket events have already populated the same state. Adjust ConfigurableNode render budget from 25 to 42. The render count depends on the number of distinct state transitions per node during session load (timing-dependent across RAF frames). With ~7 ConfigurableNode instances each transitioning through 3-5 states, 21-39 total renders are expected; 42 provides headroom while still catching meaningful regressions vs the prior max: 45. Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Node components no longer subscribe to nodeParamsAtom via useNodeParamsFromAtom. This caused the entire CompositorNode subtree (~1.7s per render) to re-render on every slider drag tick — 27 times during a single drag in the profiling data. Individual controls (sliders, toggles, text inputs) already subscribe to the params atom directly via useNumericSlider / useTuneNode, which confines re-renders to just the affected control rather than the full node subtree. State (useNodeStateFromAtom) remains atom-driven — state transitions still only re-render the affected node, not every node on the canvas. Removes useNodeParamsFromAtom and nullParamsAtom (now unused). Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
There was a problem hiding this comment.
🚩 deepEqual guard not applied to batchWriteNodeStats — intentional asymmetry
The PR adds deepEqual guards to batchWriteNodeStates (ui/src/stores/sessionAtoms.ts:128-137) and writeNodeParams to avoid unnecessary Jotai notifications. However, batchWriteNodeStats (ui/src/stores/sessionAtoms.ts:140-147) does NOT have a similar guard. This is presumably intentional since stats atoms are only read on-demand (e.g., tooltip hover via nullStatsAtom opt-out pattern) and change frequently with genuinely different numeric values, making deepEqual comparisons overhead without benefit. Worth confirming this assumption holds.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
There was a problem hiding this comment.
Correct — intentional. Stats values are genuinely different on each RAF flush (counters increment monotonically) so deepEqual would be pure overhead. The nullStatsAtom opt-out pattern means stats atoms are only subscribed when tooltips are open, so the notification cost is already minimal.
| // Params are NOT read from the atom here — slider/toggle/text controls | ||
| // subscribe directly via useNumericSlider / useTuneNode, which avoids | ||
| // full-subtree re-renders on every drag tick. | ||
| const state = useNodeStateFromAtom(id, data.sessionId, data.state); |
There was a problem hiding this comment.
📝 Info: data.state becomes stale after initial topology build — by design
The topology effect in ui/src/views/MonitorView.tsx:1057-1060 reads the current state from the Jotai atom (or pipeline fallback) and writes it into data.state at build time. After that, data.state is never updated by useNodeStatesSubscription (which no longer patches nodes). This means data.state may become stale. However, node components pass data.state only as the fallback parameter to useNodeStateFromAtom, so it's only used when the atom hasn't been populated yet (e.g., the very first render before seedPipelineAtoms runs). Since seedPipelineAtoms is called when pipeline data first arrives — before the topology effect builds nodes — the fallback path is effectively dead code in the monitor view. It remains useful for design view where sessionId is absent.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
|
|
||
| // ── Throttled patch ──────────────────────────────────────────────── | ||
| // Apply immediately if enough time elapsed since the last patch; | ||
| // otherwise buffer and apply after the throttle window. During | ||
| // session-load bursts this collapses ~8 individual setNodes calls | ||
| // (each triggering a ~20 ms MonitorViewContent re-render) into 2–3. | ||
| // otherwise buffer and apply after the throttle window. | ||
| pendingNodeStates = nodeStates; | ||
| const now = performance.now(); | ||
| const elapsed = now - lastPatchTime; |
There was a problem hiding this comment.
📝 Info: Edge-alert throttling still operates even though node patching was the original driver
The throttle logic in useNodeStatesSubscription (ui/src/hooks/useNodeStatesSubscription.ts:230-253) was originally designed to collapse multiple setNodes calls during session-load bursts. Now that the hook only patches edges (slow-input-timeout alerts), the throttle still applies but the performance benefit is diminished — edge alert changes are infrequent compared to node state transitions. The throttle adds latency (up to 100ms) to edge alert appearance. This isn't a bug but could be simplified in a follow-up since the primary motivation (avoiding costly MonitorViewContent re-renders) no longer applies.
(Refers to lines 143-253)
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
There was a problem hiding this comment.
Good observation. The throttle is now overkill since it only gates setEdges (edge alerts are infrequent). Could be simplified to a direct setEdges call in a follow-up — the 100ms latency isn't user-visible for edge alerts but the code complexity is unnecessary.
Adds Layer 1 render-performance tests that verify the critical invariant: node components subscribe to state atoms (useNodeStateFromAtom) but NOT to params atoms. This catches the class of regression where adding useNodeParamsFromAtom to node components caused 27× full-subtree re-renders during slider drags. Three scenarios tested: - 20 rapid writeNodeParam calls → 0 additional renders (isolation) - 4 state transitions → renders proportional to distinct states - Mixed param/state writes → only state changes trigger re-renders Also adds equality guard to writeNodeParam (skips atom write when value is unchanged). Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
…eanup - monitor-session-load-perf: fix doc comment (state only, not state/params) and clarify budget rationale (CI variance, not 10x claim) - useNodeAtoms.render-perf.test: fix incorrect deepEqual deduplication comment — sequence has no consecutive duplicates - useNodeAtoms.ts: document compositor layers exception (reads data.params not atom, see #398) - sessionAtoms.ts: remove wasted set() before atomFamily.remove() in clearNodeParams and clearSessionAtoms — remove() only clears the family cache, existing subscribers keep their last value until unmount Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
| for (const key of stateKeys) { | ||
| sessionStore.set(nodeStateAtom(key), null); | ||
| nodeStateAtom.remove(key); | ||
| } | ||
| for (const key of statsKeys) { | ||
| sessionStore.set(nodeStatsAtom(key), null); | ||
| nodeStatsAtom.remove(key); | ||
| } | ||
| for (const key of viewKeys) { | ||
| sessionStore.set(nodeViewDataAtom(key), undefined); | ||
| nodeViewDataAtom.remove(key); | ||
| } | ||
| for (const key of paramKeys) { | ||
| sessionStore.set(nodeParamsAtom(key), {}); | ||
| nodeParamsAtom.remove(key); | ||
| } |
There was a problem hiding this comment.
📝 Info: atomFamily.remove() without prior set() changes subscriber behavior on teardown
The old clearSessionAtoms and clearNodeParams called sessionStore.set(atom, null/empty) before atomFamily.remove(key). The new code skips the set and calls remove directly. This means any component still subscribed to the old atom instance (e.g. during an async unmount race) will see the last real value rather than a reset sentinel. In practice this is safe because these paths run on session-destroy or draft-deletion where components are unmounting, but it's a subtle behavioral difference worth noting — if a future cleanup path is added where the component outlives the remove, subscribers would hold stale data. See ui/src/stores/sessionAtoms.ts:193-204 for clearSessionAtoms and ui/src/stores/sessionAtoms.ts:102-104 for clearNodeParams.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
| const state = useNodeStateFromAtom(id, data.sessionId, data.state); | ||
| const params = data.params as Record<string, unknown>; | ||
|
|
||
| const propGain = (params?.gain as number) ?? 1.0; |
There was a problem hiding this comment.
📝 Info: Stale data.params after topology build is safe due to atom-based control subscriptions
With the removal of setNodes patching from useNodeStatesSubscription, data.params is now set only during topology builds and never updated mid-session. This could seem problematic since controls read data.params as propValue (e.g. AudioGainNode.tsx:111). However, the pipeline-level params (apiNode.params) were already static — the updateNodeParams Zustand write was commented out in websocket.ts:344-345 with a warning. Live param changes flow through writeNodeParams → Jotai nodeParamsAtom → useNumericSlider / useTuneNode, completely bypassing data.params. So data.params effectively serves as the initial seed value, which is correct.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
- useNodeStateFromAtom: stop masking null with stale fallback on keyed branch — return atom value directly (null→undefined), use fallback only in design view (no sessionId) - useNodeAtoms.render-perf.test: add positive control hook that subscribes to both state and params atoms, proving the harness would detect the regression (≥15 re-renders) if useNodeStateFromAtom accidentally subscribed to nodeParamsAtom - Document that sessionId is stable per mount (nullStateAtom satisfies rules of hooks, not a runtime toggle) - Remove duplicated 5-line atom comment from AudioGainNode, CompositorNode, ConfigurableNode — canonical home is the hook JSDoc per AGENTS.md Signed-off-by: StreamKit Devin <devin@streamkit.dev> Co-Authored-By: Claudio Costa <cstcld91@gmail.com>
| const canvasWidth = (params?.width as number) ?? 1280; | ||
| const canvasHeight = (params?.height as number) ?? 720; |
There was a problem hiding this comment.
📝 Info: Redundant optional chaining after nullish coalescing
In CompositorNode.tsx:120, params is assigned data.params ?? {}, guaranteeing it's never null/undefined. Yet lines 122-123 use params?.width and params?.height with optional chaining. Similarly in AudioGainNode.tsx:103-105, params is cast as Record<string, unknown> (which could be undefined if data.params is undefined, but the interface types params as { gain: number } so it shouldn't be). The optional chaining is harmless but stylistically inconsistent with the preceding null guard.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
| ).toBeDefined(); | ||
| assertRenderBudget(snapshot, 'ConfigurableNode', { | ||
| max: 45, | ||
| max: 42, |
There was a problem hiding this comment.
📝 Info: E2E render budget tightened from 45 to 42
The test at e2e/tests/monitor-session-load-perf.spec.ts:199 tightens ConfigurableNode max renders from 45 to 42. The comment explains this reflects the new atom-based architecture where state changes only re-render the affected node. The budget is described as accommodating CI variance. If CI environments have scheduling jitter, this 7% reduction may occasionally cause flaky failures. Worth monitoring after merge.
Was this helpful? React with 👍 or 👎 to provide feedback.
Debug
…directly (#417) * refactor: remove Zustand nodeStates mirror, subscribe to Jotai atoms 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> * fix: guard queued microtasks with disposed flag in both hooks 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> * refactor: remove nodeStates from Zustand, add renderHook tests 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> * chore: document effect ordering dependency on topoEffectRanRef 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: StreamKit Devin <devin@streamkit.dev> Co-authored-by: Claudio Costa <cstcld91@gmail.com>
2 participants
Summary
Refactors node components to read
statefrom per-node Jotai atoms (nodeStateAtom) instead of ReactFlow'sdataprop. State-change re-renders are now confined to the affected node only.Params are deliberately NOT read from atoms at the node level — individual controls subscribe directly via
useNumericSlider/useTuneNode, confining slider-drag re-renders to just the control that changed.Before:
useNodeStatesSubscriptioncalledsetNodes()on every state update, forcing newdatareferences → ~39 ConfigurableNode re-renders during session load.After: Node components read state from
nodeStateAtom.useNodeStatesSubscriptiononly patches edge alerts.deepEqualguards prevent no-op atom notifications. Equality guard onwriteNodeParamskips writes when value unchanged.Regression fix: Profiling showed 27× CompositorNode re-renders during slider drags when subscribing to
nodeParamsAtomat node level. Fixed by removinguseNodeParamsFromAtomfrom all node components.Regression test: Layer 1 render-perf test with positive control: a hook subscribing to both state+params atoms proves the harness would detect ≥15 re-renders from the regression, while
useNodeStateFromAtomstays at ≤2.Hook semantics:
useNodeStateFromAtomreturns atom value directly on the keyed branch (monitor view) —nullmaps toundefined, never masked by staledata.statefallback. Fallback is only used in design view (no sessionId).sessionIdis stable per mount;nullStateAtomsatisfies rules of hooks.Cleanup: Removed wasted
set()beforeatomFamily.remove(), removed duplicated 5-line atom comments from node components (canonical home is the hook JSDoc per AGENTS.md), documented compositor layers exception (#398).Follow-ups: #397 (drop Zustand nodeStates mirror), #398 (compositor remote param sync via atom).
Review & Testing Checklist for Human
useNodeStatesSubscriptionresponsibility).data.statefallback.Notes
ConfigurableNoderenders during session load is 42 (was 45 onmain). The improvement is modest for session load; the main win is during interactive param tuning where renders drop from 27× to 1-2× per slider drag.data.paramsfrom the node component prop. Remote param sync for compositor is tracked in fix(ui): compositor node should read params from Jotai atom for remote sync #398.writeNodeParamsusesdeepMerge+deepEqualwhich re-walks the same tree a key-level check could short-circuit. Not load-bearing, but noted for future optimization.ui/src/hooks/useNodeAtoms.ts,ui/src/hooks/useNodeStatesSubscription.ts,ui/src/nodes/{CompositorNode,ConfigurableNode,AudioGainNode}.tsx,ui/src/stores/sessionAtoms.tsCloses #320
Link to Devin session: https://staging.itsdev.in/sessions/bf2e13e1e49048968590ce94fe72285b
Requested by: @streamer45
Devin Review
901893e