From b0ee44b2f15adfdf379c2b7f1967733a66a8e090 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:54:13 +0200 Subject: [PATCH] refactor(multi-agent): add main-agent safety and llm-shaping middleware factories --- .../middleware/main_agent/context_editing.py | 50 +++++++++++++++++++ .../middleware/main_agent/dedup_hitl.py | 13 +++++ .../middleware/main_agent/doom_loop.py | 12 +++++ .../middleware/main_agent/noop_injection.py | 12 +++++ .../middleware/main_agent/repair.py | 50 +++++++++++++++++++ .../middleware/main_agent/selector.py | 39 +++++++++++++++ .../middleware/main_agent/skills.py | 39 +++++++++++++++ 7 files changed, 215 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py 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 new file mode 100644 index 000000000..e8f99933e --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py @@ -0,0 +1,50 @@ +"""Spill + clear-tool-uses passes to keep payloads under budget.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from langchain_core.tools import BaseTool + +from app.agents.multi_agent_chat.main_agent.context_prune.prune_tool_names import ( + safe_exclude_tools, +) +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ( + ClearToolUsesEdit, + SpillingContextEditingMiddleware, + SpillToBackendEdit, +) + +from ..shared.flags import enabled + + +def build_context_editing_mw( + *, + flags: AgentFeatureFlags, + max_input_tokens: int | None, + tools: Sequence[BaseTool], + backend_resolver: Any, +) -> SpillingContextEditingMiddleware | None: + if not enabled(flags, "enable_context_editing") or not max_input_tokens: + return None + spill_edit = SpillToBackendEdit( + trigger=int(max_input_tokens * 0.55), + clear_at_least=int(max_input_tokens * 0.15), + keep=5, + exclude_tools=safe_exclude_tools(tools), + clear_tool_inputs=True, + ) + clear_edit = ClearToolUsesEdit( + trigger=int(max_input_tokens * 0.55), + clear_at_least=int(max_input_tokens * 0.15), + keep=5, + exclude_tools=safe_exclude_tools(tools), + clear_tool_inputs=True, + placeholder="[cleared - older tool output trimmed for context]", + ) + return SpillingContextEditingMiddleware( + edits=[spill_edit, clear_edit], + backend_resolver=backend_resolver, + ) 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 new file mode 100644 index 000000000..66cae300b --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py @@ -0,0 +1,13 @@ +"""Drop duplicate HITL tool calls before execution.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from langchain_core.tools import BaseTool + +from app.agents.new_chat.middleware import DedupHITLToolCallsMiddleware + + +def build_dedup_hitl_mw(tools: Sequence[BaseTool]) -> DedupHITLToolCallsMiddleware: + return DedupHITLToolCallsMiddleware(agent_tools=list(tools)) 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 new file mode 100644 index 000000000..a0b294092 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py @@ -0,0 +1,12 @@ +"""Stop N identical tool calls in a row via interrupt.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import DoomLoopMiddleware + +from ..shared.flags import enabled + + +def build_doom_loop_mw(flags: AgentFeatureFlags) -> DoomLoopMiddleware | None: + return DoomLoopMiddleware(threshold=3) if enabled(flags, "enable_doom_loop") else None 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 new file mode 100644 index 000000000..6e6467ad0 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py @@ -0,0 +1,12 @@ +"""Provider-compat: append a `_noop` tool when tools=[] but history has tool calls.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import NoopInjectionMiddleware + +from ..shared.flags import enabled + + +def build_noop_injection_mw(flags: AgentFeatureFlags) -> NoopInjectionMiddleware | None: + return NoopInjectionMiddleware() if enabled(flags, "enable_compaction_v2") else None 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 new file mode 100644 index 000000000..378b61be1 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py @@ -0,0 +1,50 @@ +"""Repair miscased / unknown tool names to the registered set or invalid_tool.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from langchain_core.tools import BaseTool + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ToolCallNameRepairMiddleware + +from ..shared.flags import enabled + +# deepagents-built-in tool names the repair pass treats as known. +_DEEPAGENT_BUILTIN_TOOL_NAMES: frozenset[str] = frozenset( + { + "write_todos", + "ls", + "read_file", + "write_file", + "edit_file", + "glob", + "grep", + "execute", + "task", + "mkdir", + "cd", + "pwd", + "move_file", + "rm", + "rmdir", + "list_tree", + "execute_code", + } +) + + +def build_repair_mw( + *, + flags: AgentFeatureFlags, + tools: Sequence[BaseTool], +) -> ToolCallNameRepairMiddleware | None: + if not enabled(flags, "enable_tool_call_repair"): + return None + registered_names: set[str] = {t.name for t in tools} + registered_names |= _DEEPAGENT_BUILTIN_TOOL_NAMES + return ToolCallNameRepairMiddleware( + registered_tool_names=registered_names, + fuzzy_match_threshold=None, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py new file mode 100644 index 000000000..8e7a32be8 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py @@ -0,0 +1,39 @@ +"""LLM-based tool subset selection (only when >30 tools).""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence + +from langchain.agents.middleware import LLMToolSelectorMiddleware +from langchain_core.tools import BaseTool + +from app.agents.new_chat.feature_flags import AgentFeatureFlags + +from ..shared.flags import enabled + + +def build_selector_mw( + *, + flags: AgentFeatureFlags, + tools: Sequence[BaseTool], +) -> LLMToolSelectorMiddleware | None: + if not enabled(flags, "enable_llm_tool_selector") or len(tools) <= 30: + return None + try: + return LLMToolSelectorMiddleware( + model="openai:gpt-4o-mini", + max_tools=12, + always_include=[ + name + for name in ( + "update_memory", + "get_connected_accounts", + "scrape_webpage", + ) + if name in {t.name for t in tools} + ], + ) + except Exception: + logging.warning("LLMToolSelectorMiddleware init failed; skipping.") + return None 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 new file mode 100644 index 000000000..63a57c5a0 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py @@ -0,0 +1,39 @@ +"""Skill discovery + injection.""" + +from __future__ import annotations + +import logging + +from deepagents.middleware.skills import SkillsMiddleware + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import ( + build_skills_backend_factory, + default_skills_sources, +) + +from ..shared.flags import enabled + + +def build_skills_mw( + *, + flags: AgentFeatureFlags, + filesystem_mode: FilesystemMode, + search_space_id: int, +) -> SkillsMiddleware | None: + if not enabled(flags, "enable_skills"): + return None + try: + skills_factory = build_skills_backend_factory( + search_space_id=search_space_id + if filesystem_mode == FilesystemMode.CLOUD + else None, + ) + return SkillsMiddleware( + backend=skills_factory, + sources=default_skills_sources(), + ) + except Exception as exc: # pragma: no cover - defensive + logging.warning("SkillsMiddleware init failed; skipping: %s", exc) + return None