Merge pull request #1388 from AnishSarkar22/fix/zero-cache-stale-replica-1355

fix: zero cache stale replica & improved mentioned document chip handling

docker/docker-compose.deps-only.yml

@@ -83,7 +83,7 @@ services:
       retries: 5
 
   zero-cache:
-    image: rocicorp/zero:0.26.2
+    image: rocicorp/zero:1.4.0
     ports:
       - "${ZERO_CACHE_PORT:-4848}:4848"
     extra_hosts:

docker/docker-compose.dev.yml

@@ -179,7 +179,7 @@ services:
   #     - celery_worker
 
   zero-cache:
-    image: rocicorp/zero:0.26.2
+    image: rocicorp/zero:1.4.0
     ports:
       - "${ZERO_CACHE_PORT:-4848}:4848"
     extra_hosts:

docker/docker-compose.yml

@@ -163,7 +163,7 @@ services:
   #   restart: unless-stopped
 
   zero-cache:
-    image: rocicorp/zero:0.26.2
+    image: rocicorp/zero:1.4.0
     ports:
       - "${ZERO_CACHE_PORT:-5929}:4848"
     extra_hosts:

surfsense_backend/alembic/versions/116_create_zero_publication.py

@@ -5,6 +5,17 @@ queries via Zero, instead of replicating all tables in public schema.
 
 See: https://zero.rocicorp.dev/docs/zero-cache-config#app-publications
 
+NOTE for future migration authors: this is the ONLY migration allowed
+to use bare ``CREATE PUBLICATION``. All subsequent mutations of
+``zero_publication`` MUST use the ``COMMENT ON PUBLICATION`` bookend
+pattern wrapping an ``ALTER PUBLICATION ... SET TABLE`` -- copy the
+``upgrade()`` function from migration
+``143_force_zero_publication_resync.py`` as your starting template.
+Raw ``DROP``/``CREATE PUBLICATION`` in new migrations would
+re-introduce bug #1355 (zero-cache stuck on a stale replica snapshot
+because Zero >= 1.0's change-streamer never sees the schema-change
+event).
+
 Revision ID: 116
 Revises: 115
 """

surfsense_backend/alembic/versions/117_optimize_zero_publication_column_lists.py

@@ -17,6 +17,16 @@ IMPORTANT — before AND after running this migration:
   3. Delete / reset the zero-cache data volume
   4. Restart zero-cache  (it will do a fresh initial sync)
 
+DO NOT COPY THIS PATTERN. The ``DROP PUBLICATION`` + ``CREATE
+PUBLICATION`` dance below is the pre-#1355 anti-pattern: on Zero >=
+1.0 it does not reliably wake the zero-cache change-streamer and can
+leave the replica pinned to a stale snapshot. This file is
+grandfathered in because it has already shipped to users; new
+publication mutations MUST use the ``COMMENT ON PUBLICATION`` bookend
+pattern wrapping an ``ALTER PUBLICATION ... SET TABLE`` -- copy the
+``upgrade()`` function from migration
+``143_force_zero_publication_resync.py`` as your starting template.
+
 Revision ID: 117
 Revises: 116
 """

surfsense_backend/alembic/versions/118_add_local_folder_sync_and_versioning.py

@@ -1,5 +1,16 @@
 """Add LOCAL_FOLDER_FILE document type, folder metadata, and document_versions table
 
+DO NOT COPY THIS PATTERN. The bare ``ALTER PUBLICATION ... ADD/DROP
+TABLE`` calls below pre-date the ``COMMENT ON PUBLICATION`` bookend
+fix for bug #1355: on Zero >= 1.0 they do not reliably wake the
+zero-cache change-streamer and can leave the replica pinned to a
+stale snapshot. This file is grandfathered in because it has already
+shipped to users; new publication mutations MUST use the
+``COMMENT ON PUBLICATION`` bookend pattern wrapping an
+``ALTER PUBLICATION ... SET TABLE`` -- copy the ``upgrade()`` function
+from migration ``143_force_zero_publication_resync.py`` as your
+starting template.
+
 Revision ID: 118
 Revises: 117
 """

surfsense_backend/alembic/versions/139_add_user_to_zero_publication.py

@@ -21,6 +21,16 @@ IMPORTANT - before AND after running this migration:
   3. Delete / reset the zero-cache data volume
   4. Restart zero-cache  (it will do a fresh initial sync)
 
+DO NOT COPY THIS PATTERN. The ``DROP PUBLICATION`` + ``CREATE
+PUBLICATION`` dance below is the pre-#1355 anti-pattern: on Zero >=
+1.0 it does not reliably wake the zero-cache change-streamer and can
+leave the replica pinned to a stale snapshot. This file is
+grandfathered in because it has already shipped to users; new
+publication mutations MUST use the ``COMMENT ON PUBLICATION`` bookend
+pattern wrapping an ``ALTER PUBLICATION ... SET TABLE`` -- copy the
+``upgrade()`` function from migration
+``143_force_zero_publication_resync.py`` as your starting template.
+
 Revision ID: 139
 Revises: 138
 """

surfsense_backend/alembic/versions/140_premium_tokens_to_credit_micros.py

@@ -32,6 +32,16 @@ Skipping the zero-cache stop will deadlock at the ACCESS EXCLUSIVE LOCK on
 "user". Skipping the data-volume reset will leave IndexedDB clients seeing
 column-not-found errors from a stale catalog snapshot.
 
+DO NOT COPY THIS PATTERN. The ``DROP PUBLICATION`` + ``CREATE
+PUBLICATION`` dance below is the pre-#1355 anti-pattern: on Zero >=
+1.0 it does not reliably wake the zero-cache change-streamer and can
+leave the replica pinned to a stale snapshot. This file is
+grandfathered in because it has already shipped to users; new
+publication mutations MUST use the ``COMMENT ON PUBLICATION`` bookend
+pattern wrapping an ``ALTER PUBLICATION ... SET TABLE`` -- copy the
+``upgrade()`` function from migration
+``143_force_zero_publication_resync.py`` as your starting template.
+
 Revision ID: 140
 Revises: 139
 """

surfsense_backend/alembic/versions/143_force_zero_publication_resync.py

@@ -0,0 +1,142 @@
+"""force zero-cache to resync after upgrading to Zero >= 1.0
+
+Re-emits the current ``zero_publication`` shape using
+``ALTER PUBLICATION ... SET TABLE`` wrapped in
+``COMMENT ON PUBLICATION`` bookends. This is the publication-change
+hook documented for Zero ``>=1.0``:
+
+    https://zero.rocicorp.dev/docs/connecting-to-postgres#publication-changes
+
+Background
+----------
+Migrations 117 / 139 / 140 mutated ``zero_publication`` using
+``DROP PUBLICATION`` + ``CREATE PUBLICATION``. On Zero 0.26.2 that
+sequence did not reliably wake the zero-cache change-streamer, so
+affected installs ended up with a SQLite replica file (in the
+``surfsense-zero-cache`` volume) that was snapshotted against the
+pre-``user`` publication. The frontend Zero schema includes a
+``userTable`` query, which then failed with
+``SchemaVersionNotSupported`` and triggered the default
+``onUpdateNeeded`` -> ``location.reload()`` every WebSocket keepalive
+interval (~60s). See bug #1355.
+
+This migration emits the canonical publication shape one more time,
+this time using a pattern that fires Postgres event triggers and
+Zero's schema-change hook. With ``ZERO_AUTO_RESET=true`` (the default)
+and Zero ``>=1.0``, zero-cache responds by wiping its replica and
+doing a fresh initial sync from the corrected publication.
+
+The publication shape itself is unchanged versus migration 140 -- on
+installs whose replica is already correct, this is a no-op aside
+from the harmless event-trigger fire.
+
+Revision ID: 143
+Revises: 142
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision: str = "143"
+down_revision: str | None = "142"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+PUBLICATION_NAME = "zero_publication"
+
+# Must stay in sync with the column lists in migrations 117 / 139 / 140.
+DOCUMENT_COLS = [
+    "id",
+    "title",
+    "document_type",
+    "search_space_id",
+    "folder_id",
+    "created_by_id",
+    "status",
+    "created_at",
+    "updated_at",
+]
+
+USER_COLS = [
+    "id",
+    "pages_limit",
+    "pages_used",
+    "premium_credit_micros_limit",
+    "premium_credit_micros_used",
+]
+
+
+def _has_zero_version(conn, table: str) -> bool:
+    return (
+        conn.execute(
+            sa.text(
+                "SELECT 1 FROM information_schema.columns "
+                "WHERE table_name = :tbl AND column_name = '_0_version'"
+            ),
+            {"tbl": table},
+        ).fetchone()
+        is not None
+    )
+
+
+def _build_set_table_ddl(
+    *, documents_has_zero_ver: bool, user_has_zero_ver: bool
+) -> str:
+    doc_cols = DOCUMENT_COLS + (['"_0_version"'] if documents_has_zero_ver else [])
+    user_cols = USER_COLS + (['"_0_version"'] if user_has_zero_ver else [])
+    doc_col_list = ", ".join(doc_cols)
+    user_col_list = ", ".join(user_cols)
+    return (
+        f"ALTER PUBLICATION {PUBLICATION_NAME} SET TABLE "
+        f"notifications, "
+        f"documents ({doc_col_list}), "
+        f"folders, "
+        f"search_source_connectors, "
+        f"new_chat_messages, "
+        f"chat_comments, "
+        f"chat_session_state, "
+        f'"user" ({user_col_list})'
+    )
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+
+    exists = conn.execute(
+        sa.text("SELECT 1 FROM pg_publication WHERE pubname = :name"),
+        {"name": PUBLICATION_NAME},
+    ).fetchone()
+    if not exists:
+        return
+
+    documents_has_zero_ver = _has_zero_version(conn, "documents")
+    user_has_zero_ver = _has_zero_version(conn, "user")
+
+    # The COMMENT-ALTER-COMMENT trio MUST run in a single transaction so
+    # Zero observes them as one schema-change event. Alembic's outer
+    # transaction already covers us, but a SAVEPOINT keeps the trio
+    # atomic with asyncpg, matching the pattern used in migrations
+    # 117 / 139 / 140.
+    tx = conn.begin_nested() if conn.in_transaction() else conn.begin()
+    with tx:
+        conn.execute(
+            sa.text(f"COMMENT ON PUBLICATION {PUBLICATION_NAME} IS 'pre-143-resync'")
+        )
+        conn.execute(
+            sa.text(
+                _build_set_table_ddl(
+                    documents_has_zero_ver=documents_has_zero_ver,
+                    user_has_zero_ver=user_has_zero_ver,
+                )
+            )
+        )
+        conn.execute(
+            sa.text(f"COMMENT ON PUBLICATION {PUBLICATION_NAME} IS 'post-143-resync'")
+        )
+
+
+def downgrade() -> None:
+    """No-op. The publication shape is unchanged versus migration 140."""

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -1487,14 +1487,20 @@ async def stream_new_chat(
 
         # Resolve @-mention chips to canonical virtual paths and rewrite
         # the user-typed text so the LLM sees ``\`/documents/...\``` instead
-        # of bare ``@title``. The persisted user-message text keeps
+        # of bare ``@title``. The substitution lands in ``agent_user_query``
-        # ``@title`` so chip rendering on reload is unchanged — see
+        # ONLY — the original ``user_query`` (with ``@title`` tokens) flows
-        # ``persistence._build_user_content``.
+        # untouched into ``persist_user_turn`` below so chip rendering on
+        # reload still works (``UserTextPart`` → ``parseMentionSegments``
+        # matches ``@title``, not ``\`/documents/...\```). It also feeds
+        # the human-readable surfaces — SSE "Processing X" status, auto
+        # thread title, memory seed — which all want what the user typed.
+        # See ``persistence._build_user_content``.
         #
         # Cloud mode only: local-folder mode keeps the legacy
         # ``@title`` text path; mention support there is a follow-up
         # task because the path scheme (mount-rooted) and the picker
         # UI both need separate work.
+        agent_user_query = user_query
         accepted_folder_ids: list[int] = []
         if fs_mode == FilesystemMode.CLOUD.value and (
             mentioned_document_ids
@@ -1529,11 +1535,13 @@ async def stream_new_chat(
                 mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids,
                 mentioned_folder_ids=mentioned_folder_ids,
             )
-            user_query = substitute_in_text(user_query, resolved.token_to_path)
+            agent_user_query = substitute_in_text(user_query, resolved.token_to_path)
             accepted_folder_ids = resolved.mentioned_folder_ids
 
-        # Format the user query with context (SurfSense docs + reports only)
+        # Format the user query with context (SurfSense docs + reports only).
-        final_query = user_query
+        # Uses ``agent_user_query`` so the LLM sees backtick-wrapped paths
+        # instead of bare ``@title`` tokens.
+        final_query = agent_user_query
         context_parts = []
 
         if mentioned_surfsense_docs:
@@ -1564,7 +1572,7 @@ async def stream_new_chat(
 
         if context_parts:
             context = "\n\n".join(context_parts)
-            final_query = f"{context}\n\n<user_query>{user_query}</user_query>"
+            final_query = f"{context}\n\n<user_query>{agent_user_query}</user_query>"
 
         if visibility == ChatVisibility.SEARCH_SPACE and current_user_display_name:
             final_query = f"**[{current_user_display_name}]:** {final_query}"

surfsense_web/components/assistant-ui/inline-mention-editor.tsx

@@ -1,6 +1,7 @@
 "use client";
 
-import { Folder as FolderIcon } from "lucide-react";
+import { Folder as FolderIcon, X as XIcon } from "lucide-react";
+import type { NodeEntry, TElement } from "platejs";
 import type { PlateElementProps } from "platejs/react";
 import {
 	createPlatePlugin,
@@ -9,7 +10,16 @@ import {
 	PlateContent,
 	usePlateEditor,
 } from "platejs/react";
-import { type FC, forwardRef, useCallback, useImperativeHandle, useMemo, useRef } from "react";
+import {
+	createContext,
+	type FC,
+	forwardRef,
+	useCallback,
+	useContext,
+	useImperativeHandle,
+	useMemo,
+	useRef,
+} from "react";
 import { FOLDER_MENTION_DOCUMENT_TYPE } from "@/atoms/chat/mentioned-documents.atom";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import type { Document } from "@/contracts/types/document.types";
@@ -26,13 +36,9 @@ export interface MentionedDocument {
 }
 
 /**
- * Input shape for inserting a chip. ``kind`` defaults to ``"doc"``
+ * Input shape for inserting a chip. ``kind`` defaults to ``"doc"``.
- * when omitted so legacy callers don't have to thread the
+ * Folder chips default ``document_type`` to ``FOLDER_MENTION_DOCUMENT_TYPE``
- * discriminator. Folder callers pass ``kind: "folder"`` and the
+ * so the dedup key never collides with a doc chip sharing the same id.
- * folder ``id`` and ``title``; ``document_type`` defaults to
- * ``FOLDER_MENTION_DOCUMENT_TYPE`` inside ``insertMentionChip`` so the
- * dedup key (`kind:document_type:id`) never collides with a doc chip
- * that happens to share an id.
  */
 export type MentionChipInput = {
 	id: number;
@@ -87,12 +93,7 @@ type MentionElementNode = {
 	id: number;
 	title: string;
 	document_type?: string;
-	/**
+	/** Discriminator; defaults to ``"doc"`` for legacy nodes. */
-	 * Discriminator added so a folder chip and a doc chip with the
-	 * same id round-trip cleanly through ``getMentionedDocuments``
-	 * and the persisted ``mentioned-documents`` content part.
-	 * Defaults to ``"doc"`` for nodes that predate this field.
-	 */
 	kind?: MentionKind;
 	statusLabel?: string | null;
 	statusKind?: MentionStatusKind;
@@ -104,13 +105,22 @@ type ComposerValue = ComposerParagraph[];
 
 const MENTION_TYPE = "mention";
 const MENTION_CHIP_CLASSNAME =
-	"inline-flex h-5 items-center gap-1 mx-0.5 rounded bg-primary/10 px-1 text-xs font-bold text-primary/60 select-none align-middle leading-none";
+	"group inline-flex h-5 items-center gap-1 mx-0.5 rounded bg-primary/10 px-1 text-xs font-bold text-primary/60 select-none align-middle leading-none";
 const MENTION_CHIP_ICON_CLASSNAME = "flex items-center text-muted-foreground leading-none";
 const MENTION_CHIP_TITLE_CLASSNAME = "max-w-[120px] truncate leading-none";
 const COMPOSER_TEXT_METRICS_CLASSNAME = "text-sm leading-6";
 
 const EMPTY_VALUE: ComposerValue = [{ type: "p", children: [{ text: "" }] }];
 
+/**
+ * Lets ``MentionElement`` reach the editor's chip-removal helper so
+ * the X button and Backspace go through the same call site.
+ */
+type MentionEditorContextValue = {
+	removeChip: (docId: number, docType: string | undefined) => void;
+};
+const MentionEditorContext = createContext<MentionEditorContextValue | null>(null);
+
 const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 	attributes,
 	children,
@@ -124,17 +134,37 @@ const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 				: "text-amber-700";
 
 	const isFolder = element.kind === "folder";
+	const ctx = useContext(MentionEditorContext);
 
 	return (
 		<span {...attributes} className="inline-flex align-middle">
 			<span contentEditable={false} className={`${MENTION_CHIP_CLASSNAME} cursor-default`}>
 				<span className={MENTION_CHIP_ICON_CLASSNAME}>
+					<span className="relative flex h-3 w-3 items-center justify-center">
+						<span className="flex items-center justify-center transition-opacity group-hover:opacity-0">
 							{isFolder ? (
 								<FolderIcon className="h-3 w-3" />
 							) : (
 								getConnectorIcon(element.document_type ?? "UNKNOWN", "h-3 w-3")
 							)}
 						</span>
+						{ctx ? (
+							<button
+								type="button"
+								aria-label={`Remove mention ${element.title}`}
+								title={`Remove ${element.title}`}
+								onMouseDown={(e) => e.preventDefault()}
+								onClick={(e) => {
+									e.stopPropagation();
+									ctx.removeChip(element.id, element.document_type);
+								}}
+								className="absolute inset-0 flex items-center justify-center rounded-sm opacity-0 transition-opacity hover:text-primary focus-visible:opacity-100 focus-visible:outline-none group-hover:opacity-100"
+							>
+								<XIcon className="h-3 w-3" />
+							</button>
+						) : null}
+					</span>
+				</span>
 				<span className={MENTION_CHIP_TITLE_CLASSNAME} title={element.title}>
 					{element.title}
 				</span>
@@ -294,17 +324,16 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			value: initialText ? toValueFromText(initialText) : EMPTY_VALUE,
 		});
 
+		// Move the caret to end-of-doc and focus the editor. Falls back
+		// to DOM focus if Plate's API throws (transient unmount race).
 		const focusAtEnd = useCallback(() => {
-			const el = editableRef.current;
+			try {
-			if (!el) return;
+				editor.tf.select(editor.api.end([]));
-			el.focus();
+				editor.tf.focus();
-			const selection = window.getSelection();
+			} catch {
-			const range = document.createRange();
+				editableRef.current?.focus();
-			range.selectNodeContents(el);
+			}
-			range.collapse(false);
+		}, [editor]);
-			selection?.removeAllRanges();
-			selection?.addRange(range);
-		}, []);
 
 		const getCurrentValue = useCallback(
 			() => (editor.children as ComposerValue) ?? EMPTY_VALUE,
@@ -352,13 +381,18 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[editor, emitState]
 		);
 
+		// Insert chip + trailing space as a single ``insertNodes`` call.
+		// The chip is a void inline; ``select: true`` on it alone would
+		// land the caret inside its empty children (an unrenderable
+		// point). With the space as the last inserted node, the caret
+		// resolves to that text node and stays visible. The
+		// ``withoutNormalizing`` wrapper batches the optional trigger
+		// delete + insert into a single undo step.
 		const insertMentionChip = useCallback(
 			(mention: MentionChipInput, options?: { removeTriggerText?: boolean }) => {
 				if (typeof mention.id !== "number" || typeof mention.title !== "string") return;
 
 				const removeTriggerText = options?.removeTriggerText ?? true;
-				const current = getCurrentValue();
-				const selection = editor.selection;
 				const kind: MentionKind = mention.kind ?? "doc";
 				const document_type =
 					mention.document_type ?? (kind === "folder" ? FOLDER_MENTION_DOCUMENT_TYPE : undefined);
@@ -371,65 +405,48 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 					children: [{ text: "" }],
 				};
 
-				const cursorCtx = getCursorTextContext(current, selection);
+				editor.tf.withoutNormalizing(() => {
-				if (!cursorCtx) {
+					const selection = editor.selection;
-					const lastBlock = current[current.length - 1] ?? { type: "p", children: [{ text: "" }] };
-					const appended: ComposerValue = [
-						...current.slice(0, -1),
-						{
-							...lastBlock,
-							children: [...lastBlock.children, mentionNode, { text: " " }],
-						},
-					];
-					setValue(appended);
-					requestAnimationFrame(focusAtEnd);
-					return;
-				}
 
-				const block = current[cursorCtx.blockIndex];
+					// No active selection (focus moved to a picker) — snap
-				const currentChild = getTextNode(block.children[cursorCtx.childIndex]);
+					// to end-of-doc so the chip appends cleanly.
-				if (!currentChild) {
+					if (!selection) {
-					const children = [...block.children];
+						editor.tf.select(editor.api.end([]));
-					children.splice(cursorCtx.childIndex + 1, 0, mentionNode, { text: " " });
+					} else if (removeTriggerText) {
-					const next = [...current];
+						// Delete the in-progress "@query" so the chip stands in for it.
-					next[cursorCtx.blockIndex] = { ...block, children };
+						const cursorCtx = getCursorTextContext(getCurrentValue(), selection);
-					setValue(next as ComposerValue);
+						if (cursorCtx) {
-					requestAnimationFrame(focusAtEnd);
+							const text = cursorCtx.text;
-					return;
+							let triggerIndex = -1;
-				}
-
-				const text = currentChild.text;
-				let removeStart = cursorCtx.cursor;
-				if (removeTriggerText) {
 							for (let i = cursorCtx.cursor - 1; i >= 0; i--) {
 								if (text[i] === "@") {
-							removeStart = i;
+									triggerIndex = i;
 									break;
 								}
 								if (text[i] === " " || text[i] === "\n") break;
 							}
+							if (triggerIndex >= 0 && triggerIndex < cursorCtx.cursor) {
+								const path = [cursorCtx.blockIndex, cursorCtx.childIndex];
+								editor.tf.delete({
+									at: {
+										anchor: { path, offset: triggerIndex },
+										focus: { path, offset: cursorCtx.cursor },
+									},
+								});
+							}
+						}
 					}
 
-				const before = text.slice(0, removeStart);
+					editor.tf.insertNodes([mentionNode, { text: " " }] as unknown as TElement[], {
-				const after = text.slice(cursorCtx.cursor);
+						select: true,
-				const replacement: ComposerNode[] = [];
+					});
-				if (before.length > 0) replacement.push({ text: before });
+				});
-				replacement.push(mentionNode);
+				editor.tf.focus();
-				replacement.push({ text: ` ${after}` });
-
-				const children = [...block.children];
-				children.splice(cursorCtx.childIndex, 1, ...replacement);
-				const next = [...current];
-				next[cursorCtx.blockIndex] = { ...block, children };
-				setValue(next as ComposerValue);
-				requestAnimationFrame(focusAtEnd);
 			},
-			[editor.selection, focusAtEnd, getCurrentValue, setValue]
+			[editor, getCurrentValue]
 		);
 
-		// Backwards-compatible shim — pre-folder callers pass a doc-only
+		// Doc-only shim that routes through ``insertMentionChip``.
-		// payload; we route them through ``insertMentionChip`` with
-		// ``kind: "doc"``.
 		const insertDocumentChip = useCallback(
 			(
 				doc: Pick<Document, "id" | "title" | "document_type">,
@@ -440,26 +457,43 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[insertMentionChip]
 		);
 
+		// Remove chip(s) matching (id, document_type). Iterates in
+		// descending path order so removing one entry can't invalidate
+		// later paths. Chips are deduped today, so this typically runs
+		// at most once.
 		const removeDocumentChip = useCallback(
 			(docId: number, docType?: string) => {
-				const current = getCurrentValue();
+				const match = (n: unknown) => {
-				let changed = false;
+					if (!n || typeof n !== "object" || !("type" in n)) return false;
-				const next = current.map((block) => {
+					const node = n as MentionElementNode;
-					const children = block.children.filter((node) => {
+					if (node.type !== MENTION_TYPE) return false;
-						if (!isMentionNode(node)) return true;
+					if (node.id !== docId) return false;
-						const match =
+					return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
-							node.id === docId && (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
+				};
-						if (match) changed = true;
+
-						return !match;
+				const entries = Array.from(editor.api.nodes({ at: [], match })) as NodeEntry[];
+				if (entries.length === 0) return;
+				editor.tf.withoutNormalizing(() => {
+					for (const [, path] of entries.reverse()) {
+						editor.tf.removeNodes({ at: path });
+					}
 				});
-					return { ...block, children: children.length ? children : [{ text: "" }] };
-				});
-				if (!changed) return;
-				setValue(next as ComposerValue);
 			},
-			[getCurrentValue, setValue]
+			[editor]
 		);
 
+		// Single removal call site for Backspace and the X button so the
+		// two can never diverge (e.g. one forgetting to notify the parent).
+		const removeChip = useCallback(
+			(docId: number, docType: string | undefined) => {
+				removeDocumentChip(docId, docType);
+				onDocumentRemove?.(docId, docType);
+			},
+			[onDocumentRemove, removeDocumentChip]
+		);
+
+		// Update chip status in place via ``tf.setNodes`` so the user's
+		// selection survives backend status events arriving mid-typing.
 		const setDocumentChipStatus = useCallback(
 			(
 				docId: number,
@@ -467,31 +501,31 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				statusLabel: string | null,
 				statusKind: MentionStatusKind = "pending"
 			) => {
-				const current = getCurrentValue();
+				const match = (n: unknown) => {
-				let changed = false;
+					if (!n || typeof n !== "object" || !("type" in n)) return false;
-				const next = current.map((block) => ({
+					const node = n as MentionElementNode;
-					...block,
+					if (node.type !== MENTION_TYPE) return false;
-					children: block.children.map((node) => {
+					if (node.id !== docId) return false;
-						if (!isMentionNode(node)) return node;
+					return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
-						const sameType = (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
+				};
-						if (node.id !== docId || !sameType) return node;
+
-						changed = true;
+				editor.tf.setNodes(
-						return {
+					{
-							...node,
 						statusLabel,
 						statusKind: statusLabel ? statusKind : undefined,
-						};
+					} as Partial<TElement>,
-					}),
+					{ at: [], match }
-				}));
+				);
-				if (!changed) return;
-				setValue(next as ComposerValue);
 			},
-			[getCurrentValue, setValue]
+			[editor]
 		);
 
 		const clear = useCallback(() => {
 			setValue(EMPTY_VALUE);
-		}, [setValue]);
+			// ``tf.setValue`` wipes the selection — refocus so the caret
+			// returns after Enter-to-submit.
+			requestAnimationFrame(focusAtEnd);
+		}, [focusAtEnd, setValue]);
 
 		const setText = useCallback(
 			(text: string) => {
@@ -510,7 +544,18 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 		useImperativeHandle(
 			ref,
 			() => ({
-				focus: () => editableRef.current?.focus(),
+				// Preserve existing selection if any; otherwise seed one
+				// at end-of-doc so the contentEditable shows a caret.
+				focus: () => {
+					try {
+						if (!editor.selection) {
+							editor.tf.select(editor.api.end([]));
+						}
+						editor.tf.focus();
+					} catch {
+						editableRef.current?.focus();
+					}
+				},
 				clear,
 				setText,
 				getText,
@@ -522,6 +567,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			}),
 			[
 				clear,
+				editor,
 				getMentionedDocs,
 				getText,
 				insertMentionChip,
@@ -564,10 +610,9 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				if (!isMentionNode(prev)) return;
 
 				e.preventDefault();
-				removeDocumentChip(prev.id, prev.document_type);
+				removeChip(prev.id, prev.document_type);
-				onDocumentRemove?.(prev.id, prev.document_type);
 			},
-			[editor.selection, getCurrentValue, onDocumentRemove, onKeyDown, onSubmit, removeDocumentChip]
+			[editor.selection, getCurrentValue, onKeyDown, onSubmit, removeChip]
 		);
 
 		const editableProps = useMemo(
@@ -584,8 +629,14 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[editor, handleKeyDown, placeholder]
 		);
 
+		const mentionEditorContextValue = useMemo<MentionEditorContextValue>(
+			() => ({ removeChip }),
+			[removeChip]
+		);
+
 		return (
 			<div className="relative w-full">
+				<MentionEditorContext.Provider value={mentionEditorContextValue}>
 					<Plate
 						editor={editor}
 						onChange={({ value }) => {
@@ -604,6 +655,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 							)}
 						/>
 					</Plate>
+				</MentionEditorContext.Provider>
 			</div>
 		);
 	}

surfsense_web/components/assistant-ui/mention-chip.tsx

@@ -66,23 +66,21 @@ export function MentionChip({
 			disabled={disabled}
 			aria-label={ariaLabel ?? label}
 			className={cn(
-				"inline-flex max-w-[220px] items-center gap-1.5 rounded-md border bg-background px-2 py-0.5 align-middle text-xs font-medium text-foreground leading-5 transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
+				"inline-flex h-5 items-center gap-1 rounded bg-primary/10 px-1 align-middle text-xs font-bold text-primary/60 leading-none focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
-				isInteractive
+				isInteractive ? "cursor-pointer" : "cursor-default",
-					? "cursor-pointer hover:bg-accent hover:text-accent-foreground"
-					: "cursor-default",
 				disabled && "opacity-60",
 				className
 			)}
 		>
 			<span className="inline-flex shrink-0 text-muted-foreground">{icon}</span>
-			<span className="truncate">{label}</span>
+			<span className="max-w-[120px] truncate leading-none">{label}</span>
 		</button>
 	);
 
 	if (!tooltip) return chip;
 
 	return (
-		<Tooltip>
+		<Tooltip delayDuration={600}>
 			<TooltipTrigger asChild>{chip}</TooltipTrigger>
 			<TooltipContent side="top" className="max-w-xs break-all">
 				{tooltip}

surfsense_web/components/assistant-ui/thread.tsx

@@ -62,6 +62,7 @@ import { useDocumentUploadDialog } from "@/components/assistant-ui/document-uplo
 import {
 	InlineMentionEditor,
 	type InlineMentionEditorRef,
+	type MentionedDocument,
 } from "@/components/assistant-ui/inline-mention-editor";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { UserMessage } from "@/components/assistant-ui/user-message";
@@ -171,36 +172,24 @@ const PremiumQuotaPinnedAlert: FC = () => {
 const getTimeBasedGreeting = (user?: { display_name?: string | null; email?: string }): string => {
 	const hour = new Date().getHours();
 
-	// Extract first name: prefer display_name, fall back to email extraction
 	let firstName: string | null = null;
-
 	if (user?.display_name?.trim()) {
-		// Use display_name if available and not empty
-		// Extract first name from display_name (take first word)
 		const nameParts = user.display_name.trim().split(/\s+/);
 		firstName = nameParts[0].charAt(0).toUpperCase() + nameParts[0].slice(1).toLowerCase();
 	} else if (user?.email) {
-		// Fall back to email extraction if display_name is not available
 		firstName =
 			user.email.split("@")[0].split(".")[0].charAt(0).toUpperCase() +
 			user.email.split("@")[0].split(".")[0].slice(1);
 	}
 
-	// Array of greeting variations for each time period
 	const morningGreetings = ["Good morning", "Fresh start today", "Morning", "Hey there"];
-
 	const afternoonGreetings = ["Good afternoon", "Afternoon", "Hey there", "Hi there"];
-
 	const eveningGreetings = ["Good evening", "Evening", "Hey there", "Hi there"];
-
 	const nightGreetings = ["Good night", "Evening", "Hey there", "Winding down"];
-
 	const lateNightGreetings = ["Still up", "Night owl mode", "Up past bedtime", "Hi there"];
 
-	// Select a random greeting based on time
 	let greeting: string;
 	if (hour < 5) {
-		// Late night: midnight to 5 AM
 		greeting = lateNightGreetings[Math.floor(Math.random() * lateNightGreetings.length)];
 	} else if (hour < 12) {
 		greeting = morningGreetings[Math.floor(Math.random() * morningGreetings.length)];
@@ -209,33 +198,23 @@ const getTimeBasedGreeting = (user?: { display_name?: string | null; email?: str
 	} else if (hour < 22) {
 		greeting = eveningGreetings[Math.floor(Math.random() * eveningGreetings.length)];
 	} else {
-		// Night: 10 PM to midnight
 		greeting = nightGreetings[Math.floor(Math.random() * nightGreetings.length)];
 	}
 
-	// Add personalization with first name if available
+	return firstName ? `${greeting}, ${firstName}!` : `${greeting}!`;
-	if (firstName) {
-		return `${greeting}, ${firstName}!`;
-	}
-
-	return `${greeting}!`;
 };
 
 const ThreadWelcome: FC = () => {
 	const { data: user } = useAtomValue(currentUserAtom);
-
-	// Memoize greeting so it doesn't change on re-renders (only on user change)
 	const greeting = useMemo(() => getTimeBasedGreeting(user), [user]);
 
 	return (
 		<div className="aui-thread-welcome-root mx-auto flex w-full max-w-(--thread-max-width) grow flex-col items-center px-4 relative">
-			{/* Greeting positioned above the composer */}
 			<div className="aui-thread-welcome-message absolute bottom-[calc(50%+5rem)] left-0 right-0 flex flex-col items-center text-center">
 				<h1 className="aui-thread-welcome-message-inner text-3xl md:text-5xl select-none">
 					{greeting}
 				</h1>
 			</div>
-			{/* Composer - top edge fixed, expands downward only */}
 			<div className="w-full flex items-start justify-center absolute top-[calc(50%-3.5rem)] left-0 right-0">
 				<Composer />
 			</div>
@@ -372,7 +351,6 @@ const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDi
 };
 
 const Composer: FC = () => {
-	// Document mention state (atoms persist across component remounts)
 	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
 	const [showPromptPicker, setShowPromptPicker] = useState(false);
@@ -384,7 +362,9 @@ const Composer: FC = () => {
 	const promptPickerRef = useRef<PromptPickerRef>(null);
 	const { search_space_id, chat_id } = useParams();
 	const aui = useAui();
-	const hasAutoFocusedRef = useRef(false);
+	// Desktop-only auto-focus; on mobile, programmatic focus would
+	// summon the soft keyboard on every picker close / thread switch.
+	const isDesktop = useMediaQuery("(min-width: 640px)");
 
 	const electronAPI = useElectronAPI();
 	const [clipboardInitialText, setClipboardInitialText] = useState<string | undefined>();
@@ -404,7 +384,6 @@ const Composer: FC = () => {
 
 	const currentPlaceholder = COMPOSER_PLACEHOLDER;
 
-	// Live collaboration state
 	const { data: currentUser } = useAtomValue(currentUserAtom);
 	const { data: members } = useAtomValue(membersAtom);
 	const threadId = useMemo(() => {
@@ -418,13 +397,11 @@ const Composer: FC = () => {
 	const respondingToUserId = sessionState?.respondingToUserId ?? null;
 	const isBlockedByOtherUser = isAiResponding && respondingToUserId !== currentUser?.id;
 
-	// Sync comments for the entire thread via Zero (one subscription per thread)
+	// One Zero subscription per thread for comment sync.
 	useCommentsSync(threadId);
 
-	// Batch-prefetch comments for all assistant messages so individual useComments
+	// Batch-prefetch assistant message comments to avoid N+1 fetches.
-	// hooks never fire their own network requests (eliminates N+1 API calls).
+	// Returns a primitive string so useSyncExternalStore can compare by value.
-	// Return a primitive string from the selector so useSyncExternalStore can
-	// compare snapshots by value and avoid infinite re-render loops.
 	const assistantIdsKey = useAuiState(({ thread }) =>
 		thread.messages
 			.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
@@ -437,18 +414,17 @@ const Composer: FC = () => {
 	);
 	useBatchCommentsPreload(assistantDbMessageIds);
 
-	// Auto-focus editor on new chat page after mount
+	// Always-focused composer: refocus whenever no picker has taken
+	// over input. ``threadId`` is in the deps so the effect re-fires
+	// on thread switch (Composer instance is reused).
 	useEffect(() => {
-		if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) {
+		if (!isDesktop) return;
-			const timeoutId = setTimeout(() => {
+		if (showDocumentPopover || showPromptPicker) return;
+		void threadId;
 		editorRef.current?.focus();
-				hasAutoFocusedRef.current = true;
+	}, [isDesktop, showDocumentPopover, showPromptPicker, threadId]);
-			}, 100);
-			return () => clearTimeout(timeoutId);
-		}
-	}, [isThreadEmpty]);
 
-	// Close document picker when a slide-out panel (inbox, shared/private chats) opens
+	// Close document picker when a slide-out panel (inbox, etc.) opens.
 	useEffect(() => {
 		const handler = () => {
 			setShowDocumentPopover(false);
@@ -458,21 +434,41 @@ const Composer: FC = () => {
 		return () => window.removeEventListener(SLIDEOUT_PANEL_OPENED_EVENT, handler);
 	}, []);
 
-	// Sync editor text with assistant-ui composer runtime
+	// Sync editor text into assistant-ui's composer and mirror the chip
+	// atom from the editor's reported ``docs``. The editor is the
+	// single source of truth, so this catches every Plate deletion path
+	// (Backspace, X button, Cmd+Backspace, range-delete, cut,
+	// paste-over) without per-keybinding plumbing. The ``prev``
+	// short-circuit keeps pure-text keystrokes from churning the atom.
 	const handleEditorChange = useCallback(
-		(text: string) => {
+		(text: string, docs: MentionedDocument[]) => {
 			aui.composer().setText(text);
+			setMentionedDocuments((prev) => {
+				if (prev.length === docs.length) {
+					const editorKeys = new Set(docs.map((d) => getMentionDocKey(d)));
+					if (prev.every((d) => editorKeys.has(getMentionDocKey(d)))) {
+						return prev;
+					}
+				}
+				return docs.map<MentionedDocumentInfo>((d) => ({
+					id: d.id,
+					title: d.title,
+					// Atom requires a string; ``"UNKNOWN"`` matches the
+					// sentinel ``getMentionDocKey`` and the editor's
+					// match predicates use.
+					document_type: d.document_type ?? "UNKNOWN",
+					kind: d.kind,
+				}));
+			});
 		},
-		[aui]
+		[aui, setMentionedDocuments]
 	);
 
-	// Open document picker when @ mention is triggered
 	const handleMentionTrigger = useCallback((query: string) => {
 		setShowDocumentPopover(true);
 		setMentionQuery(query);
 	}, []);
 
-	// Close document picker and reset query
 	const handleMentionClose = useCallback(() => {
 		if (showDocumentPopover) {
 			setShowDocumentPopover(false);
@@ -480,13 +476,11 @@ const Composer: FC = () => {
 		}
 	}, [showDocumentPopover]);
 
-	// Open action picker when / is triggered
 	const handleActionTrigger = useCallback((query: string) => {
 		setShowPromptPicker(true);
 		setActionQuery(query);
 	}, []);
 
-	// Close action picker and reset query
 	const handleActionClose = useCallback(() => {
 		if (showPromptPicker) {
 			setShowPromptPicker(false);
@@ -530,7 +524,7 @@ const Composer: FC = () => {
 		[clipboardInitialText, electronAPI, aui]
 	);
 
-	// Keyboard navigation for document/action picker (arrow keys, Enter, Escape)
+	// Arrow / Enter / Escape navigation for the active picker.
 	const handleKeyDown = useCallback(
 		(e: React.KeyboardEvent) => {
 			if (showPromptPicker) {
@@ -611,7 +605,7 @@ const Composer: FC = () => {
 		(docId: number, docType?: string) => {
 			setMentionedDocuments((prev) => {
 				if (!docType) {
-					// Defensive fallback: keep UI in sync even when chip type is unavailable.
+					// Fallback when chip type is unavailable.
 					return prev.filter((doc) => doc.id !== docId);
 				}
 				const removedKey = getMentionDocKey({ id: docId, document_type: docType });
@@ -621,8 +615,7 @@ const Composer: FC = () => {
 		[setMentionedDocuments]
 	);
 
-	const handleDocumentsMention = useCallback(
+	const handleDocumentsMention = useCallback((mentions: MentionedDocumentInfo[]) => {
-		(mentions: MentionedDocumentInfo[]) => {
 		const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? [];
 		const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc)));
 
@@ -630,18 +623,14 @@ const Composer: FC = () => {
 			const key = getMentionDocKey(mention);
 			if (editorDocKeys.has(key)) continue;
 			editorRef.current?.insertMentionChip(mention);
+			// Track within the loop so a duplicate-in-batch can't double-insert.
+			editorDocKeys.add(key);
 		}
 
-			setMentionedDocuments((prev) => {
+		// Atom is reconciled by ``handleEditorChange`` via the editor's
-				const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d)));
+		// onChange — no second write path here.
-				const uniqueNew = mentions.filter((m) => !existingKeySet.has(getMentionDocKey(m)));
-				return [...prev, ...uniqueNew];
-			});
-
 		setMentionQuery("");
-		},
+	}, []);
-		[setMentionedDocuments]
-	);
 
 	useEffect(() => {
 		const editor = editorRef.current;
@@ -1212,6 +1201,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 				</div>
 			)}
 			<div className="flex items-center gap-2">
+				{isDesktop && (
 					<TooltipIconButton
 						tooltip="Capture screen"
 						type="button"
@@ -1223,6 +1213,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 					>
 						<Camera className="size-4" />
 					</TooltipIconButton>
+				)}
 				<AuiIf condition={({ thread }) => !thread.isRunning}>
 					<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
 						<TooltipIconButton
@@ -1269,12 +1260,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 	);
 };
 
-/**
+/** Friendly tool name (delegates to ``getToolDisplayName``). */
- * Friendly tool name for display in the chat UI. Delegates to the
- * shared map in ``contracts/enums/toolIcons`` so unix-style identifiers
- * (``rm``, ``ls``, ``grep`` …) and snake_cased function names render as
- * plain English (e.g. "Delete file", "List files", "Search in files").
- */
 function formatToolName(name: string): string {
 	return getToolDisplayName(name);
 }

surfsense_web/package.json

@@ -81,7 +81,7 @@
 		"@remotion/media": "^4.0.438",
 		"@remotion/player": "^4.0.438",
 		"@remotion/web-renderer": "^4.0.438",
-		"@rocicorp/zero": "^0.26.2",
+		"@rocicorp/zero": "1.4.0",
 		"@slate-serializers/html": "^2.2.3",
 		"@streamdown/code": "^1.0.2",
 		"@streamdown/math": "^1.0.2",