Abort (#1924)

* Add AbortRegistry for user-initiated generation cancellation

Introduces an AbortRegistry singleton to track and manage AbortControllers for active conversation generations, enabling user-initiated cancellation of ongoing requests. Updates text generation and endpoint logic to support abort signals, and ensures proper cleanup and partial response handling when a generation is aborted. The stop-generating endpoint now triggers aborts via the registry.

* Remove 'continue message' feature from chat flow

Eliminated the 'continue message' functionality and related code, including the ContinueBtn component, associated props, and logic in both frontend and backend. This simplifies the message handling flow and removes unused parameters and UI elements.

src/lib/buildPrompt.ts

@@ -1,7 +1,7 @@
 import type { EndpointParameters } from "./server/endpoints/endpoints";
 import type { BackendModel } from "./server/models";
 
-type buildPromptOptions = Pick<EndpointParameters, "messages" | "preprompt" | "continueMessage"> & {
+type buildPromptOptions = Pick<EndpointParameters, "messages" | "preprompt"> & {
 	model: BackendModel;
 };
 
@@ -9,7 +9,6 @@ export async function buildPrompt({
 	messages,
 	model,
 	preprompt,
-	continueMessage,
 }: buildPromptOptions): Promise<string> {
 	const filteredMessages = messages;
 
@@ -17,36 +16,18 @@ export async function buildPrompt({
 		filteredMessages[0].content = preprompt;
 	}
 
-	let prompt = model
+	const prompt = model
 		.chatPromptRender({
 			messages: filteredMessages.map((m) => ({
 				...m,
 				role: m.from,
 			})),
 			preprompt,
-			continueMessage,
 		})
 		// Not super precise, but it's truncated in the model's backend anyway
 		.split(" ")
 		.slice(-(model.parameters?.truncate ?? 0))
 		.join(" ");
 
-	if (continueMessage && model.parameters?.stop) {
-		let trimmedPrompt = prompt.trimEnd();
-		let hasRemovedStop = true;
-		while (hasRemovedStop) {
-			hasRemovedStop = false;
-			for (const stopToken of model.parameters.stop) {
-				if (trimmedPrompt.endsWith(stopToken)) {
-					trimmedPrompt = trimmedPrompt.slice(0, -stopToken.length);
-					hasRemovedStop = true;
-					break;
-				}
-			}
-			trimmedPrompt = trimmedPrompt.trimEnd();
-		}
-		prompt = trimmedPrompt;
-	}
-
 	return prompt;
 }

src/lib/components/ContinueBtn.svelte

@@ -1,18 +0,0 @@
-<script lang="ts">
-	import CarbonContinue from "~icons/carbon/continue";
-
-	interface Props {
-		classNames?: string;
-		onClick?: () => void;
-	}
-
-	let { classNames = "", onClick }: Props = $props();
-</script>
-
-<button
-	type="button"
-	onclick={onClick}
-	class="btn flex h-8 rounded-lg border bg-white px-3 py-1 text-gray-500 shadow-sm transition-all hover:bg-gray-100 dark:border-gray-600 dark:bg-gray-700 dark:text-gray-300 dark:hover:bg-gray-600 {classNames}"
->
-	<CarbonContinue class="mr-2 text-xs " /> Continue
-</button>

src/lib/components/chat/ChatWindow.svelte

@@ -14,7 +14,6 @@
 	import RetryBtn from "../RetryBtn.svelte";
 	import file2base64 from "$lib/utils/file2base64";
 	import { base } from "$app/paths";
-	import ContinueBtn from "../ContinueBtn.svelte";
 	import ChatMessage from "./ChatMessage.svelte";
 	import ScrollToBottomBtn from "../ScrollToBottomBtn.svelte";
 	import ScrollToPreviousBtn from "../ScrollToPreviousBtn.svelte";
@@ -49,7 +48,6 @@
 		onmessage?: (content: string) => void;
 		onstop?: () => void;
 		onretry?: (payload: { id: Message["id"]; content?: string }) => void;
-		oncontinue?: (payload: { id: Message["id"] }) => void;
 		onshowAlternateMsg?: (payload: { id: Message["id"] }) => void;
 	}
 
@@ -66,7 +64,6 @@
 		onmessage,
 		onstop,
 		onretry,
-		oncontinue,
 		onshowAlternateMsg,
 	}: Props = $props();
 
@@ -461,18 +458,6 @@
 								}
 							}}
 						/>
-					{:else if messages && lastMessage && lastMessage.interrupted && !isReadOnly}
-						<div class="ml-auto gap-2">
-							<ContinueBtn
-								onClick={() => {
-									if (lastMessage && lastMessage.ancestors) {
-										oncontinue?.({
-											id: lastMessage?.id,
-										});
-									}
-								}}
-							/>
-						</div>
 					{/if}
 				{/if}
 			</div>

src/lib/server/abortRegistry.ts

@@ -0,0 +1,57 @@
+import { logger } from "$lib/server/logger";
+
+/**
+ * Tracks active upstream generation requests so they can be cancelled on demand.
+ * Multiple controllers can be registered per conversation (for threaded/background runs).
+ */
+export class AbortRegistry {
+	private static instance: AbortRegistry;
+
+	private controllers = new Map<string, Set<AbortController>>();
+
+	public static getInstance(): AbortRegistry {
+		if (!AbortRegistry.instance) {
+			AbortRegistry.instance = new AbortRegistry();
+		}
+		return AbortRegistry.instance;
+	}
+
+	public register(conversationId: string, controller: AbortController) {
+		const key = conversationId.toString();
+		let set = this.controllers.get(key);
+		if (!set) {
+			set = new Set();
+			this.controllers.set(key, set);
+		}
+		set.add(controller);
+		controller.signal.addEventListener(
+			"abort",
+			() => {
+				this.unregister(key, controller);
+			},
+			{ once: true }
+		);
+	}
+
+	public abort(conversationId: string) {
+		const set = this.controllers.get(conversationId);
+		if (!set?.size) return;
+
+		logger.debug({ conversationId }, "Aborting active generation via AbortRegistry");
+		for (const controller of set) {
+			if (!controller.signal.aborted) {
+				controller.abort();
+			}
+		}
+		this.controllers.delete(conversationId);
+	}
+
+	public unregister(conversationId: string, controller: AbortController) {
+		const set = this.controllers.get(conversationId);
+		if (!set) return;
+		set.delete(controller);
+		if (set.size === 0) {
+			this.controllers.delete(conversationId);
+		}
+	}
+}

src/lib/server/endpoints/endpoints.ts

@@ -12,11 +12,11 @@ export type EndpointMessage = Omit<Message, "id">;
 export interface EndpointParameters {
 	messages: EndpointMessage[];
 	preprompt?: Conversation["preprompt"];
-	continueMessage?: boolean; // used to signal that the last message will be extended
 	generateSettings?: Partial<Model["parameters"]>;
 	isMultimodal?: boolean;
 	conversationId?: ObjectId;
 	locals: App.Locals | undefined;
+	abortSignal?: AbortSignal;
 }
 
 export type TextGenerationStreamOutputSimplified = TextGenerationStreamOutput & {

src/lib/server/endpoints/openai/endpointOai.ts

@@ -109,14 +109,13 @@ export async function endpointOai(
 		return async ({
 			messages,
 			preprompt,
-			continueMessage,
 			generateSettings,
 			conversationId,
 			locals,
+			abortSignal,
 		}) => {
 			const prompt = await buildPrompt({
 				messages,
-				continueMessage,
 				preprompt,
 				model,
 			});
@@ -141,6 +140,7 @@ export async function endpointOai(
 					"X-use-cache": "false",
 					...(locals?.token ? { Authorization: `Bearer ${locals.token}` } : {}),
 				},
+				signal: abortSignal,
 			});
 
 			return openAICompletionToTextGenerationStream(openAICompletion);
@@ -153,6 +153,7 @@ export async function endpointOai(
 			conversationId,
 			isMultimodal,
 			locals,
+			abortSignal,
 		}) => {
 			// Format messages for the chat API, handling multimodal content if supported
 			let messagesOpenAI: OpenAI.Chat.Completions.ChatCompletionMessageParam[] =
@@ -203,6 +204,7 @@ export async function endpointOai(
 							"X-use-cache": "false",
 							...(locals?.token ? { Authorization: `Bearer ${locals.token}` } : {}),
 						},
+						signal: abortSignal,
 					}
 				);
 				return openAIChatToTextGenerationStream(openChatAICompletion, () => routerMetadata);
@@ -216,6 +218,7 @@ export async function endpointOai(
 							"X-use-cache": "false",
 							...(locals?.token ? { Authorization: `Bearer ${locals.token}` } : {}),
 						},
+						signal: abortSignal,
 					}
 				);
 				return openAIChatToTextGenerationSingle(openChatAICompletion, () => routerMetadata);

src/lib/server/textGeneration/generate.ts

@@ -20,10 +20,10 @@ export async function* generate(
 		conv,
 		messages,
 		assistant,
-		isContinue,
 		promptedAt,
 		forceMultimodal,
 		locals,
+		abortController,
 	}: GenerateContext,
 	preprompt?: string
 ): AsyncIterable<MessageUpdate> {
@@ -50,16 +50,18 @@ export async function* generate(
 		};
 	}
 
-	for await (const output of await endpoint({
+	const stream = await endpoint({
 		messages,
 		preprompt,
-		continueMessage: isContinue,
 		generateSettings: assistant?.generateSettings,
 		// Allow user-level override to force multimodal
 		isMultimodal: (forceMultimodal ?? false) || model.multimodal,
 		conversationId: conv._id,
 		locals,
-	})) {
+		abortSignal: abortController.signal,
+	});
+
+	for await (const output of stream) {
 		// Check if this output contains router metadata
 		if (
 			"routerMetadata" in output &&
@@ -248,6 +250,9 @@ Do not use prefixes such as Response: or Answer: when answering to the user.`,
 
 		if (date && date > promptedAt) {
 			logger.info(`Aborting generation for conversation ${conv._id}`);
+			if (!abortController.signal.aborted) {
+				abortController.abort();
+			}
 			break;
 		}
 

src/lib/server/textGeneration/types.ts

@@ -10,11 +10,11 @@ export interface TextGenerationContext {
 	conv: Conversation;
 	messages: Message[];
 	assistant?: Pick<Assistant, "dynamicPrompt" | "generateSettings">;
-	isContinue: boolean;
 	promptedAt: Date;
 	ip: string;
 	username?: string;
 	/** Force-enable multimodal handling for endpoints that support it */
 	forceMultimodal?: boolean;
 	locals: App.Locals | undefined;
+	abortController: AbortController;
 }

src/lib/types/Template.ts

@@ -3,5 +3,4 @@ import type { Message } from "./Message";
 export type ChatTemplateInput = {
 	messages: Pick<Message, "from" | "content" | "files">[];
 	preprompt?: string;
-	continueMessage?: boolean;
 };

src/lib/utils/messageUpdates.ts

@@ -12,7 +12,6 @@ type MessageUpdateRequestOptions = {
 	inputs?: string;
 	messageId?: string;
 	isRetry: boolean;
-	isContinue: boolean;
 	files?: MessageFile[];
 };
 export async function fetchMessageUpdates(
@@ -29,7 +28,6 @@ export async function fetchMessageUpdates(
 		inputs: opts.inputs,
 		id: opts.messageId,
 		is_retry: opts.isRetry,
-		is_continue: opts.isContinue,
 	});
 
 	opts.files?.forEach((file) => {

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

@@ -131,12 +131,10 @@
 		prompt,
 		messageId = messagesPath.at(-1)?.id ?? undefined,
 		isRetry = false,
-		isContinue = false,
 	}: {
 		prompt?: string;
 		messageId?: ReturnType<typeof v4>;
 		isRetry?: boolean;
-		isContinue?: boolean;
 	}): Promise<void> {
 		try {
 			$isAborted = false;
@@ -156,13 +154,7 @@
 			let messageToWriteToId: Message["id"] | undefined = undefined;
 			// used for building the prompt, subtree of the conversation that goes from the latest message to the root
 
-			if (isContinue && messageId) {
-				if ((messages.find((msg) => msg.id === messageId)?.children?.length ?? 0) > 0) {
-					$error = "Can only continue the last message";
-				} else {
-					messageToWriteToId = messageId;
-				}
-			} else if (isRetry && messageId) {
+			if (isRetry && messageId) {
 				// two cases, if we're retrying a user message with a newPrompt set,
 				// it means we're editing a user message
 				// if we're retrying on an assistant message, newPrompt cannot be set
@@ -257,7 +249,6 @@
 					inputs: prompt,
 					messageId,
 					isRetry,
-					isContinue,
 					files: isRetry ? userMessage?.files : base64Files,
 				},
 				messageUpdatesAbortController.signal
@@ -422,25 +413,6 @@
 		messagesPath = createMessagesPath(messages, msgId);
 	}
 
-	async function onContinue(payload: { id: Message["id"] }) {
-		if (!data.shared) {
-			await writeMessage({ messageId: payload.id, isContinue: true });
-		} else {
-			await convFromShared()
-				.then(async (convId) => {
-					await goto(`${base}/conversation/${convId}`, { invalidateAll: true });
-				})
-				.then(
-					async () =>
-						await writeMessage({
-							messageId: payload.id,
-							isContinue: true,
-						})
-				)
-				.finally(() => (loading = false));
-		}
-	}
-
 	const settings = useSettingsStore();
 	let messages = $state(data.messages);
 	$effect(() => {
@@ -527,7 +499,6 @@
 	bind:files
 	onmessage={onMessage}
 	onretry={onRetry}
-	oncontinue={onContinue}
 	onshowAlternateMsg={onShowAlternateMsg}
 	onstop={async () => {
 		await fetch(`${base}/conversation/${page.params.id}/stop-generating`, {

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

@@ -23,6 +23,7 @@ import { usageLimits } from "$lib/server/usageLimits";
 import { textGeneration } from "$lib/server/textGeneration";
 import type { TextGenerationContext } from "$lib/server/textGeneration/types";
 import { logger } from "$lib/server/logger.js";
+import { AbortRegistry } from "$lib/server/abortRegistry";
 
 export async function POST({ request, locals, params, getClientAddress }) {
 	const id = z.string().parse(params.id);
@@ -148,7 +149,6 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		inputs: newPrompt,
 		id: messageId,
 		is_retry: isRetry,
-		is_continue: isContinue,
 	} = z
 		.object({
 			id: z.string().uuid().refine(isMessageId).optional(), // parent message id to append to for a normal message, or the message id for a retry/continue
@@ -159,7 +159,6 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					.transform((s) => s.replace(/\r\n/g, "\n"))
 			),
 			is_retry: z.optional(z.boolean()),
-			is_continue: z.optional(z.boolean()),
 			files: z.optional(
 				z.array(
 					z.object({
@@ -220,15 +219,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 	// used for building the prompt, subtree of the conversation that goes from the latest message to the root
 	let messagesForPrompt: Message[] = [];
 
-	if (isContinue && messageId) {
-		// if it's the last message and we continue then we build the prompt up to the last message
-		// we will strip the end tokens afterwards when the prompt is built
-		if ((conv.messages.find((msg) => msg.id === messageId)?.children?.length ?? 0) > 0) {
-			error(400, "Can only continue the last message");
-		}
-		messageToWriteToId = messageId;
-		messagesForPrompt = buildSubtree(conv, messageId);
-	} else if (isRetry && messageId) {
+	if (isRetry && messageId) {
 		// two cases, if we're retrying a user message with a newPrompt set,
 		// it means we're editing a user message
 		// if we're retrying on an assistant message, newPrompt cannot be set
@@ -331,9 +322,18 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		);
 	};
 
+	const abortRegistry = AbortRegistry.getInstance();
+
 	// we now build the stream
 	const stream = new ReadableStream({
 		async start(controller) {
+			const conversationKey = convId.toString();
+			const ctrl = new AbortController();
+			abortRegistry.register(conversationKey, ctrl);
+
+			let finalAnswerReceived = false;
+			let abortedByUser = false;
+
 			messageToWriteTo.updates ??= [];
 			async function update(event: MessageUpdate) {
 				if (!messageToWriteTo || !conv) {
@@ -372,6 +372,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 				else if (event.type === MessageUpdateType.FinalAnswer) {
 					messageToWriteTo.interrupted = event.interrupted;
 					messageToWriteTo.content = initialMessageContent + event.text;
+					finalAnswerReceived = true;
 				}
 
 				// Add file
@@ -449,7 +450,6 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					conv,
 					messages: messagesForPrompt,
 					assistant: undefined,
-					isContinue: isContinue ?? false,
 					promptedAt,
 					ip: getClientAddress(),
 					username: locals.user?.username,
@@ -460,20 +460,50 @@ export async function POST({ request, locals, params, getClientAddress }) {
 						]
 					),
 					locals,
+					abortController: ctrl,
 				};
 				// run the text generation and send updates to the client
 				for await (const event of textGeneration(ctx)) await update(event);
+				if (ctrl.signal.aborted) {
+					abortedByUser = true;
+				}
+				if (abortedByUser && !finalAnswerReceived) {
+					const partialText = messageToWriteTo.content.slice(initialMessageContent.length);
+					await update({
+						type: MessageUpdateType.FinalAnswer,
+						text: partialText,
+						interrupted: true,
+					});
+				}
 			} catch (e) {
-				hasError = true;
-				await update({
-					type: MessageUpdateType.Status,
-					status: MessageUpdateStatus.Error,
-					message: (e as Error).message,
-				});
-				logger.error(e);
+				const err = e as Error;
+				const isAbortError =
+					err?.name === "AbortError" ||
+					err?.name === "APIUserAbortError" ||
+					err?.message === "Request was aborted.";
+				if (isAbortError || ctrl.signal.aborted) {
+					abortedByUser = true;
+					logger.info({ conversationId: conversationKey }, "Generation aborted by user");
+					if (!finalAnswerReceived) {
+						const partialText = messageToWriteTo.content.slice(initialMessageContent.length);
+						await update({
+							type: MessageUpdateType.FinalAnswer,
+							text: partialText,
+							interrupted: true,
+						});
+					}
+				} else {
+					hasError = true;
+					await update({
+						type: MessageUpdateType.Status,
+						status: MessageUpdateStatus.Error,
+						message: err.message,
+					});
+					logger.error(err);
+				}
 			} finally {
 				// check if no output was generated
-				if (!hasError && messageToWriteTo.content === initialMessageContent) {
+				if (!hasError && !abortedByUser && messageToWriteTo.content === initialMessageContent) {
 					await update({
 						type: MessageUpdateType.Status,
 						status: MessageUpdateStatus.Error,
@@ -483,6 +513,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 			}
 
 			await persistConversation();
+			abortRegistry.unregister(conversationKey, ctrl);
 
 			// used to detect if cancel() is called bc of interrupt or just because the connection closes
 			doneStreaming = true;

src/routes/conversation/[id]/stop-generating/+server.ts

@@ -1,5 +1,6 @@
 import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
+import { AbortRegistry } from "$lib/server/abortRegistry";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 
@@ -18,6 +19,8 @@ export async function POST({ params, locals }) {
 		error(404, "Conversation not found");
 	}
 
+	AbortRegistry.getInstance().abort(conversationId.toString());
+
 	await collections.abortedGenerations.updateOne(
 		{ conversationId },
 		{ $set: { updatedAt: new Date() }, $setOnInsert: { createdAt: new Date() } },