fix(stream): route every agent (re)build through one helper to prevent factory drift

surfsense_backend/app/agents/multi_agent_chat/__init__.py

@@ -2,6 +2,6 @@
 
 from __future__ import annotations
 
-from .main_agent import create_surfsense_deep_agent
+from .main_agent import create_multi_agent_chat_deep_agent
 
-__all__ = ["create_surfsense_deep_agent"]
+__all__ = ["create_multi_agent_chat_deep_agent"]

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

@@ -2,6 +2,6 @@
 
 from __future__ import annotations
 
-from .runtime import create_surfsense_deep_agent
+from .runtime import create_multi_agent_chat_deep_agent
 
-__all__ = ["create_surfsense_deep_agent"]
+__all__ = ["create_multi_agent_chat_deep_agent"]

surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/__init__.py

@@ -2,6 +2,6 @@
 
 from __future__ import annotations
 
-from .factory import create_surfsense_deep_agent
+from .factory import create_multi_agent_chat_deep_agent
 
-__all__ = ["create_surfsense_deep_agent"]
+__all__ = ["create_multi_agent_chat_deep_agent"]

surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py

@@ -42,7 +42,7 @@ from .agent_cache import build_agent_with_cache
 _perf_log = get_perf_logger()
 
 
-async def create_surfsense_deep_agent(
+async def create_multi_agent_chat_deep_agent(
     llm: BaseChatModel,
     search_space_id: int,
     db_session: AsyncSession,

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -28,9 +28,7 @@ from langchain_core.messages import HumanMessage
 from sqlalchemy.future import select
 from sqlalchemy.orm import selectinload
 
-from app.agents.multi_agent_chat import (
+from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent
-    create_surfsense_deep_agent as create_multi_agent_chat,
-)
 from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent
 from app.agents.new_chat.checkpointer import get_checkpointer
 from app.agents.new_chat.context import SurfSenseContextSchema
@@ -577,6 +575,43 @@ async def _preflight_llm(llm: Any) -> None:
         )
 
 
+async def _build_main_agent_for_thread(
+    agent_factory: Any,
+    *,
+    llm: Any,
+    search_space_id: int,
+    db_session: Any,
+    connector_service: ConnectorService,
+    checkpointer: Any,
+    user_id: str | None,
+    thread_id: int | None,
+    agent_config: AgentConfig | None,
+    firecrawl_api_key: str | None,
+    thread_visibility: ChatVisibility | None,
+    filesystem_selection: FilesystemSelection | None,
+    disabled_tools: list[str] | None = None,
+    mentioned_document_ids: list[int] | None = None,
+) -> Any:
+    """Single (re)build path so the agent factory cannot drift across
+    initial build, preflight repin, and mid-stream 429 recovery for one
+    ``thread_id``: a graph swap mid-turn would corrupt checkpointer state."""
+    return await agent_factory(
+        llm=llm,
+        search_space_id=search_space_id,
+        db_session=db_session,
+        connector_service=connector_service,
+        checkpointer=checkpointer,
+        user_id=user_id,
+        thread_id=thread_id,
+        agent_config=agent_config,
+        firecrawl_api_key=firecrawl_api_key,
+        thread_visibility=thread_visibility,
+        filesystem_selection=filesystem_selection,
+        disabled_tools=disabled_tools,
+        mentioned_document_ids=mentioned_document_ids,
+    )
+
+
 async def _settle_speculative_agent_build(task: asyncio.Task[Any]) -> None:
     """Wait for a discarded speculative agent build to release shared state.
 
@@ -2767,7 +2802,7 @@ async def stream_new_chat(
 
         _t0 = time.perf_counter()
         agent_factory = (
-            create_multi_agent_chat
+            create_multi_agent_chat_deep_agent
             if use_multi_agent
             else create_surfsense_deep_agent
         )
@@ -2776,7 +2811,8 @@ async def stream_new_chat(
         # if preflight reports 429 we will discard this future and rebuild
         # against the freshly pinned config below.
         agent_build_task = asyncio.create_task(
-            agent_factory(
+            _build_main_agent_for_thread(
+                agent_factory,
                 llm=llm,
                 search_space_id=search_space_id,
                 db_session=session,
@@ -2787,9 +2823,9 @@ async def stream_new_chat(
                 agent_config=agent_config,
                 firecrawl_api_key=firecrawl_api_key,
                 thread_visibility=visibility,
+                filesystem_selection=filesystem_selection,
                 disabled_tools=disabled_tools,
                 mentioned_document_ids=mentioned_document_ids,
-                filesystem_selection=filesystem_selection,
             ),
             name="agent_build:stream_new_chat",
         )
@@ -3466,7 +3502,8 @@ async def stream_new_chat(
                 title_task = None
 
                 _t0 = time.perf_counter()
-                agent = await create_surfsense_deep_agent(
+                agent = await _build_main_agent_for_thread(
+                    agent_factory,
                     llm=llm,
                     search_space_id=search_space_id,
                     db_session=session,
@@ -3477,9 +3514,9 @@ async def stream_new_chat(
                     agent_config=agent_config,
                     firecrawl_api_key=firecrawl_api_key,
                     thread_visibility=visibility,
+                    filesystem_selection=filesystem_selection,
                     disabled_tools=disabled_tools,
                     mentioned_document_ids=mentioned_document_ids,
-                    filesystem_selection=filesystem_selection,
                 )
                 _perf_log.info(
                     "[stream_new_chat] Runtime rate-limit recovery repinned "
@@ -4130,12 +4167,13 @@ async def stream_resume_chat(
 
         _t0 = time.perf_counter()
         agent_factory = (
-            create_multi_agent_chat
+            create_multi_agent_chat_deep_agent
             if _app_config.MULTI_AGENT_CHAT_ENABLED
             else create_surfsense_deep_agent
         )
         agent_build_task = asyncio.create_task(
-            agent_factory(
+            _build_main_agent_for_thread(
+                agent_factory,
                 llm=llm,
                 search_space_id=search_space_id,
                 db_session=session,
@@ -4224,7 +4262,8 @@ async def stream_resume_chat(
                         "fallback_config_id": llm_config_id,
                     },
                 )
-                agent = await agent_factory(
+                agent = await _build_main_agent_for_thread(
+                    agent_factory,
                     llm=llm,
                     search_space_id=search_space_id,
                     db_session=session,
@@ -4409,7 +4448,8 @@ async def stream_resume_chat(
                     raise stream_exc
 
                 _t0 = time.perf_counter()
-                agent = await create_surfsense_deep_agent(
+                agent = await _build_main_agent_for_thread(
+                    agent_factory,
                     llm=llm,
                     search_space_id=search_space_id,
                     db_session=session,
@@ -4421,6 +4461,7 @@ async def stream_resume_chat(
                     firecrawl_api_key=firecrawl_api_key,
                     thread_visibility=visibility,
                     filesystem_selection=filesystem_selection,
+                    disabled_tools=disabled_tools,
                 )
                 _perf_log.info(
                     "[stream_resume] Runtime rate-limit recovery repinned "