feat: improved document, folder mentions rendering

surfsense_backend/app/tasks/chat/content_builder.py

@@ -51,7 +51,9 @@ logger = logging.getLogger(__name__)
 _MEANINGFUL_PART_TYPES: frozenset[str] = frozenset({"text", "reasoning", "tool-call"})
 
 
-def _merge_tool_part_metadata(part: dict[str, Any], metadata: dict[str, Any] | None) -> None:
+def _merge_tool_part_metadata(
+    part: dict[str, Any], metadata: dict[str, Any] | None
+) -> None:
     """Shallow-merge ``metadata`` into ``part["metadata"]``; first key wins.
 
     Used for tool-call linkage (``spanId``, ``thinkingStepId``, …): a later

surfsense_backend/app/tasks/chat/persistence.py

@@ -109,17 +109,18 @@ def _build_user_content(
         [{"type": "text", "text": "..."},
          {"type": "image", "image": "data:..."},
          {"type": "mentioned-documents", "documents": [{"id": int,
-            "title": str, "document_type": str}, ...]}]
+            "title": str, "document_type": str, "kind": "doc" | "folder"},
+            ...]}]
 
     The companion reader is
     ``app.utils.user_message_multimodal.split_persisted_user_content_parts``
     which expects exactly this shape — keep them in sync.
 
-    ``mentioned_documents``: optional list of ``{id, title, document_type}``
-    dicts. When non-empty (and a ``mentioned-documents`` part is not already
-    in some other input shape), a single ``{"type": "mentioned-documents",
-    "documents": [...]}`` part is appended. Mirrors the FE injection at
-    ``page.tsx:281-286`` (``persistUserTurn``).
+    ``mentioned_documents``: optional list of mention chip dicts. Each
+    dict may include a ``kind`` discriminator (``"doc"`` or ``"folder"``)
+    so the persisted ContentPart round-trips folder chips on reload.
+    When ``kind`` is missing we default to ``"doc"`` so legacy clients
+    that haven't migrated to the union schema still persist correctly.
     """
     parts: list[dict[str, Any]] = [{"type": "text", "text": user_query or ""}]
     for url in user_image_data_urls or ():
@@ -135,11 +136,14 @@ def _build_user_content(
             document_type = doc.get("document_type")
             if doc_id is None or title is None or document_type is None:
                 continue
+            kind_raw = doc.get("kind", "doc")
+            kind = kind_raw if kind_raw in ("doc", "folder") else "doc"
             normalized.append(
                 {
                     "id": doc_id,
                     "title": str(title),
                     "document_type": str(document_type),
+                    "kind": kind,
                 }
             )
         if normalized:

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -43,6 +43,7 @@ from app.agents.new_chat.memory_extraction import (
     extract_and_save_memory,
     extract_and_save_team_memory,
 )
+from app.agents.new_chat.mention_resolver import resolve_mentions, substitute_in_text
 from app.agents.new_chat.middleware.busy_mutex import (
     end_turn,
     get_cancel_state,
@@ -929,6 +930,7 @@ async def stream_new_chat(
     llm_config_id: int = -1,
     mentioned_document_ids: list[int] | None = None,
     mentioned_surfsense_doc_ids: list[int] | None = None,
+    mentioned_folder_ids: list[int] | None = None,
     mentioned_documents: list[dict[str, Any]] | None = None,
     checkpoint_id: str | None = None,
     needs_history_bootstrap: bool = False,
@@ -958,6 +960,7 @@ async def stream_new_chat(
         needs_history_bootstrap: If True, load message history from DB (for cloned chats)
         mentioned_document_ids: Optional list of document IDs mentioned with @ in the chat
         mentioned_surfsense_doc_ids: Optional list of SurfSense doc IDs mentioned with @ in the chat
+        mentioned_folder_ids: Optional list of knowledge-base folder IDs mentioned with @ (cloud mode)
         checkpoint_id: Optional checkpoint ID to rewind/fork from (for edit/reload operations)
 
     Yields:
@@ -1502,6 +1505,53 @@ async def stream_new_chat(
         )
         recent_reports = list(recent_reports_result.scalars().all())
 
+        # 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
+        # ``@title`` so chip rendering on reload is unchanged — 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.
+        accepted_folder_ids: list[int] = []
+        if fs_mode == FilesystemMode.CLOUD.value and (
+            mentioned_document_ids
+            or mentioned_surfsense_doc_ids
+            or mentioned_folder_ids
+            or mentioned_documents
+        ):
+            from app.schemas.new_chat import (
+                MentionedDocumentInfo as _MentionedDocumentInfo,
+            )
+
+            chip_objs: list[_MentionedDocumentInfo] | None = None
+            if mentioned_documents:
+                chip_objs = []
+                for raw in mentioned_documents:
+                    if isinstance(raw, _MentionedDocumentInfo):
+                        chip_objs.append(raw)
+                        continue
+                    try:
+                        chip_objs.append(_MentionedDocumentInfo.model_validate(raw))
+                    except Exception:
+                        logger.debug(
+                            "stream_new_chat: dropping malformed mention chip %r",
+                            raw,
+                        )
+
+            resolved = await resolve_mentions(
+                session,
+                search_space_id=search_space_id,
+                mentioned_documents=chip_objs,
+                mentioned_document_ids=mentioned_document_ids,
+                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)
+            accepted_folder_ids = resolved.mentioned_folder_ids
+
         # Format the user query with context (SurfSense docs + reports only)
         final_query = user_query
         context_parts = []
@@ -1901,6 +1951,9 @@ async def stream_new_chat(
         runtime_context = SurfSenseContextSchema(
             search_space_id=search_space_id,
             mentioned_document_ids=list(mentioned_document_ids or []),
+            mentioned_folder_ids=list(
+                accepted_folder_ids or mentioned_folder_ids or []
+            ),
             request_id=request_id,
             turn_id=stream_result.turn_id,
         )

surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py

@@ -26,9 +26,7 @@ def handle_report_progress(
         return None, last_active_step_items
 
     phase = data.get("phase", "")
-    topic_items = [
-        item for item in last_active_step_items if item.startswith("Topic:")
-    ]
+    topic_items = [item for item in last_active_step_items if item.startswith("Topic:")]
 
     if phase in ("revising_section", "adding_section"):
         plan_items = [
@@ -56,7 +54,9 @@ def handle_report_progress(
     return frame, new_items
 
 
-def handle_document_created(data: dict[str, Any], *, streaming_service: Any) -> str | None:
+def handle_document_created(
+    data: dict[str, Any], *, streaming_service: Any
+) -> str | None:
     if not data.get("id"):
         return None
     return streaming_service.format_data(

surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py

@@ -13,7 +13,9 @@ from app.tasks.chat.streaming.handlers.tools import (
 )
 from app.tasks.chat.streaming.helpers.tool_output import tool_output_has_error
 from app.tasks.chat.streaming.relay.state import AgentEventRelayState
-from app.tasks.chat.streaming.relay.task_span import clear_task_span_if_delegating_task_ended
+from app.tasks.chat.streaming.relay.task_span import (
+    clear_task_span_if_delegating_task_ended,
+)
 from app.tasks.chat.streaming.relay.thinking_step_sse import emit_thinking_step_frame
 
 
@@ -32,9 +34,7 @@ def iter_tool_end_frames(
     run_id = event.get("run_id", "")
     tool_name = event.get("name", "unknown_tool")
     raw_output = event.get("data", {}).get("output", "")
-    staged_file_path = (
-        state.file_path_by_run.pop(run_id, None) if run_id else None
-    )
+    staged_file_path = state.file_path_by_run.pop(run_id, None) if run_id else None
 
     if tool_name == "update_memory":
         state.called_update_memory = True
@@ -116,6 +116,4 @@ def iter_tool_end_frames(
     )
     yield from iter_tool_completion_emission_frames(emission_ctx)
 
-    clear_task_span_if_delegating_task_ended(
-        state, tool_name=tool_name, run_id=run_id
-    )
+    clear_task_span_if_delegating_task_ended(state, tool_name=tool_name, run_id=run_id)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/thinking.py

@@ -15,7 +15,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     return default_thinking.resolve_completed_thinking(
         tool_name, tool_output, last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/thinking.py

@@ -23,7 +23,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/thinking.py

@@ -34,7 +34,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/thinking.py

@@ -29,7 +29,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items
@@ -44,9 +46,7 @@ def resolve_completed_thinking(
         else "Report"
     )
     word_count = (
-        tool_output.get("word_count", 0)
-        if isinstance(tool_output, dict)
-        else 0
+        tool_output.get("word_count", 0) if isinstance(tool_output, dict) else 0
     )
     is_revision = (
         tool_output.get("is_revision", False)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/thinking.py

@@ -17,7 +17,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     return default_thinking.resolve_completed_thinking(
         tool_name, tool_output, last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/thinking.py

@@ -17,7 +17,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/thinking.py

@@ -21,7 +21,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/thinking.py

@@ -21,7 +21,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Editing file", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/emission.py

@@ -24,9 +24,7 @@ def iter_completion_emission_frames(
         output_text = om.group(1) if om else ""
     thread_id_str = ctx.langgraph_config.get("configurable", {}).get("thread_id", "")
 
-    for sf_match in re.finditer(
-        r"^SANDBOX_FILE:\s*(.+)$", output_text, re.MULTILINE
-    ):
+    for sf_match in re.finditer(r"^SANDBOX_FILE:\s*(.+)$", output_text, re.MULTILINE):
         fpath = sf_match.group(1).strip()
         if fpath and fpath not in ctx.stream_result.sandbox_files:
             ctx.stream_result.sandbox_files.append(fpath)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/thinking.py

@@ -22,7 +22,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/thinking.py

@@ -21,7 +21,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Searching files", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/thinking.py

@@ -25,7 +25,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Searching content", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/thinking.py

@@ -20,7 +20,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     if isinstance(tool_output, dict):
@@ -38,9 +40,7 @@ def resolve_completed_thinking(
                 paths = [str(p) for p in parsed]
         except (ValueError, SyntaxError):
             paths = [
-                line.strip()
-                for line in ls_output.strip().split("\n")
-                if line.strip()
+                line.strip() for line in ls_output.strip().split("\n") if line.strip()
             ]
         for p in paths:
             name = p.rstrip("/").split("/")[-1]

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/thinking.py

@@ -17,11 +17,15 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
     d = as_tool_input_dict(tool_input)
     p = d.get("path", "") if isinstance(tool_input, dict) else str(tool_input)
     display = p if len(p) <= 80 else "…" + p[-77:]
-    return ToolStartThinking(title="Creating folder", items=[display] if display else [])
+    return ToolStartThinking(
+        title="Creating folder", items=[display] if display else []
+    )
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Creating folder", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/thinking.py

@@ -27,7 +27,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Moving file", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/thinking.py

@@ -21,7 +21,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Reading file", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/thinking.py

@@ -22,7 +22,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Deleting file", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/thinking.py

@@ -17,11 +17,15 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
     d = as_tool_input_dict(tool_input)
     p = d.get("path", "") if isinstance(tool_input, dict) else str(tool_input)
     display = p if len(p) <= 80 else "…" + p[-77:]
-    return ToolStartThinking(title="Deleting folder", items=[display] if display else [])
+    return ToolStartThinking(
+        title="Deleting folder", items=[display] if display else []
+    )
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Deleting folder", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/thinking.py

@@ -21,7 +21,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Writing file", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/thinking.py

@@ -20,15 +20,15 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
     return ToolStartThinking(
         title="Planning tasks",
         items=(
-            [f"{todo_count} task{'s' if todo_count != 1 else ''}"]
-            if todo_count
-            else []
+            [f"{todo_count} task{'s' if todo_count != 1 else ''}"] if todo_count else []
         ),
     )
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_output, tool_name
     return ("Planning tasks", last_items)

surfsense_backend/app/tasks/chat/streaming/handlers/tools/registry.py

@@ -58,14 +58,18 @@ def _emission_module(tool_name: str) -> str:
 
 def _import_thinking(tool_name: str):
     try:
-        return importlib.import_module(f"{_BASE}.{_thinking_module(tool_name)}.thinking")
+        return importlib.import_module(
+            f"{_BASE}.{_thinking_module(tool_name)}.thinking"
+        )
     except ModuleNotFoundError:
         return importlib.import_module(f"{_BASE}.default.thinking")
 
 
 def _import_emission(tool_name: str):
     try:
-        return importlib.import_module(f"{_BASE}.{_emission_module(tool_name)}.emission")
+        return importlib.import_module(
+            f"{_BASE}.{_emission_module(tool_name)}.emission"
+        )
     except ModuleNotFoundError:
         return importlib.import_module(f"{_BASE}.default.emission")
 

surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/thinking.py

@@ -23,7 +23,9 @@ def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking
 
 
 def resolve_completed_thinking(
-    tool_name: str, tool_output: Any, last_items: list[str],
+    tool_name: str,
+    tool_output: Any,
+    last_items: list[str],
 ) -> tuple[str, list[str]]:
     del tool_name
     items = last_items

surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_search/emission.py

@@ -28,11 +28,7 @@ def iter_completion_emission_frames(
         xml,
     ):
         chunk_url, content = m.group(1).strip(), m.group(2).strip()
-        if (
-            chunk_url.startswith("http")
-            and chunk_url in citations
-            and content
-        ):
+        if chunk_url.startswith("http") and chunk_url in citations and content:
             citations[chunk_url]["snippet"] = (
                 content[:200] + "…" if len(content) > 200 else content
             )