diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/__init__.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/__init__.py
new file mode 100644
index 000000000..b5b0267ff
--- /dev/null
+++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/__init__.py
@@ -0,0 +1,13 @@
+"""End-of-turn KB persistence middleware (main-agent only)."""
+
+from .builder import build_kb_persistence_mw
+from .middleware import (
+    KnowledgeBasePersistenceMiddleware,
+    commit_staged_filesystem_state,
+)
+
+__all__ = [
+    "KnowledgeBasePersistenceMiddleware",
+    "build_kb_persistence_mw",
+    "commit_staged_filesystem_state",
+]
diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/builder.py
similarity index 89%
rename from surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence.py
rename to surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/builder.py
index 5a64d993c..7e8e06570 100644
--- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence.py
+++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/builder.py
@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from app.agents.chat.multi_agent_chat.shared.filesystem_selection import FilesystemMode
-from app.agents.chat.multi_agent_chat.shared.middleware.kb_persistence import (
+
+from .middleware import (
     KnowledgeBasePersistenceMiddleware,
 )
 
diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/kb_persistence.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/middleware.py
similarity index 100%
rename from surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/kb_persistence.py
rename to surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/middleware.py
diff --git a/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py b/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py
index 853368783..d96144bcd 100644
--- a/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py
+++ b/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py
@@ -11,10 +11,10 @@ from __future__ import annotations
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from app.agents.chat.multi_agent_chat.shared.filesystem_selection import FilesystemMode
-from app.agents.chat.multi_agent_chat.shared.middleware.kb_persistence import (
+from app.agents.chat.multi_agent_chat.main_agent.middleware.kb_persistence import (
     commit_staged_filesystem_state,
 )
+from app.agents.chat.multi_agent_chat.shared.filesystem_selection import FilesystemMode
 from app.services.new_streaming_service import VercelStreamingService
 from app.tasks.chat.streaming.contract.file_contract import (
     contract_enforcement_active,
diff --git a/surfsense_backend/tests/unit/middleware/test_kb_persistence_filesystem_parity.py b/surfsense_backend/tests/unit/middleware/test_kb_persistence_filesystem_parity.py
index b7f9b9d10..7724a4852 100644
--- a/surfsense_backend/tests/unit/middleware/test_kb_persistence_filesystem_parity.py
+++ b/surfsense_backend/tests/unit/middleware/test_kb_persistence_filesystem_parity.py
@@ -15,7 +15,7 @@ from unittest.mock import AsyncMock
 import numpy as np
 import pytest
 
-from app.agents.chat.multi_agent_chat.shared.middleware import kb_persistence
+from app.agents.chat.multi_agent_chat.main_agent.middleware.kb_persistence import middleware as kb_persistence
 from app.db import Document
 
 
diff --git a/surfsense_backend/tests/unit/middleware/test_kb_persistence_revisions.py b/surfsense_backend/tests/unit/middleware/test_kb_persistence_revisions.py
index b25f84a62..500c6cc60 100644
--- a/surfsense_backend/tests/unit/middleware/test_kb_persistence_revisions.py
+++ b/surfsense_backend/tests/unit/middleware/test_kb_persistence_revisions.py
@@ -21,7 +21,7 @@ from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
-from app.agents.chat.multi_agent_chat.shared.middleware import kb_persistence
+from app.agents.chat.multi_agent_chat.main_agent.middleware.kb_persistence import middleware as kb_persistence
 
 pytestmark = pytest.mark.unit