From 6a4dacda72bce83c6ce072e07ed928bdfe851240 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:53:49 +0200 Subject: [PATCH] refactor(multi-agent): add main-agent observability and lifecycle middleware factories --- .../middleware/main_agent/__init__.py | 0 .../middleware/main_agent/action_log.py | 36 ++++++++++++++ .../middleware/main_agent/anonymous_doc.py | 16 ++++++ .../middleware/main_agent/busy_mutex.py | 12 +++++ .../middleware/main_agent/otel.py | 12 +++++ .../middleware/main_agent/plugins.py | 49 +++++++++++++++++++ 6 files changed, 125 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py new file mode 100644 index 000000000..e69de29bb 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 new file mode 100644 index 000000000..c9f893d97 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py @@ -0,0 +1,36 @@ +"""Audit row per tool call (reversibility metadata).""" + +from __future__ import annotations + +import logging + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ActionLogMiddleware +from app.agents.new_chat.tools.registry import BUILTIN_TOOLS + +from ..shared.flags import enabled + + +def build_action_log_mw( + *, + flags: AgentFeatureFlags, + thread_id: int | None, + search_space_id: int, + user_id: str | None, +) -> ActionLogMiddleware | None: + if not enabled(flags, "enable_action_log") or thread_id is None: + return None + try: + tool_defs_by_name = {td.name: td for td in BUILTIN_TOOLS} + return ActionLogMiddleware( + thread_id=thread_id, + search_space_id=search_space_id, + user_id=user_id, + tool_definitions=tool_defs_by_name, + ) + except Exception: # pragma: no cover - defensive + logging.warning( + "ActionLogMiddleware init failed; running without it.", + exc_info=True, + ) + return None 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 new file mode 100644 index 000000000..afd54a2d3 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py @@ -0,0 +1,16 @@ +"""Anonymous document hydration from Redis (cloud only).""" + +from __future__ import annotations + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import AnonymousDocumentMiddleware + + +def build_anonymous_doc_mw( + *, + filesystem_mode: FilesystemMode, + anon_session_id: str | None, +) -> AnonymousDocumentMiddleware | None: + if filesystem_mode != FilesystemMode.CLOUD: + return None + return AnonymousDocumentMiddleware(anon_session_id=anon_session_id) 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 new file mode 100644 index 000000000..0ea53bf16 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py @@ -0,0 +1,12 @@ +"""Per-thread cooperative lock around the whole turn.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import BusyMutexMiddleware + +from ..shared.flags import enabled + + +def build_busy_mutex_mw(flags: AgentFeatureFlags) -> BusyMutexMiddleware | None: + return BusyMutexMiddleware() if enabled(flags, "enable_busy_mutex") else None 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 new file mode 100644 index 000000000..bd7516e65 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py @@ -0,0 +1,12 @@ +"""OTel spans on model and tool calls.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import OtelSpanMiddleware + +from ..shared.flags import enabled + + +def build_otel_mw(flags: AgentFeatureFlags) -> OtelSpanMiddleware | None: + return OtelSpanMiddleware() if enabled(flags, "enable_otel") else None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py new file mode 100644 index 000000000..4418e3806 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py @@ -0,0 +1,49 @@ +"""Tail-of-stack plugin slot driven by env allowlist.""" + +from __future__ import annotations + +import logging +from typing import Any + +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.plugin_loader import ( + PluginContext, + load_allowed_plugin_names_from_env, + load_plugin_middlewares, +) +from app.db import ChatVisibility + +from ..shared.flags import enabled + + +def build_plugin_middlewares( + *, + flags: AgentFeatureFlags, + search_space_id: int, + user_id: str | None, + visibility: ChatVisibility, + llm: BaseChatModel, +) -> list[Any]: + if not enabled(flags, "enable_plugin_loader"): + return [] + try: + allowed_names = load_allowed_plugin_names_from_env() + if not allowed_names: + return [] + return load_plugin_middlewares( + PluginContext.build( + search_space_id=search_space_id, + user_id=user_id, + thread_visibility=visibility, + llm=llm, + ), + allowed_plugin_names=allowed_names, + ) + except Exception: # pragma: no cover - defensive + logging.warning( + "Plugin loader failed; continuing without plugins.", + exc_info=True, + ) + return []