From a7d7155039ce0625820e2c2ab9016143ad6e04d8 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 4 Jun 2026 21:25:39 +0200 Subject: [PATCH] refactor(agents): colocate main_agent-only kernel into main_agent/ Move modules out of agents/shared/ that are consumed by a single package (main_agent), placing each next to its only consumer instead of in a "shared" grab-bag: - agent_cache.py -> main_agent/runtime/agent_cache_store.py - connector_searchable_types.py -> main_agent/runtime/ - plugin_loader.py + plugins/ -> main_agent/plugins/ - skills/ + skills_backends.py -> main_agent/skills/ - tools/invalid_tool.py -> main_agent/tools/ Drop the skills_backends re-export from the shared middleware barrel and repoint all consumers + tests. No behavior change; import-all, error-contract, and the moved tests stay green. --- .../main_agent/middleware/plugins.py | 8 ++++---- .../main_agent/middleware/skills.py | 7 ++----- .../main_agent}/plugins/__init__.py | 0 .../main_agent/plugins/loader.py} | 0 .../main_agent}/plugins/year_substituter.py | 4 ++-- .../main_agent/runtime/agent_cache.py | 12 ++++++------ .../main_agent/runtime/agent_cache_store.py} | 0 .../runtime}/connector_searchable_types.py | 0 .../main_agent/runtime/factory.py | 6 ++---- .../main_agent}/skills/__init__.py | 0 .../main_agent/skills/backends.py} | 6 +++--- .../main_agent}/skills/builtin/__init__.py | 0 .../skills/builtin/email-drafting/SKILL.md | 0 .../skills/builtin/kb-research/SKILL.md | 0 .../skills/builtin/meeting-prep/SKILL.md | 0 .../skills/builtin/report-writing/SKILL.md | 0 .../skills/builtin/slack-summary/SKILL.md | 0 .../main_agent}/tools/invalid_tool.py | 0 .../app/agents/shared/middleware/__init__.py | 10 ---------- .../agents/shared/middleware/tool_call_repair.py | 4 +++- .../unit/agents/new_chat/test_agent_cache.py | 2 +- .../unit/agents/new_chat/test_plugin_loader.py | 16 ++++++++-------- .../unit/agents/new_chat/test_skills_backends.py | 2 +- .../agents/new_chat/test_tool_call_repair.py | 2 +- 24 files changed, 33 insertions(+), 46 deletions(-) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/plugins/__init__.py (100%) rename surfsense_backend/app/agents/{shared/plugin_loader.py => multi_agent_chat/main_agent/plugins/loader.py} (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/plugins/year_substituter.py (95%) rename surfsense_backend/app/agents/{shared/agent_cache.py => multi_agent_chat/main_agent/runtime/agent_cache_store.py} (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent/runtime}/connector_searchable_types.py (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/__init__.py (100%) rename surfsense_backend/app/agents/{shared/middleware/skills_backends.py => multi_agent_chat/main_agent/skills/backends.py} (98%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/__init__.py (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/email-drafting/SKILL.md (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/kb-research/SKILL.md (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/meeting-prep/SKILL.md (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/report-writing/SKILL.md (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/skills/builtin/slack-summary/SKILL.md (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/main_agent}/tools/invalid_tool.py (100%) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/plugins.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/plugins.py index 25a791fff..6abad6765 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/plugins.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/plugins.py @@ -7,15 +7,15 @@ from typing import Any from langchain_core.language_models import BaseChatModel +from app.agents.multi_agent_chat.shared.middleware.flags import enabled from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.plugin_loader import ( +from app.db import ChatVisibility + +from ..plugins.loader import ( PluginContext, load_allowed_plugin_names_from_env, load_plugin_middlewares, ) -from app.db import ChatVisibility - -from app.agents.multi_agent_chat.shared.middleware.flags import enabled def build_plugin_middlewares( diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/skills.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/skills.py index fe415cf9e..c1d583d6f 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/skills.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/skills.py @@ -6,14 +6,11 @@ import logging from deepagents.middleware.skills import SkillsMiddleware +from app.agents.multi_agent_chat.shared.middleware.flags import enabled from app.agents.shared.feature_flags import AgentFeatureFlags from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.shared.middleware import ( - build_skills_backend_factory, - default_skills_sources, -) -from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from ..skills.backends import build_skills_backend_factory, default_skills_sources def build_skills_mw( diff --git a/surfsense_backend/app/agents/shared/plugins/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/__init__.py similarity index 100% rename from surfsense_backend/app/agents/shared/plugins/__init__.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/__init__.py diff --git a/surfsense_backend/app/agents/shared/plugin_loader.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/loader.py similarity index 100% rename from surfsense_backend/app/agents/shared/plugin_loader.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/loader.py diff --git a/surfsense_backend/app/agents/shared/plugins/year_substituter.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/year_substituter.py similarity index 95% rename from surfsense_backend/app/agents/shared/plugins/year_substituter.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/year_substituter.py index c0095ddd7..bc0ef87a1 100644 --- a/surfsense_backend/app/agents/shared/plugins/year_substituter.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/plugins/year_substituter.py @@ -17,7 +17,7 @@ Wire-up in ``pyproject.toml`` (illustrative; the in-repo plugin doesn't need this -- it's already on the import path):: [project.entry-points."surfsense.plugins"] - year_substituter = "app.agents.shared.plugins.year_substituter:make_middleware" + year_substituter = "app.agents.multi_agent_chat.main_agent.plugins.year_substituter:make_middleware" """ from __future__ import annotations @@ -34,7 +34,7 @@ if TYPE_CHECKING: # pragma: no cover - type-only from langchain_core.messages import ToolMessage from langgraph.types import Command - from app.agents.shared.plugin_loader import PluginContext + from .loader import PluginContext logger = logging.getLogger(__name__) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py index 4d726abb6..bed40856d 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py @@ -10,18 +10,18 @@ from langchain_core.language_models import BaseChatModel from langchain_core.tools import BaseTool from langgraph.types import Checkpointer -from app.agents.shared.agent_cache import ( +from app.agents.shared.feature_flags import AgentFeatureFlags +from app.agents.shared.filesystem_selection import FilesystemMode +from app.db import ChatVisibility + +from ..graph.compile_graph_sync import build_compiled_agent_graph_sync +from .agent_cache_store import ( flags_signature, get_cache, stable_hash, system_prompt_hash, tools_signature, ) -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.filesystem_selection import FilesystemMode -from app.db import ChatVisibility - -from ..graph.compile_graph_sync import build_compiled_agent_graph_sync def mcp_signature(mcp_tools_by_agent: dict[str, list[BaseTool]]) -> str: diff --git a/surfsense_backend/app/agents/shared/agent_cache.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache_store.py similarity index 100% rename from surfsense_backend/app/agents/shared/agent_cache.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache_store.py diff --git a/surfsense_backend/app/agents/shared/connector_searchable_types.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/connector_searchable_types.py similarity index 100% rename from surfsense_backend/app/agents/shared/connector_searchable_types.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/connector_searchable_types.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 8e1087357..d692a3d94 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -19,15 +19,11 @@ from app.agents.multi_agent_chat.subagents import ( from app.agents.multi_agent_chat.subagents.mcp_tools.index import ( load_mcp_tools_by_connector, ) -from app.agents.shared.connector_searchable_types import ( - map_connectors_to_searchable_types, -) from app.agents.shared.feature_flags import AgentFeatureFlags, get_flags from app.agents.shared.filesystem_backends import build_backend_resolver from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection from app.agents.shared.llm_config import AgentConfig from app.agents.shared.prompt_caching import apply_litellm_prompt_caching -from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME, invalid_tool from app.db import ChatVisibility from app.services.connector_service import ConnectorService from app.services.user_tool_allowlist import ( @@ -41,8 +37,10 @@ from ..tools import ( MAIN_AGENT_SURFSENSE_TOOL_NAMES, MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED, ) +from ..tools.invalid_tool import INVALID_TOOL_NAME, invalid_tool from ..tools.registry import build_main_agent_tools from .agent_cache import build_agent_with_cache +from .connector_searchable_types import map_connectors_to_searchable_types _perf_log = get_perf_logger() diff --git a/surfsense_backend/app/agents/shared/skills/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/__init__.py similarity index 100% rename from surfsense_backend/app/agents/shared/skills/__init__.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/__init__.py diff --git a/surfsense_backend/app/agents/shared/middleware/skills_backends.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/backends.py similarity index 98% rename from surfsense_backend/app/agents/shared/middleware/skills_backends.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/backends.py index 76a1e7f49..51528e686 100644 --- a/surfsense_backend/app/agents/shared/middleware/skills_backends.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/backends.py @@ -59,10 +59,10 @@ _MAX_SKILL_FILE_SIZE = 10 * 1024 * 1024 def _default_builtin_root() -> Path: """Return the absolute path to the bundled builtin skills directory. - Located at ``app/agents/shared/skills/builtin/`` relative to this module - (this module lives at ``app/agents/shared/middleware/skills_backends.py``). + Located at ``builtin/`` next to this module (this module lives at + ``app/agents/multi_agent_chat/main_agent/skills/backends.py``). """ - return (Path(__file__).resolve().parent.parent / "skills" / "builtin").resolve() + return (Path(__file__).resolve().parent / "builtin").resolve() class BuiltinSkillsBackend(BackendProtocol): diff --git a/surfsense_backend/app/agents/shared/skills/builtin/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/__init__.py similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/__init__.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/__init__.py diff --git a/surfsense_backend/app/agents/shared/skills/builtin/email-drafting/SKILL.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/email-drafting/SKILL.md similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/email-drafting/SKILL.md rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/email-drafting/SKILL.md diff --git a/surfsense_backend/app/agents/shared/skills/builtin/kb-research/SKILL.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/kb-research/SKILL.md similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/kb-research/SKILL.md rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/kb-research/SKILL.md diff --git a/surfsense_backend/app/agents/shared/skills/builtin/meeting-prep/SKILL.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/meeting-prep/SKILL.md similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/meeting-prep/SKILL.md rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/meeting-prep/SKILL.md diff --git a/surfsense_backend/app/agents/shared/skills/builtin/report-writing/SKILL.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/report-writing/SKILL.md similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/report-writing/SKILL.md rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/report-writing/SKILL.md diff --git a/surfsense_backend/app/agents/shared/skills/builtin/slack-summary/SKILL.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/slack-summary/SKILL.md similarity index 100% rename from surfsense_backend/app/agents/shared/skills/builtin/slack-summary/SKILL.md rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/skills/builtin/slack-summary/SKILL.md diff --git a/surfsense_backend/app/agents/shared/tools/invalid_tool.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/tools/invalid_tool.py similarity index 100% rename from surfsense_backend/app/agents/shared/tools/invalid_tool.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/tools/invalid_tool.py diff --git a/surfsense_backend/app/agents/shared/middleware/__init__.py b/surfsense_backend/app/agents/shared/middleware/__init__.py index 234a7ee29..fb6eacfdb 100644 --- a/surfsense_backend/app/agents/shared/middleware/__init__.py +++ b/surfsense_backend/app/agents/shared/middleware/__init__.py @@ -45,12 +45,6 @@ from app.agents.shared.middleware.noop_injection import NoopInjectionMiddleware from app.agents.shared.middleware.otel_span import OtelSpanMiddleware from app.agents.shared.middleware.permission import PermissionMiddleware from app.agents.shared.middleware.retry_after import RetryAfterMiddleware -from app.agents.shared.middleware.skills_backends import ( - BuiltinSkillsBackend, - SearchSpaceSkillsBackend, - build_skills_backend_factory, - default_skills_sources, -) from app.agents.shared.middleware.tool_call_repair import ( ToolCallNameRepairMiddleware, ) @@ -58,7 +52,6 @@ from app.agents.shared.middleware.tool_call_repair import ( __all__ = [ "ActionLogMiddleware", "AnonymousDocumentMiddleware", - "BuiltinSkillsBackend", "BusyMutexMiddleware", "ClearToolUsesEdit", "DedupHITLToolCallsMiddleware", @@ -74,14 +67,11 @@ __all__ = [ "OtelSpanMiddleware", "PermissionMiddleware", "RetryAfterMiddleware", - "SearchSpaceSkillsBackend", "SpillToBackendEdit", "SpillingContextEditingMiddleware", "SurfSenseCompactionMiddleware", "ToolCallNameRepairMiddleware", "ToolDefinition", - "build_skills_backend_factory", "commit_staged_filesystem_state", "create_surfsense_compaction_middleware", - "default_skills_sources", ] diff --git a/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py b/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py index 7a0e7c4c6..93130ab30 100644 --- a/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py +++ b/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py @@ -120,7 +120,9 @@ class ToolCallNameRepairMiddleware( # Stage 2 — invalid fallback # Local import keeps the middleware module import-light and avoids any # tools <-> middleware import-order coupling at module scope. - from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME + from app.agents.multi_agent_chat.main_agent.tools.invalid_tool import ( + INVALID_TOOL_NAME, + ) if INVALID_TOOL_NAME in registered: original_args = call.get("args") or {} diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_agent_cache.py b/surfsense_backend/tests/unit/agents/new_chat/test_agent_cache.py index 48eb86a2c..d50d7b91c 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_agent_cache.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_agent_cache.py @@ -16,7 +16,7 @@ from dataclasses import dataclass import pytest -from app.agents.shared.agent_cache import ( +from app.agents.multi_agent_chat.main_agent.runtime.agent_cache_store import ( flags_signature, reload_for_tests, stable_hash, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_plugin_loader.py b/surfsense_backend/tests/unit/agents/new_chat/test_plugin_loader.py index fa7ec223b..2a1b4c51b 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_plugin_loader.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_plugin_loader.py @@ -6,13 +6,13 @@ from unittest.mock import MagicMock, patch from langchain.agents.middleware import AgentMiddleware -from app.agents.shared.plugin_loader import ( +from app.agents.multi_agent_chat.main_agent.plugins.loader import ( PLUGIN_ENTRY_POINT_GROUP, PluginContext, load_allowed_plugin_names_from_env, load_plugin_middlewares, ) -from app.agents.shared.plugins.year_substituter import ( +from app.agents.multi_agent_chat.main_agent.plugins.year_substituter import ( _YearSubstituterMiddleware, make_middleware as year_substituter_factory, ) @@ -66,7 +66,7 @@ class TestPluginLoaderBasics: ep = _FakeEntryPoint("dangerous_plugin", factory) with patch( - "app.agents.shared.plugin_loader.entry_points", + "app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=[ep], ): result = load_plugin_middlewares( @@ -78,7 +78,7 @@ class TestPluginLoaderBasics: def test_loads_allowlisted_plugin(self) -> None: ep = _FakeEntryPoint("year_substituter", year_substituter_factory) with patch( - "app.agents.shared.plugin_loader.entry_points", + "app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=[ep], ): result = load_plugin_middlewares( @@ -95,7 +95,7 @@ class TestPluginLoaderIsolation: ep = _FakeEntryPoint("buggy", crashing_factory) with patch( - "app.agents.shared.plugin_loader.entry_points", + "app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=[ep], ): result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"buggy"}) @@ -107,7 +107,7 @@ class TestPluginLoaderIsolation: ep = _FakeEntryPoint("liar", bad_factory) with patch( - "app.agents.shared.plugin_loader.entry_points", + "app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=[ep], ): result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"liar"}) @@ -121,7 +121,7 @@ class TestPluginLoaderIsolation: raise ImportError("cannot import") with patch( - "app.agents.shared.plugin_loader.entry_points", + "app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=[_BrokenEP()], ): result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"broken"}) @@ -137,7 +137,7 @@ class TestPluginLoaderIsolation: _FakeEntryPoint("crashing", crashing_factory), _FakeEntryPoint("ok", year_substituter_factory), ] - with patch("app.agents.shared.plugin_loader.entry_points", return_value=eps): + with patch("app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=eps): result = load_plugin_middlewares( _ctx(), allowed_plugin_names={"crashing", "ok"} ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_skills_backends.py b/surfsense_backend/tests/unit/agents/new_chat/test_skills_backends.py index b49cdfa1d..477b4f7fc 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_skills_backends.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_skills_backends.py @@ -7,7 +7,7 @@ from pathlib import Path import pytest -from app.agents.shared.middleware.skills_backends import ( +from app.agents.multi_agent_chat.main_agent.skills.backends import ( SKILLS_BUILTIN_PREFIX, SKILLS_SPACE_PREFIX, BuiltinSkillsBackend, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py b/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py index 068d8415b..eb2704546 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py @@ -5,10 +5,10 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage +from app.agents.multi_agent_chat.main_agent.tools.invalid_tool import INVALID_TOOL_NAME from app.agents.shared.middleware.tool_call_repair import ( ToolCallNameRepairMiddleware, ) -from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME pytestmark = pytest.mark.unit