From 390dc9307fac6733df8848b043cac7d7fabd1365 Mon Sep 17 00:00:00 2001
From: CREDO23 <bakerathierry@gmail.com>
Date: Tue, 5 May 2026 20:53:59 +0200
Subject: [PATCH] refactor(multi-agent): add main-agent knowledge middleware
 factories

---
 .../middleware/main_agent/kb_persistence.py   | 23 ++++++++++++++++
 .../main_agent/knowledge_priority.py          | 27 +++++++++++++++++++
 .../middleware/main_agent/knowledge_tree.py   | 23 ++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py
 create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py
 create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py

diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py
new file mode 100644
index 000000000..4b27581e7
--- /dev/null
+++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py
@@ -0,0 +1,23 @@
+"""Commit staged cloud filesystem mutations to Postgres at end of turn."""
+
+from __future__ import annotations
+
+from app.agents.new_chat.filesystem_selection import FilesystemMode
+from app.agents.new_chat.middleware import KnowledgeBasePersistenceMiddleware
+
+
+def build_kb_persistence_mw(
+    *,
+    filesystem_mode: FilesystemMode,
+    search_space_id: int,
+    user_id: str | None,
+    thread_id: int | None,
+) -> KnowledgeBasePersistenceMiddleware | None:
+    if filesystem_mode != FilesystemMode.CLOUD:
+        return None
+    return KnowledgeBasePersistenceMiddleware(
+        search_space_id=search_space_id,
+        created_by_id=user_id,
+        filesystem_mode=filesystem_mode,
+        thread_id=thread_id,
+    )
diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py
new file mode 100644
index 000000000..395d2a7af
--- /dev/null
+++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py
@@ -0,0 +1,27 @@
+"""KB priority planner: <priority_documents> injection."""
+
+from __future__ import annotations
+
+from langchain_core.language_models import BaseChatModel
+
+from app.agents.new_chat.filesystem_selection import FilesystemMode
+from app.agents.new_chat.middleware import KnowledgePriorityMiddleware
+
+
+def build_knowledge_priority_mw(
+    *,
+    llm: BaseChatModel,
+    search_space_id: int,
+    filesystem_mode: FilesystemMode,
+    available_connectors: list[str] | None,
+    available_document_types: list[str] | None,
+    mentioned_document_ids: list[int] | None,
+) -> KnowledgePriorityMiddleware:
+    return KnowledgePriorityMiddleware(
+        llm=llm,
+        search_space_id=search_space_id,
+        filesystem_mode=filesystem_mode,
+        available_connectors=available_connectors,
+        available_document_types=available_document_types,
+        mentioned_document_ids=mentioned_document_ids,
+    )
diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py
new file mode 100644
index 000000000..404082401
--- /dev/null
+++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py
@@ -0,0 +1,23 @@
+"""<workspace_tree> injection (cloud only)."""
+
+from __future__ import annotations
+
+from langchain_core.language_models import BaseChatModel
+
+from app.agents.new_chat.filesystem_selection import FilesystemMode
+from app.agents.new_chat.middleware import KnowledgeTreeMiddleware
+
+
+def build_knowledge_tree_mw(
+    *,
+    filesystem_mode: FilesystemMode,
+    search_space_id: int,
+    llm: BaseChatModel,
+) -> KnowledgeTreeMiddleware | None:
+    if filesystem_mode != FilesystemMode.CLOUD:
+        return None
+    return KnowledgeTreeMiddleware(
+        search_space_id=search_space_id,
+        filesystem_mode=filesystem_mode,
+        llm=llm,
+    )