Add background generation polling for conversation updates (#1882)

* Add background generation polling for conversations

Introduces a BackgroundGenerationPoller Svelte component and store to track and poll background conversation generations when users navigate away during loading. Updates conversation streaming server logic to persist messages if the client detaches, ensuring conversation state is saved. Integrates poller into layout and conversation page, improving reliability of background message generation and UI updates.

* Improve background generation polling and loading state

* Refactor backgroundGenerations store to use $state

* Update +server.ts

* Harden background generation polling

* Format updated polling flows

src/lib/components/BackgroundGenerationPoller.svelte

@@ -0,0 +1,154 @@
+<script lang="ts">
+	import { browser, dev } from "$app/environment";
+	import { invalidate } from "$app/navigation";
+
+	import {
+		type BackgroundGeneration,
+		backgroundGenerationEntries,
+		removeBackgroundGeneration,
+	} from "$lib/stores/backgroundGenerations";
+	import { handleResponse, useAPIClient } from "$lib/APIClient";
+	import { UrlDependency } from "$lib/types/UrlDependency";
+	import { MessageUpdateStatus, MessageUpdateType } from "$lib/types/MessageUpdate";
+	import type { Message } from "$lib/types/Message";
+
+	const POLL_INTERVAL_MS = 1000;
+
+	const client = useAPIClient();
+	const pollers = new Map<string, () => void>();
+	const inflight = new Set<string>();
+	const assistantSnapshots = new Map<string, string>();
+	const failureCounts = new Map<string, number>();
+
+	$effect.root(() => {
+		if (!browser) {
+			pollers.clear();
+			return;
+		}
+
+		let destroyed = false;
+
+		const stopPoller = (id: string) => {
+			const stop = pollers.get(id);
+			if (!stop) return;
+
+			stop();
+			pollers.delete(id);
+			inflight.delete(id);
+			assistantSnapshots.delete(id);
+			failureCounts.delete(id);
+		};
+
+		const pollOnce = async (id: string) => {
+			if (destroyed || inflight.has(id)) return;
+			inflight.add(id);
+			if (dev) {
+				console.log("background generation poll", id);
+			}
+
+			try {
+				const response = await client.conversations({ id }).get();
+				const conversation = handleResponse(response);
+				const messages = conversation?.messages ?? [];
+				const lastAssistant = [...messages]
+					.reverse()
+					.find((message: Message) => message.from === "assistant");
+
+				const hasFinalAnswer =
+					lastAssistant?.updates?.some((update) => update.type === MessageUpdateType.FinalAnswer) ??
+					false;
+				const hasError =
+					lastAssistant?.updates?.some(
+						(update) =>
+							update.type === MessageUpdateType.Status &&
+							update.status === MessageUpdateStatus.Error
+					) ?? false;
+
+				const snapshot = lastAssistant
+					? JSON.stringify({
+							id: lastAssistant.id,
+							updatedAt: lastAssistant.updatedAt,
+							contentLength: lastAssistant.content?.length ?? 0,
+							updatesLength: lastAssistant.updates?.length ?? 0,
+						})
+					: "__none__";
+				const previousSnapshot = assistantSnapshots.get(id);
+				let shouldInvalidateConversation = false;
+
+				if (lastAssistant) {
+					assistantSnapshots.set(id, snapshot);
+					if (snapshot !== previousSnapshot) {
+						shouldInvalidateConversation = true;
+					}
+				} else if (assistantSnapshots.has(id)) {
+					assistantSnapshots.delete(id);
+					shouldInvalidateConversation = true;
+				}
+
+				if (lastAssistant && (hasFinalAnswer || hasError)) {
+					removeBackgroundGeneration(id);
+					assistantSnapshots.delete(id);
+					failureCounts.delete(id);
+					shouldInvalidateConversation = true;
+					await invalidate(UrlDependency.ConversationList);
+				}
+
+				if (shouldInvalidateConversation) {
+					await invalidate(UrlDependency.Conversation);
+				}
+
+				failureCounts.delete(id);
+			} catch (err) {
+				console.error("Background generation poll failed", err);
+				const failures = (failureCounts.get(id) ?? 0) + 1;
+				failureCounts.set(id, failures);
+				if (failures >= 3) {
+					removeBackgroundGeneration(id);
+					assistantSnapshots.delete(id);
+					failureCounts.delete(id);
+					await invalidate(UrlDependency.ConversationList);
+				}
+			} finally {
+				inflight.delete(id);
+			}
+		};
+
+		const startPoller = (entry: BackgroundGeneration) => {
+			if (pollers.has(entry.id)) return;
+
+			const intervalId = setInterval(() => {
+				void pollOnce(entry.id);
+			}, POLL_INTERVAL_MS);
+
+			pollers.set(entry.id, () => clearInterval(intervalId));
+			void pollOnce(entry.id);
+		};
+
+		$effect(() => {
+			const entries = backgroundGenerationEntries;
+
+			if (destroyed) return;
+
+			const activeIds = new Set(entries.map((entry) => entry.id));
+
+			for (const id of pollers.keys()) {
+				if (!activeIds.has(id)) {
+					stopPoller(id);
+				}
+			}
+
+			for (const entry of entries) {
+				startPoller(entry);
+			}
+		});
+
+		return () => {
+			destroyed = true;
+			for (const stop of pollers.values()) stop();
+			pollers.clear();
+			inflight.clear();
+			assistantSnapshots.clear();
+			failureCounts.clear();
+		};
+	});
+</script>

src/lib/stores/backgroundGenerations.svelte.ts

@@ -0,0 +1,32 @@
+export type BackgroundGeneration = {
+	id: string;
+	startedAt: number;
+};
+
+export const backgroundGenerationEntries = $state<BackgroundGeneration[]>([]);
+
+export function addBackgroundGeneration(entry: BackgroundGeneration) {
+	const index = backgroundGenerationEntries.findIndex(({ id }) => id === entry.id);
+
+	if (index === -1) {
+		backgroundGenerationEntries.push(entry);
+		return;
+	}
+
+	backgroundGenerationEntries[index] = entry;
+}
+
+export function removeBackgroundGeneration(id: string) {
+	const index = backgroundGenerationEntries.findIndex((entry) => entry.id === id);
+	if (index === -1) return;
+
+	backgroundGenerationEntries.splice(index, 1);
+}
+
+export function clearBackgroundGenerations() {
+	backgroundGenerationEntries.length = 0;
+}
+
+export function hasBackgroundGeneration(id: string) {
+	return backgroundGenerationEntries.some((entry) => entry.id === id);
+}

src/lib/stores/backgroundGenerations.ts

@@ -0,0 +1 @@
+export * from "./backgroundGenerations.svelte";

src/routes/+layout.svelte

@@ -23,6 +23,7 @@
 	import { isAborted } from "$lib/stores/isAborted";
 	import IconShare from "$lib/components/icons/IconShare.svelte";
 	import { shareModal } from "$lib/stores/shareModal";
+	import BackgroundGenerationPoller from "$lib/components/BackgroundGenerationPoller.svelte";
 
 	let { data = $bindable(), children } = $props();
 
@@ -237,6 +238,8 @@
 	<OverloadedModal onClose={() => (overloadedModalOpen = false)} />
 {/if}
 
+<BackgroundGenerationPoller />
+
 <div
 	class="fixed grid h-full w-screen grid-cols-1 grid-rows-[auto,1fr] overflow-hidden text-smd {!isNavCollapsed
 		? 'md:grid-cols-[290px,1fr]'

src/routes/conversation/[id]/+page.svelte

@@ -22,7 +22,11 @@
 	import type { v4 } from "uuid";
 	import { useSettingsStore } from "$lib/stores/settings.js";
 	import { browser } from "$app/environment";
-
+	import {
+		addBackgroundGeneration,
+		hasBackgroundGeneration,
+		removeBackgroundGeneration,
+	} from "$lib/stores/backgroundGenerations";
 	import type { TreeNode, TreeId } from "$lib/utils/tree/tree";
 	import "katex/dist/katex.min.css";
 	import { updateDebouncer } from "$lib/utils/updates.js";
@@ -443,6 +447,33 @@
 		messages = data.messages;
 	});
 
+	function isConversationStreaming(msgs: Message[]): boolean {
+		const lastAssistant = [...msgs].reverse().find((msg) => msg.from === "assistant");
+		if (!lastAssistant) return false;
+		const hasFinalAnswer =
+			lastAssistant.updates?.some((update) => update.type === MessageUpdateType.FinalAnswer) ??
+			false;
+		const hasError =
+			lastAssistant.updates?.some(
+				(update) =>
+					update.type === MessageUpdateType.Status && update.status === MessageUpdateStatus.Error
+			) ?? false;
+		return !hasFinalAnswer && !hasError;
+	}
+
+	$effect(() => {
+		const streaming = isConversationStreaming(messages);
+		if (streaming) {
+			loading = true;
+		} else if (!pending) {
+			loading = false;
+		}
+
+		if (!streaming && browser) {
+			removeBackgroundGeneration(page.params.id);
+		}
+	});
+
 	// create a linear list of `messagesPath` from `messages` that is a tree of threaded messages
 	let messagesPath = $derived(createMessagesPath(messages));
 	let messagesAlternatives = $derived(createMessagesAlternatives(messages));
@@ -453,10 +484,26 @@
 		}
 	});
 
-	beforeNavigate(() => {
-		if (page.params.id) {
-			$isAborted = true;
-			loading = false;
+	beforeNavigate((navigation) => {
+		if (!page.params.id) return;
+
+		const navigatingAway =
+			navigation.to?.route.id !== page.route.id || navigation.to?.params?.id !== page.params.id;
+
+		if (loading && navigatingAway) {
+			addBackgroundGeneration({ id: page.params.id, startedAt: Date.now() });
+		}
+
+		$isAborted = true;
+		loading = false;
+	});
+
+	onMount(() => {
+		const hasBackgroundEntry = hasBackgroundGeneration(page.params.id);
+		const streaming = isConversationStreaming(messages);
+		if (hasBackgroundEntry && streaming) {
+			addBackgroundGeneration({ id: page.params.id, startedAt: Date.now() });
+			loading = true;
 		}
 	});
 

src/routes/conversation/[id]/+server.ts

@@ -320,9 +320,17 @@ export async function POST({ request, locals, params, getClientAddress }) {
 	);
 
 	let doneStreaming = false;
+	let clientDetached = false;
 
 	let lastTokenTimestamp: undefined | Date = undefined;
 
+	const persistConversation = async () => {
+		await collections.conversations.updateOne(
+			{ _id: convId },
+			{ $set: { messages: conv.messages, title: conv.title, updatedAt: new Date() } }
+		);
+	};
+
 	// we now build the stream
 	const stream = new ReadableStream({
 		async start(controller) {
@@ -406,21 +414,31 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					event = { ...event, token: event.token.padEnd(16, "\0") };
 				}
 
-				// Send the update to the client
-				controller.enqueue(JSON.stringify(event) + "\n");
+				messageToWriteTo.updatedAt = new Date();
+
+				const enqueueUpdate = async () => {
+					if (clientDetached) return;
+					try {
+						controller.enqueue(JSON.stringify(event) + "\n");
+						if (event.type === MessageUpdateType.FinalAnswer) {
+							controller.enqueue(" ".repeat(4096));
+						}
+					} catch (err) {
+						clientDetached = true;
+						logger.info(
+							{ conversationId: convId.toString() },
+							"Client detached during message streaming"
+						);
+					}
+				};
+
+				await enqueueUpdate();
 
-				// Send 4096 of spaces to make sure the browser doesn't blocking buffer that holding the response
-				if (event.type === MessageUpdateType.FinalAnswer) {
-					controller.enqueue(" ".repeat(4096));
+				if (clientDetached) {
+					await persistConversation();
 				}
 			}
 
-			await collections.conversations.updateOne(
-				{ _id: convId },
-				{ $set: { title: conv.title, updatedAt: new Date() } }
-			);
-			messageToWriteTo.updatedAt = new Date();
-
 			let hasError = false;
 			const initialMessageContent = messageToWriteTo.content;
 
@@ -463,22 +481,18 @@ export async function POST({ request, locals, params, getClientAddress }) {
 				}
 			}
 
-			await collections.conversations.updateOne(
-				{ _id: convId },
-				{ $set: { messages: conv.messages, title: conv?.title, updatedAt: new Date() } }
-			);
+			await persistConversation();
 
 			// used to detect if cancel() is called bc of interrupt or just because the connection closes
 			doneStreaming = true;
-
-			controller.close();
+			if (!clientDetached) {
+				controller.close();
+			}
 		},
 		async cancel() {
 			if (doneStreaming) return;
-			await collections.conversations.updateOne(
-				{ _id: convId },
-				{ $set: { messages: conv.messages, title: conv.title, updatedAt: new Date() } }
-			);
+			clientDetached = true;
+			await persistConversation();
 		},
 	});