From 227983a104990e9edb4622fca65b98feb1542482 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 4 Jun 2026 13:00:41 +0200 Subject: [PATCH] refactor(agents): move middleware package to app/agents/shared (slice 5c) Relocate the entire new_chat/middleware/ package to the shared kernel as one cohesive unit (it is live shared infrastructure: the multi-agent stack wraps nearly every middleware via multi_agent_chat/middleware/main_agent/*, and anonymous_agent consumes it too). Flip 69 live importers across both the package-path and submodule-path forms. Shims left for the frozen single-agent stack: a package __init__ re-export plus submodule shims for permission, skills_backends, and scoped_model_fallback (the three imported via submodule path by chat_deepagent/subagents). Cycle break: importing shared.middleware previously reached back into new_chat.tools at module load, which dragged in new_chat.__init__ -> chat_deepagent -> the middleware shim -> half-initialized shared.middleware. Made action_log's ToolDefinition import TYPE_CHECKING-only and tool_call_repair's INVALID_TOOL_NAME import function-local. These tools-package back-edges fully resolve in slice 6. Asset note: skills_backends._default_builtin_root now walks to app/agents/new_chat/skills/builtin (the skills/ tree migrates in slice 7). --- .../middleware/main_agent/action_log.py | 2 +- .../middleware/main_agent/anonymous_doc.py | 2 +- .../middleware/main_agent/busy_mutex.py | 2 +- .../middleware/main_agent/context_editing.py | 2 +- .../middleware/main_agent/dedup_hitl.py | 2 +- .../middleware/main_agent/doom_loop.py | 2 +- .../middleware/main_agent/kb_persistence.py | 2 +- .../main_agent/knowledge_priority.py | 2 +- .../middleware/main_agent/knowledge_tree.py | 2 +- .../middleware/main_agent/noop_injection.py | 2 +- .../middleware/main_agent/otel.py | 2 +- .../middleware/main_agent/repair.py | 2 +- .../middleware/main_agent/skills.py | 2 +- .../middleware/shared/compaction.py | 2 +- .../filesystem/middleware/path_resolution.py | 2 +- .../filesystem/tools/edit_file/index.py | 2 +- .../filesystem/tools/list_tree/index.py | 2 +- .../shared/filesystem/tools/ls/index.py | 2 +- .../filesystem/tools/move_file/helpers.py | 2 +- .../filesystem/tools/read_file/index.py | 2 +- .../shared/filesystem/tools/rm/helpers.py | 2 +- .../shared/filesystem/tools/rmdir/helpers.py | 2 +- .../shared/kb_context_projection.py | 2 +- .../middleware/shared/memory.py | 2 +- .../middleware/shared/resilience/bundle.py | 4 +- .../middleware/shared/resilience/fallback.py | 2 +- .../middleware/shared/resilience/retry.py | 2 +- .../app/agents/new_chat/anonymous_agent.py | 2 +- .../agents/new_chat/filesystem_backends.py | 4 +- .../agents/new_chat/middleware/__init__.py | 66 +-- .../agents/new_chat/middleware/permission.py | 424 +---------------- .../middleware/scoped_model_fallback.py | 114 +---- .../new_chat/middleware/skills_backends.py | 336 +------------- .../app/agents/new_chat/tools/mcp_tool.py | 2 +- .../app/agents/new_chat/tools/registry.py | 2 +- .../app/agents/shared/middleware/__init__.py | 87 ++++ .../middleware/action_log.py | 6 +- .../middleware/anonymous_document.py | 0 .../middleware/busy_mutex.py | 0 .../middleware/compaction.py | 0 .../middleware/context_editing.py | 0 .../middleware/dedup_tool_calls.py | 0 .../middleware/doom_loop.py | 0 .../middleware/file_intent.py | 0 .../middleware/filesystem.py | 4 +- .../middleware/flatten_system.py | 0 .../middleware/kb_persistence.py | 0 .../middleware/kb_postgres_backend.py | 0 .../middleware/knowledge_search.py | 2 +- .../middleware/knowledge_tree.py | 0 .../middleware/local_folder_backend.py | 0 .../middleware/memory_injection.py | 0 .../multi_root_local_folder_backend.py | 2 +- .../middleware/noop_injection.py | 0 .../middleware/otel_span.py | 0 .../agents/shared/middleware/permission.py | 427 ++++++++++++++++++ .../middleware/retry_after.py | 0 .../middleware/scoped_model_fallback.py | 111 +++++ .../shared/middleware/skills_backends.py | 344 ++++++++++++++ .../middleware/tool_call_repair.py | 7 +- .../app/agents/shared/receipt.py | 2 +- .../app/routes/new_chat_routes.py | 2 +- .../app/tasks/chat/stream_new_chat.py | 4 +- .../tasks/chat/streaming/agent/event_loop.py | 2 +- .../tasks/chat/streaming/errors/classifier.py | 2 +- .../streaming/flows/new_chat/orchestrator.py | 2 +- .../flows/resume_chat/orchestrator.py | 2 +- .../flows/shared/rate_limit_recovery.py | 2 +- .../test_knowledge_search_date_filters.py | 6 +- .../unit/agents/new_chat/test_action_log.py | 8 +- .../unit/agents/new_chat/test_busy_mutex.py | 2 +- .../unit/agents/new_chat/test_compaction.py | 4 +- .../agents/new_chat/test_context_editing.py | 2 +- .../agents/new_chat/test_dedup_tool_calls.py | 6 +- .../new_chat/test_desktop_safety_rules.py | 2 +- .../unit/agents/new_chat/test_doom_loop.py | 2 +- .../agents/new_chat/test_flatten_system.py | 2 +- .../agents/new_chat/test_noop_injection.py | 2 +- .../unit/agents/new_chat/test_otel_span.py | 10 +- .../new_chat/test_permission_middleware.py | 2 +- .../unit/agents/new_chat/test_retry_after.py | 2 +- .../agents/new_chat/test_rm_rmdir_cloud.py | 4 +- .../agents/new_chat/test_skills_backends.py | 2 +- .../new_chat/test_specialized_subagents.py | 4 +- .../agents/new_chat/test_tool_call_repair.py | 2 +- .../middleware/test_dedup_hitl_tool_calls.py | 2 +- .../middleware/test_file_intent_middleware.py | 2 +- .../middleware/test_filesystem_backends.py | 2 +- .../middleware/test_filesystem_middleware.py | 2 +- .../test_filesystem_verification.py | 4 +- .../test_kb_persistence_filesystem_parity.py | 2 +- .../test_kb_persistence_revisions.py | 2 +- .../unit/middleware/test_knowledge_search.py | 28 +- .../unit/middleware/test_knowledge_tree.py | 2 +- .../middleware/test_local_folder_backend.py | 2 +- .../test_multi_root_local_folder_backend.py | 2 +- .../chat/streaming/test_stage_1_parity.py | 2 +- .../unit/test_stream_new_chat_contract.py | 2 +- 98 files changed, 1131 insertions(+), 999 deletions(-) create mode 100644 surfsense_backend/app/agents/shared/middleware/__init__.py rename surfsense_backend/app/agents/{new_chat => shared}/middleware/action_log.py (97%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/anonymous_document.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/busy_mutex.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/compaction.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/context_editing.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/dedup_tool_calls.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/doom_loop.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/file_intent.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/filesystem.py (99%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/flatten_system.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/kb_persistence.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/kb_postgres_backend.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/knowledge_search.py (99%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/knowledge_tree.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/local_folder_backend.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/memory_injection.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/multi_root_local_folder_backend.py (99%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/noop_injection.py (100%) rename surfsense_backend/app/agents/{new_chat => shared}/middleware/otel_span.py (100%) create mode 100644 surfsense_backend/app/agents/shared/middleware/permission.py rename surfsense_backend/app/agents/{new_chat => shared}/middleware/retry_after.py (100%) create mode 100644 surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py create mode 100644 surfsense_backend/app/agents/shared/middleware/skills_backends.py rename surfsense_backend/app/agents/{new_chat => shared}/middleware/tool_call_repair.py (95%) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py index d1fa31512..66758de8f 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import ActionLogMiddleware +from app.agents.shared.middleware import ActionLogMiddleware from app.agents.new_chat.tools.registry import BUILTIN_TOOLS from ..shared.flags import enabled diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py index 1361bde20..40090320e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import AnonymousDocumentMiddleware +from app.agents.shared.middleware import AnonymousDocumentMiddleware def build_anonymous_doc_mw( diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py index 56147d850..54f82526c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import BusyMutexMiddleware +from app.agents.shared.middleware import BusyMutexMiddleware from ..shared.flags import enabled diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py index 82fca9f14..d5188891d 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py @@ -11,7 +11,7 @@ from app.agents.multi_agent_chat.main_agent.context_prune.prune_tool_names impor safe_exclude_tools, ) from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import ( +from app.agents.shared.middleware import ( ClearToolUsesEdit, SpillingContextEditingMiddleware, SpillToBackendEdit, diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py index 66cae300b..f5536bca9 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from langchain_core.tools import BaseTool -from app.agents.new_chat.middleware import DedupHITLToolCallsMiddleware +from app.agents.shared.middleware import DedupHITLToolCallsMiddleware def build_dedup_hitl_mw(tools: Sequence[BaseTool]) -> DedupHITLToolCallsMiddleware: diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py index f1a82c206..bd380bca0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import DoomLoopMiddleware +from app.agents.shared.middleware import DoomLoopMiddleware from ..shared.flags import enabled 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 index afe72347b..4b4c8a9bb 100644 --- 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 @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import KnowledgeBasePersistenceMiddleware +from app.agents.shared.middleware import KnowledgeBasePersistenceMiddleware def build_kb_persistence_mw( 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 index 2e5d90a62..3d95a7d20 100644 --- 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 @@ -5,7 +5,7 @@ from __future__ import annotations from langchain_core.language_models import BaseChatModel from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import KnowledgePriorityMiddleware +from app.agents.shared.middleware import KnowledgePriorityMiddleware from app.services.llm_service import get_planner_llm 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 index 965a7e744..4f6a19805 100644 --- 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 @@ -5,7 +5,7 @@ from __future__ import annotations from langchain_core.language_models import BaseChatModel from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import KnowledgeTreeMiddleware +from app.agents.shared.middleware import KnowledgeTreeMiddleware def build_knowledge_tree_mw( diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py index f5371d0e1..0226da710 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import NoopInjectionMiddleware +from app.agents.shared.middleware import NoopInjectionMiddleware from ..shared.flags import enabled diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py index 73b04672b..bb655a9b0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import OtelSpanMiddleware +from app.agents.shared.middleware import OtelSpanMiddleware from ..shared.flags import enabled diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py index e4ecdd0ed..503dc6c96 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py @@ -7,7 +7,7 @@ from collections.abc import Sequence from langchain_core.tools import BaseTool from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import ToolCallNameRepairMiddleware +from app.agents.shared.middleware import ToolCallNameRepairMiddleware from ..shared.flags import enabled diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py index c13a37d6e..a7bde3bac 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py @@ -8,7 +8,7 @@ from deepagents.middleware.skills import SkillsMiddleware from app.agents.shared.feature_flags import AgentFeatureFlags from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import ( +from app.agents.shared.middleware import ( build_skills_backend_factory, default_skills_sources, ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py index b59e7d2c4..b5ca1afb4 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py @@ -7,7 +7,7 @@ from typing import Any from deepagents.backends import StateBackend from langchain_core.language_models import BaseChatModel -from app.agents.new_chat.middleware import create_surfsense_compaction_middleware +from app.agents.shared.middleware import create_surfsense_compaction_middleware def build_compaction_mw(llm: BaseChatModel) -> Any: diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/middleware/path_resolution.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/middleware/path_resolution.py index d01da8074..1eb062a11 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/middleware/path_resolution.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/middleware/path_resolution.py @@ -9,7 +9,7 @@ from langchain.tools import ToolRuntime from app.agents.shared.filesystem_selection import FilesystemMode from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/edit_file/index.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/edit_file/index.py index 675d36077..0a3009db3 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/edit_file/index.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/edit_file/index.py @@ -12,7 +12,7 @@ from langchain_core.tools import BaseTool, StructuredTool from langgraph.types import Command from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/list_tree/index.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/list_tree/index.py index ca15d0160..c629175a5 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/list_tree/index.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/list_tree/index.py @@ -10,7 +10,7 @@ from langchain.tools import ToolRuntime from langchain_core.tools import BaseTool, StructuredTool from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_list_target_path diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/ls/index.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/ls/index.py index cec881e8b..305a8233a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/ls/index.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/ls/index.py @@ -9,7 +9,7 @@ from langchain.tools import ToolRuntime from langchain_core.tools import BaseTool, StructuredTool from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import paginate_listing +from app.agents.shared.middleware.kb_postgres_backend import paginate_listing from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_list_target_path diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/move_file/helpers.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/move_file/helpers.py index b82e825cb..9df446bdb 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/move_file/helpers.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/move_file/helpers.py @@ -9,7 +9,7 @@ from langchain_core.messages import ToolMessage from langgraph.types import Command from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from app.agents.shared.path_resolver import DOCUMENTS_ROOT from app.agents.shared.state_reducers import _CLEAR diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/read_file/index.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/read_file/index.py index e922ec417..784d00b6c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/read_file/index.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/read_file/index.py @@ -11,7 +11,7 @@ from langchain_core.tools import BaseTool, StructuredTool from langgraph.types import Command from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_relative diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rm/helpers.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rm/helpers.py index 82706a6b2..e990208e0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rm/helpers.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rm/helpers.py @@ -13,7 +13,7 @@ from langchain_core.messages import ToolMessage from langgraph.types import Command from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from app.agents.shared.path_resolver import DOCUMENTS_ROOT from app.agents.shared.state_reducers import _CLEAR diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rmdir/helpers.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rmdir/helpers.py index 79b3c3488..b07009792 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rmdir/helpers.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem/tools/rmdir/helpers.py @@ -14,7 +14,7 @@ from langchain_core.messages import ToolMessage from langgraph.types import Command from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend from app.agents.shared.path_resolver import DOCUMENTS_ROOT from app.agents.shared.state_reducers import _CLEAR diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/kb_context_projection.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/kb_context_projection.py index dcb984be6..f1bc618af 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/kb_context_projection.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/kb_context_projection.py @@ -10,7 +10,7 @@ from langchain_core.messages import SystemMessage from langgraph.runtime import Runtime from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.knowledge_search import _render_priority_message +from app.agents.shared.middleware.knowledge_search import _render_priority_message from app.utils.perf import get_perf_logger _perf_log = get_perf_logger() diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py index 9316b3e21..3e1df1783 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py @@ -2,7 +2,7 @@ from __future__ import annotations -from app.agents.new_chat.middleware import MemoryInjectionMiddleware +from app.agents.shared.middleware import MemoryInjectionMiddleware from app.db import ChatVisibility diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py index 5940135a8..13354aff6 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py @@ -11,8 +11,8 @@ from langchain.agents.middleware import ( ) from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import RetryAfterMiddleware -from app.agents.new_chat.middleware.scoped_model_fallback import ( +from app.agents.shared.middleware import RetryAfterMiddleware +from app.agents.shared.middleware.scoped_model_fallback import ( ScopedModelFallbackMiddleware, ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py index 1146ae887..3bd8987f5 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware.scoped_model_fallback import ( +from app.agents.shared.middleware.scoped_model_fallback import ( ScopedModelFallbackMiddleware, ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py index 2b08f0cbb..02d15078e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py @@ -3,7 +3,7 @@ from __future__ import annotations from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware import RetryAfterMiddleware +from app.agents.shared.middleware import RetryAfterMiddleware from ..flags import enabled diff --git a/surfsense_backend/app/agents/new_chat/anonymous_agent.py b/surfsense_backend/app/agents/new_chat/anonymous_agent.py index c783d9a45..0e9f70d78 100644 --- a/surfsense_backend/app/agents/new_chat/anonymous_agent.py +++ b/surfsense_backend/app/agents/new_chat/anonymous_agent.py @@ -28,7 +28,7 @@ from langchain_core.language_models import BaseChatModel from langgraph.types import Checkpointer from app.agents.new_chat.context import SurfSenseContextSchema -from app.agents.new_chat.middleware import ( +from app.agents.shared.middleware import ( RetryAfterMiddleware, create_surfsense_compaction_middleware, ) diff --git a/surfsense_backend/app/agents/new_chat/filesystem_backends.py b/surfsense_backend/app/agents/new_chat/filesystem_backends.py index 781266a00..c20eaea5f 100644 --- a/surfsense_backend/app/agents/new_chat/filesystem_backends.py +++ b/surfsense_backend/app/agents/new_chat/filesystem_backends.py @@ -10,8 +10,8 @@ from deepagents.backends.state import StateBackend from langgraph.prebuilt.tool_node import ToolRuntime from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) diff --git a/surfsense_backend/app/agents/new_chat/middleware/__init__.py b/surfsense_backend/app/agents/new_chat/middleware/__init__.py index 6742bd8de..5a7324e2e 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/__init__.py +++ b/surfsense_backend/app/agents/new_chat/middleware/__init__.py @@ -1,58 +1,40 @@ -"""Middleware components for the SurfSense new chat agent.""" +"""Backward-compatible shim package. -from app.agents.new_chat.middleware.action_log import ActionLogMiddleware -from app.agents.new_chat.middleware.anonymous_document import ( +The agent middleware now lives in the shared kernel at +``app.agents.shared.middleware``. This package re-exports it so frozen +single-agent code (``chat_deepagent`` and ``subagents/*``) keeps working +until that stack is retired. +""" + +from app.agents.shared.middleware import ( + ActionLogMiddleware, AnonymousDocumentMiddleware, -) -from app.agents.new_chat.middleware.busy_mutex import BusyMutexMiddleware -from app.agents.new_chat.middleware.compaction import ( - SurfSenseCompactionMiddleware, - create_surfsense_compaction_middleware, -) -from app.agents.new_chat.middleware.context_editing import ( + BuiltinSkillsBackend, + BusyMutexMiddleware, ClearToolUsesEdit, - SpillingContextEditingMiddleware, - SpillToBackendEdit, -) -from app.agents.new_chat.middleware.dedup_tool_calls import ( DedupHITLToolCallsMiddleware, -) -from app.agents.new_chat.middleware.doom_loop import DoomLoopMiddleware -from app.agents.new_chat.middleware.file_intent import ( + DoomLoopMiddleware, FileIntentMiddleware, -) -from app.agents.new_chat.middleware.filesystem import ( - SurfSenseFilesystemMiddleware, -) -from app.agents.new_chat.middleware.flatten_system import ( FlattenSystemMessageMiddleware, -) -from app.agents.new_chat.middleware.kb_persistence import ( KnowledgeBasePersistenceMiddleware, - commit_staged_filesystem_state, -) -from app.agents.new_chat.middleware.knowledge_search import ( KnowledgeBaseSearchMiddleware, KnowledgePriorityMiddleware, -) -from app.agents.new_chat.middleware.knowledge_tree import ( KnowledgeTreeMiddleware, -) -from app.agents.new_chat.middleware.memory_injection import ( MemoryInjectionMiddleware, -) -from app.agents.new_chat.middleware.noop_injection import NoopInjectionMiddleware -from app.agents.new_chat.middleware.otel_span import OtelSpanMiddleware -from app.agents.new_chat.middleware.permission import PermissionMiddleware -from app.agents.new_chat.middleware.retry_after import RetryAfterMiddleware -from app.agents.new_chat.middleware.skills_backends import ( - BuiltinSkillsBackend, + NoopInjectionMiddleware, + OtelSpanMiddleware, + PermissionMiddleware, + RetryAfterMiddleware, SearchSpaceSkillsBackend, - build_skills_backend_factory, - default_skills_sources, -) -from app.agents.new_chat.middleware.tool_call_repair import ( + SpillingContextEditingMiddleware, + SpillToBackendEdit, + SurfSenseCompactionMiddleware, + SurfSenseFilesystemMiddleware, ToolCallNameRepairMiddleware, + build_skills_backend_factory, + commit_staged_filesystem_state, + create_surfsense_compaction_middleware, + default_skills_sources, ) __all__ = [ diff --git a/surfsense_backend/app/agents/new_chat/middleware/permission.py b/surfsense_backend/app/agents/new_chat/middleware/permission.py index 8601a3296..e8395cbdb 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/permission.py +++ b/surfsense_backend/app/agents/new_chat/middleware/permission.py @@ -1,424 +1,14 @@ -""" -PermissionMiddleware — pattern-based allow/deny/ask with HITL fallback. +"""Backward-compatible shim. -LangChain's :class:`HumanInTheLoopMiddleware` only supports a static -"this tool always asks" decision per tool. There's no rule-based -allow/deny/ask layered ruleset, no glob patterns, no per-search-space or -per-thread overrides, and no auto-deny synthesis. - -This middleware ports OpenCode's ``packages/opencode/src/permission/index.ts`` -ruleset model on top of SurfSense's existing ``interrupt({type, action, -context})`` payload shape (see ``app/agents/new_chat/tools/hitl.py``) so -the frontend keeps working unchanged. - -Operation: -1. ``aafter_model`` inspects the latest ``AIMessage.tool_calls``. -2. For each call, the middleware builds a list of ``patterns`` (the - tool name plus any tool-specific patterns from the resolver). It - evaluates each pattern against the layered rulesets and aggregates - the results: ``deny`` > ``ask`` > ``allow``. -3. On ``deny``: replaces the call with a synthetic ``ToolMessage`` - containing a :class:`StreamingError`. -4. On ``ask``: raises a SurfSense-style ``interrupt(...)``. Both the legacy - SurfSense shape and LangChain HITL ``{"decisions": [{"type": ...}]}`` - replies are accepted via :func:`_normalize_permission_decision`. - - ``once``: proceed. - - ``approve_always``: also persist allow rules for ``request.always`` patterns. - - ``reject`` w/o feedback: raise :class:`RejectedError`. - - ``reject`` w/ feedback: raise :class:`CorrectedError`. -5. On ``allow``: proceed unchanged. - -The middleware also performs a *pre-model* tool-filter step (the -``before_model`` hook) so globally denied tools are stripped from the -exposed tool list before the model gets to see them. This mirrors -OpenCode's ``Permission.disabled`` and dramatically reduces the chance -the model emits a deny-only call. +Moved to ``app.agents.shared.middleware.permission``. Re-exported here for the +frozen single-agent stack (``chat_deepagent``/``subagents``). """ -from __future__ import annotations - -import logging -from collections.abc import Callable -from typing import Any - -from langchain.agents.middleware.types import ( - AgentMiddleware, - AgentState, - ContextT, +from app.agents.shared.middleware.permission import ( + PatternResolver, + PermissionMiddleware, + _normalize_permission_decision, ) -from langchain_core.messages import AIMessage, ToolMessage -from langgraph.runtime import Runtime -from langgraph.types import interrupt - -from app.agents.shared.errors import ( - CorrectedError, - RejectedError, - StreamingError, -) -from app.agents.shared.permissions import ( - Rule, - Ruleset, - aggregate_action, - evaluate_many, -) -from app.observability import metrics as ot_metrics, otel as ot - -logger = logging.getLogger(__name__) - - -# Mapping ``tool_name -> resolver`` that converts ``args`` to a list of -# patterns to evaluate. The first pattern is conventionally the bare -# tool name; later entries narrow down to specific resources. -PatternResolver = Callable[[dict[str, Any]], list[str]] - - -def _default_pattern_resolver(name: str) -> PatternResolver: - def _resolve(args: dict[str, Any]) -> list[str]: - # Bare name covers the default catch-all; primary-arg fallbacks - # are best added per-tool by callers. - del args - return [name] - - return _resolve - - -# Translation from the LangChain HITL envelope (what ``stream_resume_chat`` -# sends) to SurfSense's legacy ``decision_type`` shape. ``edit`` keeps the -# original tool args — tools needing argument edits should use -# ``request_approval`` from ``app/agents/new_chat/tools/hitl.py``. -_LC_TYPE_TO_PERMISSION_DECISION: dict[str, str] = { - "approve": "once", - "reject": "reject", - "edit": "once", - "approve_always": "approve_always", -} - - -def _normalize_permission_decision(decision: Any) -> dict[str, Any]: - """Coerce any accepted reply shape into ``{"decision_type": ..., "feedback"?}``. - - Falls back to ``reject`` (with a warning) on unrecognized payloads so the - middleware fails closed. - """ - if isinstance(decision, str): - return {"decision_type": decision} - if not isinstance(decision, dict): - logger.warning( - "Unrecognized permission resume value (%s); treating as reject", - type(decision).__name__, - ) - return {"decision_type": "reject"} - - if decision.get("decision_type"): - return decision - - payload: dict[str, Any] = decision - decisions = decision.get("decisions") - if isinstance(decisions, list) and decisions: - first = decisions[0] - if isinstance(first, dict): - payload = first - - raw_type = payload.get("type") or payload.get("decision_type") - if not raw_type: - logger.warning( - "Permission resume missing decision type (keys=%s); treating as reject", - list(payload.keys()), - ) - return {"decision_type": "reject"} - - raw_type = str(raw_type).lower() - mapped = _LC_TYPE_TO_PERMISSION_DECISION.get(raw_type) - if mapped is None: - # Tolerate legacy values arriving without ``decision_type`` wrapping. - if raw_type in {"once", "approve_always", "reject"}: - mapped = raw_type - else: - logger.warning( - "Unknown permission decision type %r; treating as reject", raw_type - ) - mapped = "reject" - - if raw_type == "edit": - logger.warning( - "Permission middleware received an 'edit' decision; original args " - "kept (edits not merged here)." - ) - - out: dict[str, Any] = {"decision_type": mapped} - feedback = payload.get("feedback") or payload.get("message") - if isinstance(feedback, str) and feedback.strip(): - out["feedback"] = feedback - return out - - -class PermissionMiddleware(AgentMiddleware): # type: ignore[type-arg] - """Allow/deny/ask layer over the agent's tool calls. - - Args: - rulesets: Layered rulesets to evaluate. Earlier entries are - overridden by later ones (last-match-wins). Typical layering: - ``defaults < global < space < thread < runtime_approved``. - pattern_resolvers: Optional per-tool callables that return a list - of patterns to evaluate. When a tool isn't listed, the bare - tool name is used as the only pattern. - runtime_ruleset: Mutable :class:`Ruleset` that the middleware - extends in-place when the user replies ``"approve_always"`` to - an ask interrupt. Reused across all calls in the same agent - instance so newly-allowed rules apply to subsequent calls. - always_emit_interrupt_payload: If True, every ask uses the - SurfSense interrupt wire format (default). Set False to - disable interrupts and treat ``ask`` as ``deny`` for - non-interactive deployments. - """ - - tools = () - - def __init__( - self, - *, - rulesets: list[Ruleset] | None = None, - pattern_resolvers: dict[str, PatternResolver] | None = None, - runtime_ruleset: Ruleset | None = None, - always_emit_interrupt_payload: bool = True, - ) -> None: - super().__init__() - self._static_rulesets: list[Ruleset] = list(rulesets or []) - self._pattern_resolvers: dict[str, PatternResolver] = dict( - pattern_resolvers or {} - ) - self._runtime_ruleset: Ruleset = runtime_ruleset or Ruleset( - origin="runtime_approved" - ) - self._emit_interrupt = always_emit_interrupt_payload - - # ------------------------------------------------------------------ - # Tool-filter step (mirrors OpenCode's ``Permission.disabled``) - # ------------------------------------------------------------------ - - def _globally_denied(self, tool_name: str) -> bool: - """Return True if a deny rule with no narrowing pattern matches.""" - rules = evaluate_many(tool_name, ["*"], *self._all_rulesets()) - return aggregate_action(rules) == "deny" - - def _all_rulesets(self) -> list[Ruleset]: - return [*self._static_rulesets, self._runtime_ruleset] - - # NOTE: ``before_model`` filtering of the tools list is left to the - # agent factory. This middleware only blocks at execution time — and - # only via the rule-evaluator path, not by mutating ``request.tools``. - # Mutating ``request.tools`` per-call would invalidate provider - # prompt-cache prefixes (see Operational risks: prompt-cache regression). - - # ------------------------------------------------------------------ - # Tool-call evaluation - # ------------------------------------------------------------------ - - def _resolve_patterns(self, tool_name: str, args: dict[str, Any]) -> list[str]: - resolver = self._pattern_resolvers.get( - tool_name, _default_pattern_resolver(tool_name) - ) - try: - patterns = resolver(args or {}) - except Exception: - logger.exception( - "Pattern resolver for %s raised; using bare name", tool_name - ) - patterns = [tool_name] - if not patterns: - patterns = [tool_name] - return patterns - - def _evaluate( - self, tool_name: str, args: dict[str, Any] - ) -> tuple[str, list[str], list[Rule]]: - patterns = self._resolve_patterns(tool_name, args) - rules = evaluate_many(tool_name, patterns, *self._all_rulesets()) - action = aggregate_action(rules) - return action, patterns, rules - - # ------------------------------------------------------------------ - # HITL ask flow — SurfSense wire format - # ------------------------------------------------------------------ - - def _raise_interrupt( - self, - *, - tool_name: str, - args: dict[str, Any], - patterns: list[str], - rules: list[Rule], - ) -> dict[str, Any]: - """Block on user approval via SurfSense's ``interrupt`` shape.""" - if not self._emit_interrupt: - return {"decision_type": "reject"} - - # ``params`` (NOT ``args``) is what SurfSense's streaming - # normalizer forwards. Other fields move into ``context``. - payload = { - "type": "permission_ask", - "action": {"tool": tool_name, "params": args or {}}, - "context": { - "patterns": patterns, - "rules": [ - { - "permission": r.permission, - "pattern": r.pattern, - "action": r.action, - } - for r in rules - ], - # Rules of thumb for the frontend: surface the patterns - # the user can promote to "approve_always" with a single reply. - "always": patterns, - }, - } - # Open ``permission.asked`` + ``interrupt.raised`` OTel spans - # (no-op when OTel is disabled) so dashboards can correlate - # "we asked X" with "interrupt was actually delivered". - with ( - ot.permission_asked_span( - permission=tool_name, - pattern=patterns[0] if patterns else None, - extra={"permission.patterns": list(patterns)}, - ), - ot.interrupt_span(interrupt_type="permission_ask"), - ): - ot_metrics.record_permission_ask(permission=tool_name) - ot_metrics.record_interrupt(interrupt_type="permission_ask") - decision = interrupt(payload) - return _normalize_permission_decision(decision) - - def _persist_always(self, tool_name: str, patterns: list[str]) -> None: - """Promote ``approve_always`` reply into runtime allow rules. - - Persistence to ``agent_permission_rules`` is done by the - streaming layer (``stream_new_chat``) once it observes the - ``approve_always`` reply — the middleware just keeps an - in-memory copy so subsequent calls in the same stream see the rule. - """ - for pattern in patterns: - self._runtime_ruleset.rules.append( - Rule(permission=tool_name, pattern=pattern, action="allow") - ) - - # ------------------------------------------------------------------ - # Synthesizing deny -> ToolMessage - # ------------------------------------------------------------------ - - @staticmethod - def _deny_message( - tool_call: dict[str, Any], - rule: Rule, - ) -> ToolMessage: - err = StreamingError( - code="permission_denied", - retryable=False, - suggestion=( - f"rule permission={rule.permission!r} pattern={rule.pattern!r} " - f"blocked this call" - ), - ) - return ToolMessage( - content=( - f"Permission denied: rule {rule.permission}/{rule.pattern} " - f"blocked tool {tool_call.get('name')!r}." - ), - tool_call_id=tool_call.get("id") or "", - name=tool_call.get("name"), - status="error", - additional_kwargs={"error": err.model_dump()}, - ) - - # ------------------------------------------------------------------ - # The hook: aafter_model - # ------------------------------------------------------------------ - - def _process( - self, - state: AgentState, - runtime: Runtime[Any], - ) -> dict[str, Any] | None: - del runtime # unused - messages = state.get("messages") or [] - if not messages: - return None - last = messages[-1] - if not isinstance(last, AIMessage) or not last.tool_calls: - return None - - deny_messages: list[ToolMessage] = [] - kept_calls: list[dict[str, Any]] = [] - any_change = False - - for raw in last.tool_calls: - call = ( - dict(raw) - if isinstance(raw, dict) - else { - "name": getattr(raw, "name", None), - "args": getattr(raw, "args", {}), - "id": getattr(raw, "id", None), - "type": "tool_call", - } - ) - name = call.get("name") or "" - args = call.get("args") or {} - action, patterns, rules = self._evaluate(name, args) - - if action == "deny": - # Find the deny rule for the suggestion text - deny_rule = next((r for r in rules if r.action == "deny"), rules[0]) - deny_messages.append(self._deny_message(call, deny_rule)) - any_change = True - continue - - if action == "ask": - decision = self._raise_interrupt( - tool_name=name, args=args, patterns=patterns, rules=rules - ) - kind = str(decision.get("decision_type") or "reject").lower() - if kind == "once": - kept_calls.append(call) - elif kind == "approve_always": - self._persist_always(name, patterns) - kept_calls.append(call) - elif kind == "reject": - feedback = decision.get("feedback") - if isinstance(feedback, str) and feedback.strip(): - raise CorrectedError(feedback, tool=name) - raise RejectedError( - tool=name, pattern=patterns[0] if patterns else None - ) - else: - logger.warning( - "Unknown permission decision %r; treating as reject", kind - ) - raise RejectedError(tool=name) - continue - - # allow - kept_calls.append(call) - - if not any_change and len(kept_calls) == len(last.tool_calls): - return None - - updated = last.model_copy(update={"tool_calls": kept_calls}) - result_messages: list[Any] = [updated] - if deny_messages: - result_messages.extend(deny_messages) - return {"messages": result_messages} - - def after_model( # type: ignore[override] - self, state: AgentState, runtime: Runtime[ContextT] - ) -> dict[str, Any] | None: - return self._process(state, runtime) - - async def aafter_model( # type: ignore[override] - self, state: AgentState, runtime: Runtime[ContextT] - ) -> dict[str, Any] | None: - return self._process(state, runtime) - __all__ = [ "PatternResolver", diff --git a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py index 0294e2839..d1f4fffef 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py +++ b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py @@ -1,111 +1,11 @@ -"""Fallback only on provider/network errors; let programming bugs raise.""" +"""Backward-compatible shim. -from __future__ import annotations +Moved to ``app.agents.shared.middleware.scoped_model_fallback``. Re-exported here +for the frozen single-agent stack (``chat_deepagent``). +""" -from typing import TYPE_CHECKING, Any - -from langchain.agents.middleware import ModelFallbackMiddleware - -from app.observability import metrics as ot_metrics, otel as ot - -if TYPE_CHECKING: - from collections.abc import Awaitable, Callable - - from langchain.agents.middleware.types import ModelRequest, ModelResponse - from langchain_core.messages import AIMessage - - -# Matched by class name across the MRO so we don't have to import every -# provider SDK (openai/anthropic/google/...). Extend as new providers ship. -_FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( - { - "RateLimitError", - "APIStatusError", - "InternalServerError", - "ServiceUnavailableError", - "BadGatewayError", - "GatewayTimeoutError", - "APIConnectionError", - "APITimeoutError", - "ConnectError", - "ConnectTimeout", - "ReadTimeout", - "RemoteProtocolError", - "TimeoutError", - "TimeoutException", - } +from app.agents.shared.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, ) - -def _is_fallback_eligible(exc: BaseException) -> bool: - return any(cls.__name__ in _FALLBACK_ELIGIBLE_NAMES for cls in type(exc).__mro__) - - -class ScopedModelFallbackMiddleware(ModelFallbackMiddleware): - """Re-raise non-provider exceptions instead of walking the fallback chain.""" - - def wrap_model_call( # type: ignore[override] - self, - request: ModelRequest[Any], - handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], - ) -> ModelResponse[Any] | AIMessage: - last_exception: Exception - try: - return handler(request) - except Exception as e: - if not _is_fallback_eligible(e): - raise - last_exception = e - - for attempt, fallback_model in enumerate(self.models, start=1): - ot.add_event( - "model.fallback", - { - "fallback.attempt": attempt, - "fallback.from": attempt - 1, - "fallback.to": attempt, - "fallback.reason": ot_metrics.categorize_exception(last_exception), - }, - ) - try: - return handler(request.override(model=fallback_model)) - except Exception as e: - if not _is_fallback_eligible(e): - raise - last_exception = e - continue - - raise last_exception - - async def awrap_model_call( # type: ignore[override] - self, - request: ModelRequest[Any], - handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]], - ) -> ModelResponse[Any] | AIMessage: - last_exception: Exception - try: - return await handler(request) - except Exception as e: - if not _is_fallback_eligible(e): - raise - last_exception = e - - for attempt, fallback_model in enumerate(self.models, start=1): - ot.add_event( - "model.fallback", - { - "fallback.attempt": attempt, - "fallback.from": attempt - 1, - "fallback.to": attempt, - "fallback.reason": ot_metrics.categorize_exception(last_exception), - }, - ) - try: - return await handler(request.override(model=fallback_model)) - except Exception as e: - if not _is_fallback_eligible(e): - raise - last_exception = e - continue - - raise last_exception +__all__ = ["ScopedModelFallbackMiddleware"] diff --git a/surfsense_backend/app/agents/new_chat/middleware/skills_backends.py b/surfsense_backend/app/agents/new_chat/middleware/skills_backends.py index dad22db50..37d0c6680 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/skills_backends.py +++ b/surfsense_backend/app/agents/new_chat/middleware/skills_backends.py @@ -1,333 +1,17 @@ -"""Skills backends for SurfSense. +"""Backward-compatible shim. -Implements two minimal :class:`deepagents.backends.protocol.BackendProtocol` -subclasses tailored for use with :class:`deepagents.middleware.skills.SkillsMiddleware`. - -The middleware only needs four methods to load skills from a backend: - -* ``ls_info`` / ``als_info`` — list directories under a source path. -* ``download_files`` / ``adownload_files`` — fetch ``SKILL.md`` bytes. - -Other ``BackendProtocol`` methods (``read``/``write``/``edit``/``grep_raw`` …) -default to ``NotImplementedError`` from the base class. They are never reached -by the skills middleware because skill content is rendered into the system -prompt at agent build time, not edited at runtime. - -Two backends are provided: - -* :class:`BuiltinSkillsBackend` — disk-backed read of bundled skills from - ``app/agents/new_chat/skills/builtin/``. -* :class:`SearchSpaceSkillsBackend` — a thin read-only wrapper over - :class:`KBPostgresBackend` that filters notes under the privileged folder - ``/documents/_skills/``. - -Both backends are intentionally read-only: skill authoring happens out of band -(via filesystem or a search-space-admin route), so we never expose -``write`` / ``edit`` / ``upload_files``. The base class' ``NotImplementedError`` -gives a clean failure mode if anything tries. +Moved to ``app.agents.shared.middleware.skills_backends``. Re-exported here for +the frozen single-agent stack (``subagents/config``). """ -from __future__ import annotations - -import contextlib -import logging -from collections.abc import Callable -from dataclasses import replace -from pathlib import Path -from typing import TYPE_CHECKING - -from deepagents.backends.composite import CompositeBackend -from deepagents.backends.protocol import ( - BackendProtocol, - FileDownloadResponse, - FileInfo, +from app.agents.shared.middleware.skills_backends import ( + SKILLS_BUILTIN_PREFIX, + SKILLS_SPACE_PREFIX, + BuiltinSkillsBackend, + SearchSpaceSkillsBackend, + build_skills_backend_factory, + default_skills_sources, ) -from deepagents.backends.state import StateBackend - -if TYPE_CHECKING: - from langchain.tools import ToolRuntime - - from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend - -logger = logging.getLogger(__name__) - - -# Limit per Agent Skills spec; matches deepagents.middleware.skills.MAX_SKILL_FILE_SIZE. -_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/new_chat/skills/builtin/`` relative to this module. - """ - return (Path(__file__).resolve().parent.parent / "skills" / "builtin").resolve() - - -class BuiltinSkillsBackend(BackendProtocol): - """Read-only disk-backed skills source. - - Maps a virtual ``/skills/builtin/`` namespace onto a directory on local disk, - where each skill is its own subdirectory containing a ``SKILL.md`` file:: - - //SKILL.md - - The middleware calls :meth:`als_info` with the source path and expects a - ``list[FileInfo]`` whose ``is_dir=True`` entries are descended into. Then it - calls :meth:`adownload_files` with the synthesized ``SKILL.md`` paths and - parses YAML frontmatter from the returned ``content`` bytes. - - Mounting under :class:`~deepagents.backends.composite.CompositeBackend` at - prefix ``/skills/builtin/`` means the middleware can issue paths like - ``/skills/builtin/kb-research/SKILL.md`` which the composite strips down to - ``/kb-research/SKILL.md`` before forwarding here. We treat any leading - slash as anchoring at :attr:`root`. - """ - - def __init__(self, root: Path | str | None = None) -> None: - self.root: Path = Path(root).resolve() if root else _default_builtin_root() - if not self.root.exists(): - logger.info( - "BuiltinSkillsBackend root %s does not exist; skills will be empty.", - self.root, - ) - - def _resolve(self, path: str) -> Path: - """Resolve a virtual posix path under :attr:`root`, refusing escapes.""" - bare = path.lstrip("/") - candidate = (self.root / bare).resolve() if bare else self.root - # Refuse symlink/.. traversal that escapes the root. - try: - candidate.relative_to(self.root) - except ValueError as exc: - raise ValueError(f"path {path!r} escapes builtin skills root") from exc - return candidate - - def ls_info(self, path: str) -> list[FileInfo]: - try: - target = self._resolve(path) - except ValueError as exc: - logger.warning("BuiltinSkillsBackend.ls_info refused: %s", exc) - return [] - if not target.exists() or not target.is_dir(): - return [] - - infos: list[FileInfo] = [] - # Build virtual paths anchored at "/" because CompositeBackend already - # stripped the route prefix before calling us. - target_virtual = ( - "/" - if target == self.root - else ("/" + str(target.relative_to(self.root)).replace("\\", "/")) - ) - for child in sorted(target.iterdir()): - if child.name == "__pycache__" or child.name.startswith("."): - continue - child_virtual = ( - target_virtual.rstrip("/") + "/" + child.name - if target_virtual != "/" - else "/" + child.name - ) - info: FileInfo = { - "path": child_virtual, - "is_dir": child.is_dir(), - } - if child.is_file(): - with contextlib.suppress(OSError): # pragma: no cover - defensive - info["size"] = child.stat().st_size - infos.append(info) - return infos - - def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: - responses: list[FileDownloadResponse] = [] - for p in paths: - try: - target = self._resolve(p) - except ValueError: - responses.append(FileDownloadResponse(path=p, error="invalid_path")) - continue - if not target.exists(): - responses.append(FileDownloadResponse(path=p, error="file_not_found")) - continue - if target.is_dir(): - responses.append(FileDownloadResponse(path=p, error="is_directory")) - continue - try: - # Hard cap to avoid loading rogue mega-files into memory. - size = target.stat().st_size - if size > _MAX_SKILL_FILE_SIZE: - logger.warning( - "Builtin skill file %s exceeds %d bytes; truncating.", - target, - _MAX_SKILL_FILE_SIZE, - ) - with target.open("rb") as fh: - content = fh.read(_MAX_SKILL_FILE_SIZE) - else: - content = target.read_bytes() - except PermissionError: - responses.append( - FileDownloadResponse(path=p, error="permission_denied") - ) - continue - except OSError as exc: # pragma: no cover - defensive - logger.warning("Builtin skill read failed %s: %s", target, exc) - responses.append(FileDownloadResponse(path=p, error="file_not_found")) - continue - responses.append(FileDownloadResponse(path=p, content=content, error=None)) - return responses - - -class SearchSpaceSkillsBackend(BackendProtocol): - """Read-only view of search-space-authored skills. - - Wraps a :class:`KBPostgresBackend` and only ever reads under the privileged - folder ``/documents/_skills/`` (configurable). The folder is intended to be - writable only by search-space admins; this backend never writes. - - The skills middleware expects a layout like:: - - ///SKILL.md - - But the KB stores documents like ``/documents/_skills//SKILL.md``. - We expose the inner namespace by remapping each path. When mounted under - :class:`CompositeBackend` at prefix ``/skills/space/`` the paths the - middleware sees become ``/skills/space//SKILL.md``; the composite - strips ``/skills/space/`` and hands us ``//SKILL.md``, which we - rewrite to ``/documents/_skills//SKILL.md`` before forwarding to the - KB. - - No new database table is needed: the privileged folder convention is - enforced server-side outside of this class. We intentionally swallow any - write/edit attempts (the base class raises ``NotImplementedError``). - """ - - DEFAULT_KB_ROOT: str = "/documents/_skills" - - def __init__( - self, - kb_backend: KBPostgresBackend, - *, - kb_root: str = DEFAULT_KB_ROOT, - ) -> None: - self._kb = kb_backend - # Normalize trailing slash off so we can join cleanly. - self._kb_root = kb_root.rstrip("/") or "/" - - def _to_kb(self, path: str) -> str: - """Rewrite a virtual path into the underlying KB namespace.""" - bare = path.lstrip("/") - if not bare: - return self._kb_root - return f"{self._kb_root}/{bare}" - - def _from_kb(self, kb_path: str) -> str: - """Rewrite a KB path back into our virtual namespace.""" - if not kb_path.startswith(self._kb_root): - return kb_path # pragma: no cover - defensive - rel = kb_path[len(self._kb_root) :] - return rel if rel.startswith("/") else "/" + rel - - def ls_info(self, path: str) -> list[FileInfo]: - # KBPostgresBackend exposes only the async API meaningfully; the sync - # path falls back to ``asyncio.to_thread(...)`` in the base class. We - # keep this stub to satisfy abstract resolution; the middleware calls - # ``als_info``. - raise NotImplementedError("SearchSpaceSkillsBackend is async-only") - - async def als_info(self, path: str) -> list[FileInfo]: - kb_path = self._to_kb(path) - try: - infos = await self._kb.als_info(kb_path) - except Exception as exc: # pragma: no cover - defensive - logger.warning("SearchSpaceSkillsBackend.als_info failed: %s", exc) - return [] - remapped: list[FileInfo] = [] - for info in infos: - kb_p = info.get("path", "") - if not kb_p.startswith(self._kb_root): - continue - remapped.append({**info, "path": self._from_kb(kb_p)}) - return remapped - - def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: - raise NotImplementedError("SearchSpaceSkillsBackend is async-only") - - async def adownload_files(self, paths: list[str]) -> list[FileDownloadResponse]: - kb_paths = [self._to_kb(p) for p in paths] - responses = await self._kb.adownload_files(kb_paths) - # Re-map response paths back to the virtual namespace so the middleware - # correlates them to the input list correctly. - remapped: list[FileDownloadResponse] = [] - for original, resp in zip(paths, responses, strict=True): - remapped.append(replace(resp, path=original)) - return remapped - - -SKILLS_BUILTIN_PREFIX = "/skills/builtin/" -SKILLS_SPACE_PREFIX = "/skills/space/" - - -def build_skills_backend_factory( - *, - builtin_root: Path | str | None = None, - search_space_id: int | None = None, -) -> Callable[[ToolRuntime], BackendProtocol]: - """Return a runtime-aware factory for the skills :class:`CompositeBackend`. - - When ``search_space_id`` is provided the composite includes a - :class:`SearchSpaceSkillsBackend` route at ``/skills/space/`` over a fresh - per-runtime :class:`KBPostgresBackend`, mirroring how - :func:`build_backend_resolver` constructs the main filesystem backend. - - When ``search_space_id`` is ``None`` (e.g., desktop-local mode or unit - tests) only the bundled :class:`BuiltinSkillsBackend` is exposed. - - Returning a factory rather than a fixed instance is intentional: the - underlying KB backend depends on per-call ``ToolRuntime`` state - (``staged_dirs``, ``files`` cache, runtime config), so a single shared - instance cannot serve multiple concurrent agent runs. - """ - builtin = BuiltinSkillsBackend(builtin_root) - - if search_space_id is None: - - def _factory_builtin_only(runtime: ToolRuntime) -> BackendProtocol: - # Default StateBackend is intentionally inert: any path outside the - # ``/skills/builtin/`` route resolves to an empty per-runtime state - # so the SkillsMiddleware can iterate sources without raising. - return CompositeBackend( - default=StateBackend(runtime), - routes={SKILLS_BUILTIN_PREFIX: builtin}, - ) - - return _factory_builtin_only - - def _factory_with_space(runtime: ToolRuntime) -> BackendProtocol: - # Imported lazily to avoid a hard dependency at module import time: - # ``KBPostgresBackend`` pulls in DB models, which are unnecessary for - # the unit-tested builtin path. - from app.agents.new_chat.middleware.kb_postgres_backend import ( - KBPostgresBackend, - ) - - kb = KBPostgresBackend(search_space_id, runtime) - space = SearchSpaceSkillsBackend(kb) - return CompositeBackend( - default=StateBackend(runtime), - routes={ - SKILLS_BUILTIN_PREFIX: builtin, - SKILLS_SPACE_PREFIX: space, - }, - ) - - return _factory_with_space - - -def default_skills_sources() -> list[str]: - """Return the canonical source list for SkillsMiddleware (built-in then space).""" - return [SKILLS_BUILTIN_PREFIX, SKILLS_SPACE_PREFIX] - __all__ = [ "SKILLS_BUILTIN_PREFIX", diff --git a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py index 6c4cfb6be..8bef19050 100644 --- a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py +++ b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py @@ -33,7 +33,7 @@ from sqlalchemy import cast, select from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.middleware.dedup_tool_calls import dedup_key_full_args +from app.agents.shared.middleware.dedup_tool_calls import dedup_key_full_args from app.agents.new_chat.tools.hitl import request_approval from app.agents.new_chat.tools.mcp_client import MCPClient from app.agents.new_chat.tools.mcp_tools_cache import ( diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 6f011e372..9b1944aa5 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/surfsense_backend/app/agents/new_chat/tools/registry.py @@ -43,7 +43,7 @@ from typing import Any from langchain_core.tools import BaseTool -from app.agents.new_chat.middleware.dedup_tool_calls import ( +from app.agents.shared.middleware.dedup_tool_calls import ( wrap_dedup_key_by_arg_name, ) from app.db import ChatVisibility diff --git a/surfsense_backend/app/agents/shared/middleware/__init__.py b/surfsense_backend/app/agents/shared/middleware/__init__.py new file mode 100644 index 000000000..9ca2f3960 --- /dev/null +++ b/surfsense_backend/app/agents/shared/middleware/__init__.py @@ -0,0 +1,87 @@ +"""Middleware components for the SurfSense new chat agent.""" + +from app.agents.shared.middleware.action_log import ActionLogMiddleware +from app.agents.shared.middleware.anonymous_document import ( + AnonymousDocumentMiddleware, +) +from app.agents.shared.middleware.busy_mutex import BusyMutexMiddleware +from app.agents.shared.middleware.compaction import ( + SurfSenseCompactionMiddleware, + create_surfsense_compaction_middleware, +) +from app.agents.shared.middleware.context_editing import ( + ClearToolUsesEdit, + SpillingContextEditingMiddleware, + SpillToBackendEdit, +) +from app.agents.shared.middleware.dedup_tool_calls import ( + DedupHITLToolCallsMiddleware, +) +from app.agents.shared.middleware.doom_loop import DoomLoopMiddleware +from app.agents.shared.middleware.file_intent import ( + FileIntentMiddleware, +) +from app.agents.shared.middleware.filesystem import ( + SurfSenseFilesystemMiddleware, +) +from app.agents.shared.middleware.flatten_system import ( + FlattenSystemMessageMiddleware, +) +from app.agents.shared.middleware.kb_persistence import ( + KnowledgeBasePersistenceMiddleware, + commit_staged_filesystem_state, +) +from app.agents.shared.middleware.knowledge_search import ( + KnowledgeBaseSearchMiddleware, + KnowledgePriorityMiddleware, +) +from app.agents.shared.middleware.knowledge_tree import ( + KnowledgeTreeMiddleware, +) +from app.agents.shared.middleware.memory_injection import ( + MemoryInjectionMiddleware, +) +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, +) + +__all__ = [ + "ActionLogMiddleware", + "AnonymousDocumentMiddleware", + "BuiltinSkillsBackend", + "BusyMutexMiddleware", + "ClearToolUsesEdit", + "DedupHITLToolCallsMiddleware", + "DoomLoopMiddleware", + "FileIntentMiddleware", + "FlattenSystemMessageMiddleware", + "KnowledgeBasePersistenceMiddleware", + "KnowledgeBaseSearchMiddleware", + "KnowledgePriorityMiddleware", + "KnowledgeTreeMiddleware", + "MemoryInjectionMiddleware", + "NoopInjectionMiddleware", + "OtelSpanMiddleware", + "PermissionMiddleware", + "RetryAfterMiddleware", + "SearchSpaceSkillsBackend", + "SpillToBackendEdit", + "SpillingContextEditingMiddleware", + "SurfSenseCompactionMiddleware", + "SurfSenseFilesystemMiddleware", + "ToolCallNameRepairMiddleware", + "build_skills_backend_factory", + "commit_staged_filesystem_state", + "create_surfsense_compaction_middleware", + "default_skills_sources", +] diff --git a/surfsense_backend/app/agents/new_chat/middleware/action_log.py b/surfsense_backend/app/agents/shared/middleware/action_log.py similarity index 97% rename from surfsense_backend/app/agents/new_chat/middleware/action_log.py rename to surfsense_backend/app/agents/shared/middleware/action_log.py index ff6c9c53e..bba790c06 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/action_log.py +++ b/surfsense_backend/app/agents/shared/middleware/action_log.py @@ -34,12 +34,16 @@ from langchain_core.callbacks import adispatch_custom_event from langchain_core.messages import ToolMessage from app.agents.shared.feature_flags import get_flags -from app.agents.new_chat.tools.registry import ToolDefinition if TYPE_CHECKING: # pragma: no cover - type-only from langchain.agents.middleware.types import ToolCallRequest from langgraph.types import Command + # Type-only import: keeping it lazy avoids a module-load cycle through the + # frozen single-agent package (new_chat.__init__ -> chat_deepagent -> + # middleware shim). Resolves to app.agents.shared.tools once tools migrate. + from app.agents.new_chat.tools.registry import ToolDefinition + logger = logging.getLogger(__name__) diff --git a/surfsense_backend/app/agents/new_chat/middleware/anonymous_document.py b/surfsense_backend/app/agents/shared/middleware/anonymous_document.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/anonymous_document.py rename to surfsense_backend/app/agents/shared/middleware/anonymous_document.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/busy_mutex.py b/surfsense_backend/app/agents/shared/middleware/busy_mutex.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/busy_mutex.py rename to surfsense_backend/app/agents/shared/middleware/busy_mutex.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/compaction.py b/surfsense_backend/app/agents/shared/middleware/compaction.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/compaction.py rename to surfsense_backend/app/agents/shared/middleware/compaction.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/context_editing.py b/surfsense_backend/app/agents/shared/middleware/context_editing.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/context_editing.py rename to surfsense_backend/app/agents/shared/middleware/context_editing.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py b/surfsense_backend/app/agents/shared/middleware/dedup_tool_calls.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py rename to surfsense_backend/app/agents/shared/middleware/dedup_tool_calls.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/doom_loop.py b/surfsense_backend/app/agents/shared/middleware/doom_loop.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/doom_loop.py rename to surfsense_backend/app/agents/shared/middleware/doom_loop.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/file_intent.py b/surfsense_backend/app/agents/shared/middleware/file_intent.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/file_intent.py rename to surfsense_backend/app/agents/shared/middleware/file_intent.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/filesystem.py b/surfsense_backend/app/agents/shared/middleware/filesystem.py similarity index 99% rename from surfsense_backend/app/agents/new_chat/middleware/filesystem.py rename to surfsense_backend/app/agents/shared/middleware/filesystem.py index befd69da4..7968d76ba 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/filesystem.py +++ b/surfsense_backend/app/agents/shared/middleware/filesystem.py @@ -48,11 +48,11 @@ from langgraph.types import Command from app.agents.shared.filesystem_selection import FilesystemMode from app.agents.shared.filesystem_state import SurfSenseFilesystemState -from app.agents.new_chat.middleware.kb_postgres_backend import ( +from app.agents.shared.middleware.kb_postgres_backend import ( KBPostgresBackend, paginate_listing, ) -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) from app.agents.shared.path_resolver import DOCUMENTS_ROOT diff --git a/surfsense_backend/app/agents/new_chat/middleware/flatten_system.py b/surfsense_backend/app/agents/shared/middleware/flatten_system.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/flatten_system.py rename to surfsense_backend/app/agents/shared/middleware/flatten_system.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/kb_persistence.py b/surfsense_backend/app/agents/shared/middleware/kb_persistence.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/kb_persistence.py rename to surfsense_backend/app/agents/shared/middleware/kb_persistence.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/kb_postgres_backend.py b/surfsense_backend/app/agents/shared/middleware/kb_postgres_backend.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/kb_postgres_backend.py rename to surfsense_backend/app/agents/shared/middleware/kb_postgres_backend.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py b/surfsense_backend/app/agents/shared/middleware/knowledge_search.py similarity index 99% rename from surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py rename to surfsense_backend/app/agents/shared/middleware/knowledge_search.py index 9d5aebc60..9fbfc2a3c 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py +++ b/surfsense_backend/app/agents/shared/middleware/knowledge_search.py @@ -634,7 +634,7 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg] if not flags.enable_kb_planner_runnable or flags.disable_new_agent_stack: return None - from app.agents.new_chat.middleware.retry_after import RetryAfterMiddleware + from app.agents.shared.middleware.retry_after import RetryAfterMiddleware try: self._planner = create_agent( diff --git a/surfsense_backend/app/agents/new_chat/middleware/knowledge_tree.py b/surfsense_backend/app/agents/shared/middleware/knowledge_tree.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/knowledge_tree.py rename to surfsense_backend/app/agents/shared/middleware/knowledge_tree.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/local_folder_backend.py b/surfsense_backend/app/agents/shared/middleware/local_folder_backend.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/local_folder_backend.py rename to surfsense_backend/app/agents/shared/middleware/local_folder_backend.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/memory_injection.py b/surfsense_backend/app/agents/shared/middleware/memory_injection.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/memory_injection.py rename to surfsense_backend/app/agents/shared/middleware/memory_injection.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/multi_root_local_folder_backend.py b/surfsense_backend/app/agents/shared/middleware/multi_root_local_folder_backend.py similarity index 99% rename from surfsense_backend/app/agents/new_chat/middleware/multi_root_local_folder_backend.py rename to surfsense_backend/app/agents/shared/middleware/multi_root_local_folder_backend.py index a5add6248..220890fea 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/multi_root_local_folder_backend.py +++ b/surfsense_backend/app/agents/shared/middleware/multi_root_local_folder_backend.py @@ -15,7 +15,7 @@ from deepagents.backends.protocol import ( WriteResult, ) -from app.agents.new_chat.middleware.local_folder_backend import LocalFolderBackend +from app.agents.shared.middleware.local_folder_backend import LocalFolderBackend _INVALID_PATH = "invalid_path" _FILE_NOT_FOUND = "file_not_found" diff --git a/surfsense_backend/app/agents/new_chat/middleware/noop_injection.py b/surfsense_backend/app/agents/shared/middleware/noop_injection.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/noop_injection.py rename to surfsense_backend/app/agents/shared/middleware/noop_injection.py diff --git a/surfsense_backend/app/agents/new_chat/middleware/otel_span.py b/surfsense_backend/app/agents/shared/middleware/otel_span.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/otel_span.py rename to surfsense_backend/app/agents/shared/middleware/otel_span.py diff --git a/surfsense_backend/app/agents/shared/middleware/permission.py b/surfsense_backend/app/agents/shared/middleware/permission.py new file mode 100644 index 000000000..8601a3296 --- /dev/null +++ b/surfsense_backend/app/agents/shared/middleware/permission.py @@ -0,0 +1,427 @@ +""" +PermissionMiddleware — pattern-based allow/deny/ask with HITL fallback. + +LangChain's :class:`HumanInTheLoopMiddleware` only supports a static +"this tool always asks" decision per tool. There's no rule-based +allow/deny/ask layered ruleset, no glob patterns, no per-search-space or +per-thread overrides, and no auto-deny synthesis. + +This middleware ports OpenCode's ``packages/opencode/src/permission/index.ts`` +ruleset model on top of SurfSense's existing ``interrupt({type, action, +context})`` payload shape (see ``app/agents/new_chat/tools/hitl.py``) so +the frontend keeps working unchanged. + +Operation: +1. ``aafter_model`` inspects the latest ``AIMessage.tool_calls``. +2. For each call, the middleware builds a list of ``patterns`` (the + tool name plus any tool-specific patterns from the resolver). It + evaluates each pattern against the layered rulesets and aggregates + the results: ``deny`` > ``ask`` > ``allow``. +3. On ``deny``: replaces the call with a synthetic ``ToolMessage`` + containing a :class:`StreamingError`. +4. On ``ask``: raises a SurfSense-style ``interrupt(...)``. Both the legacy + SurfSense shape and LangChain HITL ``{"decisions": [{"type": ...}]}`` + replies are accepted via :func:`_normalize_permission_decision`. + - ``once``: proceed. + - ``approve_always``: also persist allow rules for ``request.always`` patterns. + - ``reject`` w/o feedback: raise :class:`RejectedError`. + - ``reject`` w/ feedback: raise :class:`CorrectedError`. +5. On ``allow``: proceed unchanged. + +The middleware also performs a *pre-model* tool-filter step (the +``before_model`` hook) so globally denied tools are stripped from the +exposed tool list before the model gets to see them. This mirrors +OpenCode's ``Permission.disabled`` and dramatically reduces the chance +the model emits a deny-only call. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from langchain.agents.middleware.types import ( + AgentMiddleware, + AgentState, + ContextT, +) +from langchain_core.messages import AIMessage, ToolMessage +from langgraph.runtime import Runtime +from langgraph.types import interrupt + +from app.agents.shared.errors import ( + CorrectedError, + RejectedError, + StreamingError, +) +from app.agents.shared.permissions import ( + Rule, + Ruleset, + aggregate_action, + evaluate_many, +) +from app.observability import metrics as ot_metrics, otel as ot + +logger = logging.getLogger(__name__) + + +# Mapping ``tool_name -> resolver`` that converts ``args`` to a list of +# patterns to evaluate. The first pattern is conventionally the bare +# tool name; later entries narrow down to specific resources. +PatternResolver = Callable[[dict[str, Any]], list[str]] + + +def _default_pattern_resolver(name: str) -> PatternResolver: + def _resolve(args: dict[str, Any]) -> list[str]: + # Bare name covers the default catch-all; primary-arg fallbacks + # are best added per-tool by callers. + del args + return [name] + + return _resolve + + +# Translation from the LangChain HITL envelope (what ``stream_resume_chat`` +# sends) to SurfSense's legacy ``decision_type`` shape. ``edit`` keeps the +# original tool args — tools needing argument edits should use +# ``request_approval`` from ``app/agents/new_chat/tools/hitl.py``. +_LC_TYPE_TO_PERMISSION_DECISION: dict[str, str] = { + "approve": "once", + "reject": "reject", + "edit": "once", + "approve_always": "approve_always", +} + + +def _normalize_permission_decision(decision: Any) -> dict[str, Any]: + """Coerce any accepted reply shape into ``{"decision_type": ..., "feedback"?}``. + + Falls back to ``reject`` (with a warning) on unrecognized payloads so the + middleware fails closed. + """ + if isinstance(decision, str): + return {"decision_type": decision} + if not isinstance(decision, dict): + logger.warning( + "Unrecognized permission resume value (%s); treating as reject", + type(decision).__name__, + ) + return {"decision_type": "reject"} + + if decision.get("decision_type"): + return decision + + payload: dict[str, Any] = decision + decisions = decision.get("decisions") + if isinstance(decisions, list) and decisions: + first = decisions[0] + if isinstance(first, dict): + payload = first + + raw_type = payload.get("type") or payload.get("decision_type") + if not raw_type: + logger.warning( + "Permission resume missing decision type (keys=%s); treating as reject", + list(payload.keys()), + ) + return {"decision_type": "reject"} + + raw_type = str(raw_type).lower() + mapped = _LC_TYPE_TO_PERMISSION_DECISION.get(raw_type) + if mapped is None: + # Tolerate legacy values arriving without ``decision_type`` wrapping. + if raw_type in {"once", "approve_always", "reject"}: + mapped = raw_type + else: + logger.warning( + "Unknown permission decision type %r; treating as reject", raw_type + ) + mapped = "reject" + + if raw_type == "edit": + logger.warning( + "Permission middleware received an 'edit' decision; original args " + "kept (edits not merged here)." + ) + + out: dict[str, Any] = {"decision_type": mapped} + feedback = payload.get("feedback") or payload.get("message") + if isinstance(feedback, str) and feedback.strip(): + out["feedback"] = feedback + return out + + +class PermissionMiddleware(AgentMiddleware): # type: ignore[type-arg] + """Allow/deny/ask layer over the agent's tool calls. + + Args: + rulesets: Layered rulesets to evaluate. Earlier entries are + overridden by later ones (last-match-wins). Typical layering: + ``defaults < global < space < thread < runtime_approved``. + pattern_resolvers: Optional per-tool callables that return a list + of patterns to evaluate. When a tool isn't listed, the bare + tool name is used as the only pattern. + runtime_ruleset: Mutable :class:`Ruleset` that the middleware + extends in-place when the user replies ``"approve_always"`` to + an ask interrupt. Reused across all calls in the same agent + instance so newly-allowed rules apply to subsequent calls. + always_emit_interrupt_payload: If True, every ask uses the + SurfSense interrupt wire format (default). Set False to + disable interrupts and treat ``ask`` as ``deny`` for + non-interactive deployments. + """ + + tools = () + + def __init__( + self, + *, + rulesets: list[Ruleset] | None = None, + pattern_resolvers: dict[str, PatternResolver] | None = None, + runtime_ruleset: Ruleset | None = None, + always_emit_interrupt_payload: bool = True, + ) -> None: + super().__init__() + self._static_rulesets: list[Ruleset] = list(rulesets or []) + self._pattern_resolvers: dict[str, PatternResolver] = dict( + pattern_resolvers or {} + ) + self._runtime_ruleset: Ruleset = runtime_ruleset or Ruleset( + origin="runtime_approved" + ) + self._emit_interrupt = always_emit_interrupt_payload + + # ------------------------------------------------------------------ + # Tool-filter step (mirrors OpenCode's ``Permission.disabled``) + # ------------------------------------------------------------------ + + def _globally_denied(self, tool_name: str) -> bool: + """Return True if a deny rule with no narrowing pattern matches.""" + rules = evaluate_many(tool_name, ["*"], *self._all_rulesets()) + return aggregate_action(rules) == "deny" + + def _all_rulesets(self) -> list[Ruleset]: + return [*self._static_rulesets, self._runtime_ruleset] + + # NOTE: ``before_model`` filtering of the tools list is left to the + # agent factory. This middleware only blocks at execution time — and + # only via the rule-evaluator path, not by mutating ``request.tools``. + # Mutating ``request.tools`` per-call would invalidate provider + # prompt-cache prefixes (see Operational risks: prompt-cache regression). + + # ------------------------------------------------------------------ + # Tool-call evaluation + # ------------------------------------------------------------------ + + def _resolve_patterns(self, tool_name: str, args: dict[str, Any]) -> list[str]: + resolver = self._pattern_resolvers.get( + tool_name, _default_pattern_resolver(tool_name) + ) + try: + patterns = resolver(args or {}) + except Exception: + logger.exception( + "Pattern resolver for %s raised; using bare name", tool_name + ) + patterns = [tool_name] + if not patterns: + patterns = [tool_name] + return patterns + + def _evaluate( + self, tool_name: str, args: dict[str, Any] + ) -> tuple[str, list[str], list[Rule]]: + patterns = self._resolve_patterns(tool_name, args) + rules = evaluate_many(tool_name, patterns, *self._all_rulesets()) + action = aggregate_action(rules) + return action, patterns, rules + + # ------------------------------------------------------------------ + # HITL ask flow — SurfSense wire format + # ------------------------------------------------------------------ + + def _raise_interrupt( + self, + *, + tool_name: str, + args: dict[str, Any], + patterns: list[str], + rules: list[Rule], + ) -> dict[str, Any]: + """Block on user approval via SurfSense's ``interrupt`` shape.""" + if not self._emit_interrupt: + return {"decision_type": "reject"} + + # ``params`` (NOT ``args``) is what SurfSense's streaming + # normalizer forwards. Other fields move into ``context``. + payload = { + "type": "permission_ask", + "action": {"tool": tool_name, "params": args or {}}, + "context": { + "patterns": patterns, + "rules": [ + { + "permission": r.permission, + "pattern": r.pattern, + "action": r.action, + } + for r in rules + ], + # Rules of thumb for the frontend: surface the patterns + # the user can promote to "approve_always" with a single reply. + "always": patterns, + }, + } + # Open ``permission.asked`` + ``interrupt.raised`` OTel spans + # (no-op when OTel is disabled) so dashboards can correlate + # "we asked X" with "interrupt was actually delivered". + with ( + ot.permission_asked_span( + permission=tool_name, + pattern=patterns[0] if patterns else None, + extra={"permission.patterns": list(patterns)}, + ), + ot.interrupt_span(interrupt_type="permission_ask"), + ): + ot_metrics.record_permission_ask(permission=tool_name) + ot_metrics.record_interrupt(interrupt_type="permission_ask") + decision = interrupt(payload) + return _normalize_permission_decision(decision) + + def _persist_always(self, tool_name: str, patterns: list[str]) -> None: + """Promote ``approve_always`` reply into runtime allow rules. + + Persistence to ``agent_permission_rules`` is done by the + streaming layer (``stream_new_chat``) once it observes the + ``approve_always`` reply — the middleware just keeps an + in-memory copy so subsequent calls in the same stream see the rule. + """ + for pattern in patterns: + self._runtime_ruleset.rules.append( + Rule(permission=tool_name, pattern=pattern, action="allow") + ) + + # ------------------------------------------------------------------ + # Synthesizing deny -> ToolMessage + # ------------------------------------------------------------------ + + @staticmethod + def _deny_message( + tool_call: dict[str, Any], + rule: Rule, + ) -> ToolMessage: + err = StreamingError( + code="permission_denied", + retryable=False, + suggestion=( + f"rule permission={rule.permission!r} pattern={rule.pattern!r} " + f"blocked this call" + ), + ) + return ToolMessage( + content=( + f"Permission denied: rule {rule.permission}/{rule.pattern} " + f"blocked tool {tool_call.get('name')!r}." + ), + tool_call_id=tool_call.get("id") or "", + name=tool_call.get("name"), + status="error", + additional_kwargs={"error": err.model_dump()}, + ) + + # ------------------------------------------------------------------ + # The hook: aafter_model + # ------------------------------------------------------------------ + + def _process( + self, + state: AgentState, + runtime: Runtime[Any], + ) -> dict[str, Any] | None: + del runtime # unused + messages = state.get("messages") or [] + if not messages: + return None + last = messages[-1] + if not isinstance(last, AIMessage) or not last.tool_calls: + return None + + deny_messages: list[ToolMessage] = [] + kept_calls: list[dict[str, Any]] = [] + any_change = False + + for raw in last.tool_calls: + call = ( + dict(raw) + if isinstance(raw, dict) + else { + "name": getattr(raw, "name", None), + "args": getattr(raw, "args", {}), + "id": getattr(raw, "id", None), + "type": "tool_call", + } + ) + name = call.get("name") or "" + args = call.get("args") or {} + action, patterns, rules = self._evaluate(name, args) + + if action == "deny": + # Find the deny rule for the suggestion text + deny_rule = next((r for r in rules if r.action == "deny"), rules[0]) + deny_messages.append(self._deny_message(call, deny_rule)) + any_change = True + continue + + if action == "ask": + decision = self._raise_interrupt( + tool_name=name, args=args, patterns=patterns, rules=rules + ) + kind = str(decision.get("decision_type") or "reject").lower() + if kind == "once": + kept_calls.append(call) + elif kind == "approve_always": + self._persist_always(name, patterns) + kept_calls.append(call) + elif kind == "reject": + feedback = decision.get("feedback") + if isinstance(feedback, str) and feedback.strip(): + raise CorrectedError(feedback, tool=name) + raise RejectedError( + tool=name, pattern=patterns[0] if patterns else None + ) + else: + logger.warning( + "Unknown permission decision %r; treating as reject", kind + ) + raise RejectedError(tool=name) + continue + + # allow + kept_calls.append(call) + + if not any_change and len(kept_calls) == len(last.tool_calls): + return None + + updated = last.model_copy(update={"tool_calls": kept_calls}) + result_messages: list[Any] = [updated] + if deny_messages: + result_messages.extend(deny_messages) + return {"messages": result_messages} + + def after_model( # type: ignore[override] + self, state: AgentState, runtime: Runtime[ContextT] + ) -> dict[str, Any] | None: + return self._process(state, runtime) + + async def aafter_model( # type: ignore[override] + self, state: AgentState, runtime: Runtime[ContextT] + ) -> dict[str, Any] | None: + return self._process(state, runtime) + + +__all__ = [ + "PatternResolver", + "PermissionMiddleware", + "_normalize_permission_decision", +] diff --git a/surfsense_backend/app/agents/new_chat/middleware/retry_after.py b/surfsense_backend/app/agents/shared/middleware/retry_after.py similarity index 100% rename from surfsense_backend/app/agents/new_chat/middleware/retry_after.py rename to surfsense_backend/app/agents/shared/middleware/retry_after.py diff --git a/surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py new file mode 100644 index 000000000..0294e2839 --- /dev/null +++ b/surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py @@ -0,0 +1,111 @@ +"""Fallback only on provider/network errors; let programming bugs raise.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from langchain.agents.middleware import ModelFallbackMiddleware + +from app.observability import metrics as ot_metrics, otel as ot + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from langchain.agents.middleware.types import ModelRequest, ModelResponse + from langchain_core.messages import AIMessage + + +# Matched by class name across the MRO so we don't have to import every +# provider SDK (openai/anthropic/google/...). Extend as new providers ship. +_FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( + { + "RateLimitError", + "APIStatusError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GatewayTimeoutError", + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "RemoteProtocolError", + "TimeoutError", + "TimeoutException", + } +) + + +def _is_fallback_eligible(exc: BaseException) -> bool: + return any(cls.__name__ in _FALLBACK_ELIGIBLE_NAMES for cls in type(exc).__mro__) + + +class ScopedModelFallbackMiddleware(ModelFallbackMiddleware): + """Re-raise non-provider exceptions instead of walking the fallback chain.""" + + def wrap_model_call( # type: ignore[override] + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], + ) -> ModelResponse[Any] | AIMessage: + last_exception: Exception + try: + return handler(request) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + + for attempt, fallback_model in enumerate(self.models, start=1): + ot.add_event( + "model.fallback", + { + "fallback.attempt": attempt, + "fallback.from": attempt - 1, + "fallback.to": attempt, + "fallback.reason": ot_metrics.categorize_exception(last_exception), + }, + ) + try: + return handler(request.override(model=fallback_model)) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + continue + + raise last_exception + + async def awrap_model_call( # type: ignore[override] + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]], + ) -> ModelResponse[Any] | AIMessage: + last_exception: Exception + try: + return await handler(request) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + + for attempt, fallback_model in enumerate(self.models, start=1): + ot.add_event( + "model.fallback", + { + "fallback.attempt": attempt, + "fallback.from": attempt - 1, + "fallback.to": attempt, + "fallback.reason": ot_metrics.categorize_exception(last_exception), + }, + ) + try: + return await handler(request.override(model=fallback_model)) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + continue + + raise last_exception diff --git a/surfsense_backend/app/agents/shared/middleware/skills_backends.py b/surfsense_backend/app/agents/shared/middleware/skills_backends.py new file mode 100644 index 000000000..091926627 --- /dev/null +++ b/surfsense_backend/app/agents/shared/middleware/skills_backends.py @@ -0,0 +1,344 @@ +"""Skills backends for SurfSense. + +Implements two minimal :class:`deepagents.backends.protocol.BackendProtocol` +subclasses tailored for use with :class:`deepagents.middleware.skills.SkillsMiddleware`. + +The middleware only needs four methods to load skills from a backend: + +* ``ls_info`` / ``als_info`` — list directories under a source path. +* ``download_files`` / ``adownload_files`` — fetch ``SKILL.md`` bytes. + +Other ``BackendProtocol`` methods (``read``/``write``/``edit``/``grep_raw`` …) +default to ``NotImplementedError`` from the base class. They are never reached +by the skills middleware because skill content is rendered into the system +prompt at agent build time, not edited at runtime. + +Two backends are provided: + +* :class:`BuiltinSkillsBackend` — disk-backed read of bundled skills from + ``app/agents/new_chat/skills/builtin/``. +* :class:`SearchSpaceSkillsBackend` — a thin read-only wrapper over + :class:`KBPostgresBackend` that filters notes under the privileged folder + ``/documents/_skills/``. + +Both backends are intentionally read-only: skill authoring happens out of band +(via filesystem or a search-space-admin route), so we never expose +``write`` / ``edit`` / ``upload_files``. The base class' ``NotImplementedError`` +gives a clean failure mode if anything tries. +""" + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import Callable +from dataclasses import replace +from pathlib import Path +from typing import TYPE_CHECKING + +from deepagents.backends.composite import CompositeBackend +from deepagents.backends.protocol import ( + BackendProtocol, + FileDownloadResponse, + FileInfo, +) +from deepagents.backends.state import StateBackend + +if TYPE_CHECKING: + from langchain.tools import ToolRuntime + + from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend + +logger = logging.getLogger(__name__) + + +# Limit per Agent Skills spec; matches deepagents.middleware.skills.MAX_SKILL_FILE_SIZE. +_MAX_SKILL_FILE_SIZE = 10 * 1024 * 1024 + + +def _default_builtin_root() -> Path: + """Return the absolute path to the bundled builtin skills directory. + + The skill assets still live at ``app/agents/new_chat/skills/builtin/`` (the + ``skills/`` tree migrates to the shared kernel in a later slice). This module + now lives under ``app/agents/shared/middleware/``, so we walk up to + ``app/agents/`` and back into ``new_chat/skills/builtin``. Once skills move, + this becomes ``Path(__file__).resolve().parent.parent / "skills" / "builtin"``. + """ + agents_dir = Path(__file__).resolve().parent.parent.parent + return (agents_dir / "new_chat" / "skills" / "builtin").resolve() + + +class BuiltinSkillsBackend(BackendProtocol): + """Read-only disk-backed skills source. + + Maps a virtual ``/skills/builtin/`` namespace onto a directory on local disk, + where each skill is its own subdirectory containing a ``SKILL.md`` file:: + + //SKILL.md + + The middleware calls :meth:`als_info` with the source path and expects a + ``list[FileInfo]`` whose ``is_dir=True`` entries are descended into. Then it + calls :meth:`adownload_files` with the synthesized ``SKILL.md`` paths and + parses YAML frontmatter from the returned ``content`` bytes. + + Mounting under :class:`~deepagents.backends.composite.CompositeBackend` at + prefix ``/skills/builtin/`` means the middleware can issue paths like + ``/skills/builtin/kb-research/SKILL.md`` which the composite strips down to + ``/kb-research/SKILL.md`` before forwarding here. We treat any leading + slash as anchoring at :attr:`root`. + """ + + def __init__(self, root: Path | str | None = None) -> None: + self.root: Path = Path(root).resolve() if root else _default_builtin_root() + if not self.root.exists(): + logger.info( + "BuiltinSkillsBackend root %s does not exist; skills will be empty.", + self.root, + ) + + def _resolve(self, path: str) -> Path: + """Resolve a virtual posix path under :attr:`root`, refusing escapes.""" + bare = path.lstrip("/") + candidate = (self.root / bare).resolve() if bare else self.root + # Refuse symlink/.. traversal that escapes the root. + try: + candidate.relative_to(self.root) + except ValueError as exc: + raise ValueError(f"path {path!r} escapes builtin skills root") from exc + return candidate + + def ls_info(self, path: str) -> list[FileInfo]: + try: + target = self._resolve(path) + except ValueError as exc: + logger.warning("BuiltinSkillsBackend.ls_info refused: %s", exc) + return [] + if not target.exists() or not target.is_dir(): + return [] + + infos: list[FileInfo] = [] + # Build virtual paths anchored at "/" because CompositeBackend already + # stripped the route prefix before calling us. + target_virtual = ( + "/" + if target == self.root + else ("/" + str(target.relative_to(self.root)).replace("\\", "/")) + ) + for child in sorted(target.iterdir()): + if child.name == "__pycache__" or child.name.startswith("."): + continue + child_virtual = ( + target_virtual.rstrip("/") + "/" + child.name + if target_virtual != "/" + else "/" + child.name + ) + info: FileInfo = { + "path": child_virtual, + "is_dir": child.is_dir(), + } + if child.is_file(): + with contextlib.suppress(OSError): # pragma: no cover - defensive + info["size"] = child.stat().st_size + infos.append(info) + return infos + + def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: + responses: list[FileDownloadResponse] = [] + for p in paths: + try: + target = self._resolve(p) + except ValueError: + responses.append(FileDownloadResponse(path=p, error="invalid_path")) + continue + if not target.exists(): + responses.append(FileDownloadResponse(path=p, error="file_not_found")) + continue + if target.is_dir(): + responses.append(FileDownloadResponse(path=p, error="is_directory")) + continue + try: + # Hard cap to avoid loading rogue mega-files into memory. + size = target.stat().st_size + if size > _MAX_SKILL_FILE_SIZE: + logger.warning( + "Builtin skill file %s exceeds %d bytes; truncating.", + target, + _MAX_SKILL_FILE_SIZE, + ) + with target.open("rb") as fh: + content = fh.read(_MAX_SKILL_FILE_SIZE) + else: + content = target.read_bytes() + except PermissionError: + responses.append( + FileDownloadResponse(path=p, error="permission_denied") + ) + continue + except OSError as exc: # pragma: no cover - defensive + logger.warning("Builtin skill read failed %s: %s", target, exc) + responses.append(FileDownloadResponse(path=p, error="file_not_found")) + continue + responses.append(FileDownloadResponse(path=p, content=content, error=None)) + return responses + + +class SearchSpaceSkillsBackend(BackendProtocol): + """Read-only view of search-space-authored skills. + + Wraps a :class:`KBPostgresBackend` and only ever reads under the privileged + folder ``/documents/_skills/`` (configurable). The folder is intended to be + writable only by search-space admins; this backend never writes. + + The skills middleware expects a layout like:: + + ///SKILL.md + + But the KB stores documents like ``/documents/_skills//SKILL.md``. + We expose the inner namespace by remapping each path. When mounted under + :class:`CompositeBackend` at prefix ``/skills/space/`` the paths the + middleware sees become ``/skills/space//SKILL.md``; the composite + strips ``/skills/space/`` and hands us ``//SKILL.md``, which we + rewrite to ``/documents/_skills//SKILL.md`` before forwarding to the + KB. + + No new database table is needed: the privileged folder convention is + enforced server-side outside of this class. We intentionally swallow any + write/edit attempts (the base class raises ``NotImplementedError``). + """ + + DEFAULT_KB_ROOT: str = "/documents/_skills" + + def __init__( + self, + kb_backend: KBPostgresBackend, + *, + kb_root: str = DEFAULT_KB_ROOT, + ) -> None: + self._kb = kb_backend + # Normalize trailing slash off so we can join cleanly. + self._kb_root = kb_root.rstrip("/") or "/" + + def _to_kb(self, path: str) -> str: + """Rewrite a virtual path into the underlying KB namespace.""" + bare = path.lstrip("/") + if not bare: + return self._kb_root + return f"{self._kb_root}/{bare}" + + def _from_kb(self, kb_path: str) -> str: + """Rewrite a KB path back into our virtual namespace.""" + if not kb_path.startswith(self._kb_root): + return kb_path # pragma: no cover - defensive + rel = kb_path[len(self._kb_root) :] + return rel if rel.startswith("/") else "/" + rel + + def ls_info(self, path: str) -> list[FileInfo]: + # KBPostgresBackend exposes only the async API meaningfully; the sync + # path falls back to ``asyncio.to_thread(...)`` in the base class. We + # keep this stub to satisfy abstract resolution; the middleware calls + # ``als_info``. + raise NotImplementedError("SearchSpaceSkillsBackend is async-only") + + async def als_info(self, path: str) -> list[FileInfo]: + kb_path = self._to_kb(path) + try: + infos = await self._kb.als_info(kb_path) + except Exception as exc: # pragma: no cover - defensive + logger.warning("SearchSpaceSkillsBackend.als_info failed: %s", exc) + return [] + remapped: list[FileInfo] = [] + for info in infos: + kb_p = info.get("path", "") + if not kb_p.startswith(self._kb_root): + continue + remapped.append({**info, "path": self._from_kb(kb_p)}) + return remapped + + def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: + raise NotImplementedError("SearchSpaceSkillsBackend is async-only") + + async def adownload_files(self, paths: list[str]) -> list[FileDownloadResponse]: + kb_paths = [self._to_kb(p) for p in paths] + responses = await self._kb.adownload_files(kb_paths) + # Re-map response paths back to the virtual namespace so the middleware + # correlates them to the input list correctly. + remapped: list[FileDownloadResponse] = [] + for original, resp in zip(paths, responses, strict=True): + remapped.append(replace(resp, path=original)) + return remapped + + +SKILLS_BUILTIN_PREFIX = "/skills/builtin/" +SKILLS_SPACE_PREFIX = "/skills/space/" + + +def build_skills_backend_factory( + *, + builtin_root: Path | str | None = None, + search_space_id: int | None = None, +) -> Callable[[ToolRuntime], BackendProtocol]: + """Return a runtime-aware factory for the skills :class:`CompositeBackend`. + + When ``search_space_id`` is provided the composite includes a + :class:`SearchSpaceSkillsBackend` route at ``/skills/space/`` over a fresh + per-runtime :class:`KBPostgresBackend`, mirroring how + :func:`build_backend_resolver` constructs the main filesystem backend. + + When ``search_space_id`` is ``None`` (e.g., desktop-local mode or unit + tests) only the bundled :class:`BuiltinSkillsBackend` is exposed. + + Returning a factory rather than a fixed instance is intentional: the + underlying KB backend depends on per-call ``ToolRuntime`` state + (``staged_dirs``, ``files`` cache, runtime config), so a single shared + instance cannot serve multiple concurrent agent runs. + """ + builtin = BuiltinSkillsBackend(builtin_root) + + if search_space_id is None: + + def _factory_builtin_only(runtime: ToolRuntime) -> BackendProtocol: + # Default StateBackend is intentionally inert: any path outside the + # ``/skills/builtin/`` route resolves to an empty per-runtime state + # so the SkillsMiddleware can iterate sources without raising. + return CompositeBackend( + default=StateBackend(runtime), + routes={SKILLS_BUILTIN_PREFIX: builtin}, + ) + + return _factory_builtin_only + + def _factory_with_space(runtime: ToolRuntime) -> BackendProtocol: + # Imported lazily to avoid a hard dependency at module import time: + # ``KBPostgresBackend`` pulls in DB models, which are unnecessary for + # the unit-tested builtin path. + from app.agents.shared.middleware.kb_postgres_backend import ( + KBPostgresBackend, + ) + + kb = KBPostgresBackend(search_space_id, runtime) + space = SearchSpaceSkillsBackend(kb) + return CompositeBackend( + default=StateBackend(runtime), + routes={ + SKILLS_BUILTIN_PREFIX: builtin, + SKILLS_SPACE_PREFIX: space, + }, + ) + + return _factory_with_space + + +def default_skills_sources() -> list[str]: + """Return the canonical source list for SkillsMiddleware (built-in then space).""" + return [SKILLS_BUILTIN_PREFIX, SKILLS_SPACE_PREFIX] + + +__all__ = [ + "SKILLS_BUILTIN_PREFIX", + "SKILLS_SPACE_PREFIX", + "BuiltinSkillsBackend", + "SearchSpaceSkillsBackend", + "build_skills_backend_factory", + "default_skills_sources", +] diff --git a/surfsense_backend/app/agents/new_chat/middleware/tool_call_repair.py b/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py similarity index 95% rename from surfsense_backend/app/agents/new_chat/middleware/tool_call_repair.py rename to surfsense_backend/app/agents/shared/middleware/tool_call_repair.py index 9f81a168b..96154e7ab 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/tool_call_repair.py +++ b/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py @@ -34,8 +34,6 @@ from langchain.agents.middleware.types import ( from langchain_core.messages import AIMessage from langgraph.runtime import Runtime -from app.agents.new_chat.tools.invalid_tool import INVALID_TOOL_NAME - logger = logging.getLogger(__name__) @@ -120,6 +118,11 @@ class ToolCallNameRepairMiddleware( return call # Stage 2 — invalid fallback + # Local import avoids a module-load cycle through the frozen single-agent + # package (new_chat.__init__ -> chat_deepagent -> middleware shim). + # Resolves to app.agents.shared.tools once tools migrate. + from app.agents.new_chat.tools.invalid_tool import INVALID_TOOL_NAME + if INVALID_TOOL_NAME in registered: original_args = call.get("args") or {} error_msg = ( diff --git a/surfsense_backend/app/agents/shared/receipt.py b/surfsense_backend/app/agents/shared/receipt.py index 6f30067ee..b1318fb90 100644 --- a/surfsense_backend/app/agents/shared/receipt.py +++ b/surfsense_backend/app/agents/shared/receipt.py @@ -23,7 +23,7 @@ the receipt into the parent's ``receipts`` state via the append reducer. The KB write path is the one exception: file-tool calls cannot emit a durable receipt because the actual DB writes happen end-of-turn inside -:class:`app.agents.new_chat.middleware.kb_persistence.KnowledgeBasePersistenceMiddleware`. +:class:`app.agents.shared.middleware.kb_persistence.KnowledgeBasePersistenceMiddleware`. KB tools therefore emit a *provisional* receipt with ``status="pending"``; the persistence middleware flips it to ``"success"`` or ``"failed"`` before returning control to the parent. diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index aa83b6847..967b88e7a 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -30,7 +30,7 @@ from app.agents.shared.filesystem_selection import ( FilesystemSelection, LocalFilesystemMount, ) -from app.agents.new_chat.middleware.busy_mutex import ( +from app.agents.shared.middleware.busy_mutex import ( get_cancel_state, is_cancel_requested, manager, diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 323ea2736..727f1f9ad 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -40,12 +40,12 @@ from app.agents.shared.llm_config import ( load_global_llm_config_by_id, ) from app.agents.shared.mention_resolver import resolve_mentions, substitute_in_text -from app.agents.new_chat.middleware.busy_mutex import ( +from app.agents.shared.middleware.busy_mutex import ( end_turn, get_cancel_state, is_cancel_requested, ) -from app.agents.new_chat.middleware.kb_persistence import ( +from app.agents.shared.middleware.kb_persistence import ( commit_staged_filesystem_state, ) from app.db import ( 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 4cba8fdad..0cbfdc160 100644 --- a/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py +++ b/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py @@ -12,7 +12,7 @@ from collections.abc import AsyncGenerator from typing import Any from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware.kb_persistence import ( +from app.agents.shared.middleware.kb_persistence import ( commit_staged_filesystem_state, ) from app.services.new_streaming_service import VercelStreamingService diff --git a/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py b/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py index a95a104b2..0baae627b 100644 --- a/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py +++ b/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py @@ -8,7 +8,7 @@ import time from typing import Any, Literal from app.agents.shared.errors import BusyError -from app.agents.new_chat.middleware.busy_mutex import ( +from app.agents.shared.middleware.busy_mutex import ( get_cancel_state, is_cancel_requested, ) diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py index 984115e88..f71e18770 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py @@ -32,7 +32,7 @@ import anyio from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection -from app.agents.new_chat.middleware.busy_mutex import end_turn +from app.agents.shared.middleware.busy_mutex import end_turn from app.config import config as _app_config from app.db import ChatVisibility, async_session_maker from app.observability import otel as ot diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py index ff5f5c2f5..6ed6fa166 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py @@ -26,7 +26,7 @@ import anyio from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection -from app.agents.new_chat.middleware.busy_mutex import end_turn +from app.agents.shared.middleware.busy_mutex import end_turn from app.config import config as _app_config from app.db import ChatVisibility, async_session_maker from app.observability import otel as ot diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/shared/rate_limit_recovery.py b/surfsense_backend/app/tasks/chat/streaming/flows/shared/rate_limit_recovery.py index 6b3857594..dd0f050d7 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/shared/rate_limit_recovery.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/shared/rate_limit_recovery.py @@ -17,7 +17,7 @@ from typing import Literal from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.middleware.busy_mutex import end_turn +from app.agents.shared.middleware.busy_mutex import end_turn from app.observability import otel as ot from app.services.auto_model_pin_service import ( mark_runtime_cooldown, diff --git a/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py b/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py index 910d882a7..9a911cc44 100644 --- a/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py +++ b/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py @@ -8,7 +8,7 @@ from datetime import UTC, datetime, timedelta import numpy as np import pytest -from app.agents.new_chat.middleware.knowledge_search import search_knowledge_base +from app.agents.shared.middleware.knowledge_search import search_knowledge_base from .conftest import DUMMY_EMBEDDING @@ -27,11 +27,11 @@ async def test_search_knowledge_base_applies_date_filters( yield db_session monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.shielded_async_session", + "app.agents.shared.middleware.knowledge_search.shielded_async_session", fake_shielded_async_session, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.embed_texts", + "app.agents.shared.middleware.knowledge_search.embed_texts", lambda texts: [np.array(DUMMY_EMBEDDING) for _ in texts], ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py b/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py index 7772a38d4..5e3955bf1 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py @@ -11,7 +11,7 @@ from langchain_core.messages import ToolMessage from langchain_core.tools import tool from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.new_chat.middleware.action_log import ActionLogMiddleware +from app.agents.shared.middleware.action_log import ActionLogMiddleware from app.agents.new_chat.tools.registry import ToolDefinition @@ -58,7 +58,7 @@ def _disabled_flags() -> AgentFeatureFlags: def patch_get_flags(): def _patch(flags: AgentFeatureFlags): return patch( - "app.agents.new_chat.middleware.action_log.get_flags", + "app.agents.shared.middleware.action_log.get_flags", return_value=flags, ) @@ -360,7 +360,7 @@ class TestActionLogDispatch: patch_get_flags(_enabled_flags()), patch("app.db.shielded_async_session", side_effect=lambda: factory()), patch( - "app.agents.new_chat.middleware.action_log.adispatch_custom_event", + "app.agents.shared.middleware.action_log.adispatch_custom_event", dispatch_mock, ), ): @@ -395,7 +395,7 @@ class TestActionLogDispatch: patch_get_flags(_enabled_flags()), patch("app.db.shielded_async_session", side_effect=_exploding_session), patch( - "app.agents.new_chat.middleware.action_log.adispatch_custom_event", + "app.agents.shared.middleware.action_log.adispatch_custom_event", dispatch_mock, ), ): diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_busy_mutex.py b/surfsense_backend/tests/unit/agents/new_chat/test_busy_mutex.py index 0c5f41bfd..b460c3e65 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_busy_mutex.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_busy_mutex.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest from app.agents.shared.errors import BusyError -from app.agents.new_chat.middleware.busy_mutex import ( +from app.agents.shared.middleware.busy_mutex import ( BusyMutexMiddleware, end_turn, get_cancel_event, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_compaction.py b/surfsense_backend/tests/unit/agents/new_chat/test_compaction.py index c6d4cc452..cf33f2260 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_compaction.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_compaction.py @@ -10,7 +10,7 @@ from langchain_core.messages import ( ToolMessage, ) -from app.agents.new_chat.middleware.compaction import ( +from app.agents.shared.middleware.compaction import ( PROTECTED_SYSTEM_PREFIXES, _is_protected_system_message, _sanitize_message_content, @@ -72,7 +72,7 @@ class TestPartitionMessages: # SurfSenseCompactionMiddleware without a real model, but the # override path needs ``_lc_helper`` to delegate to. We mock # that with a simple slicing partitioner equivalent to the real one. - from app.agents.new_chat.middleware.compaction import ( + from app.agents.shared.middleware.compaction import ( SurfSenseCompactionMiddleware, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py index ba2246413..fbb32713e 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py @@ -7,7 +7,7 @@ from typing import Any import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage -from app.agents.new_chat.middleware.context_editing import ( +from app.agents.shared.middleware.context_editing import ( SpillToBackendEdit, _build_spill_placeholder, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_dedup_tool_calls.py b/surfsense_backend/tests/unit/agents/new_chat/test_dedup_tool_calls.py index 61d9b499f..b0a3b2e00 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_dedup_tool_calls.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_dedup_tool_calls.py @@ -6,7 +6,7 @@ import pytest from langchain_core.messages import AIMessage from langchain_core.tools import StructuredTool -from app.agents.new_chat.middleware.dedup_tool_calls import ( +from app.agents.shared.middleware.dedup_tool_calls import ( DedupHITLToolCallsMiddleware, ) @@ -137,7 +137,7 @@ def test_full_args_dedup_keeps_distinct_calls_sharing_a_field() -> None: With :func:`dedup_key_full_args` only fully identical arg dicts dedup. """ - from app.agents.new_chat.middleware.dedup_tool_calls import dedup_key_full_args + from app.agents.shared.middleware.dedup_tool_calls import dedup_key_full_args tool = _make_tool("createJiraIssue", dedup_key=dedup_key_full_args) mw = DedupHITLToolCallsMiddleware(agent_tools=[tool]) @@ -179,7 +179,7 @@ def test_full_args_dedup_keeps_distinct_calls_sharing_a_field() -> None: def test_full_args_dedup_drops_only_exact_duplicates() -> None: - from app.agents.new_chat.middleware.dedup_tool_calls import dedup_key_full_args + from app.agents.shared.middleware.dedup_tool_calls import dedup_key_full_args tool = _make_tool("createJiraIssue", dedup_key=dedup_key_full_args) mw = DedupHITLToolCallsMiddleware(agent_tools=[tool]) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_desktop_safety_rules.py b/surfsense_backend/tests/unit/agents/new_chat/test_desktop_safety_rules.py index b513d68d8..e387e53f0 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_desktop_safety_rules.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_desktop_safety_rules.py @@ -10,7 +10,7 @@ from __future__ import annotations import pytest -from app.agents.new_chat.middleware.permission import PermissionMiddleware +from app.agents.shared.middleware.permission import PermissionMiddleware from app.agents.shared.permissions import ( Rule, Ruleset, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py b/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py index 802129bf6..0bb338628 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage -from app.agents.new_chat.middleware.doom_loop import DoomLoopMiddleware, _signature +from app.agents.shared.middleware.doom_loop import DoomLoopMiddleware, _signature pytestmark = pytest.mark.unit diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_flatten_system.py b/surfsense_backend/tests/unit/agents/new_chat/test_flatten_system.py index 6c323d920..f38d1ebc2 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_flatten_system.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_flatten_system.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock import pytest from langchain_core.messages import HumanMessage, SystemMessage -from app.agents.new_chat.middleware.flatten_system import ( +from app.agents.shared.middleware.flatten_system import ( FlattenSystemMessageMiddleware, _flatten_text_blocks, _flattened_request, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py b/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py index 346271f4b..92c4ba1b3 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage, HumanMessage -from app.agents.new_chat.middleware.noop_injection import ( +from app.agents.shared.middleware.noop_injection import ( NOOP_TOOL_NAME, NoopInjectionMiddleware, _last_ai_has_tool_calls, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py b/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py index dc59c6dac..0e924210d 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py @@ -8,7 +8,7 @@ from unittest.mock import MagicMock import pytest from langchain_core.messages import AIMessage, ToolMessage -from app.agents.new_chat.middleware.otel_span import ( +from app.agents.shared.middleware.otel_span import ( OtelSpanMiddleware, _annotate_model_response, _annotate_tool_result, @@ -206,13 +206,13 @@ class TestMiddlewareIntegration: duration_calls: list[dict[str, Any]] = [] token_calls: list[dict[str, Any]] = [] monkeypatch.setattr( - "app.agents.new_chat.middleware.otel_span.ot_metrics.record_model_call_duration", + "app.agents.shared.middleware.otel_span.ot_metrics.record_model_call_duration", lambda duration_ms, **attrs: duration_calls.append( {"duration_ms": duration_ms, **attrs} ), ) monkeypatch.setattr( - "app.agents.new_chat.middleware.otel_span.ot_metrics.record_model_token_usage", + "app.agents.shared.middleware.otel_span.ot_metrics.record_model_token_usage", lambda **attrs: token_calls.append(attrs), ) @@ -257,11 +257,11 @@ class TestMiddlewareIntegration: errors: list[str] = [] monkeypatch.setattr( - "app.agents.new_chat.middleware.otel_span.ot_metrics.record_tool_call_error", + "app.agents.shared.middleware.otel_span.ot_metrics.record_tool_call_error", lambda *, tool_name: errors.append(tool_name), ) monkeypatch.setattr( - "app.agents.new_chat.middleware.otel_span.ot_metrics.record_tool_call_duration", + "app.agents.shared.middleware.otel_span.ot_metrics.record_tool_call_duration", lambda *args, **kwargs: None, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py b/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py index faf27328f..e84ed4f61 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py @@ -6,7 +6,7 @@ import pytest from langchain_core.messages import AIMessage, ToolMessage from app.agents.shared.errors import CorrectedError, RejectedError -from app.agents.new_chat.middleware.permission import ( +from app.agents.shared.middleware.permission import ( PermissionMiddleware, _normalize_permission_decision, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_retry_after.py b/surfsense_backend/tests/unit/agents/new_chat/test_retry_after.py index d23fd693b..b5890f65e 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_retry_after.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_retry_after.py @@ -4,7 +4,7 @@ from __future__ import annotations import pytest -from app.agents.new_chat.middleware.retry_after import ( +from app.agents.shared.middleware.retry_after import ( RetryAfterMiddleware, _extract_retry_after_seconds, _is_non_retryable, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_rm_rmdir_cloud.py b/surfsense_backend/tests/unit/agents/new_chat/test_rm_rmdir_cloud.py index aa08ab59f..4f0d4c48c 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_rm_rmdir_cloud.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_rm_rmdir_cloud.py @@ -21,8 +21,8 @@ from unittest.mock import AsyncMock import pytest from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware.filesystem import SurfSenseFilesystemMiddleware -from app.agents.new_chat.middleware.kb_postgres_backend import KBPostgresBackend +from app.agents.shared.middleware.filesystem import SurfSenseFilesystemMiddleware +from app.agents.shared.middleware.kb_postgres_backend import KBPostgresBackend pytestmark = pytest.mark.unit 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 eb9cf396c..b49cdfa1d 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.new_chat.middleware.skills_backends import ( +from app.agents.shared.middleware.skills_backends import ( SKILLS_BUILTIN_PREFIX, SKILLS_SPACE_PREFIX, BuiltinSkillsBackend, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_specialized_subagents.py b/surfsense_backend/tests/unit/agents/new_chat/test_specialized_subagents.py index 7259c49f8..79d517d9a 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_specialized_subagents.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_specialized_subagents.py @@ -4,7 +4,7 @@ from __future__ import annotations from langchain_core.tools import tool -from app.agents.new_chat.middleware.permission import PermissionMiddleware +from app.agents.shared.middleware.permission import PermissionMiddleware from app.agents.new_chat.subagents import ( build_connector_negotiator_subagent, build_explore_subagent, @@ -140,7 +140,7 @@ class TestExploreSubagent: def test_includes_dedup_and_patch_middleware(self) -> None: from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware - from app.agents.new_chat.middleware import DedupHITLToolCallsMiddleware + from app.agents.shared.middleware import DedupHITLToolCallsMiddleware spec = build_explore_subagent(tools=ALL_TOOLS) types = {type(m) for m in spec["middleware"]} # type: ignore[index] 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 e02a04774..0cd338ce3 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,7 +5,7 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage -from app.agents.new_chat.middleware.tool_call_repair import ( +from app.agents.shared.middleware.tool_call_repair import ( ToolCallNameRepairMiddleware, ) from app.agents.new_chat.tools.invalid_tool import INVALID_TOOL_NAME diff --git a/surfsense_backend/tests/unit/middleware/test_dedup_hitl_tool_calls.py b/surfsense_backend/tests/unit/middleware/test_dedup_hitl_tool_calls.py index 467ba6d5f..aa4bab204 100644 --- a/surfsense_backend/tests/unit/middleware/test_dedup_hitl_tool_calls.py +++ b/surfsense_backend/tests/unit/middleware/test_dedup_hitl_tool_calls.py @@ -2,7 +2,7 @@ import pytest from langchain_core.messages import AIMessage from langchain_core.tools import StructuredTool -from app.agents.new_chat.middleware.dedup_tool_calls import ( +from app.agents.shared.middleware.dedup_tool_calls import ( DedupHITLToolCallsMiddleware, wrap_dedup_key_by_arg_name, ) diff --git a/surfsense_backend/tests/unit/middleware/test_file_intent_middleware.py b/surfsense_backend/tests/unit/middleware/test_file_intent_middleware.py index 7fd3fe4a7..e1d522201 100644 --- a/surfsense_backend/tests/unit/middleware/test_file_intent_middleware.py +++ b/surfsense_backend/tests/unit/middleware/test_file_intent_middleware.py @@ -1,7 +1,7 @@ import pytest from langchain_core.messages import AIMessage, HumanMessage -from app.agents.new_chat.middleware.file_intent import ( +from app.agents.shared.middleware.file_intent import ( FileIntentMiddleware, FileOperationIntent, _fallback_path, diff --git a/surfsense_backend/tests/unit/middleware/test_filesystem_backends.py b/surfsense_backend/tests/unit/middleware/test_filesystem_backends.py index db16342f6..f0d7ab3e9 100644 --- a/surfsense_backend/tests/unit/middleware/test_filesystem_backends.py +++ b/surfsense_backend/tests/unit/middleware/test_filesystem_backends.py @@ -9,7 +9,7 @@ from app.agents.shared.filesystem_selection import ( FilesystemSelection, LocalFilesystemMount, ) -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) diff --git a/surfsense_backend/tests/unit/middleware/test_filesystem_middleware.py b/surfsense_backend/tests/unit/middleware/test_filesystem_middleware.py index 7b53948b3..482d01b7a 100644 --- a/surfsense_backend/tests/unit/middleware/test_filesystem_middleware.py +++ b/surfsense_backend/tests/unit/middleware/test_filesystem_middleware.py @@ -19,7 +19,7 @@ from types import SimpleNamespace import pytest from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware.filesystem import ( +from app.agents.shared.middleware.filesystem import ( SurfSenseFilesystemMiddleware, _build_filesystem_system_prompt, _build_tool_descriptions, diff --git a/surfsense_backend/tests/unit/middleware/test_filesystem_verification.py b/surfsense_backend/tests/unit/middleware/test_filesystem_verification.py index 889683fd2..80306a801 100644 --- a/surfsense_backend/tests/unit/middleware/test_filesystem_verification.py +++ b/surfsense_backend/tests/unit/middleware/test_filesystem_verification.py @@ -3,8 +3,8 @@ from pathlib import Path import pytest from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware.filesystem import SurfSenseFilesystemMiddleware -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.filesystem import SurfSenseFilesystemMiddleware +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) 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 ef95434bf..907f5080b 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.new_chat.middleware import kb_persistence +from app.agents.shared.middleware import 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 feca23d27..99abd7baa 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.new_chat.middleware import kb_persistence +from app.agents.shared.middleware import kb_persistence pytestmark = pytest.mark.unit diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py index 39ae70bad..72f54f400 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py @@ -6,7 +6,7 @@ import pytest from langchain_core.messages import AIMessage, HumanMessage from app.agents.shared.document_xml import build_document_xml as _build_document_xml -from app.agents.new_chat.middleware.knowledge_search import ( +from app.agents.shared.middleware.knowledge_search import ( KBSearchPlan, KnowledgeBaseSearchMiddleware, _normalize_optional_date_range, @@ -258,7 +258,7 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -301,7 +301,7 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -330,7 +330,7 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -375,11 +375,11 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.browse_recent_documents", + "app.agents.shared.middleware.knowledge_search.browse_recent_documents", fake_browse_recent_documents, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -422,11 +422,11 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.browse_recent_documents", + "app.agents.shared.middleware.knowledge_search.browse_recent_documents", fake_browse_recent_documents, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -549,11 +549,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.fetch_mentioned_documents", + "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -597,11 +597,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.fetch_mentioned_documents", + "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) @@ -642,11 +642,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.fetch_mentioned_documents", + "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", + "app.agents.shared.middleware.knowledge_search.search_knowledge_base", fake_search_knowledge_base, ) diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py b/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py index 741c2d980..87403c14a 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py @@ -9,7 +9,7 @@ contract cannot silently regress. from __future__ import annotations -from app.agents.new_chat.middleware.knowledge_tree import KnowledgeTreeMiddleware +from app.agents.shared.middleware.knowledge_tree import KnowledgeTreeMiddleware from app.agents.shared.path_resolver import DOCUMENTS_ROOT diff --git a/surfsense_backend/tests/unit/middleware/test_local_folder_backend.py b/surfsense_backend/tests/unit/middleware/test_local_folder_backend.py index 6e81ecf8e..2111e2f66 100644 --- a/surfsense_backend/tests/unit/middleware/test_local_folder_backend.py +++ b/surfsense_backend/tests/unit/middleware/test_local_folder_backend.py @@ -2,7 +2,7 @@ from pathlib import Path import pytest -from app.agents.new_chat.middleware.local_folder_backend import LocalFolderBackend +from app.agents.shared.middleware.local_folder_backend import LocalFolderBackend pytestmark = pytest.mark.unit diff --git a/surfsense_backend/tests/unit/middleware/test_multi_root_local_folder_backend.py b/surfsense_backend/tests/unit/middleware/test_multi_root_local_folder_backend.py index 43a671178..2176698ec 100644 --- a/surfsense_backend/tests/unit/middleware/test_multi_root_local_folder_backend.py +++ b/surfsense_backend/tests/unit/middleware/test_multi_root_local_folder_backend.py @@ -2,7 +2,7 @@ from pathlib import Path import pytest -from app.agents.new_chat.middleware.multi_root_local_folder_backend import ( +from app.agents.shared.middleware.multi_root_local_folder_backend import ( MultiRootLocalFolderBackend, ) diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py index 20e421eb5..a4bd1d56c 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py @@ -17,7 +17,7 @@ from typing import Any import pytest from app.agents.shared.errors import BusyError -from app.agents.new_chat.middleware.busy_mutex import request_cancel, reset_cancel +from app.agents.shared.middleware.busy_mutex import request_cancel, reset_cancel from app.tasks.chat.stream_new_chat import ( _classify_stream_exception as old_classify, _emit_stream_terminal_error as old_emit_terminal_error, diff --git a/surfsense_backend/tests/unit/test_stream_new_chat_contract.py b/surfsense_backend/tests/unit/test_stream_new_chat_contract.py index 794db41a7..9b29fdd6a 100644 --- a/surfsense_backend/tests/unit/test_stream_new_chat_contract.py +++ b/surfsense_backend/tests/unit/test_stream_new_chat_contract.py @@ -8,7 +8,7 @@ import pytest import app.tasks.chat.stream_new_chat as stream_new_chat_module from app.agents.shared.errors import BusyError -from app.agents.new_chat.middleware.busy_mutex import request_cancel, reset_cancel +from app.agents.shared.middleware.busy_mutex import request_cancel, reset_cancel from app.tasks.chat.stream_new_chat import ( StreamResult, _classify_stream_exception,