fix(stream): add keep alive stream event

src/lib/server/textGeneration/index.ts

@@ -21,15 +21,31 @@ import type { TextGenerationContext } from "./types";
 import type { ToolResult } from "$lib/types/Tool";
 import { toolHasName } from "../tools/utils";
 
+async function* keepAlive(done: AbortSignal): AsyncGenerator<MessageUpdate, undefined, undefined> {
+	while (!done.aborted) {
+		yield {
+			type: MessageUpdateType.Status,
+			status: MessageUpdateStatus.KeepAlive,
+		};
+		await new Promise((resolve) => setTimeout(resolve, 5000));
+	}
+}
+
 export async function* textGeneration(ctx: TextGenerationContext) {
-	yield* mergeAsyncGenerators([
-		textGenerationWithoutTitle(ctx),
-		generateTitleForConversation(ctx.conv),
-	]);
+	const done = new AbortController();
+
+	const titleGen = generateTitleForConversation(ctx.conv);
+	const textGen = textGenerationWithoutTitle(ctx, done);
+	const keepAliveGen = keepAlive(done.signal);
+
+	// keep alive until textGen is done
+
+	yield* mergeAsyncGenerators([titleGen, textGen, keepAliveGen]);
 }
 
 async function* textGenerationWithoutTitle(
-	ctx: TextGenerationContext
+	ctx: TextGenerationContext,
+	done: AbortController
 ): AsyncGenerator<MessageUpdate, undefined, undefined> {
 	yield {
 		type: MessageUpdateType.Status,
@@ -69,4 +85,5 @@ async function* textGenerationWithoutTitle(
 
 	const processedMessages = await preprocessMessages(messages, webSearchResult, convId);
 	yield* generate({ ...ctx, messages: processedMessages }, toolResults, preprompt);
+	done.abort();
 }

src/lib/types/MessageUpdate.ts

@@ -25,6 +25,7 @@ export enum MessageUpdateStatus {
 	Started = "started",
 	Error = "error",
 	Finished = "finished",
+	KeepAlive = "keepAlive",
 }
 export interface MessageStatusUpdate {
 	type: MessageUpdateType.Status;

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

@@ -220,11 +220,6 @@
 					messageUpdatesAbortController.abort();
 					return;
 				}
-				if (update.type === "finalAnswer") {
-					loading = false;
-					pending = false;
-					break;
-				}
 
 				// Remove null characters added due to remote keylogging prevention
 				// See server code for more details

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

@@ -409,7 +409,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 				controller.enqueue(JSON.stringify(event) + "\n");
 
 				// Send 4096 of spaces to make sure the browser doesn't blocking buffer that holding the response
-				if (event.type === "finalAnswer") {
+				if (event.type === MessageUpdateType.FinalAnswer) {
 					controller.enqueue(" ".repeat(4096));
 				}
 			}