mirror of
https://github.com/catlog22/Claude-Code-Workflow.git
synced 2026-02-05 01:50:27 +08:00
refactor: optimize useWebSocket hook by consolidating store references and improving handler stability
This commit is contained in:
catlog22
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user