Refactor multi-agent supervisor prompts and wiring; thread citations into streaming.

surfsense_backend/app/agents/multi_agent_chat/core/bindings/__init__.py

@@ -1,4 +1,4 @@
-"""Search-space / DB kwargs shared by ``new_chat`` tool factories (distinct from ``expert_agent.connectors`` integrations)."""
+"""Search-space / DB kwargs shared by main-chat tool factories (distinct from ``expert_agent.connectors`` integrations)."""
 
 from app.agents.multi_agent_chat.core.bindings.binding import connector_binding
 

surfsense_backend/app/agents/multi_agent_chat/core/bindings/binding.py

@@ -1,4 +1,4 @@
-"""Shared kwargs dict for ``new_chat`` tool factories (DB session + search space + user)."""
+"""Shared kwargs dict for main-chat tool factories (DB session + search space + user)."""
 
 from __future__ import annotations
 

surfsense_backend/app/agents/multi_agent_chat/core/mcp_partition.py

@@ -1,4 +1,4 @@
-"""Partition MCP tools onto multi-agent expert routes without modifying ``new_chat``.
+"""Partition MCP tools onto multi-agent expert routes (read-only; does not change the MCP loader).
 
 Uses the same connector discovery shape as ``load_mcp_tools`` (copied query below). Tools come from
 ``app.agents.new_chat.tools.mcp_tool.load_mcp_tools``; routing uses metadata already set there:
@@ -61,7 +61,7 @@ async def fetch_mcp_connector_metadata_maps(
 ) -> tuple[dict[int, str], dict[str, str]]:
     """Read-only copy of connector discovery used alongside ``load_mcp_tools``.
 
-    Same filter as ``new_chat.tools.mcp_tool.load_mcp_tools`` (connectors with ``server_config``).
+    Same filter as :func:`app.agents.new_chat.tools.mcp_tool.load_mcp_tools` (connectors with ``server_config``).
     """
     result = await session.execute(
         select(SearchSourceConnector).filter(
@@ -90,7 +90,7 @@ def partition_mcp_tools_by_expert_route(
 ) -> dict[str, list[BaseTool]]:
     """Bucket MCP tools by expert route key. Supervisor never receives raw MCP tools.
 
-    Same inclusion rule as ``new_chat.tools.registry.build_tools_async``: all tools returned by
+    Same inclusion rule as :func:`app.agents.new_chat.tools.registry.build_tools_async`: all tools returned by
     ``load_mcp_tools`` are partitioned — connector availability for **registry** builtins is handled via
     ``get_connector_gated_tools`` / routing gates; MCP tools are not pre-filtered by inventory here.
     """

surfsense_backend/app/agents/multi_agent_chat/core/registry/__init__.py

@@ -1,4 +1,4 @@
-"""``new_chat`` tool registry grouping + dependency bundles for domain slices."""
+"""Main chat tool registry grouping + dependency bundles for domain slices."""
 
 from app.agents.multi_agent_chat.core.registry.categories import (
     REGISTRY_ROUTING_CATEGORY_KEYS,

surfsense_backend/app/agents/multi_agent_chat/core/registry/dependencies.py

@@ -1,4 +1,4 @@
-"""Dependency dict for :func:`app.agents.new_chat.tools.registry.build_tools` in multi-agent graphs."""
+"""Dependency dict for :func:`app.agents.new_chat.tools.registry.build_tools` on expert subgraphs."""
 
 from __future__ import annotations
 

surfsense_backend/app/agents/multi_agent_chat/core/registry/subset.py

@@ -1,4 +1,4 @@
-"""Build :mod:`new_chat` registry tool subsets for multi-agent domain slices."""
+"""Build registry tool subsets (``app.agents.new_chat.tools.registry``) for multi-agent domain slices."""
 
 from __future__ import annotations
 

surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/deliverables/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_deliverables_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``deliverables`` category."""
+    """Registry-backed tools for the ``deliverables`` category."""
     return build_registry_tools_for_category(dependencies, "deliverables")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/memory/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_memory_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``memory`` category."""
+    """Registry-backed tools for the ``memory`` category."""
     return build_registry_tools_for_category(dependencies, "memory")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/research/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_research_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``research`` category."""
+    """Registry-backed tools for the ``research`` category."""
     return build_registry_tools_for_category(dependencies, "research")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/calendar/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_calendar_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``calendar`` category."""
+    """Registry-backed tools for the ``calendar`` category."""
     return build_registry_tools_for_category(dependencies, "calendar")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/confluence/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_confluence_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``confluence`` category."""
+    """Registry-backed tools for the ``confluence`` category."""
     return build_registry_tools_for_category(dependencies, "confluence")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/discord/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_discord_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``discord`` category."""
+    """Registry-backed tools for the ``discord`` category."""
     return build_registry_tools_for_category(dependencies, "discord")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/dropbox/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_dropbox_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``dropbox`` category."""
+    """Registry-backed tools for the ``dropbox`` category."""
     return build_registry_tools_for_category(dependencies, "dropbox")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/gmail/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_gmail_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``gmail`` category."""
+    """Registry-backed tools for the ``gmail`` category."""
     return build_registry_tools_for_category(dependencies, "gmail")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/google_drive/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_google_drive_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``google_drive`` category."""
+    """Registry-backed tools for the ``google_drive`` category."""
     return build_registry_tools_for_category(dependencies, "google_drive")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/luma/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_luma_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``luma`` category."""
+    """Registry-backed tools for the ``luma`` category."""
     return build_registry_tools_for_category(dependencies, "luma")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/notion/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_notion_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``notion`` category."""
+    """Registry-backed tools for the ``notion`` category."""
     return build_registry_tools_for_category(dependencies, "notion")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/onedrive/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_onedrive_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``onedrive`` category."""
+    """Registry-backed tools for the ``onedrive`` category."""
     return build_registry_tools_for_category(dependencies, "onedrive")

surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/teams/slice_tools.py

@@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c
 
 
 def build_teams_tools(dependencies: dict[str, Any]) -> list[BaseTool]:
-    """Tools from ``new_chat`` registry: ``teams`` category."""
+    """Registry-backed tools for the ``teams`` category."""
     return build_registry_tools_for_category(dependencies, "teams")

surfsense_backend/app/agents/multi_agent_chat/integration/create_multi_agent_chat.py

@@ -1,4 +1,4 @@
-"""Single entry: SurfSense connectors + multi-agent stack → compiled supervisor graph."""
+"""Build the multi-agent supervisor graph: MCP partition, registry, routing tools, optional SurfSense middleware."""
 
 from __future__ import annotations
 
@@ -11,14 +11,6 @@ from langchain_core.tools import BaseTool
 from langgraph.types import Checkpointer
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.agents.new_chat.chat_deepagent import _map_connectors_to_searchable_types
-from app.agents.new_chat.context import SurfSenseContextSchema
-from app.agents.new_chat.feature_flags import get_flags
-from app.agents.new_chat.filesystem_backends import build_backend_resolver
-from app.agents.new_chat.filesystem_selection import FilesystemSelection
-from app.agents.new_chat.tools.mcp_tool import load_mcp_tools
-from app.db import ChatVisibility
-
 from app.agents.multi_agent_chat.core.mcp_partition import (
     fetch_mcp_connector_metadata_maps,
     partition_mcp_tools_by_expert_route,
@@ -27,14 +19,95 @@ from app.agents.multi_agent_chat.core.registry.dependencies import (
     build_registry_dependencies,
     coerce_thread_id_for_registry,
 )
-from app.agents.multi_agent_chat.middleware.supervisor_stack import build_supervisor_middleware_stack
-from app.agents.multi_agent_chat.routing.supervisor_routing import build_supervisor_routing_tools
+from app.agents.multi_agent_chat.middleware.supervisor_stack import (
+    build_supervisor_middleware_stack,
+)
+from app.agents.multi_agent_chat.routing.supervisor_routing import (
+    build_supervisor_routing_tools,
+)
 from app.agents.multi_agent_chat.supervisor import build_supervisor_agent
+from app.agents.new_chat.chat_deepagent import _map_connectors_to_searchable_types
+from app.agents.new_chat.context import SurfSenseContextSchema
+from app.agents.new_chat.feature_flags import get_flags
+from app.agents.new_chat.filesystem_backends import build_backend_resolver
+from app.agents.new_chat.filesystem_selection import FilesystemSelection
+from app.agents.new_chat.tools.mcp_tool import load_mcp_tools
+from app.db import ChatVisibility
 
 logger = logging.getLogger(__name__)
 
 
-def _compile_supervisor_chat_blocking(
+async def _discover_connectors_and_doc_types(
+    *,
+    connector_service: Any | None,
+    search_space_id: int,
+    available_connectors: list[str] | None,
+    available_document_types: list[str] | None,
+) -> tuple[list[str] | None, list[str] | None]:
+    """Fill connector / document-type lists from ``connector_service`` when callers omit them."""
+    connectors = available_connectors
+    doc_types = available_document_types
+    if connector_service is None:
+        return connectors, doc_types
+    try:
+        if connectors is None:
+            raw = await connector_service.get_available_connectors(search_space_id)
+            if raw:
+                connectors = _map_connectors_to_searchable_types(raw)
+        if doc_types is None:
+            doc_types = await connector_service.get_available_document_types(search_space_id)
+    except Exception as exc:
+        logger.warning("Failed to discover available connectors/document types: %s", exc)
+    return connectors, doc_types
+
+
+async def _mcp_tools_by_expert_route(
+    *,
+    db_session: AsyncSession,
+    search_space_id: int,
+) -> dict[str, list[BaseTool]] | None:
+    mcp_flat = await load_mcp_tools(db_session, search_space_id)
+    id_map, name_map = await fetch_mcp_connector_metadata_maps(db_session, search_space_id)
+    return partition_mcp_tools_by_expert_route(mcp_flat, id_map, name_map)
+
+
+def _make_supervisor_routing_tools(
+    llm: BaseChatModel,
+    *,
+    db_session: AsyncSession,
+    search_space_id: int,
+    user_id: str,
+    thread_id: str | int | None,
+    firecrawl_api_key: str | None,
+    connector_service: Any | None,
+    available_connectors: list[str] | None,
+    available_document_types: list[str] | None,
+    thread_visibility: ChatVisibility,
+    mcp_tools_by_route: dict[str, list[BaseTool]] | None,
+) -> list[BaseTool]:
+    registry_dependencies = build_registry_dependencies(
+        db_session=db_session,
+        search_space_id=search_space_id,
+        user_id=user_id,
+        thread_id=thread_id,
+        llm=llm,
+        firecrawl_api_key=firecrawl_api_key,
+        connector_service=connector_service,
+        available_connectors=available_connectors,
+        available_document_types=available_document_types,
+        thread_visibility=thread_visibility,
+    )
+    return build_supervisor_routing_tools(
+        llm,
+        registry_dependencies=registry_dependencies,
+        include_deliverables=coerce_thread_id_for_registry(thread_id) is not None,
+        mcp_tools_by_route=mcp_tools_by_route,
+        available_connectors=available_connectors,
+        thread_visibility=thread_visibility,
+    )
+
+
+def _compile_supervisor_agent_sync(
     *,
     llm: BaseChatModel,
     routing_tools: list[BaseTool],
@@ -43,16 +116,16 @@ def _compile_supervisor_chat_blocking(
     filesystem_mode: Any,
     search_space_id: int,
     user_id: str,
-    thread_id: str | None,
+    thread_id: str | int | None,
     thread_visibility: ChatVisibility,
     anon_session_id: str | None,
     available_connectors: list[str] | None,
     available_document_types: list[str] | None,
     mentioned_document_ids: list[int] | None,
     max_input_tokens: int | None,
+    citations_enabled: bool,
 ) -> Any:
-    """CPU-heavy: middleware assembly + ``create_agent`` (runs in a worker thread)."""
-    flags = get_flags()
+    """CPU-heavy: middleware stack + ``create_agent`` (intended for ``asyncio.to_thread``)."""
     middleware = build_supervisor_middleware_stack(
         llm=llm,
         tools=routing_tools,
@@ -67,7 +140,7 @@ def _compile_supervisor_chat_blocking(
         available_document_types=available_document_types,
         mentioned_document_ids=mentioned_document_ids,
         max_input_tokens=max_input_tokens,
-        flags=flags,
+        flags=get_flags(),
     )
     return build_supervisor_agent(
         llm,
@@ -76,6 +149,7 @@ def _compile_supervisor_chat_blocking(
         thread_visibility=thread_visibility,
         middleware=middleware,
         context_schema=SurfSenseContextSchema,
+        citations_enabled=citations_enabled,
     )
 
 
@@ -98,16 +172,17 @@ async def create_multi_agent_chat(
     mentioned_document_ids: list[int] | None = None,
     max_input_tokens: int | None = None,
     surfsense_stack: bool = True,
+    citations_enabled: bool | None = None,
 ):
-    """Build the full multi-agent chat graph (supervisor + domain subgraphs via routing tools).
+    """Build the full multi-agent chat graph (supervisor + expert subgraphs via routing tools).
 
     **Builtins** (:mod:`expert_agent.builtins`): registry-grouped **categories** (research, memory, deliverables).
     **Connectors** (:mod:`expert_agent.connectors`): **vendor integrations** — one subgraph per route in
     ``TOOL_NAMES_BY_CATEGORY`` (e.g. calendar, confluence, discord, dropbox, gmail, google_drive, luma, notion, onedrive, teams).
 
-    MCP tools from ``new_chat`` (``load_mcp_tools``) are partitioned inside this package and attached only
+    MCP tools (via ``load_mcp_tools``) are partitioned inside this package and attached only
     to the matching expert subgraphs — not to the supervisor tool list as raw MCP calls. Inclusion matches
-    ``new_chat.tools.registry.build_tools_async``: all tools returned by ``load_mcp_tools`` are merged
+    ``app.agents.new_chat.tools.registry.build_tools_async``: all tools returned by ``load_mcp_tools`` are merged
     after partitioning (no extra inventory filter on MCP). Connector routing uses ``available_connectors``:
     pass explicitly, or provide ``connector_service`` so lists are resolved like
     ``create_surfsense_deep_agent`` (``get_available_connectors`` → searchable types).
@@ -115,57 +190,38 @@ async def create_multi_agent_chat(
     Deliverables (thread-scoped reports, podcasts, etc.) are registered only when ``thread_id`` is set.
 
     When ``surfsense_stack`` is true (default), the supervisor uses the same SurfSense middleware shell as
-    ``new_chat`` (KB priority/tree, filesystem, compaction, permissions, etc.) except ``SubAgentMiddleware`` /
-    ``task``, since experts are separate graphs behind routing tools. Graph compilation runs in
-    ``asyncio.to_thread`` so heavy CPU work does not block the event loop.
+    the main single-agent chat (KB priority/tree, filesystem, compaction, permissions, etc.) except
+    ``SubAgentMiddleware`` / ``task``, since experts are separate graphs behind routing tools. Graph
+    compilation runs in ``asyncio.to_thread`` so heavy CPU work does not block the event loop.
+
+    ``citations_enabled``: when ``None``, defaults to ``True`` (same default as ``AgentConfig`` / main chat).
     """
-    resolved_connectors = available_connectors
-    resolved_doc_types = available_document_types
-    if connector_service is not None:
-        try:
-            if resolved_connectors is None:
-                connector_types = await connector_service.get_available_connectors(
-                    search_space_id
-                )
-                if connector_types:
-                    resolved_connectors = _map_connectors_to_searchable_types(
-                        connector_types
-                    )
-            if resolved_doc_types is None:
-                resolved_doc_types = (
-                    await connector_service.get_available_document_types(search_space_id)
-                )
-        except Exception as exc:
-            logger.warning(
-                "Failed to discover available connectors/document types: %s",
-                exc,
-            )
+    citations = True if citations_enabled is None else citations_enabled
+    connectors, doc_types = await _discover_connectors_and_doc_types(
+        connector_service=connector_service,
+        search_space_id=search_space_id,
+        available_connectors=available_connectors,
+        available_document_types=available_document_types,
+    )
 
-    mcp_tools_by_route: dict[str, list[BaseTool]] | None = None
+    mcp_by_route: dict[str, list[BaseTool]] | None = None
     if include_mcp_tools:
-        mcp_flat = await load_mcp_tools(db_session, search_space_id)
-        id_map, name_map = await fetch_mcp_connector_metadata_maps(db_session, search_space_id)
-        mcp_tools_by_route = partition_mcp_tools_by_expert_route(mcp_flat, id_map, name_map)
+        mcp_by_route = await _mcp_tools_by_expert_route(
+            db_session=db_session, search_space_id=search_space_id
+        )
 
-    registry_dependencies = build_registry_dependencies(
+    routing_tools = _make_supervisor_routing_tools(
+        llm,
         db_session=db_session,
         search_space_id=search_space_id,
         user_id=user_id,
         thread_id=thread_id,
-        llm=llm,
         firecrawl_api_key=firecrawl_api_key,
         connector_service=connector_service,
-        available_connectors=resolved_connectors,
-        available_document_types=resolved_doc_types,
-        thread_visibility=thread_visibility,
-    )
-    routing_tools = build_supervisor_routing_tools(
-        llm,
-        registry_dependencies=registry_dependencies,
-        include_deliverables=coerce_thread_id_for_registry(thread_id) is not None,
-        mcp_tools_by_route=mcp_tools_by_route,
-        available_connectors=resolved_connectors,
+        available_connectors=connectors,
+        available_document_types=doc_types,
         thread_visibility=thread_visibility,
+        mcp_tools_by_route=mcp_by_route,
     )
 
     fs_sel = filesystem_selection or FilesystemSelection()
@@ -177,10 +233,11 @@ async def create_multi_agent_chat(
             tools=routing_tools,
             checkpointer=checkpointer,
             thread_visibility=thread_visibility,
+            citations_enabled=citations,
         )
 
     return await asyncio.to_thread(
-        _compile_supervisor_chat_blocking,
+        _compile_supervisor_agent_sync,
         llm=llm,
         routing_tools=routing_tools,
         checkpointer=checkpointer,
@@ -191,8 +248,9 @@ async def create_multi_agent_chat(
         thread_id=thread_id,
         thread_visibility=thread_visibility,
         anon_session_id=anon_session_id,
-        available_connectors=resolved_connectors,
-        available_document_types=resolved_doc_types,
+        available_connectors=connectors,
+        available_document_types=doc_types,
         mentioned_document_ids=mentioned_document_ids,
         max_input_tokens=max_input_tokens,
+        citations_enabled=citations,
     )

surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py

@@ -1,4 +1,4 @@
-"""SurfSense supervisor middleware (parity with ``new_chat`` main agent, minus subagents)."""
+"""SurfSense supervisor middleware (parity with the main single-agent chat, minus subagents)."""
 
 from app.agents.multi_agent_chat.middleware.supervisor_stack import (
     build_supervisor_middleware_stack,

surfsense_backend/app/agents/multi_agent_chat/middleware/supervisor_stack.py

@@ -1,4 +1,4 @@
-"""Supervisor middleware stack matching ``new_chat`` main agent (no ``SubAgentMiddleware`` / ``task``)."""
+"""Supervisor middleware stack matching the main single-agent chat (no ``SubAgentMiddleware`` / ``task``)."""
 
 from __future__ import annotations
 

surfsense_backend/app/agents/multi_agent_chat/routing/route_connector_gate.py

@@ -1,4 +1,4 @@
-"""Gate supervisor routing tools by connected searchable connector types (aligned with ``new_chat`` KB).
+"""Gate supervisor routing tools by connected searchable connector types (aligned with main chat KB).
 
 When ``available_connectors`` is ``None``, all routes are emitted (caller did not pass an inventory).
 

surfsense_backend/app/agents/multi_agent_chat/routing/supervisor_routing.py

@@ -116,7 +116,7 @@ def build_supervisor_routing_tools(
 
     ``mcp_tools_by_route`` maps route keys to MCP tools merged into the matching expert subgraph.
 
-    When ``available_connectors`` is set (searchable connector strings, same shape as ``new_chat``),
+    When ``available_connectors`` is set (searchable connector strings, same shape as the main chat agent),
     a connector-backed route is registered only if its required searchable connector type is available.
     """
     if registry_dependencies is None:

surfsense_backend/app/agents/multi_agent_chat/supervisor/graph.py

@@ -1,80 +1,18 @@
-"""Compile the supervisor agent graph (supervisor prompt + caller-supplied routing tools)."""
+"""Compile the supervisor agent graph (LangChain ``create_agent`` + caller routing tools)."""
 
 from __future__ import annotations
 
 from collections.abc import Sequence
 from typing import Any
 
-import app.agents.multi_agent_chat.supervisor as supervisor_pkg
-
 from langchain.agents import create_agent
 from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import BaseTool
 from langgraph.types import Checkpointer
 
-from app.agents.multi_agent_chat.core.prompts import read_prompt_md
-
-_BUILTIN_SPECIALISTS: frozenset[str] = frozenset({"research", "memory", "deliverables"})
-_SPECIALIST_CAPABILITIES: dict[str, str] = {
-    "research": "external research: web lookup, source gathering, and SurfSense documentation help.",
-    "memory": "save durable long-lived memory items.",
-    "deliverables": "deliverables and shareable artifacts: reports, podcasts, video presentations, resumes, and images.",
-    "gmail": "email inbox actions: search/read emails, draft updates, send messages, and trash emails.",
-    "calendar": "scheduling actions: check availability, inspect events, create events, and update events.",
-    "google_drive": "Drive file/document actions: locate files, inspect content, and manage files/folders.",
-    "notion": "Notion page actions: create pages, update content, and delete pages.",
-    "confluence": "Confluence page actions: find/read pages and create/update pages.",
-    "dropbox": "Dropbox file storage actions: browse folders, read files, and manage file content.",
-    "onedrive": "OneDrive file storage actions: browse folders, read files, and manage file content.",
-    "discord": "Discord communication actions: read channels/threads and post replies.",
-    "teams": "Microsoft Teams communication actions: read channels/threads and post replies.",
-    "luma": "Luma event actions: list events, inspect event details, and create events.",
-    "linear": "Linear workflow actions: search/update issues and inspect projects/cycles.",
-    "jira": "Jira workflow actions: search/update issues and manage workflow transitions.",
-    "clickup": "ClickUp workflow actions: find/update tasks and lists.",
-    "airtable": "Airtable data actions: locate bases/tables and create/read/update records.",
-    "slack": "Slack communication actions: read channel/thread history and post replies.",
-    # generic_mcp specialist intentionally disabled for now.
-    # "generic_mcp": "handle tasks through user-defined custom app integration tools not covered above.",
-}
-_SPECIALIST_ORDER: tuple[str, ...] = tuple(_SPECIALIST_CAPABILITIES.keys())
-
-
-def _memory_capability_for_visibility(thread_visibility: Any | None) -> str:
-    vis = str(getattr(thread_visibility, "value", thread_visibility)).upper()
-    if vis == "SEARCH_SPACE":
-        return "team memory actions: save shared team preferences, conventions, and long-lived team facts."
-    return "user memory actions: save personal preferences, instructions, and long-lived user facts."
-
-
-def _render_available_specialists_list(
-    tools: Sequence[BaseTool],
-    *,
-    thread_visibility: Any | None,
-) -> str:
-    available_names = {
-        tool.name for tool in tools if isinstance(getattr(tool, "name", None), str)
-    }
-    capabilities = dict(_SPECIALIST_CAPABILITIES)
-    capabilities["memory"] = _memory_capability_for_visibility(thread_visibility)
-    lines: list[str] = []
-    for name in _SPECIALIST_ORDER:
-        if name in _BUILTIN_SPECIALISTS or name in available_names:
-            capability = capabilities[name]
-            lines.append(f"- {name}: {capability}")
-    return "\n".join(lines)
-
-
-def _render_supervisor_prompt(
-    template: str,
-    tools: Sequence[BaseTool],
-    *,
-    thread_visibility: Any | None,
-) -> str:
-    specialist_list = _render_available_specialists_list(
-        tools, thread_visibility=thread_visibility
-    )
-    return template.replace("{{AVAILABLE_SPECIALISTS_LIST}}", specialist_list)
+from app.agents.multi_agent_chat.supervisor.prompt_assembly import (
+    build_supervisor_system_prompt,
+)
 
 
 def build_supervisor_agent(
@@ -85,13 +23,13 @@ def build_supervisor_agent(
     thread_visibility: Any | None = None,
     middleware: Sequence[Any] | None = None,
     context_schema: Any | None = None,
+    citations_enabled: bool = True,
 ):
     """Compile the supervisor **agent** (graph). ``tools`` = output of ``build_supervisor_routing_tools``."""
-    template = read_prompt_md(supervisor_pkg.__name__, "supervisor_prompt")
-    system_prompt = _render_supervisor_prompt(
-        template,
+    system_prompt = build_supervisor_system_prompt(
         tools,
         thread_visibility=thread_visibility,
+        citations_enabled=citations_enabled,
     )
     kwargs: dict[str, Any] = {
         "system_prompt": system_prompt,

surfsense_backend/app/agents/multi_agent_chat/supervisor/prompt_assembly.py

@@ -0,0 +1,128 @@
+"""Supervisor system prompt: template load, shared agent-identity injection, specialist list."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import UTC, datetime
+from typing import Any
+
+from langchain_core.tools import BaseTool
+
+import app.agents.multi_agent_chat.supervisor as supervisor_pkg
+from app.agents.multi_agent_chat.core.prompts import read_prompt_md
+from app.agents.new_chat.prompts.composer import _build_citation_block, _read_fragment
+from app.db import ChatVisibility
+
+_MEMORY_SPECIALIST_PHRASE = "invoke the **memory** specialist"
+
+_BUILTIN_SPECIALISTS: frozenset[str] = frozenset({"research", "memory", "deliverables"})
+_SPECIALIST_CAPABILITIES: dict[str, str] = {
+    "research": "external research: web lookup, source gathering, and SurfSense documentation help.",
+    "memory": "save durable long-lived memory items.",
+    "deliverables": "deliverables and shareable artifacts: reports, podcasts, video presentations, resumes, and images.",
+    "gmail": "email inbox actions: search/read emails, draft updates, send messages, and trash emails.",
+    "calendar": "scheduling actions: check availability, inspect events, create events, and update events.",
+    "google_drive": "Drive file/document actions: locate files, inspect content, and manage files/folders.",
+    "notion": "Notion page actions: create pages, update content, and delete pages.",
+    "confluence": "Confluence page actions: find/read pages and create/update pages.",
+    "dropbox": "Dropbox file storage actions: browse folders, read files, and manage file content.",
+    "onedrive": "OneDrive file storage actions: browse folders, read files, and manage file content.",
+    "discord": "Discord communication actions: read channels/threads and post replies.",
+    "teams": "Microsoft Teams communication actions: read channels/threads and post replies.",
+    "luma": "Luma event actions: list events, inspect event details, and create events.",
+    "linear": "Linear workflow actions: search/update issues and inspect projects/cycles.",
+    "jira": "Jira workflow actions: search/update issues and manage workflow transitions.",
+    "clickup": "ClickUp workflow actions: find/update tasks and lists.",
+    "airtable": "Airtable data actions: locate bases/tables and create/read/update records.",
+    "slack": "Slack communication actions: read channel/thread history and post replies.",
+    # generic_mcp specialist intentionally disabled for now.
+    # "generic_mcp": "handle tasks through user-defined custom app integration tools not covered above.",
+}
+_SPECIALIST_ORDER: tuple[str, ...] = tuple(_SPECIALIST_CAPABILITIES.keys())
+
+
+def _normalize_chat_visibility(thread_visibility: Any | None) -> ChatVisibility:
+    if thread_visibility is None:
+        return ChatVisibility.PRIVATE
+    if thread_visibility == ChatVisibility.SEARCH_SPACE:
+        return ChatVisibility.SEARCH_SPACE
+    raw = getattr(thread_visibility, "value", thread_visibility)
+    if str(raw).upper() == "SEARCH_SPACE":
+        return ChatVisibility.SEARCH_SPACE
+    return ChatVisibility.PRIVATE
+
+
+def _identity_fragment_key(thread_visibility: Any | None) -> str:
+    """``private`` / ``team`` suffix for ``agent_*`` and ``memory_protocol_*`` fragments."""
+    return (
+        "team"
+        if _normalize_chat_visibility(thread_visibility) == ChatVisibility.SEARCH_SPACE
+        else "private"
+    )
+
+
+def _compose_identity_memory_citations(
+    *,
+    thread_visibility: Any | None,
+    citations_enabled: bool,
+) -> str:
+    """Main-chat identity, memory protocol, and citation fragments (supervisor slice only)."""
+    key = _identity_fragment_key(thread_visibility)
+    today = datetime.now(UTC).date().isoformat()
+
+    intro = _read_fragment(f"base/agent_{key}.md")
+    if intro:
+        intro = intro.format(resolved_today=today)
+
+    memory = _read_fragment(f"base/memory_protocol_{key}.md").replace(
+        "call update_memory",
+        _MEMORY_SPECIALIST_PHRASE,
+    )
+    tail = (
+        f"<system_instruction>\n{memory}\n\n</system_instruction>\n"
+        + _build_citation_block(citations_enabled)
+    )
+    return "\n\n".join(part for part in (intro.strip(), tail.strip()) if part)
+
+
+def _memory_specialist_capability(thread_visibility: Any | None) -> str:
+    vis = str(getattr(thread_visibility, "value", thread_visibility)).upper()
+    if vis == "SEARCH_SPACE":
+        return "team memory actions: save shared team preferences, conventions, and long-lived team facts."
+    return "user memory actions: save personal preferences, instructions, and long-lived user facts."
+
+
+def _specialists_markdown(
+    tools: Sequence[BaseTool],
+    *,
+    thread_visibility: Any | None,
+) -> str:
+    available_names = {
+        tool.name for tool in tools if isinstance(getattr(tool, "name", None), str)
+    }
+    capabilities = dict(_SPECIALIST_CAPABILITIES)
+    capabilities["memory"] = _memory_specialist_capability(thread_visibility)
+    lines: list[str] = []
+    for name in _SPECIALIST_ORDER:
+        if name in _BUILTIN_SPECIALISTS or name in available_names:
+            lines.append(f"- {name}: {capabilities[name]}")
+    return "\n".join(lines)
+
+
+def build_supervisor_system_prompt(
+    tools: Sequence[BaseTool],
+    *,
+    thread_visibility: Any | None,
+    citations_enabled: bool,
+) -> str:
+    """Load ``supervisor_prompt.md`` and fill placeholders."""
+    template = read_prompt_md(supervisor_pkg.__name__, "supervisor_prompt")
+    specialists = _specialists_markdown(tools, thread_visibility=thread_visibility)
+    injected = _compose_identity_memory_citations(
+        thread_visibility=thread_visibility,
+        citations_enabled=citations_enabled,
+    )
+    return template.replace("{{AVAILABLE_SPECIALISTS_LIST}}", specialists).replace(
+        "{{SUPERVISOR_BASE_INJECTION}}",
+        injected,
+    )

surfsense_backend/app/agents/multi_agent_chat/supervisor/supervisor_prompt.md

@@ -1,9 +1,8 @@
-You are SurfSense's multi-agent supervisor.
+{{SUPERVISOR_BASE_INJECTION}}
 
-<role>
-Your job is to decide whether to answer directly or delegate to one or more specialists.
-You optimize for correctness, low confusion, and minimal unnecessary delegation.
-</role>
+<supervisor_role>
+In this **multi-agent** session you also **coordinate specialists** (listed below): call a specialist only when their domain matches the need; give each call a compact, outcome-focused task; merge structured results into one clear user-facing reply. When you can satisfy the turn with your own tools and reasoning, do so without delegating.
+</supervisor_role>
 
 <available_specialists>
 Use only the specialists listed below.

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -1581,6 +1581,7 @@ async def stream_new_chat(
                 thread_visibility=visibility,
                 filesystem_selection=filesystem_selection,
                 mentioned_document_ids=mentioned_document_ids,
+                citations_enabled=agent_config.citations_enabled,
             )
         else:
             agent = await create_surfsense_deep_agent(
@@ -2305,6 +2306,7 @@ async def stream_resume_chat(
                 connector_service=connector_service,
                 thread_visibility=visibility,
                 filesystem_selection=filesystem_selection,
+                citations_enabled=agent_config.citations_enabled,
             )
         else:
             agent = await create_surfsense_deep_agent(