Skip to content

Commit 7b2ac46

Browse files
author
catlog22
committed
refactor: streamline store access in useWebSocket by consolidating state retrieval into getStoreState function
1 parent 346c87a commit 7b2ac46

1 file changed

Lines changed: 72 additions & 100 deletions

File tree

‎ccw/frontend/src/hooks/useWebSocket.ts‎

Lines changed: 72 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,42 @@ const RECONNECT_DELAY_BASE = 1000; // 1 second
2121
const RECONNECT_DELAY_MAX = 30000; // 30 seconds
2222
const RECONNECT_DELAY_MULTIPLIER = 1.5;
2323

24+
// Access store state/actions via getState() - avoids calling hooks in callbacks/effects
25+
// This is the zustand-recommended pattern for non-rendering store access
26+
function getStoreState() {
27+
const notification = useNotificationStore.getState();
28+
const execution = useExecutionStore.getState();
29+
const flow = useFlowStore.getState();
30+
const cliStream = useCliStreamStore.getState();
31+
const coordinator = useCoordinatorStore.getState();
32+
return {
33+
// Notification store
34+
setWsStatus: notification.setWsStatus,
35+
setWsLastMessage: notification.setWsLastMessage,
36+
incrementReconnectAttempts: notification.incrementReconnectAttempts,
37+
resetReconnectAttempts: notification.resetReconnectAttempts,
38+
addA2UINotification: notification.addA2UINotification,
39+
// Execution store
40+
setExecutionStatus: execution.setExecutionStatus,
41+
setNodeStarted: execution.setNodeStarted,
42+
setNodeCompleted: execution.setNodeCompleted,
43+
setNodeFailed: execution.setNodeFailed,
44+
addLog: execution.addLog,
45+
completeExecution: execution.completeExecution,
46+
currentExecution: execution.currentExecution,
47+
// Flow store
48+
updateNode: flow.updateNode,
49+
// CLI stream store
50+
addOutput: cliStream.addOutput,
51+
// Coordinator store
52+
updateNodeStatus: coordinator.updateNodeStatus,
53+
addCoordinatorLog: coordinator.addLog,
54+
setActiveQuestion: coordinator.setActiveQuestion,
55+
markExecutionComplete: coordinator.markExecutionComplete,
56+
coordinatorExecutionId: coordinator.currentExecutionId,
57+
};
58+
}
59+
2460
interface UseWebSocketOptions {
2561
enabled?: boolean;
2662
onMessage?: (message: OrchestratorWebSocketMessage) => void;
@@ -40,74 +76,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
4076
const reconnectDelayRef = useRef(RECONNECT_DELAY_BASE);
4177
const mountedRef = useRef(true);
4278

43-
// Store refs to prevent handler recreation - use useRef to keep stable references
44-
const storeRefs = useRef({
45-
// Notification store
46-
setWsStatus: useNotificationStore((state) => state.setWsStatus),
47-
setWsLastMessage: useNotificationStore((state) => state.setWsLastMessage),
48-
incrementReconnectAttempts: useNotificationStore((state) => state.incrementReconnectAttempts),
49-
resetReconnectAttempts: useNotificationStore((state) => state.resetReconnectAttempts),
50-
addA2UINotification: useNotificationStore((state) => state.addA2UINotification),
51-
52-
// Execution store
53-
setExecutionStatus: useExecutionStore((state) => state.setExecutionStatus),
54-
setNodeStarted: useExecutionStore((state) => state.setNodeStarted),
55-
setNodeCompleted: useExecutionStore((state) => state.setNodeCompleted),
56-
setNodeFailed: useExecutionStore((state) => state.setNodeFailed),
57-
addLog: useExecutionStore((state) => state.addLog),
58-
completeExecution: useExecutionStore((state) => state.completeExecution),
59-
currentExecution: useExecutionStore((state) => state.currentExecution),
60-
61-
// Flow store
62-
updateNode: useFlowStore((state) => state.updateNode),
63-
64-
// CLI stream store
65-
addOutput: useCliStreamStore((state) => state.addOutput),
66-
67-
// Coordinator store
68-
updateNodeStatus: useCoordinatorStore((state) => state.updateNodeStatus),
69-
addCoordinatorLog: useCoordinatorStore((state) => state.addLog),
70-
setActiveQuestion: useCoordinatorStore((state) => state.setActiveQuestion),
71-
markExecutionComplete: useCoordinatorStore((state) => state.markExecutionComplete),
72-
coordinatorExecutionId: useCoordinatorStore((state) => state.currentExecutionId),
73-
});
74-
75-
// Update refs periodically to ensure they have fresh store references
76-
useEffect(() => {
77-
storeRefs.current = {
78-
// Notification store
79-
setWsStatus: useNotificationStore((state) => state.setWsStatus),
80-
setWsLastMessage: useNotificationStore((state) => state.setWsLastMessage),
81-
incrementReconnectAttempts: useNotificationStore((state) => state.incrementReconnectAttempts),
82-
resetReconnectAttempts: useNotificationStore((state) => state.resetReconnectAttempts),
83-
addA2UINotification: useNotificationStore((state) => state.addA2UINotification),
84-
85-
// Execution store
86-
setExecutionStatus: useExecutionStore((state) => state.setExecutionStatus),
87-
setNodeStarted: useExecutionStore((state) => state.setNodeStarted),
88-
setNodeCompleted: useExecutionStore((state) => state.setNodeCompleted),
89-
setNodeFailed: useExecutionStore((state) => state.setNodeFailed),
90-
addLog: useExecutionStore((state) => state.addLog),
91-
completeExecution: useExecutionStore((state) => state.completeExecution),
92-
currentExecution: useExecutionStore((state) => state.currentExecution),
93-
94-
// Flow store
95-
updateNode: useFlowStore((state) => state.updateNode),
96-
97-
// CLI stream store
98-
addOutput: useCliStreamStore((state) => state.addOutput),
99-
100-
// Coordinator store
101-
updateNodeStatus: useCoordinatorStore((state) => state.updateNodeStatus),
102-
addCoordinatorLog: useCoordinatorStore((state) => state.addLog),
103-
setActiveQuestion: useCoordinatorStore((state) => state.setActiveQuestion),
104-
markExecutionComplete: useCoordinatorStore((state) => state.markExecutionComplete),
105-
coordinatorExecutionId: useCoordinatorStore((state) => state.currentExecutionId),
106-
};
107-
}); // Run on every render to keep refs fresh
108-
10979
// Handle incoming WebSocket messages
110-
// Note: Using refs via storeRefs to prevent handler recreation on every store change
11180
const handleMessage = useCallback(
11281
(event: MessageEvent) => {
11382
// Guard against state updates after unmount
@@ -117,9 +86,10 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
11786

11887
try {
11988
const data = JSON.parse(event.data);
89+
const stores = getStoreState();
12090

12191
// Store last message for debugging
122-
storeRefs.current.setWsLastMessage(data);
92+
stores.setWsLastMessage(data);
12393

12494
// Handle CLI messages
12595
if (data.type?.startsWith('CLI_')) {
@@ -128,7 +98,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
12898
const { executionId, tool, mode, timestamp } = data.payload;
12999

130100
// Add system message for CLI start
131-
storeRefs.current.addOutput(executionId, {
101+
stores.addOutput(executionId, {
132102
type: 'system',
133103
content: `[${new Date(timestamp).toLocaleTimeString()}] CLI execution started: ${tool} (${mode || 'default'} mode)`,
134104
timestamp: Date.now(),
@@ -157,7 +127,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
157127
lines.forEach((line: string) => {
158128
// Add non-empty lines, or single line if that's all we have
159129
if (line.trim() || lines.length === 1) {
160-
storeRefs.current.addOutput(executionId, {
130+
stores.addOutput(executionId, {
161131
type: unitType as any,
162132
content: line,
163133
timestamp: Date.now(),
@@ -173,7 +143,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
173143
const statusText = success ? 'completed successfully' : 'failed';
174144
const durationText = duration ? ` (${duration}ms)` : '';
175145

176-
storeRefs.current.addOutput(executionId, {
146+
stores.addOutput(executionId, {
177147
type: 'system',
178148
content: `[${new Date().toLocaleTimeString()}] CLI execution ${statusText}${durationText}`,
179149
timestamp: Date.now(),
@@ -188,7 +158,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
188158
if (data.type === 'a2ui-surface') {
189159
const parsed = SurfaceUpdateSchema.safeParse(data.payload);
190160
if (parsed.success) {
191-
storeRefs.current.addA2UINotification(parsed.data, 'Interactive UI');
161+
stores.addA2UINotification(parsed.data, 'Interactive UI');
192162
} else {
193163
console.warn('[WebSocket] Invalid A2UI surface:', parsed.error.issues);
194164
}
@@ -197,7 +167,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
197167

198168
// Handle Coordinator messages
199169
if (data.type?.startsWith('COORDINATOR_')) {
200-
const { coordinatorExecutionId } = storeRefs.current;
170+
const { coordinatorExecutionId } = stores;
201171
// Only process messages for current coordinator execution
202172
if (coordinatorExecutionId && data.executionId !== coordinatorExecutionId) {
203173
return;
@@ -208,26 +178,26 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
208178
case 'COORDINATOR_STATE_UPDATE':
209179
// Check for completion
210180
if (data.status === 'completed') {
211-
storeRefs.current.markExecutionComplete(true);
181+
stores.markExecutionComplete(true);
212182
} else if (data.status === 'failed') {
213-
storeRefs.current.markExecutionComplete(false);
183+
stores.markExecutionComplete(false);
214184
}
215185
break;
216186

217187
case 'COORDINATOR_COMMAND_STARTED':
218-
storeRefs.current.updateNodeStatus(data.nodeId, 'running');
188+
stores.updateNodeStatus(data.nodeId, 'running');
219189
break;
220190

221191
case 'COORDINATOR_COMMAND_COMPLETED':
222-
storeRefs.current.updateNodeStatus(data.nodeId, 'completed', data.result);
192+
stores.updateNodeStatus(data.nodeId, 'completed', data.result);
223193
break;
224194

225195
case 'COORDINATOR_COMMAND_FAILED':
226-
storeRefs.current.updateNodeStatus(data.nodeId, 'failed', undefined, data.error);
196+
stores.updateNodeStatus(data.nodeId, 'failed', undefined, data.error);
227197
break;
228198

229199
case 'COORDINATOR_LOG_ENTRY':
230-
storeRefs.current.addCoordinatorLog(
200+
stores.addCoordinatorLog(
231201
data.log.message,
232202
data.log.level,
233203
data.log.nodeId,
@@ -236,7 +206,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
236206
break;
237207

238208
case 'COORDINATOR_QUESTION_ASKED':
239-
storeRefs.current.setActiveQuestion(data.question);
209+
stores.setActiveQuestion(data.question);
240210
break;
241211

242212
case 'COORDINATOR_ANSWER_RECEIVED':
@@ -262,47 +232,47 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
262232
const message = parsed.data as OrchestratorWebSocketMessage;
263233

264234
// Only process messages for current execution
265-
const { currentExecution } = storeRefs.current;
235+
const { currentExecution } = stores;
266236
if (currentExecution && message.execId !== currentExecution.execId) {
267237
return;
268238
}
269239

270240
// Dispatch to execution store based on message type
271241
switch (message.type) {
272242
case 'ORCHESTRATOR_STATE_UPDATE':
273-
storeRefs.current.setExecutionStatus(message.status, message.currentNodeId);
243+
stores.setExecutionStatus(message.status, message.currentNodeId);
274244
// Check for completion
275245
if (message.status === 'completed' || message.status === 'failed') {
276-
storeRefs.current.completeExecution(message.status);
246+
stores.completeExecution(message.status);
277247
}
278248
break;
279249

280250
case 'ORCHESTRATOR_NODE_STARTED':
281-
storeRefs.current.setNodeStarted(message.nodeId);
251+
stores.setNodeStarted(message.nodeId);
282252
// Update canvas node status
283-
storeRefs.current.updateNode(message.nodeId, { executionStatus: 'running' });
253+
stores.updateNode(message.nodeId, { executionStatus: 'running' });
284254
break;
285255

286256
case 'ORCHESTRATOR_NODE_COMPLETED':
287-
storeRefs.current.setNodeCompleted(message.nodeId, message.result);
257+
stores.setNodeCompleted(message.nodeId, message.result);
288258
// Update canvas node status
289-
storeRefs.current.updateNode(message.nodeId, {
259+
stores.updateNode(message.nodeId, {
290260
executionStatus: 'completed',
291261
executionResult: message.result,
292262
});
293263
break;
294264

295265
case 'ORCHESTRATOR_NODE_FAILED':
296-
storeRefs.current.setNodeFailed(message.nodeId, message.error);
266+
stores.setNodeFailed(message.nodeId, message.error);
297267
// Update canvas node status
298-
storeRefs.current.updateNode(message.nodeId, {
268+
stores.updateNode(message.nodeId, {
299269
executionStatus: 'failed',
300270
executionError: message.error,
301271
});
302272
break;
303273

304274
case 'ORCHESTRATOR_LOG':
305-
storeRefs.current.addLog(message.log as ExecutionLog);
275+
stores.addLog(message.log as ExecutionLog);
306276
break;
307277
}
308278

@@ -312,7 +282,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
312282
console.error('[WebSocket] Failed to parse message:', error);
313283
}
314284
},
315-
[onMessage] // Only dependency is onMessage, all other functions accessed via refs
285+
[onMessage] // Only dependency is onMessage, store access via getState()
316286
);
317287

318288
// Connect to WebSocket
@@ -329,8 +299,9 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
329299
const delay = reconnectDelayRef.current;
330300
console.log(`[WebSocket] Reconnecting in ${delay}ms...`);
331301

332-
storeRefs.current.setWsStatus('reconnecting');
333-
storeRefs.current.incrementReconnectAttempts();
302+
const stores = getStoreState();
303+
stores.setWsStatus('reconnecting');
304+
stores.incrementReconnectAttempts();
334305

335306
reconnectTimeoutRef.current = setTimeout(() => {
336307
connectRef.current?.();
@@ -341,7 +312,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
341312
reconnectDelayRef.current * RECONNECT_DELAY_MULTIPLIER,
342313
RECONNECT_DELAY_MAX
343314
);
344-
}, []); // No dependencies - uses connectRef
315+
}, []); // No dependencies - uses connectRef and getStoreState()
345316

346317
const connect = useCallback(() => {
347318
if (!enabled) return;
@@ -351,34 +322,35 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
351322
const wsUrl = `${protocol}//${window.location.host}/ws`;
352323

353324
try {
354-
storeRefs.current.setWsStatus('connecting');
325+
getStoreState().setWsStatus('connecting');
355326

356327
const ws = new WebSocket(wsUrl);
357328
wsRef.current = ws;
358329

359330
ws.onopen = () => {
360331
console.log('[WebSocket] Connected');
361-
storeRefs.current.setWsStatus('connected');
362-
storeRefs.current.resetReconnectAttempts();
332+
const s = getStoreState();
333+
s.setWsStatus('connected');
334+
s.resetReconnectAttempts();
363335
reconnectDelayRef.current = RECONNECT_DELAY_BASE;
364336
};
365337

366338
ws.onmessage = handleMessage;
367339

368340
ws.onclose = () => {
369341
console.log('[WebSocket] Disconnected');
370-
storeRefs.current.setWsStatus('disconnected');
342+
getStoreState().setWsStatus('disconnected');
371343
wsRef.current = null;
372344
scheduleReconnect();
373345
};
374346

375347
ws.onerror = (error) => {
376348
console.error('[WebSocket] Error:', error);
377-
storeRefs.current.setWsStatus('error');
349+
getStoreState().setWsStatus('error');
378350
};
379351
} catch (error) {
380352
console.error('[WebSocket] Failed to connect:', error);
381-
storeRefs.current.setWsStatus('error');
353+
getStoreState().setWsStatus('error');
382354
scheduleReconnect();
383355
}
384356
}, [enabled, handleMessage, scheduleReconnect]);

0 commit comments

Comments
 (0)