catlog22/Claude-Code-Workflow
1
0
Fork 0
mirror of https://github.com/catlog22/Claude-Code-Workflow.git synced 2026-02-05 01:50:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: optimize useWebSocket hook by consolidating store references and improving handler stability

Browse Source
This commit is contained in:
catlog22
2026-02-04 17:29:30 +08:00
parent e260a3f77b
commit 346c87a706
2 changed files with 341 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

248
ccw/frontend/src/hooks/useWebSocket.ts
View File

@@ -40,36 +40,74 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const reconnectDelayRef = useRef(RECONNECT_DELAY_BASE);
const mountedRef = useRef(true);
// Notification store for connection status
const setWsStatus = useNotificationStore((state) => state.setWsStatus);
const setWsLastMessage = useNotificationStore((state) => state.setWsLastMessage);
const incrementReconnectAttempts = useNotificationStore((state) => state.incrementReconnectAttempts);
const resetReconnectAttempts = useNotificationStore((state) => state.resetReconnectAttempts);
const addA2UINotification = useNotificationStore((state) => state.addA2UINotification);
// Store refs to prevent handler recreation - use useRef to keep stable references
const storeRefs = useRef({
// Notification store
setWsStatus: useNotificationStore((state) => state.setWsStatus),
setWsLastMessage: useNotificationStore((state) => state.setWsLastMessage),
incrementReconnectAttempts: useNotificationStore((state) => state.incrementReconnectAttempts),
resetReconnectAttempts: useNotificationStore((state) => state.resetReconnectAttempts),
addA2UINotification: useNotificationStore((state) => state.addA2UINotification),
// Execution store for state updates
const setExecutionStatus = useExecutionStore((state) => state.setExecutionStatus);
const setNodeStarted = useExecutionStore((state) => state.setNodeStarted);
const setNodeCompleted = useExecutionStore((state) => state.setNodeCompleted);
const setNodeFailed = useExecutionStore((state) => state.setNodeFailed);
const addLog = useExecutionStore((state) => state.addLog);
const completeExecution = useExecutionStore((state) => state.completeExecution);
const currentExecution = useExecutionStore((state) => state.currentExecution);
// Execution store
setExecutionStatus: useExecutionStore((state) => state.setExecutionStatus),
setNodeStarted: useExecutionStore((state) => state.setNodeStarted),
setNodeCompleted: useExecutionStore((state) => state.setNodeCompleted),
setNodeFailed: useExecutionStore((state) => state.setNodeFailed),
addLog: useExecutionStore((state) => state.addLog),
completeExecution: useExecutionStore((state) => state.completeExecution),
currentExecution: useExecutionStore((state) => state.currentExecution),
// Flow store for node status updates on canvas
const updateNode = useFlowStore((state) => state.updateNode);
// Flow store
updateNode: useFlowStore((state) => state.updateNode),
// CLI stream store for CLI output handling
const addOutput = useCliStreamStore((state) => state.addOutput);
// CLI stream store
addOutput: useCliStreamStore((state) => state.addOutput),
// Coordinator store for coordinator state updates
const updateNodeStatus = useCoordinatorStore((state) => state.updateNodeStatus);
const addCoordinatorLog = useCoordinatorStore((state) => state.addLog);
const setActiveQuestion = useCoordinatorStore((state) => state.setActiveQuestion);
const markExecutionComplete = useCoordinatorStore((state) => state.markExecutionComplete);
const coordinatorExecutionId = useCoordinatorStore((state) => state.currentExecutionId);
// Coordinator store
updateNodeStatus: useCoordinatorStore((state) => state.updateNodeStatus),
addCoordinatorLog: useCoordinatorStore((state) => state.addLog),
setActiveQuestion: useCoordinatorStore((state) => state.setActiveQuestion),
markExecutionComplete: useCoordinatorStore((state) => state.markExecutionComplete),
coordinatorExecutionId: useCoordinatorStore((state) => state.currentExecutionId),
});
// Update refs periodically to ensure they have fresh store references
useEffect(() => {
storeRefs.current = {
// Notification store
setWsStatus: useNotificationStore((state) => state.setWsStatus),
setWsLastMessage: useNotificationStore((state) => state.setWsLastMessage),
incrementReconnectAttempts: useNotificationStore((state) => state.incrementReconnectAttempts),
resetReconnectAttempts: useNotificationStore((state) => state.resetReconnectAttempts),
addA2UINotification: useNotificationStore((state) => state.addA2UINotification),
// Execution store
setExecutionStatus: useExecutionStore((state) => state.setExecutionStatus),
setNodeStarted: useExecutionStore((state) => state.setNodeStarted),
setNodeCompleted: useExecutionStore((state) => state.setNodeCompleted),
setNodeFailed: useExecutionStore((state) => state.setNodeFailed),
addLog: useExecutionStore((state) => state.addLog),
completeExecution: useExecutionStore((state) => state.completeExecution),
currentExecution: useExecutionStore((state) => state.currentExecution),
// Flow store
updateNode: useFlowStore((state) => state.updateNode),
// CLI stream store
addOutput: useCliStreamStore((state) => state.addOutput),
// Coordinator store
updateNodeStatus: useCoordinatorStore((state) => state.updateNodeStatus),
addCoordinatorLog: useCoordinatorStore((state) => state.addLog),
setActiveQuestion: useCoordinatorStore((state) => state.setActiveQuestion),
markExecutionComplete: useCoordinatorStore((state) => state.markExecutionComplete),
coordinatorExecutionId: useCoordinatorStore((state) => state.currentExecutionId),
};
}); // Run on every render to keep refs fresh
// Handle incoming WebSocket messages
// Note: Using refs via storeRefs to prevent handler recreation on every store change
const handleMessage = useCallback(
(event: MessageEvent) => {
// Guard against state updates after unmount
@@ -81,7 +119,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const data = JSON.parse(event.data);
// Store last message for debugging
setWsLastMessage(data);
storeRefs.current.setWsLastMessage(data);
// Handle CLI messages
if (data.type?.startsWith('CLI_')) {
@@ -90,7 +128,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const { executionId, tool, mode, timestamp } = data.payload;
// Add system message for CLI start
addOutput(executionId, {
storeRefs.current.addOutput(executionId, {
type: 'system',
content: `[${new Date(timestamp).toLocaleTimeString()}] CLI execution started: ${tool} (${mode || 'default'} mode)`,
timestamp: Date.now(),
@@ -119,7 +157,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
lines.forEach((line: string) => {
// Add non-empty lines, or single line if that's all we have
if (line.trim() || lines.length === 1) {
addOutput(executionId, {
storeRefs.current.addOutput(executionId, {
type: unitType as any,
content: line,
timestamp: Date.now(),
@@ -135,7 +173,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const statusText = success ? 'completed successfully' : 'failed';
const durationText = duration ? ` (${duration}ms)` : '';
addOutput(executionId, {
storeRefs.current.addOutput(executionId, {
type: 'system',
content: `[${new Date().toLocaleTimeString()}] CLI execution ${statusText}${durationText}`,
timestamp: Date.now(),
@@ -150,7 +188,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
if (data.type === 'a2ui-surface') {
const parsed = SurfaceUpdateSchema.safeParse(data.payload);
if (parsed.success) {
addA2UINotification(parsed.data, 'Interactive UI');
storeRefs.current.addA2UINotification(parsed.data, 'Interactive UI');
} else {
console.warn('[WebSocket] Invalid A2UI surface:', parsed.error.issues);
}
@@ -159,6 +197,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
// Handle Coordinator messages
if (data.type?.startsWith('COORDINATOR_')) {
const { coordinatorExecutionId } = storeRefs.current;
// Only process messages for current coordinator execution
if (coordinatorExecutionId && data.executionId !== coordinatorExecutionId) {
return;
@@ -169,26 +208,26 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
case 'COORDINATOR_STATE_UPDATE':
// Check for completion
if (data.status === 'completed') {
markExecutionComplete(true);
storeRefs.current.markExecutionComplete(true);
} else if (data.status === 'failed') {
markExecutionComplete(false);
storeRefs.current.markExecutionComplete(false);
}
break;
case 'COORDINATOR_COMMAND_STARTED':
updateNodeStatus(data.nodeId, 'running');
storeRefs.current.updateNodeStatus(data.nodeId, 'running');
break;
case 'COORDINATOR_COMMAND_COMPLETED':
updateNodeStatus(data.nodeId, 'completed', data.result);
storeRefs.current.updateNodeStatus(data.nodeId, 'completed', data.result);
break;
case 'COORDINATOR_COMMAND_FAILED':
updateNodeStatus(data.nodeId, 'failed', undefined, data.error);
storeRefs.current.updateNodeStatus(data.nodeId, 'failed', undefined, data.error);
break;
case 'COORDINATOR_LOG_ENTRY':
addCoordinatorLog(
storeRefs.current.addCoordinatorLog(
data.log.message,
data.log.level,
data.log.nodeId,
@@ -197,7 +236,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
break;
case 'COORDINATOR_QUESTION_ASKED':
setActiveQuestion(data.question);
storeRefs.current.setActiveQuestion(data.question);
break;
case 'COORDINATOR_ANSWER_RECEIVED':
@@ -223,6 +262,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const message = parsed.data as OrchestratorWebSocketMessage;
// Only process messages for current execution
const { currentExecution } = storeRefs.current;
if (currentExecution && message.execId !== currentExecution.execId) {
return;
}
@@ -230,39 +270,39 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
// Dispatch to execution store based on message type
switch (message.type) {
case 'ORCHESTRATOR_STATE_UPDATE':
setExecutionStatus(message.status, message.currentNodeId);
storeRefs.current.setExecutionStatus(message.status, message.currentNodeId);
// Check for completion
if (message.status === 'completed' || message.status === 'failed') {
completeExecution(message.status);
storeRefs.current.completeExecution(message.status);
}
break;
case 'ORCHESTRATOR_NODE_STARTED':
setNodeStarted(message.nodeId);
storeRefs.current.setNodeStarted(message.nodeId);
// Update canvas node status
updateNode(message.nodeId, { executionStatus: 'running' });
storeRefs.current.updateNode(message.nodeId, { executionStatus: 'running' });
break;
case 'ORCHESTRATOR_NODE_COMPLETED':
setNodeCompleted(message.nodeId, message.result);
storeRefs.current.setNodeCompleted(message.nodeId, message.result);
// Update canvas node status
updateNode(message.nodeId, {
storeRefs.current.updateNode(message.nodeId, {
executionStatus: 'completed',
executionResult: message.result,
});
break;
case 'ORCHESTRATOR_NODE_FAILED':
setNodeFailed(message.nodeId, message.error);
storeRefs.current.setNodeFailed(message.nodeId, message.error);
// Update canvas node status
updateNode(message.nodeId, {
storeRefs.current.updateNode(message.nodeId, {
executionStatus: 'failed',
executionError: message.error,
});
break;
case 'ORCHESTRATOR_LOG':
addLog(message.log as ExecutionLog);
storeRefs.current.addLog(message.log as ExecutionLog);
break;
}
@@ -272,69 +312,15 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
console.error('[WebSocket] Failed to parse message:', error);
}
},
[
currentExecution,
coordinatorExecutionId,
setWsLastMessage,
setExecutionStatus,
setNodeStarted,
setNodeCompleted,
setNodeFailed,
addLog,
completeExecution,
updateNode,
addOutput,
addA2UINotification,
updateNodeStatus,
addCoordinatorLog,
setActiveQuestion,
markExecutionComplete,
onMessage,
]
[onMessage] // Only dependency is onMessage, all other functions accessed via refs
);
// Connect to WebSocket
const connect = useCallback(() => {
if (!enabled) return;
// Construct WebSocket URL
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
const wsUrl = `${protocol}//${window.location.host}/ws`;
try {
setWsStatus('connecting');
const ws = new WebSocket(wsUrl);
wsRef.current = ws;
ws.onopen = () => {
console.log('[WebSocket] Connected');
setWsStatus('connected');
resetReconnectAttempts();
reconnectDelayRef.current = RECONNECT_DELAY_BASE;
};
ws.onmessage = handleMessage;
ws.onclose = () => {
console.log('[WebSocket] Disconnected');
setWsStatus('disconnected');
wsRef.current = null;
scheduleReconnect();
};
ws.onerror = (error) => {
console.error('[WebSocket] Error:', error);
setWsStatus('error');
};
} catch (error) {
console.error('[WebSocket] Failed to connect:', error);
setWsStatus('error');
scheduleReconnect();
}
}, [enabled, handleMessage, setWsStatus, resetReconnectAttempts]);
// Use ref to avoid circular dependency with scheduleReconnect
const connectRef = useRef<(() => void) | null>(null);
// Schedule reconnection with exponential backoff
// Define this first to avoid circular dependency
const scheduleReconnect = useCallback(() => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
@@ -343,11 +329,11 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
const delay = reconnectDelayRef.current;
console.log(`[WebSocket] Reconnecting in ${delay}ms...`);
setWsStatus('reconnecting');
incrementReconnectAttempts();
storeRefs.current.setWsStatus('reconnecting');
storeRefs.current.incrementReconnectAttempts();
reconnectTimeoutRef.current = setTimeout(() => {
connect();
connectRef.current?.();
}, delay);
// Increase delay for next attempt (exponential backoff)
@@ -355,7 +341,50 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
reconnectDelayRef.current * RECONNECT_DELAY_MULTIPLIER,
RECONNECT_DELAY_MAX
);
}, [connect, setWsStatus, incrementReconnectAttempts]);
}, []); // No dependencies - uses connectRef
const connect = useCallback(() => {
if (!enabled) return;
// Construct WebSocket URL
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
const wsUrl = `${protocol}//${window.location.host}/ws`;
try {
storeRefs.current.setWsStatus('connecting');
const ws = new WebSocket(wsUrl);
wsRef.current = ws;
ws.onopen = () => {
console.log('[WebSocket] Connected');
storeRefs.current.setWsStatus('connected');
storeRefs.current.resetReconnectAttempts();
reconnectDelayRef.current = RECONNECT_DELAY_BASE;
};
ws.onmessage = handleMessage;
ws.onclose = () => {
console.log('[WebSocket] Disconnected');
storeRefs.current.setWsStatus('disconnected');
wsRef.current = null;
scheduleReconnect();
};
ws.onerror = (error) => {
console.error('[WebSocket] Error:', error);
storeRefs.current.setWsStatus('error');
};
} catch (error) {
console.error('[WebSocket] Failed to connect:', error);
storeRefs.current.setWsStatus('error');
scheduleReconnect();
}
}, [enabled, handleMessage, scheduleReconnect]);
// Update connect ref after connect is defined
connectRef.current = connect;
// Send message through WebSocket
const send = useCallback((message: unknown) => {
@@ -367,13 +396,14 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
}, []);
// Manual reconnect
// Use connectRef to avoid depending on connect
const reconnect = useCallback(() => {
if (wsRef.current) {
wsRef.current.close();
}
reconnectDelayRef.current = RECONNECT_DELAY_BASE;
connect();
}, [connect]);
connectRef.current?.();
}, []); // No dependencies - uses connectRef
// Check connection status
const isConnected = wsRef.current?.readyState === WebSocket.OPEN;