Merge upstream/dev into feature/multi-agent

surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/compile_graph_sync.py

@@ -11,7 +11,6 @@ from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import BaseTool
 from langgraph.types import Checkpointer
 
-from .middleware import build_main_agent_deepagent_middleware
 from app.agents.multi_agent_chat.subagents.shared.permissions import (
     ToolsPermissions,
 )
@@ -20,6 +19,8 @@ from app.agents.new_chat.feature_flags import AgentFeatureFlags
 from app.agents.new_chat.filesystem_selection import FilesystemMode
 from app.db import ChatVisibility
 
+from .middleware import build_main_agent_deepagent_middleware
+
 
 def build_compiled_agent_graph_sync(
     *,

surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/task_tool.py

@@ -31,8 +31,8 @@ from .propagation import (
 from .resume import (
     build_resume_command,
     fan_out_decisions_to_match,
-    hitlrequest_action_count,
     get_first_pending_subagent_interrupt,
+    hitlrequest_action_count,
 )
 
 logger = logging.getLogger(__name__)
@@ -51,7 +51,9 @@ def build_task_tool_with_parent_config(
     )
 
     if task_description is None:
-        description = TASK_TOOL_DESCRIPTION.format(available_agents=subagent_description_str)
+        description = TASK_TOOL_DESCRIPTION.format(
+            available_agents=subagent_description_str
+        )
     elif "{available_agents}" in task_description:
         description = task_description.format(available_agents=subagent_description_str)
     else:
@@ -90,11 +92,11 @@ def build_task_tool_with_parent_config(
     def task(
         description: Annotated[
             str,
-            "A detailed description of the task for the subagent to perform autonomously. Include all necessary context and specify the expected output format.",  # noqa: E501
+            "A detailed description of the task for the subagent to perform autonomously. Include all necessary context and specify the expected output format.",
         ],
         subagent_type: Annotated[
             str,
-            "The type of subagent to use. Must be one of the available agent types listed in the tool description.",  # noqa: E501
+            "The type of subagent to use. Must be one of the available agent types listed in the tool description.",
         ],
         runtime: ToolRuntime,
     ) -> str | Command:
@@ -119,7 +121,9 @@ def build_task_tool_with_parent_config(
         if callable(get_state):
             try:
                 snapshot = get_state(sub_config)
-                pending_id, pending_value = get_first_pending_subagent_interrupt(snapshot)
+                pending_id, pending_value = get_first_pending_subagent_interrupt(
+                    snapshot
+                )
             except Exception:
                 # Fail loud if a resume is queued: silent fallback would
                 # replay the original interrupt to the user.
@@ -158,11 +162,11 @@ def build_task_tool_with_parent_config(
     async def atask(
         description: Annotated[
             str,
-            "A detailed description of the task for the subagent to perform autonomously. Include all necessary context and specify the expected output format.",  # noqa: E501
+            "A detailed description of the task for the subagent to perform autonomously. Include all necessary context and specify the expected output format.",
         ],
         subagent_type: Annotated[
             str,
-            "The type of subagent to use. Must be one of the available agent types listed in the tool description.",  # noqa: E501
+            "The type of subagent to use. Must be one of the available agent types listed in the tool description.",
         ],
         runtime: ToolRuntime,
     ) -> str | Command:
@@ -186,7 +190,9 @@ def build_task_tool_with_parent_config(
         if callable(aget_state):
             try:
                 snapshot = await aget_state(sub_config)
-                pending_id, pending_value = get_first_pending_subagent_interrupt(snapshot)
+                pending_id, pending_value = get_first_pending_subagent_interrupt(
+                    snapshot
+                )
             except Exception:
                 if has_surfsense_resume(runtime):
                     logger.exception(

surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py

@@ -23,7 +23,6 @@ from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import BaseTool
 from langgraph.types import Checkpointer
 
-from ...context_prune.prune_tool_names import safe_exclude_tools
 from app.agents.multi_agent_chat.subagents import (
     build_subagents,
     get_subagents_to_exclude,
@@ -66,6 +65,7 @@ from app.agents.new_chat.plugin_loader import (
 from app.agents.new_chat.tools.registry import BUILTIN_TOOLS
 from app.db import ChatVisibility
 
+from ...context_prune.prune_tool_names import safe_exclude_tools
 from .checkpointed_subagent_middleware import SurfSenseCheckpointedSubAgentMiddleware
 
 

surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py

@@ -14,8 +14,10 @@ from langchain_core.tools import BaseTool
 from langgraph.types import Checkpointer
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from ..graph.compile_graph_sync import build_compiled_agent_graph_sync
-from ..tools import MAIN_AGENT_SURFSENSE_TOOL_NAMES, MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED
+from app.agents.multi_agent_chat.subagents import (
+    get_subagents_to_exclude,
+    main_prompt_registry_subagent_lines,
+)
 from app.agents.multi_agent_chat.subagents.mcp_tools.index import (
     load_mcp_tools_by_connector,
 )
@@ -24,17 +26,19 @@ from app.agents.new_chat.feature_flags import AgentFeatureFlags, get_flags
 from app.agents.new_chat.filesystem_backends import build_backend_resolver
 from app.agents.new_chat.filesystem_selection import FilesystemMode, FilesystemSelection
 from app.agents.new_chat.llm_config import AgentConfig
-from app.agents.multi_agent_chat.subagents import (
-    get_subagents_to_exclude,
-    main_prompt_registry_subagent_lines,
-)
-from ..system_prompt import build_main_agent_system_prompt
 from app.agents.new_chat.tools.invalid_tool import INVALID_TOOL_NAME, invalid_tool
 from app.agents.new_chat.tools.registry import build_tools_async
 from app.db import ChatVisibility
 from app.services.connector_service import ConnectorService
 from app.utils.perf import get_perf_logger
 
+from ..graph.compile_graph_sync import build_compiled_agent_graph_sync
+from ..system_prompt import build_main_agent_system_prompt
+from ..tools import (
+    MAIN_AGENT_SURFSENSE_TOOL_NAMES,
+    MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED,
+)
+
 _perf_log = get_perf_logger()
 
 

surfsense_backend/app/agents/multi_agent_chat/main_agent/tools/__init__.py

@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 
-from .index import MAIN_AGENT_SURFSENSE_TOOL_NAMES, MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED
+from .index import (
+    MAIN_AGENT_SURFSENSE_TOOL_NAMES,
+    MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED,
+)
 
 __all__ = ["MAIN_AGENT_SURFSENSE_TOOL_NAMES", "MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED"]

surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/deliverables/tools/index.py

@@ -13,7 +13,9 @@ from .resume import create_generate_resume_tool
 from .video_presentation import create_generate_video_presentation_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     resolved_dependencies = {**(dependencies or {}), **kwargs}
     podcast = create_generate_podcast_tool(
         search_space_id=resolved_dependencies["search_space_id"],

surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/tools/index.py

@@ -10,7 +10,9 @@ from app.db import ChatVisibility
 from .update_memory import create_update_memory_tool, create_update_team_memory_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     resolved_dependencies = {**(dependencies or {}), **kwargs}
     if resolved_dependencies.get("thread_visibility") == ChatVisibility.SEARCH_SPACE:
         mem = create_update_team_memory_tool(
@@ -18,7 +20,10 @@ def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) ->
             db_session=resolved_dependencies["db_session"],
             llm=resolved_dependencies.get("llm"),
         )
-        return {"allow": [{"name": getattr(mem, "name", "") or "", "tool": mem}], "ask": []}
+        return {
+            "allow": [{"name": getattr(mem, "name", "") or "", "tool": mem}],
+            "ask": [],
+        }
     mem = create_update_memory_tool(
         user_id=resolved_dependencies["user_id"],
         db_session=resolved_dependencies["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/research/tools/index.py

@@ -11,14 +11,20 @@ from .search_surfsense_docs import create_search_surfsense_docs_tool
 from .web_search import create_web_search_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     resolved_dependencies = {**(dependencies or {}), **kwargs}
     web = create_web_search_tool(
         search_space_id=resolved_dependencies.get("search_space_id"),
         available_connectors=resolved_dependencies.get("available_connectors"),
     )
-    scrape = create_scrape_webpage_tool(firecrawl_api_key=resolved_dependencies.get("firecrawl_api_key"))
-    docs = create_search_surfsense_docs_tool(db_session=resolved_dependencies["db_session"])
+    scrape = create_scrape_webpage_tool(
+        firecrawl_api_key=resolved_dependencies.get("firecrawl_api_key")
+    )
+    docs = create_search_surfsense_docs_tool(
+        db_session=resolved_dependencies["db_session"]
+    )
     return {
         "allow": [
             {"name": getattr(web, "name", "") or "", "tool": web},

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/airtable/tools/index.py

@@ -7,6 +7,8 @@ from app.agents.multi_agent_chat.subagents.shared.permissions import (
 )
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     _ = {**(dependencies or {}), **kwargs}
     return {"allow": [], "ask": []}

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/index.py

@@ -12,7 +12,9 @@ from .search_events import create_search_calendar_events_tool
 from .update_event import create_update_calendar_event_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     resolved_dependencies = {**(dependencies or {}), **kwargs}
     session_dependencies = {
         "db_session": resolved_dependencies["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/clickup/tools/index.py

@@ -7,6 +7,8 @@ from app.agents.multi_agent_chat.subagents.shared.permissions import (
 )
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     _ = {**(dependencies or {}), **kwargs}
     return {"allow": [], "ask": []}

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/confluence/tools/index.py

@@ -11,7 +11,9 @@ from .delete_page import create_delete_confluence_page_tool
 from .update_page import create_update_confluence_page_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     resolved_dependencies = {**(dependencies or {}), **kwargs}
     session_dependencies = {
         "db_session": resolved_dependencies["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/discord/tools/index.py

@@ -11,7 +11,9 @@ from .read_messages import create_read_discord_messages_tool
 from .send_message import create_send_discord_message_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/dropbox/tools/index.py

@@ -10,7 +10,9 @@ from .create_file import create_create_dropbox_file_tool
 from .trash_file import create_delete_dropbox_file_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/index.py

@@ -14,7 +14,9 @@ from .trash_email import create_trash_gmail_email_tool
 from .update_draft import create_update_gmail_draft_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/index.py

@@ -10,7 +10,9 @@ from .create_file import create_create_google_drive_file_tool
 from .trash_file import create_delete_google_drive_file_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/jira/tools/index.py

@@ -11,7 +11,9 @@ from .delete_issue import create_delete_jira_issue_tool
 from .update_issue import create_update_jira_issue_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/linear/tools/index.py

@@ -11,7 +11,9 @@ from .delete_issue import create_delete_linear_issue_tool
 from .update_issue import create_update_linear_issue_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/luma/tools/index.py

@@ -11,7 +11,9 @@ from .list_events import create_list_luma_events_tool
 from .read_event import create_read_luma_event_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/notion/tools/index.py

@@ -11,7 +11,9 @@ from .delete_page import create_delete_notion_page_tool
 from .update_page import create_update_notion_page_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/onedrive/tools/index.py

@@ -10,7 +10,9 @@ from .create_file import create_create_onedrive_file_tool
 from .trash_file import create_delete_onedrive_file_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/slack/tools/index.py

@@ -7,6 +7,8 @@ from app.agents.multi_agent_chat.subagents.shared.permissions import (
 )
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     _ = {**(dependencies or {}), **kwargs}
     return {"allow": [], "ask": []}

surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/teams/tools/index.py

@@ -11,7 +11,9 @@ from .read_messages import create_read_teams_messages_tool
 from .send_message import create_send_teams_message_tool
 
 
-def load_tools(*, dependencies: dict[str, Any] | None = None, **kwargs: Any) -> ToolsPermissions:
+def load_tools(
+    *, dependencies: dict[str, Any] | None = None, **kwargs: Any
+) -> ToolsPermissions:
     d = {**(dependencies or {}), **kwargs}
     common = {
         "db_session": d["db_session"],

surfsense_backend/app/agents/multi_agent_chat/subagents/mcp_tools/index.py

@@ -31,6 +31,7 @@ logger = logging.getLogger(__name__)
 
 ## Helper functions for fetching connector metadata maps
 
+
 async def fetch_mcp_connector_metadata_maps(
     session: AsyncSession,
     search_space_id: int,
@@ -58,6 +59,7 @@ async def fetch_mcp_connector_metadata_maps(
 
 ## Helper functions for partitioning tools by connector agent
 
+
 def partition_mcp_tools_by_connector(
     tools: Sequence[BaseTool],
     connector_id_to_type: dict[int, str],
@@ -104,8 +106,10 @@ def partition_mcp_tools_by_connector(
 
     return dict(buckets)
 
+
 ## Helper functions for splitting tools by permissions
 
+
 def _get_mcp_tool_name(tool: BaseTool) -> str:
     meta: dict[str, Any] = getattr(tool, "metadata", None) or {}
     orig = meta.get("mcp_original_tool_name")
@@ -139,6 +143,7 @@ def _split_tools_by_permissions(
 
 ## Main function to load MCP tools and split them by permissions for each connector agent
 
+
 async def load_mcp_tools_by_connector(
     session: AsyncSession,
     search_space_id: int,
@@ -148,9 +153,7 @@ async def load_mcp_tools_by_connector(
     Pass ``bypass_internal_hitl=True`` so the subagent's
     ``HumanInTheLoopMiddleware`` is the single HITL gate.
     """
-    flat = await load_mcp_tools(
-        session, search_space_id, bypass_internal_hitl=True
-    )
+    flat = await load_mcp_tools(session, search_space_id, bypass_internal_hitl=True)
     id_map, name_map = await fetch_mcp_connector_metadata_maps(session, search_space_id)
     buckets = partition_mcp_tools_by_connector(flat, id_map, name_map)
     return {

surfsense_backend/app/agents/multi_agent_chat/subagents/registry.py

@@ -8,6 +8,9 @@ from typing import Any, Protocol
 from deepagents import SubAgent
 from langchain_core.language_models import BaseChatModel
 
+from app.agents.multi_agent_chat.constants import (
+    SUBAGENT_TO_REQUIRED_CONNECTOR_MAP,
+)
 from app.agents.multi_agent_chat.subagents.builtins.deliverables.agent import (
     build_subagent as build_deliverables_subagent,
 )
@@ -62,9 +65,6 @@ from app.agents.multi_agent_chat.subagents.connectors.slack.agent import (
 from app.agents.multi_agent_chat.subagents.connectors.teams.agent import (
     build_subagent as build_teams_subagent,
 )
-from app.agents.multi_agent_chat.constants import (
-    SUBAGENT_TO_REQUIRED_CONNECTOR_MAP,
-)
 from app.agents.multi_agent_chat.subagents.shared.md_file_reader import (
     read_md_file,
 )
@@ -105,6 +105,7 @@ SUBAGENT_BUILDERS_BY_NAME: dict[str, SubagentBuilder] = {
     "teams": build_teams_subagent,
 }
 
+
 def _route_resource_package(builder: SubagentBuilder) -> str:
     mod = builder.__module__
     return mod[: -len(".agent")] if mod.endswith(".agent") else mod.rsplit(".", 1)[0]