fix: WebSocket race condition in session hooks and chat input performance (#119)

This commit is contained in:
João Salavisa
2026-01-28 22:53:35 +00:00
committed by GitHub
parent 7f0592cccd
commit a97c35904f
9 changed files with 69 additions and 42 deletions
+38 -30
View File
@@ -1,5 +1,6 @@
'use client';
import { useMemo, useCallback } from 'react';
import CodeMirror from '@uiw/react-codemirror';
import { markdown } from '@codemirror/lang-markdown';
import { EditorView, keymap, placeholder as cmPlaceholder } from '@codemirror/view';
@@ -37,40 +38,47 @@ export function TaskChatInput({
const submitKey = shortcutToCodeMirrorKeybinding(SHORTCUTS.SUBMIT);
const completionSource = useChatCompletions(sessionId);
const handleSubmitKey = () => {
// Memoize the submit handler - include onSubmit in deps since it's now stable from parent
const handleSubmitKey = useCallback(() => {
onSubmit();
return true;
};
}, [onSubmit]);
const submitKeymap = Prec.highest(
keymap.of([
{ key: submitKey, run: handleSubmitKey, preventDefault: true },
{ key: 'Mod-Enter', run: handleSubmitKey, preventDefault: true },
])
);
const handleKeydown = useCallback((event: KeyboardEvent) => {
if (matchesShortcut(event, SHORTCUTS.SUBMIT)) {
event.preventDefault();
event.stopPropagation();
onSubmit();
return true;
}
return false;
}, [onSubmit]);
const extensions: Extension[] = [
EditorView.lineWrapping,
chatEditorTheme,
EditorView.domEventHandlers({
keydown: (event) => {
if (matchesShortcut(event, SHORTCUTS.SUBMIT)) {
event.preventDefault();
event.stopPropagation();
onSubmit();
return true;
}
return false;
},
}),
history(),
closeBrackets(),
markdown(),
cmPlaceholder(placeholder ?? ''),
autocompletion({ override: [completionSource], aboveCursor: true }),
submitKeymap,
keymap.of([...defaultKeymap, ...historyKeymap, ...closeBracketsKeymap]),
];
// Memoize extensions to prevent CodeMirror reconfiguration on every render
// Only recreate when dependencies actually change
const extensions = useMemo<Extension[]>(() => {
const submitKeymap = Prec.highest(
keymap.of([
{ key: submitKey, run: handleSubmitKey, preventDefault: true },
{ key: 'Mod-Enter', run: handleSubmitKey, preventDefault: true },
])
);
return [
EditorView.lineWrapping,
chatEditorTheme,
EditorView.domEventHandlers({
keydown: handleKeydown,
}),
history(),
closeBrackets(),
markdown(),
cmPlaceholder(placeholder ?? ''),
autocompletion({ override: [completionSource], aboveCursor: true }),
submitKeymap,
keymap.of([...defaultKeymap, ...historyKeymap, ...closeBracketsKeymap]),
];
}, [placeholder, completionSource, submitKey, handleSubmitKey, handleKeydown]);
return (
<div
+3 -3
View File
@@ -135,7 +135,7 @@ export function TaskChatPanel({
const showApproveButton =
session?.review_status != null && session.review_status !== 'approved' && !isWorking;
const handleSubmit = async (event?: FormEvent) => {
const handleSubmit = useCallback(async (event?: FormEvent) => {
event?.preventDefault();
const trimmed = messageInput.trim();
if (!trimmed || isSending) return;
@@ -156,7 +156,7 @@ export function TaskChatPanel({
} finally {
setIsSending(false);
}
};
}, [messageInput, isSending, planModeEnabled, onSend, handleSendMessage]);
return (
@@ -182,7 +182,7 @@ export function TaskChatPanel({
<TaskChatInput
value={messageInput}
onChange={setMessageInput}
onSubmit={() => handleSubmit()}
onSubmit={handleSubmit}
placeholder={
agentMessageCount > 0
? 'Continue working on this task...'
@@ -9,13 +9,15 @@ export function useSessionAgentctl(sessionId: string | null) {
const status = useAppStore((state) =>
sessionId ? state.sessionAgentctl.itemsBySessionId[sessionId] : undefined
);
const connectionStatus = useAppStore((state) => state.connection.status);
useEffect(() => {
if (!session?.id) return;
if (connectionStatus !== 'connected') return;
const client = getWebSocketClient();
if (!client) return;
return client.subscribeSession(session.id);
}, [session?.id]);
}, [session?.id, connectionStatus]);
return {
status: status?.status ?? 'starting',
@@ -16,6 +16,7 @@ export function useSessionCommits(sessionId: string | null) {
);
const setSessionCommits = useAppStore((state) => state.setSessionCommits);
const setSessionCommitsLoading = useAppStore((state) => state.setSessionCommitsLoading);
const connectionStatus = useAppStore((state) => state.connection.status);
const fetchCommits = useCallback(async () => {
if (!sessionId) return;
@@ -42,10 +43,11 @@ export function useSessionCommits(sessionId: string | null) {
// Fetch commits on mount
useEffect(() => {
if (connectionStatus !== 'connected') return;
if (sessionId && !commits) {
fetchCommits();
}
}, [sessionId, commits, fetchCommits]);
}, [sessionId, commits, fetchCommits, connectionStatus]);
return {
commits: commits ?? [],
@@ -31,6 +31,7 @@ export function useSessionContextWindow(sessionId: string | null): ContextWindow
sessionId ? state.taskSessions.items[sessionId] : undefined
);
const setContextWindow = useAppStore((state) => state.setContextWindow);
const connectionStatus = useAppStore((state) => state.connection.status);
// Populate context window from session metadata if not already in store
useEffect(() => {
@@ -59,6 +60,7 @@ export function useSessionContextWindow(sessionId: string | null): ContextWindow
// Subscribe to session updates via WebSocket
useEffect(() => {
if (!sessionId) return;
if (connectionStatus !== 'connected') return;
const client = getWebSocketClient();
if (client) {
const unsubscribe = client.subscribeSession(sessionId);
@@ -67,7 +69,7 @@ export function useSessionContextWindow(sessionId: string | null): ContextWindow
// Don't clear context window on cleanup - keep it cached
};
}
}, [sessionId]);
}, [sessionId, connectionStatus]);
return contextWindow;
}
@@ -19,6 +19,7 @@ export function useSessionGitSnapshots(sessionId: string | null, options?: { lim
);
const setGitSnapshots = useAppStore((state) => state.setGitSnapshots);
const setGitSnapshotsLoading = useAppStore((state) => state.setGitSnapshotsLoading);
const connectionStatus = useAppStore((state) => state.connection.status);
const fetchSnapshots = useCallback(async () => {
if (!sessionId) return;
@@ -48,10 +49,11 @@ export function useSessionGitSnapshots(sessionId: string | null, options?: { lim
// Fetch snapshots on mount
useEffect(() => {
if (connectionStatus !== 'connected') return;
if (sessionId && !snapshots) {
fetchSnapshots();
}
}, [sessionId, snapshots, fetchSnapshots]);
}, [sessionId, snapshots, fetchSnapshots, connectionStatus]);
return {
snapshots: snapshots ?? [],
@@ -14,13 +14,17 @@ export function useSessionGitStatus(sessionId: string | null) {
sessionId ? state.gitStatus.bySessionId[sessionId] : undefined
);
const setGitStatus = useAppStore((state) => state.setGitStatus);
const connectionStatus = useAppStore((state) => state.connection.status);
const prevSessionIdRef = useRef<string | null>(null);
// Fetch initial status on mount if not already loaded
// Also refetch when switching to a different session
// Also refetch when switching to a different session or when WebSocket connects
useEffect(() => {
if (!sessionId) return;
// Wait for WebSocket to be connected before fetching
if (connectionStatus !== 'connected') return;
// Detect session change to force refetch
const sessionChanged = prevSessionIdRef.current !== null &&
prevSessionIdRef.current !== sessionId;
@@ -87,11 +91,15 @@ export function useSessionGitStatus(sessionId: string | null) {
};
fetchSnapshot();
}, [sessionId, gitStatus, setGitStatus]);
}, [sessionId, gitStatus, setGitStatus, connectionStatus]);
// Subscribe to session updates to receive git status via WebSocket
useEffect(() => {
if (!sessionId) return;
// Wait for WebSocket to be connected before subscribing
if (connectionStatus !== 'connected') return;
const client = getWebSocketClient();
if (client) {
const unsubscribe = client.subscribeSession(sessionId);
@@ -100,7 +108,7 @@ export function useSessionGitStatus(sessionId: string | null) {
// Don't clear git status on cleanup - keep it cached for when user switches back
};
}
}, [sessionId]);
}, [sessionId, connectionStatus]);
return gitStatus;
}
@@ -134,13 +134,14 @@ export function useSessionMessages(
useEffect(() => {
if (!taskSessionId) return;
if (connectionStatus !== 'connected') return;
const client = getWebSocketClient();
if (!client) return;
const unsubscribe = client.subscribeSession(taskSessionId);
return () => {
unsubscribe();
};
}, [taskSessionId]);
}, [taskSessionId, connectionStatus]);
useEffect(() => {
if (!taskSessionId || !taskSessionState) return;
@@ -12,6 +12,7 @@ export function useSession(sessionId: string | null): UseSessionResult {
const session = useAppStore((state) =>
sessionId ? state.taskSessions.items[sessionId] ?? null : null
);
const connectionStatus = useAppStore((state) => state.connection.status);
const isActive = useMemo(() => {
if (!session?.state) return false;
@@ -19,6 +20,7 @@ export function useSession(sessionId: string | null): UseSessionResult {
}, [session?.state]);
useEffect(() => {
if (connectionStatus !== 'connected') return;
if (!session?.id) return;
const client = getWebSocketClient();
if (!client) return;
@@ -26,7 +28,7 @@ export function useSession(sessionId: string | null): UseSessionResult {
return () => {
unsubscribe();
};
}, [session?.id]);
}, [session?.id, connectionStatus]);
return { session, isActive };
}