fix: Fix aborting chat streaming

Improves the chat stream abortion process by ensuring that partial responses are saved before the abort signal is sent.

This avoids a race condition where the onError callback could clear the streaming state before the partial response is saved. Additionally, the stream reading loop and callbacks are now checked for abort signals to prevent further processing after abortion.

Author: allozaur

tools/server/webui/src/lib/services/chat.ts

@@ -194,14 +194,16 @@ export class ChatService {
 			}
 
 			if (stream) {
-				return this.handleStreamResponse(
+				await this.handleStreamResponse(
 					response,
 					onChunk,
 					onComplete,
 					onError,
 					options.onReasoningChunk,
-					conversationId
+					conversationId,
+					abortController.signal
 				);
+				return;
 			} else {
 				return this.handleNonStreamResponse(response, onComplete, onError);
 			}
@@ -264,7 +266,8 @@ export class ChatService {
 		) => void,
 		onError?: (error: Error) => void,
 		onReasoningChunk?: (chunk: string) => void,
-		conversationId?: string
+		conversationId?: string,
+		abortSignal?: AbortSignal
 	): Promise<void> {
 		const reader = response.body?.getReader();
 
@@ -282,14 +285,29 @@ export class ChatService {
 		try {
 			let chunk = '';
 			while (true) {
+				// Check if we've been aborted before reading more data
+				if (abortSignal?.aborted) {
+					break;
+				}
+
 				const { done, value } = await reader.read();
 				if (done) break;
 
+				// Check again after async read
+				if (abortSignal?.aborted) {
+					break;
+				}
+
 				chunk += decoder.decode(value, { stream: true });
 				const lines = chunk.split('\n');
 				chunk = lines.pop() || ''; // Save incomplete line for next read
 
 				for (const line of lines) {
+					// Check abort signal before processing each line
+					if (abortSignal?.aborted) {
+						break;
+					}
+
 					if (line.startsWith('data: ')) {
 						const data = line.slice(6);
 						if (data === '[DONE]') {
@@ -317,19 +335,35 @@ export class ChatService {
 							if (content) {
 								hasReceivedData = true;
 								aggregatedContent += content;
-								onChunk?.(content);
+								// Only call callback if not aborted
+								if (!abortSignal?.aborted) {
+									onChunk?.(content);
+								}
 							}
 
 							if (reasoningContent) {
 								hasReceivedData = true;
 								fullReasoningContent += reasoningContent;
-								onReasoningChunk?.(reasoningContent);
+								// Only call callback if not aborted
+								if (!abortSignal?.aborted) {
+									onReasoningChunk?.(reasoningContent);
+								}
 							}
 						} catch (e) {
 							console.error('Error parsing JSON chunk:', e);
 						}
 					}
 				}
+
+				// Break outer loop if aborted during line processing
+				if (abortSignal?.aborted) {
+					break;
+				}
+			}
+
+			// Don't call onComplete if we've been aborted
+			if (abortSignal?.aborted) {
+				return;
 			}
 
 			if (streamFinished) {

tools/server/webui/src/lib/stores/chat.svelte.ts

@@ -685,16 +685,24 @@ class ChatStore {
 	 * Stops the current message generation
 	 * Aborts ongoing requests and saves partial response if available
 	 */
-	stopGeneration(): void {
+	async stopGeneration(): Promise<void> {
+		if (!this.activeConversation) return;
+
+		const convId = this.activeConversation.id;
+
+		// Save partial response BEFORE aborting to avoid race condition
+		// where onError callback clears the streaming state
+		await this.savePartialResponseIfNeeded(convId);
+
+		// Stop streaming and abort the request
 		slotsService.stopStreaming();
-		chatService.abort();
-		this.savePartialResponseIfNeeded();
+		chatService.abort(convId);
 
-		// Clear all conversation loading and streaming states
-		this.conversationLoadingStates.clear();
-		this.conversationStreamingStates.clear();
-		this.isLoading = false;
-		this.currentResponse = '';
+		// Clear loading and streaming states for the active conversation only
+		// Note: onError callback will also clear these, but that's okay - idempotent operations
+		this.setConversationLoading(convId, false);
+		this.clearConversationStreaming(convId);
+		slotsService.clearConversationState(convId);
 	}
 
 	/**
@@ -718,12 +726,26 @@ class ChatStore {
 	 * Saves partial response if generation was interrupted
 	 * Preserves user's partial content and timing data when generation is stopped early
 	 */
-	private async savePartialResponseIfNeeded(): Promise<void> {
-		if (!this.currentResponse.trim() || !this.activeMessages.length) {
+	private async savePartialResponseIfNeeded(convId?: string): Promise<void> {
+		// Use provided conversation ID or active conversation
+		const conversationId = convId || this.activeConversation?.id;
+		if (!conversationId) return;
+
+		// Get the streaming state for this conversation
+		const streamingState = this.conversationStreamingStates.get(conversationId);
+		if (!streamingState || !streamingState.response.trim()) {
 			return;
 		}
 
-		const lastMessage = this.activeMessages[this.activeMessages.length - 1];
+		// Get messages for this conversation
+		const messages =
+			conversationId === this.activeConversation?.id
+				? this.activeMessages
+				: await DatabaseStore.getConversationMessages(conversationId);
+
+		if (!messages.length) return;
+
+		const lastMessage = messages[messages.length - 1];
 
 		if (lastMessage && lastMessage.role === 'assistant') {
 			try {
@@ -732,7 +754,7 @@ class ChatStore {
 					thinking?: string;
 					timings?: ChatMessageTimings;
 				} = {
-					content: this.currentResponse
+					content: streamingState.response
 				};
 
 				if (lastMessage.thinking?.trim()) {