From 0081b627e9e02bfc66d587321565ba343edab388 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 5 Jun 2026 14:11:55 +0200 Subject: [PATCH] refactor(agents): move kb_persistence middleware into main_agent (owner) The KB-persistence impl lived in shared/middleware/ but no subagent uses it -- consumers are the main_agent builder and the boundary event loop. Colocate with its owner using the folder-per-middleware shape; __init__ re-exports the public surface. Tests that reached module internals now alias the .middleware submodule. main_agent/middleware/kb_persistence.py -> kb_persistence/builder.py shared/middleware/kb_persistence.py -> kb_persistence/middleware.py --- .../middleware/kb_persistence/__init__.py | 13 +++++++++++++ .../builder.py} | 3 ++- .../middleware/kb_persistence/middleware.py} | 0 .../app/tasks/chat/streaming/agent/event_loop.py | 4 ++-- .../test_kb_persistence_filesystem_parity.py | 2 +- .../middleware/test_kb_persistence_revisions.py | 2 +- 6 files changed, 19 insertions(+), 5 deletions(-) create mode 100644 surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/kb_persistence/__init__.py rename surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/{kb_persistence.py => kb_persistence/builder.py} (89%) rename surfsense_backend/app/agents/chat/multi_agent_chat/{shared/middleware/kb_persistence.py => main_agent/middleware/kb_persistence/middleware.py} (100%) 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