Merge branch 'refactor/core-messaging'

src/lib/types/MessageContext.ts

@@ -0,0 +1,41 @@
+import type { Message } from "./Message";
+import type { v4 } from "uuid";
+import type { goto, invalidate } from "$app/navigation";
+
+export interface WriteMessageContext {
+	page: { params: { id: string } };
+	messages: Message[];
+	messagesPath: Message[];
+	data: { rootMessageId: string };
+	files: File[];
+	settings: {
+		disableStream: boolean;
+		personas?: Array<{ id: string; name: string }>;
+	};
+	isAborted: () => boolean;
+	branchState: {
+		messageId: string;
+		personaId: string;
+		personaName: string;
+	} | null;
+
+	setLoading: (val: boolean) => void;
+	setPending: (val: boolean) => void;
+	setFiles: (val: File[]) => void;
+	setError: (val: string) => void;
+	setIsAborted: (val: boolean) => void;
+	setTitleUpdate: (val: { title: string; convId: string }) => void;
+	onTitleUpdate?: (title: string) => void;
+	onMessageCreated?: (id: string) => void;
+	updateBranchState: (val: unknown) => void;
+	invalidate: typeof invalidate;
+	goto: typeof goto;
+}
+
+export interface WriteMessageParams {
+	prompt?: string;
+	messageId?: ReturnType<typeof v4>;
+	isRetry?: boolean;
+	isContinue?: boolean;
+	personaId?: string;
+}

src/lib/utils/message/ConversationTreeManager.ts

@@ -0,0 +1,268 @@
+import type { Message, MessageFile } from "$lib/types/Message";
+import { MessageRole } from "$lib/types/Message";
+import { addChildren } from "$lib/utils/tree/addChildren";
+import { addSibling } from "$lib/utils/tree/addSibling";
+import type { WriteMessageContext, WriteMessageParams } from "$lib/types/MessageContext";
+
+export class ConversationTreeManager {
+	constructor(private ctx: WriteMessageContext) {}
+
+	public prepareMessageForWrite(
+		params: WriteMessageParams,
+		base64Files: MessageFile[] = []
+	): {
+		messageToWriteToId: string;
+		navigateToMessageId: string | null;
+	} {
+		const {
+			prompt,
+			messageId = this.ctx.messagesPath.at(-1)?.id ?? undefined,
+			isRetry = false,
+			isContinue = false,
+			personaId,
+		} = params;
+
+		let messageToWriteToId: string | undefined;
+		let navigateToMessageId: string | null = null;
+
+		if (isContinue && messageId) {
+			const msg = this.ctx.messages.find((m) => m.id === messageId);
+			if ((msg?.children?.length ?? 0) > 0) {
+				throw new Error("Can only continue the last message");
+			}
+			messageToWriteToId = messageId;
+		} else if (isRetry && messageId) {
+			const messageToRetry = this.ctx.messages.find((m) => m.id === messageId);
+			if (!messageToRetry) {
+				throw new Error("Message not found");
+			}
+
+			if (messageToRetry.from === MessageRole.User && prompt !== undefined) {
+				const newUserMessageId = addSibling(
+					{
+						messages: this.ctx.messages,
+						rootMessageId: this.ctx.data.rootMessageId,
+					},
+					{
+						from: MessageRole.User,
+						content: prompt,
+						files: messageToRetry.files,
+						...(messageToRetry.branchedFrom && {
+							branchedFrom: messageToRetry.branchedFrom,
+						}),
+					},
+					messageId
+				);
+
+				const targetPersonaId =
+					this.ctx.branchState?.personaId || messageToRetry.branchedFrom?.personaId;
+				const initialResponses: Message["personaResponses"] = [];
+				if (targetPersonaId) {
+					const persona = this.ctx.settings.personas?.find((p) => p.id === targetPersonaId);
+					initialResponses.push({
+						personaId: targetPersonaId,
+						personaName: this.ctx.branchState?.personaName || persona?.name || targetPersonaId,
+						content: "",
+					});
+				}
+
+				messageToWriteToId = addChildren(
+					{
+						messages: this.ctx.messages,
+						rootMessageId: this.ctx.data.rootMessageId,
+					},
+					{
+						from: MessageRole.Assistant,
+						content: "",
+						personaResponses: initialResponses,
+						...((this.ctx.branchState || messageToRetry.branchedFrom) && {
+							branchedFrom: this.ctx.branchState
+								? {
+										messageId: this.ctx.branchState.messageId,
+										personaId: this.ctx.branchState.personaId,
+									}
+								: messageToRetry.branchedFrom,
+						}),
+					},
+					newUserMessageId
+				);
+
+				if (messageToRetry.branchedFrom && !this.ctx.branchState) {
+					const persona = this.ctx.settings.personas?.find(
+						(p) => p.id === messageToRetry.branchedFrom?.personaId
+					);
+					this.ctx.updateBranchState({
+						messageId: messageToRetry.branchedFrom.messageId,
+						personaId: messageToRetry.branchedFrom.personaId,
+						personaName: persona?.name || messageToRetry.branchedFrom.personaId,
+					});
+					navigateToMessageId = newUserMessageId;
+				}
+				this.ctx.onMessageCreated?.(messageToWriteToId);
+			} else if (messageToRetry.from === MessageRole.User && prompt === undefined) {
+				messageToWriteToId = addChildren(
+					{
+						messages: this.ctx.messages,
+						rootMessageId: this.ctx.data.rootMessageId,
+					},
+					{
+						from: MessageRole.Assistant,
+						content: "",
+						personaResponses: [],
+						...(this.ctx.branchState && {
+							branchedFrom: {
+								messageId: this.ctx.branchState.messageId,
+								personaId: this.ctx.branchState.personaId,
+							},
+						}),
+					},
+					messageId
+				);
+				navigateToMessageId = messageToWriteToId;
+				this.ctx.onMessageCreated?.(messageToWriteToId);
+			} else if (messageToRetry.from === MessageRole.Assistant) {
+				let initialPersonaResponses: Message["personaResponses"] = [];
+				if (personaId && messageToRetry.personaResponses) {
+					initialPersonaResponses = messageToRetry.personaResponses.map((p) => {
+						if (p.personaId === personaId) {
+							return {
+								...p,
+								content: "",
+								interrupted: undefined,
+								reasoning: undefined,
+								updates: undefined,
+								routerMetadata: undefined,
+							};
+						}
+						// Defensive copy using structuredClone
+						return structuredClone(p);
+					});
+				}
+
+				messageToWriteToId = addSibling(
+					{
+						messages: this.ctx.messages,
+						rootMessageId: this.ctx.data.rootMessageId,
+					},
+					{
+						from: MessageRole.Assistant,
+						content: "",
+						personaResponses: initialPersonaResponses,
+						...((this.ctx.branchState || messageToRetry.branchedFrom) && {
+							branchedFrom: this.ctx.branchState
+								? {
+										messageId: this.ctx.branchState.messageId,
+										personaId: this.ctx.branchState.personaId,
+									}
+								: messageToRetry.branchedFrom,
+						}),
+					},
+					messageId
+				);
+
+				if (messageToRetry.branchedFrom && !this.ctx.branchState) {
+					const persona = this.ctx.settings.personas?.find(
+						(p) => p.id === messageToRetry.branchedFrom?.personaId
+					);
+					this.ctx.updateBranchState({
+						messageId: messageToRetry.branchedFrom.messageId,
+						personaId: messageToRetry.branchedFrom.personaId,
+						personaName: persona?.name || messageToRetry.branchedFrom.personaId,
+					});
+					navigateToMessageId = messageToWriteToId;
+				}
+				this.ctx.onMessageCreated?.(messageToWriteToId);
+			}
+		} else {
+			// New message
+			const newUserMessageId = addChildren(
+				{
+					messages: this.ctx.messages,
+					rootMessageId: this.ctx.data.rootMessageId,
+				},
+				{
+					from: MessageRole.User,
+					content: prompt ?? "",
+					files: base64Files,
+					...(this.ctx.branchState && {
+						branchedFrom: {
+							messageId: this.ctx.branchState.messageId,
+							personaId: this.ctx.branchState.personaId,
+						},
+					}),
+				},
+				messageId
+			);
+
+			if (!this.ctx.data.rootMessageId) {
+				this.ctx.data.rootMessageId = newUserMessageId;
+			}
+
+			messageToWriteToId = addChildren(
+				{
+					messages: this.ctx.messages,
+					rootMessageId: this.ctx.data.rootMessageId,
+				},
+				{
+					from: MessageRole.Assistant,
+					content: "",
+					personaResponses: [],
+					...(this.ctx.branchState && {
+						branchedFrom: {
+							messageId: this.ctx.branchState.messageId,
+							personaId: this.ctx.branchState.personaId,
+						},
+					}),
+				},
+				newUserMessageId
+			);
+
+			this.ctx.onMessageCreated?.(messageToWriteToId);
+		}
+
+		if (!messageToWriteToId) {
+			throw new Error("Failed to determine message ID to write to");
+		}
+
+		return { messageToWriteToId, navigateToMessageId };
+	}
+
+	/**
+	 * Safely updates a message ID in the tree, ensuring parent linkage is maintained.
+	 */
+	public syncMessageId(oldId: string, newId: string): void {
+		const message = this.ctx.messages.find((m) => m.id === oldId || m.id === newId);
+		if (!message) return;
+
+		// If ID is already updated, just return
+		if (message.id === newId) return;
+
+		message.id = newId;
+
+		if (message.ancestors && message.ancestors.length > 0) {
+			const parentId = message.ancestors[message.ancestors.length - 1];
+			const parent = this.ctx.messages.find((m) => m.id === parentId);
+
+			if (parent) {
+				if (!parent.children) parent.children = [];
+
+				const childIndex = parent.children.indexOf(oldId);
+				if (childIndex !== -1) {
+					parent.children[childIndex] = newId;
+				} else {
+					// Fallback: append if not found
+					console.warn(
+						`[TreeManager] Parent ${parentId} missing child ${oldId}, appending ${newId}`
+					);
+					parent.children.push(newId);
+				}
+			} else {
+				console.error(
+					`[TreeManager] Parent ${parentId} not found for message ${oldId} -> ${newId}`
+				);
+			}
+		}
+
+		console.debug(`[TreeManager] Synced message ID: ${oldId} -> ${newId}`);
+	}
+}

src/lib/utils/message/MessageStreamHandler.ts

@@ -0,0 +1,228 @@
+import {
+	MessageReasoningUpdateType,
+	MessageUpdateStatus,
+	MessageUpdateType,
+} from "$lib/types/MessageUpdate";
+import { fetchMessageUpdates } from "$lib/utils/messageUpdates";
+import { updateDebouncer } from "$lib/utils/updates.js";
+import type { WriteMessageContext } from "$lib/types/MessageContext";
+import type { ConversationTreeManager } from "./ConversationTreeManager";
+
+export class MessageStreamHandler {
+	constructor(
+		private ctx: WriteMessageContext,
+		private treeManager: ConversationTreeManager
+	) {}
+
+	public async handleStream(
+		conversationId: string,
+		params: {
+			base: string;
+			prompt?: string;
+			messageId?: string;
+			isRetry: boolean;
+			isContinue: boolean;
+			// eslint-disable-next-line @typescript-eslint/no-explicit-any
+			files?: any[];
+			personaId?: string;
+		},
+		messageToWriteToId: string
+	) {
+		const messageToWriteTo = this.ctx.messages.find((m) => m.id === messageToWriteToId);
+		if (!messageToWriteTo) {
+			throw new Error("Message to write to not found");
+		}
+
+		const messageUpdatesAbortController = new AbortController();
+
+		const messageUpdatesIterator = await fetchMessageUpdates(
+			conversationId,
+			{
+				base: params.base,
+				inputs: params.prompt,
+				messageId: params.messageId,
+				isRetry: params.isRetry,
+				isContinue: params.isContinue,
+				files: params.files,
+				personaId: params.personaId,
+				branchedFrom: this.ctx.branchState
+					? {
+							messageId: this.ctx.branchState.messageId,
+							personaId: this.ctx.branchState.personaId,
+						}
+					: undefined,
+			},
+			messageUpdatesAbortController.signal
+		).catch((err) => {
+			this.ctx.setError(err.message);
+		});
+
+		if (messageUpdatesIterator === undefined) return;
+
+		this.ctx.setFiles([]);
+		let buffer = "";
+		let lastUpdateTime = new Date();
+
+		let reasoningBuffer = "";
+		let reasoningLastUpdate = new Date();
+
+		const personaBuffers = new Map<string, string>();
+		const personaLastUpdateTimes = new Map<string, Date>();
+
+		for await (const update of messageUpdatesIterator) {
+			if (this.ctx.isAborted()) {
+				messageUpdatesAbortController.abort();
+				return;
+			}
+
+			if (update.type === MessageUpdateType.Stream) {
+				update.token = update.token.replaceAll("\0", "");
+			}
+
+			const isHighFrequencyUpdate =
+				(update.type === MessageUpdateType.Reasoning &&
+					update.subtype === MessageReasoningUpdateType.Stream) ||
+				update.type === MessageUpdateType.Stream ||
+				update.type === MessageUpdateType.Persona ||
+				(update.type === MessageUpdateType.Status &&
+					update.status === MessageUpdateStatus.KeepAlive);
+
+			if (!isHighFrequencyUpdate) {
+				messageToWriteTo.updates = [...(messageToWriteTo.updates ?? []), update];
+			}
+			const currentTime = new Date();
+
+			if (update.type === MessageUpdateType.PersonaInit) {
+				const newResponses = update.personas.map((p) => ({
+					personaId: p.personaId,
+					personaName: p.personaName,
+					personaOccupation: p.personaOccupation,
+					personaStance: p.personaStance,
+					content: "",
+				}));
+
+				if (!messageToWriteTo.personaResponses) {
+					messageToWriteTo.personaResponses = newResponses;
+				} else {
+					// Merge with existing personas (preserving those not in the update)
+					for (const newRes of newResponses) {
+						const existingIdx = messageToWriteTo.personaResponses.findIndex(
+							(p) => p.personaId === newRes.personaId
+						);
+						if (existingIdx !== -1) {
+							// Update existing persona in place
+							messageToWriteTo.personaResponses[existingIdx] = {
+								...messageToWriteTo.personaResponses[existingIdx],
+								...newRes,
+								content: messageToWriteTo.personaResponses[existingIdx].content || newRes.content,
+							};
+						} else {
+							messageToWriteTo.personaResponses.push(newRes);
+						}
+					}
+				}
+			} else if (update.type === MessageUpdateType.Persona) {
+				if (!messageToWriteTo.personaResponses) {
+					messageToWriteTo.personaResponses = [];
+				}
+
+				let personaResponse = messageToWriteTo.personaResponses.find(
+					(pr) => pr.personaId === update.personaId
+				);
+				if (!personaResponse) {
+					personaResponse = {
+						personaId: update.personaId,
+						personaName: update.personaName,
+						personaOccupation: update.personaOccupation,
+						personaStance: update.personaStance,
+						content: "",
+					};
+					messageToWriteTo.personaResponses.push(personaResponse);
+				}
+
+				if (update.updateType === "stream" && update.token && !this.ctx.settings.disableStream) {
+					const personaBuffer = personaBuffers.get(update.personaId) || "";
+					const newBuffer = personaBuffer + update.token;
+					personaBuffers.set(update.personaId, newBuffer);
+
+					const lastUpdate = personaLastUpdateTimes.get(update.personaId) || new Date(0);
+					if (currentTime.getTime() - lastUpdate.getTime() > updateDebouncer.maxUpdateTime) {
+						personaResponse.content += newBuffer;
+						personaBuffers.set(update.personaId, "");
+						personaLastUpdateTimes.set(update.personaId, currentTime);
+					}
+					this.ctx.setPending(false);
+				} else if (update.updateType === "finalAnswer" && update.text) {
+					personaResponse.content = update.text;
+					personaResponse.interrupted = update.interrupted;
+				} else if (update.updateType === "routerMetadata" && update.route && update.model) {
+					personaResponse.routerMetadata = {
+						route: update.route,
+						model: update.model,
+					};
+				} else if (update.updateType === "status" && update.error) {
+					personaResponse.interrupted = true;
+					personaResponse.content = personaResponse.content || `Error: ${update.error}`;
+				}
+			} else if (update.type === MessageUpdateType.Stream && !this.ctx.settings.disableStream) {
+				buffer += update.token;
+				if (currentTime.getTime() - lastUpdateTime.getTime() > updateDebouncer.maxUpdateTime) {
+					messageToWriteTo.content += buffer;
+					buffer = "";
+					lastUpdateTime = currentTime;
+				}
+				this.ctx.setPending(false);
+			} else if (
+				update.type === MessageUpdateType.Status &&
+				update.status === MessageUpdateStatus.Error
+			) {
+				this.ctx.setError(update.message ?? "An error has occurred");
+			} else if (
+				update.type === MessageUpdateType.Status &&
+				update.status === MessageUpdateStatus.Started &&
+				update.messageId
+			) {
+				if (messageToWriteTo.id !== update.messageId) {
+					// Use TreeManager to safely update the ID and parent links
+					const oldId = messageToWriteTo.id;
+					this.treeManager.syncMessageId(oldId, update.messageId);
+				}
+			} else if (update.type === MessageUpdateType.Title) {
+				this.ctx.setTitleUpdate({
+					title: update.title,
+					convId: conversationId,
+				});
+				this.ctx.onTitleUpdate?.(update.title);
+			} else if (update.type === MessageUpdateType.File) {
+				messageToWriteTo.files = [
+					...(messageToWriteTo.files ?? []),
+					{ type: "hash", value: update.sha, mime: update.mime, name: update.name },
+				];
+			} else if (update.type === MessageUpdateType.Reasoning) {
+				if (!messageToWriteTo.reasoning) {
+					messageToWriteTo.reasoning = "";
+				}
+				if (update.subtype === MessageReasoningUpdateType.Stream) {
+					reasoningBuffer += update.token;
+					if (
+						currentTime.getTime() - reasoningLastUpdate.getTime() >
+						updateDebouncer.maxUpdateTime
+					) {
+						messageToWriteTo.reasoning += reasoningBuffer;
+						reasoningBuffer = "";
+						reasoningLastUpdate = currentTime;
+					}
+				}
+			} else if (update.type === MessageUpdateType.RouterMetadata) {
+				messageToWriteTo.routerMetadata = {
+					route: update.route,
+					model: update.model,
+				};
+			} else if (update.type === MessageUpdateType.FinalAnswer) {
+				messageToWriteTo.content = update.text;
+				messageToWriteTo.interrupted = update.interrupted;
+				this.ctx.setPending(false);
+			}
+		}
+	}
+}