tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Separate SSE streaming from POST work submission (#7834)

Browse Source
This commit is contained in:
jh-block
2026-03-18 18:00:06 +01:00
committed by GitHub
parent c2d10cea9b
commit 8d72344602
15 changed files with 1830 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/src/api/index.ts
+2 -2
View File
File diff suppressed because one or more lines are too long
ui/desktop/src/api/sdk.gen.ts
+21 -1
View File
File diff suppressed because one or more lines are too long
ui/desktop/src/api/types.gen.ts
+109
View File
@@ -59,6 +59,10 @@ export type CallToolResponse = {
structuredContent?: unknown;
};
export type CancelRequest = {
request_id: string;
};
export type ChatRequest = {
/**
* Override the server's conversation history. Only use this when you need absolute control
@@ -696,6 +700,9 @@ export type MessageEvent = {
} | {
conversation: Conversation;
type: 'UpdateConversation';
} | {
request_ids: Array<string>;
type: 'ActiveRequests';
} | {
type: 'Ping';
};
@@ -1255,6 +1262,21 @@ export type SessionListResponse = {
sessions: Array<Session>;
};
export type SessionReplyRequest = {
override_conversation?: Array<Message> | null;
recipe_name?: string | null;
recipe_version?: string | null;
/**
* Client-generated UUIDv7 identifying this request.
*/
request_id: string;
user_message: Message;
};
export type SessionReplyResponse = {
request_id: string;
};
export type SessionType = 'user' | 'scheduled' | 'sub_agent' | 'hidden' | 'terminal' | 'gateway';
export type SessionsQuery = {
@@ -4129,6 +4151,93 @@ export type SearchSessionsResponses = {
export type SearchSessionsResponse = SearchSessionsResponses[keyof SearchSessionsResponses];
export type SessionCancelData = {
body: CancelRequest;
path: {
/**
* Session ID
*/
id: string;
};
query?: never;
url: '/sessions/{id}/cancel';
};
export type SessionCancelResponses = {
/**
* Cancellation accepted
*/
200: unknown;
};
export type SessionEventsData = {
body?: never;
path: {
/**
* Session ID
*/
id: string;
};
query?: never;
url: '/sessions/{id}/events';
};
export type SessionEventsErrors = {
/**
* Session not found
*/
404: unknown;
};
export type SessionEventsResponses = {
/**
* SSE event stream
*/
200: MessageEvent;
};
export type SessionEventsResponse = SessionEventsResponses[keyof SessionEventsResponses];
export type SessionReplyData = {
body: SessionReplyRequest;
path: {
/**
* Session ID
*/
id: string;
};
query?: never;
url: '/sessions/{id}/reply';
};
export type SessionReplyErrors = {
/**
* Invalid request
*/
400: unknown;
/**
* Session not found
*/
404: unknown;
/**
* Agent not initialized
*/
424: unknown;
/**
* Internal server error
*/
500: unknown;
};
export type SessionReplyResponses = {
/**
* Request accepted
*/
200: SessionReplyResponse;
};
export type SessionReplyResponse2 = SessionReplyResponses[keyof SessionReplyResponses];
export type DeleteSessionData = {
body?: never;
path: {
ui/desktop/src/hooks/useChatStream.ts
+384 -198
View File
@@ -1,15 +1,15 @@
import { useCallback, useEffect, useMemo, useReducer, useRef } from 'react';
import { v7 as uuidv7 } from 'uuid';
import { AppEvents } from '../constants/events';
import { ChatState } from '../types/chatState';
import { toastError } from '../toasts';
import {
getSession,
Message,
MessageEvent,
reply,
resumeAgent,
Session,
sessionCancel,
sessionReply,
TokenState,
updateFromSession,
updateSessionUserRecipeValues,
@@ -27,6 +27,7 @@ import {
import { errorMessage } from '../utils/conversionUtils';
import { showExtensionLoadResults } from '../utils/extensionErrorUtils';
import { maybeHandlePlatformEvent } from '../utils/platform_events';
import { useSessionEvents, type SessionEvent } from './useSessionEvents';
const resultsCache = new Map<string, { messages: Message[]; session: Session }>();
@@ -217,14 +218,17 @@ function prefersReducedMotion(): boolean {
const REDUCED_MOTION_BATCH_INTERVAL = 1000;
async function streamFromResponse(
stream: AsyncIterable<MessageEvent>,
/**
* Creates an event processor that handles individual SSE events for a request.
* Returns an unsubscribe function and a handler to process events.
*/
function createEventProcessor(
initialMessages: Message[],
dispatch: React.Dispatch<StreamAction>,
onFinish: (error?: string) => void,
sessionId: string,
signal?: AbortController['signal']
): Promise<void> {
onReloadNeeded?: () => void,
) {
let currentMessages = initialMessages;
const reduceMotion = prefersReducedMotion();
let latestTokenState: TokenState | null = null;
@@ -266,88 +270,80 @@ async function streamFromResponse(
}
};
try {
for await (const event of stream) {
switch (event.type) {
case 'Message': {
const msg = event.message;
currentMessages = pushMessage(currentMessages, msg);
// Returns true if the event is terminal (Finish or Error)
const processEvent = (event: SessionEvent): boolean => {
switch (event.type) {
case 'Message': {
const msg = (event as Record<string, unknown>).message as Message;
const tokenState = (event as Record<string, unknown>).token_state as TokenState;
currentMessages = pushMessage(currentMessages, msg);
const hasToolConfirmation = msg.content.some(
(content) =>
content.type === 'actionRequired' && content.data.actionType === 'toolConfirmation'
);
const hasToolConfirmation = msg.content.some(
(content) =>
content.type === 'actionRequired' && content.data.actionType === 'toolConfirmation'
);
const hasElicitation = msg.content.some(
(content) =>
content.type === 'actionRequired' && content.data.actionType === 'elicitation'
);
const hasElicitation = msg.content.some(
(content) =>
content.type === 'actionRequired' && content.data.actionType === 'elicitation'
);
if (hasToolConfirmation || hasElicitation) {
maybeUpdateUI(event.token_state, ChatState.WaitingForUserInput, true);
} else if (getCompactingMessage(msg)) {
maybeUpdateUI(event.token_state, ChatState.Compacting);
} else if (getThinkingMessage(msg)) {
maybeUpdateUI(event.token_state, ChatState.Thinking);
} else {
maybeUpdateUI(event.token_state, ChatState.Streaming);
}
break;
if (hasToolConfirmation || hasElicitation) {
maybeUpdateUI(tokenState, ChatState.WaitingForUserInput, true);
} else if (getCompactingMessage(msg)) {
maybeUpdateUI(tokenState, ChatState.Compacting);
} else if (getThinkingMessage(msg)) {
maybeUpdateUI(tokenState, ChatState.Thinking);
} else {
maybeUpdateUI(tokenState, ChatState.Streaming);
}
case 'Error': {
flushBatchedUpdates();
onFinish('Stream error: ' + event.error);
return;
}
case 'Finish': {
flushBatchedUpdates();
onFinish();
return;
}
case 'ModelChange': {
break;
}
case 'UpdateConversation': {
currentMessages = event.conversation;
if (!reduceMotion) {
dispatch({ type: 'SET_MESSAGES', payload: event.conversation });
} else {
hasPendingUpdate = true;
}
break;
}
case 'Notification': {
dispatch({ type: 'ADD_NOTIFICATION', payload: event as NotificationEvent });
maybeHandlePlatformEvent(event.message, sessionId);
break;
}
case 'Ping':
break;
return false;
}
case 'Error': {
flushBatchedUpdates();
const errorMsg = String((event as Record<string, unknown>).error ?? '');
if (errorMsg.includes('too far behind') && onReloadNeeded) {
// Server indicated we missed events — end streaming without setting
// an error (which would show a blocking error screen), then reload
// the full conversation so the UI reflects the actual state.
onFinish();
onReloadNeeded();
} else {
onFinish('Stream error: ' + errorMsg);
}
return true;
}
case 'Finish': {
flushBatchedUpdates();
onFinish();
return true;
}
case 'ModelChange': {
return false;
}
case 'UpdateConversation': {
const conversation = (event as Record<string, unknown>).conversation as Message[];
currentMessages = conversation;
if (!reduceMotion) {
dispatch({ type: 'SET_MESSAGES', payload: conversation });
} else {
hasPendingUpdate = true;
}
return false;
}
case 'Notification': {
dispatch({ type: 'ADD_NOTIFICATION', payload: event as unknown as NotificationEvent });
maybeHandlePlatformEvent((event as Record<string, unknown>).message, sessionId);
return false;
}
case 'Ping':
return false;
default:
return false;
}
};
// If we reach here, the stream ended without a Finish or Error event.
// This happens when the connection drops and retries are exhausted — the
// generator exits its loop without yielding a terminal event. We call
// onFinish() without an error to keep the conversation visible (passing
// an error would trigger a full-page error screen via sessionLoadError),
// then show a toast so the user knows the response may be incomplete.
// If the signal was aborted, the user intentionally stopped streaming,
// so we skip the toast.
flushBatchedUpdates();
onFinish();
if (!signal?.aborted) {
toastError({
title: 'Connection lost',
msg: 'The response may be incomplete. You can try sending your message again.',
});
}
} catch (error) {
flushBatchedUpdates();
if (error instanceof Error && error.name !== 'AbortError') {
onFinish('Stream error: ' + errorMessage(error));
}
}
return processEvent;
}
export function useChatStream({
@@ -357,14 +353,26 @@ export function useChatStream({
}: UseChatStreamProps): UseChatStreamReturn {
const [state, dispatch] = useReducer(streamReducer, initialState);
// Refs for values needed in callbacks without causing re-renders
const abortControllerRef = useRef<AbortController | null>(null);
// Long-lived SSE connection for this session
const { addListener, setActiveRequestsHandler } = useSessionEvents(sessionId);
// Track the active request for cancellation (includes the session that started it)
const activeRequestIdRef = useRef<string | null>(null);
const activeRequestSessionIdRef = useRef<string | null>(null);
const activeAbortRef = useRef<AbortController | null>(null);
const activeUnsubscribeRef = useRef<(() => void) | null>(null);
const lastInteractionTimeRef = useRef<number>(Date.now());
// When ActiveRequests fires before resumeAgent populates messages (cold mount),
// defer the reattach until the session is loaded so the event processor has
// the full conversation history. Events are buffered in the meantime.
const pendingReattachRequestIdRef = useRef<string | null>(null);
const pendingReattachBufferRef = useRef<SessionEvent[]>([]);
const namePollingRef = useRef<ReturnType<typeof setTimeout> | null>(null);
// Ref to access latest state in callbacks (avoids stale closures)
const stateRef = useRef(state);
stateRef.current = state;
const doReattachRef = useRef<((requestId: string, messages: Message[]) => void) | null>(null);
useEffect(() => {
return () => {
@@ -383,6 +391,10 @@ export function useChatStream({
const onFinish = useCallback(
async (error?: string): Promise<void> => {
// Note: SSE listener/ref cleanup is handled by the terminal-event
// handler in each listener closure (which guards on requestId) so
// that overlapping requests don't clobber each other's state.
if (namePollingRef.current) {
clearTimeout(namePollingRef.current);
namePollingRef.current = null;
@@ -404,13 +416,10 @@ export function useChatStream({
}
// Refresh session name after each reply for the first 3 user messages
// The backend regenerates the name after each of the first 3 user messages
// to refine it as more context becomes available
if (!error && sessionId) {
const currentState = stateRef.current;
const userMessageCount = currentState.messages.filter((m) => m.role === 'user').length;
// Only refresh for the first 3 user messages
if (userMessageCount <= 3) {
try {
const response = await getSession({
@@ -424,7 +433,6 @@ export function useChatStream({
? { ...currentState.session, name: response.data.name }
: undefined,
});
// Notify sidebar of the name change
window.dispatchEvent(
new CustomEvent(AppEvents.SESSION_RENAMED, {
detail: { sessionId, newName: response.data.name },
@@ -432,7 +440,6 @@ export function useChatStream({
);
}
} catch (refreshError) {
// Silently fail - this is a nice-to-have feature
console.warn('Failed to refresh session name:', refreshError);
}
}
@@ -443,6 +450,208 @@ export function useChatStream({
[onStreamFinish, sessionId]
);
// Reload the full conversation from the server, e.g. after the SSE
// stream indicates the client fell too far behind the replay buffer.
const reloadConversation = useCallback(() => {
getSession({
path: { session_id: sessionId },
throwOnError: true,
}).then((response) => {
const session = response.data as Session;
if (session?.conversation) {
dispatch({ type: 'SET_MESSAGES', payload: session.conversation });
}
}).catch((e) => {
console.warn('Failed to reload conversation after buffer overflow:', e);
});
}, [sessionId]);
// Perform the actual reattach: wire up an event processor and listener
// for a request that is already in-flight on the server.
const doReattach = useCallback(
(requestId: string, messages: Message[]) => {
activeRequestIdRef.current = requestId;
activeRequestSessionIdRef.current = sessionId;
pendingReattachRequestIdRef.current = null;
dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Streaming });
dispatch({ type: 'SET_SESSION_LOAD_ERROR', payload: undefined });
const processEvent = createEventProcessor(
messages,
dispatch,
onFinish,
sessionId,
reloadConversation,
);
// Replay any events that were buffered during cold-mount wait
const buffered = pendingReattachBufferRef.current;
pendingReattachBufferRef.current = [];
let finished = false;
for (const event of buffered) {
if (processEvent(event)) {
finished = true;
break;
}
}
if (finished) {
// The reply already completed while we were waiting for session load.
// Clean up — the buffering listener will be replaced below but the
// old one captured into activeUnsubscribeRef should be removed.
if (activeUnsubscribeRef.current) {
activeUnsubscribeRef.current();
activeUnsubscribeRef.current = null;
}
activeRequestIdRef.current = null;
activeRequestSessionIdRef.current = null;
return;
}
// Replace the buffering listener with a real processing listener
if (activeUnsubscribeRef.current) {
activeUnsubscribeRef.current();
}
const unsubscribe = addListener(requestId, (event) => {
const isTerminal = processEvent(event);
if (isTerminal) {
unsubscribe();
if (activeRequestIdRef.current === requestId) {
activeUnsubscribeRef.current = null;
activeRequestIdRef.current = null;
activeRequestSessionIdRef.current = null;
}
}
});
activeUnsubscribeRef.current = unsubscribe;
},
[sessionId, addListener, onFinish, reloadConversation],
);
doReattachRef.current = doReattach;
// Reattach to in-flight replies discovered via the SSE ActiveRequests event.
// This handles the case where the chat view remounts while a reply is still
// running on the server — the new hook instance picks up the existing request
// and starts processing its events.
useEffect(() => {
setActiveRequestsHandler((requestIds: string[]) => {
// Only reattach if we don't already have an active request
if (activeRequestIdRef.current) return;
if (requestIds.length === 0) return;
// Reattach to the first (most recent) active request.
// Multiple concurrent requests per session aren't supported in the UI.
const requestId = requestIds[0];
const currentMessages = stateRef.current.messages;
if (currentMessages.length === 0) {
// Cold mount: resumeAgent hasn't populated messages yet.
// Defer event processing until session load completes so the
// processor starts with the full conversation history.
// Register a buffering listener NOW so replayed events aren't
// lost while we wait.
pendingReattachRequestIdRef.current = requestId;
pendingReattachBufferRef.current = [];
activeRequestIdRef.current = requestId;
activeRequestSessionIdRef.current = sessionId;
dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Streaming });
dispatch({ type: 'SET_SESSION_LOAD_ERROR', payload: undefined });
const unsubscribe = addListener(requestId, (event) => {
pendingReattachBufferRef.current.push(event);
});
activeUnsubscribeRef.current = unsubscribe;
return;
}
doReattach(requestId, currentMessages);
});
return () => {
setActiveRequestsHandler(null);
};
}, [sessionId, addListener, onFinish, reloadConversation, setActiveRequestsHandler, doReattach]);
/**
* Submit a message via the new POST+SSE pattern.
* 1. Generate request_id
* 2. Register SSE listener BEFORE POST (no race condition)
* 3. POST to /sessions/{id}/reply
* 4. Events arrive on the long-lived SSE connection
*/
const submitToSession = useCallback(
async (
targetSessionId: string,
userMessage: Message,
currentMessages: Message[],
overrideConversation?: Message[],
recipeName?: string,
recipeVersion?: string,
) => {
const requestId = uuidv7();
const abortController = new AbortController();
activeRequestIdRef.current = requestId;
activeRequestSessionIdRef.current = targetSessionId;
activeAbortRef.current = abortController;
// Create event processor and register listener BEFORE the POST
const processEvent = createEventProcessor(
currentMessages,
dispatch,
onFinish,
targetSessionId,
reloadConversation,
);
const unsubscribe = addListener(requestId, (event) => {
const isTerminal = processEvent(event);
if (isTerminal) {
unsubscribe();
// Only clear global refs if this request is still the active one.
// A newer request may have already replaced them.
if (activeRequestIdRef.current === requestId) {
activeUnsubscribeRef.current = null;
activeRequestIdRef.current = null;
activeRequestSessionIdRef.current = null;
activeAbortRef.current = null;
}
}
});
activeUnsubscribeRef.current = unsubscribe;
try {
await sessionReply({
path: { id: targetSessionId },
body: {
request_id: requestId,
user_message: userMessage,
override_conversation: overrideConversation,
recipe_name: recipeName,
recipe_version: recipeVersion,
},
signal: abortController.signal,
throwOnError: true,
});
} catch (error) {
// Abort is expected when stopStreaming races with the POST
if (abortController.signal.aborted) return;
// POST failed — clean up listener and report error.
// Only clear global refs if this request is still the active one;
// a newer request may have already replaced them.
unsubscribe();
if (activeRequestIdRef.current === requestId) {
activeUnsubscribeRef.current = null;
activeRequestIdRef.current = null;
activeRequestSessionIdRef.current = null;
activeAbortRef.current = null;
}
onFinish('Submit error: ' + errorMessage(error));
}
},
[addListener, onFinish, reloadConversation]
);
// Load session on mount or sessionId change
useEffect(() => {
if (!sessionId) return;
@@ -494,12 +703,38 @@ export function useChatStream({
showExtensionLoadResults(extensionResults);
window.dispatchEvent(new CustomEvent(AppEvents.SESSION_EXTENSIONS_LOADED));
dispatch({
type: 'SESSION_LOADED',
payload: {
session: loadedSession!,
messages: loadedSession?.conversation || [],
tokenState: {
const pendingRequestId = pendingReattachRequestIdRef.current;
const reattachedToActiveRequest = activeRequestIdRef.current !== null;
if (pendingRequestId) {
// Cold-mount reattach: ActiveRequests arrived before resumeAgent
// returned. Load session state first, then complete the reattach
// with the full conversation so the event processor has context.
dispatch({
type: 'SESSION_LOADED',
payload: {
session: loadedSession!,
messages: loadedSession?.conversation || [],
tokenState: {
inputTokens: loadedSession?.input_tokens ?? 0,
outputTokens: loadedSession?.output_tokens ?? 0,
totalTokens: loadedSession?.total_tokens ?? 0,
accumulatedInputTokens: loadedSession?.accumulated_input_tokens ?? 0,
accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0,
accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0,
},
},
});
// Now complete the deferred reattach with the loaded messages
doReattachRef.current?.(pendingRequestId, loadedSession?.conversation || []);
} else if (reattachedToActiveRequest) {
// ActiveRequests already wired up an event processor with existing
// messages — only load session metadata, don't overwrite messages
// with the stale DB snapshot.
dispatch({ type: 'SET_SESSION', payload: loadedSession });
dispatch({
type: 'SET_TOKEN_STATE',
payload: {
inputTokens: loadedSession?.input_tokens ?? 0,
outputTokens: loadedSession?.output_tokens ?? 0,
totalTokens: loadedSession?.total_tokens ?? 0,
@@ -507,8 +742,24 @@ export function useChatStream({
accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0,
accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0,
},
},
});
});
} else {
dispatch({
type: 'SESSION_LOADED',
payload: {
session: loadedSession!,
messages: loadedSession?.conversation || [],
tokenState: {
inputTokens: loadedSession?.input_tokens ?? 0,
outputTokens: loadedSession?.output_tokens ?? 0,
totalTokens: loadedSession?.total_tokens ?? 0,
accumulatedInputTokens: loadedSession?.accumulated_input_tokens ?? 0,
accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0,
accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0,
},
},
});
}
listApps({
throwOnError: true,
@@ -535,7 +786,6 @@ export function useChatStream({
const { msg: userMessage, images } = input;
const currentState = stateRef.current;
// Guard: Don't submit if session hasn't been loaded yet
if (!currentState.session || currentState.chatState === ChatState.LoadingConversation) {
return;
}
@@ -543,7 +793,6 @@ export function useChatStream({
const hasExistingMessages = currentState.messages.length > 0;
const hasNewMessage = userMessage.trim().length > 0 || images.length > 0;
// Don't submit if there's no message and no conversation to continue
if (!hasNewMessage && !hasExistingMessages) {
return;
}
@@ -554,10 +803,8 @@ export function useChatStream({
if (!hasExistingMessages && hasNewMessage) {
window.dispatchEvent(new CustomEvent(AppEvents.SESSION_CREATED));
// Start polling for session name update during streaming
// The backend generates the name in parallel with the response
const pollForName = async (attempts = 0) => {
if (attempts >= 20) return; // Max 20 attempts (10 seconds)
if (attempts >= 20) return;
try {
const response = await getSession({
@@ -568,7 +815,6 @@ export function useChatStream({
const currentName = currentState.session?.name;
const newName = response.data?.name;
// Check if name has changed from the initial name
if (newName && newName !== currentName) {
dispatch({
type: 'SET_SESSION',
@@ -581,13 +827,12 @@ export function useChatStream({
detail: { sessionId, newName },
})
);
return; // Stop polling once name is updated
return;
}
} catch {
// Silently continue polling
}
// Continue polling if still streaming
const latestState = stateRef.current;
if (
latestState.chatState === ChatState.Streaming ||
@@ -598,7 +843,6 @@ export function useChatStream({
}
};
// Start polling after a short delay to give backend time to start name generation
namePollingRef.current = setTimeout(() => pollForName(0), 1000);
}
@@ -614,38 +858,10 @@ export function useChatStream({
}
dispatch({ type: 'START_STREAMING' });
abortControllerRef.current = new AbortController();
try {
const { stream } = await reply({
body: {
session_id: sessionId,
user_message: newMessage,
},
throwOnError: true,
signal: abortControllerRef.current.signal,
sseMaxRetryAttempts: 0,
});
await streamFromResponse(
stream,
currentMessages,
dispatch,
onFinish,
sessionId,
abortControllerRef.current.signal
);
} catch (error) {
// AbortError is expected when user stops streaming
if (error instanceof Error && error.name === 'AbortError') {
// Silently handle abort
} else {
// Unexpected error during fetch setup (streamFromResponse handles its own errors)
onFinish('Submit error: ' + errorMessage(error));
}
}
await submitToSession(sessionId, newMessage, currentMessages);
},
[sessionId, onFinish]
[sessionId, submitToSession]
);
const submitElicitationResponse = useCallback(
@@ -663,36 +879,10 @@ export function useChatStream({
dispatch({ type: 'SET_MESSAGES', payload: currentMessages });
dispatch({ type: 'START_STREAMING' });
abortControllerRef.current = new AbortController();
try {
const { stream } = await reply({
body: {
session_id: sessionId,
user_message: responseMessage,
},
throwOnError: true,
signal: abortControllerRef.current.signal,
sseMaxRetryAttempts: 0,
});
await streamFromResponse(
stream,
currentMessages,
dispatch,
onFinish,
sessionId,
abortControllerRef.current.signal
);
} catch (error) {
if (error instanceof Error && error.name === 'AbortError') {
// Silently handle abort
} else {
onFinish('Submit error: ' + errorMessage(error));
}
}
await submitToSession(sessionId, responseMessage, currentMessages);
},
[sessionId, onFinish]
[sessionId, submitToSession]
);
const setRecipeUserParams = useCallback(
@@ -709,7 +899,6 @@ export function useChatStream({
},
throwOnError: true,
});
// TODO(Douwe): get this from the server instead of emulating it here
dispatch({
type: 'SET_SESSION',
payload: {
@@ -728,9 +917,6 @@ export function useChatStream({
);
useEffect(() => {
// This should happen on the server when the session is loaded or changed
// use session.id to support changing of sessions rather than depending on the
// stable sessionId.
if (state.session) {
updateFromSession({
body: {
@@ -742,7 +928,34 @@ export function useChatStream({
}, [state.session]);
const stopStreaming = useCallback(() => {
abortControllerRef.current?.abort();
const requestId = activeRequestIdRef.current;
const requestSessionId = activeRequestSessionIdRef.current;
// Abort the in-flight POST so the reply never starts if cancel wins the race
if (activeAbortRef.current) {
activeAbortRef.current.abort();
activeAbortRef.current = null;
}
if (requestId && requestSessionId) {
// Cancel against the session that originally started the request,
// not the current sessionId (which may have changed if user navigated).
sessionCancel({
path: { id: requestSessionId },
body: { request_id: requestId },
}).catch((e) => {
console.warn('Failed to cancel request:', e);
});
}
// Clean up listener
if (activeUnsubscribeRef.current) {
activeUnsubscribeRef.current();
activeUnsubscribeRef.current = null;
}
activeRequestIdRef.current = null;
activeRequestSessionIdRef.current = null;
dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Idle });
lastInteractionTimeRef.current = Date.now();
}, []);
@@ -810,34 +1023,7 @@ export function useChatStream({
dispatch({ type: 'SET_MESSAGES', payload: messagesForUI });
dispatch({ type: 'START_STREAMING' });
abortControllerRef.current = new AbortController();
try {
const { stream } = await reply({
body: {
session_id: targetSessionId,
user_message: updatedUserMessage,
},
throwOnError: true,
signal: abortControllerRef.current.signal,
sseMaxRetryAttempts: 0,
});
await streamFromResponse(
stream,
messagesForUI,
dispatch,
onFinish,
targetSessionId,
abortControllerRef.current.signal
);
} catch (error) {
if (error instanceof Error && error.name === 'AbortError') {
dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Idle });
} else {
throw error;
}
}
await submitToSession(targetSessionId, updatedUserMessage, messagesForUI);
} else {
await handleSubmit({ msg: newContent, images: [] });
}
@@ -853,7 +1039,7 @@ export function useChatStream({
});
}
},
[sessionId, handleSubmit, onFinish]
[sessionId, handleSubmit, submitToSession]
);
const setChatState = useCallback((newState: ChatState) => {
ui/desktop/src/hooks/useSessionEvents.ts
+196
View File
@@ -0,0 +1,196 @@
import { useEffect, useRef, useState, useCallback } from 'react';
import { sessionEvents, type MessageEvent } from '../api';
/**
* An SSE event with an optional request_id (added by the server at the
* SSE framing layer, not part of the generated MessageEvent type).
*/
export type SessionEvent = MessageEvent & {
request_id?: string;
/** Chat-level request UUID used for routing events to the correct handler. */
chat_request_id?: string;
};
type EventHandler = (event: SessionEvent) => void;
type ActiveRequestsHandler = (requestIds: string[]) => void;
export function useSessionEvents(sessionId: string) {
const listenersRef = useRef(new Map<string, Set<EventHandler>>());
const activeRequestsHandlerRef = useRef<ActiveRequestsHandler | null>(null);
const abortRef = useRef<AbortController | null>(null);
const [connected, setConnected] = useState(false);
useEffect(() => {
if (!sessionId) return;
const abortController = new AbortController();
abortRef.current = abortController;
(async () => {
let retryDelay = 500;
const MAX_RETRY_DELAY = 10_000;
const MAX_CONSECUTIVE_ERRORS = 10;
let consecutiveErrors = 0;
let lastEventId: string | undefined;
while (!abortController.signal.aborted) {
try {
const { stream } = await sessionEvents({
path: { id: sessionId },
signal: abortController.signal,
headers: lastEventId ? { 'Last-Event-ID': lastEventId } : undefined,
// Disable the inner retry loop so errors surface to our outer
// loop which tracks consecutive failures and notifies listeners.
sseMaxRetryAttempts: 1,
onSseEvent: (event) => {
if (event.id) {
lastEventId = event.id;
}
},
});
let receivedEvent = false;
for await (const event of stream) {
if (abortController.signal.aborted) break;
// Only mark as connected after the first real event arrives,
// since the HTTP request doesn't happen until iteration starts.
if (!receivedEvent) {
receivedEvent = true;
setConnected(true);
retryDelay = 500;
consecutiveErrors = 0;
}
// The server adds chat_request_id (the chat UUID) and request_id
// to the JSON at the SSE framing layer. Route using chat_request_id
// so that Notification events (which carry their own MCP tool-call
// request_id) still reach the correct handler.
const sessionEvent = event as SessionEvent;
const routingId = sessionEvent.chat_request_id ?? sessionEvent.request_id;
// ActiveRequests events notify the client about in-flight requests
// it can reattach to (e.g. after a remount).
if (sessionEvent.type === 'ActiveRequests') {
const ids = (sessionEvent as unknown as { request_ids: string[] }).request_ids;
activeRequestsHandlerRef.current?.(ids);
continue;
}
// Server-level errors without a request ID (e.g. "client too far
// behind") affect all active listeners — broadcast to everyone.
if (!routingId && sessionEvent.type === 'Error') {
for (const [id, handlers] of listenersRef.current) {
for (const handler of handlers) {
handler({ ...sessionEvent, request_id: id, chat_request_id: id });
}
}
} else if (routingId) {
const handlers = listenersRef.current.get(routingId);
if (handlers) {
for (const handler of handlers) {
handler(sessionEvent);
}
}
}
}
// Stream ended. Reconnect unless we were intentionally aborted.
if (abortController.signal.aborted) break;
setConnected(false);
// If the stream ended without delivering any events, the connection
// likely failed silently (e.g. 404 with sseMaxRetryAttempts: 1).
// Treat it as an error so backoff and error counting apply.
if (!receivedEvent) {
consecutiveErrors++;
console.warn(
`SSE stream ended with no events (${consecutiveErrors}/${MAX_CONSECUTIVE_ERRORS})`
);
if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) {
console.error('SSE reconnect limit reached, notifying active listeners');
const errorEvent: SessionEvent = {
type: 'Error',
error: 'Lost connection to server',
} as SessionEvent;
for (const [routingId, handlers] of listenersRef.current) {
for (const handler of handlers) {
handler({ ...errorEvent, request_id: routingId, chat_request_id: routingId });
}
}
consecutiveErrors = 0;
}
await new Promise((r) => setTimeout(r, retryDelay));
retryDelay = Math.min(retryDelay * 2, MAX_RETRY_DELAY);
}
} catch (error) {
if (abortController.signal.aborted) break;
consecutiveErrors++;
console.warn(
`SSE connection error (${consecutiveErrors}/${MAX_CONSECUTIVE_ERRORS}), reconnecting:`,
error,
);
setConnected(false);
if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) {
console.error('SSE reconnect limit reached, notifying active listeners');
// Send an error event to all active listeners so they can
// transition out of streaming state. Reset the counter so
// the loop keeps reconnecting for future requests.
const errorEvent: SessionEvent = {
type: 'Error',
error: 'Lost connection to server',
} as SessionEvent;
for (const [routingId, handlers] of listenersRef.current) {
for (const handler of handlers) {
handler({ ...errorEvent, request_id: routingId, chat_request_id: routingId });
}
}
consecutiveErrors = 0;
}
// Back off before retrying
await new Promise((r) => setTimeout(r, retryDelay));
retryDelay = Math.min(retryDelay * 2, MAX_RETRY_DELAY);
}
}
setConnected(false);
})();
const listeners = listenersRef.current;
return () => {
abortController.abort();
abortRef.current = null;
listeners.clear();
setConnected(false);
};
}, [sessionId]);
const addListener = useCallback(
(requestId: string, handler: EventHandler): (() => void) => {
if (!listenersRef.current.has(requestId)) {
listenersRef.current.set(requestId, new Set());
}
listenersRef.current.get(requestId)!.add(handler);
return () => {
const set = listenersRef.current.get(requestId);
if (set) {
set.delete(handler);
if (set.size === 0) {
listenersRef.current.delete(requestId);
}
}
};
},
[]
);
const setActiveRequestsHandler = useCallback((handler: ActiveRequestsHandler | null) => {
activeRequestsHandlerRef.current = handler;
}, []);
return { connected, addListener, setActiveRequestsHandler };
}