fix: WebSocket race condition in session hooks and chat input performance (#119)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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 };
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user