mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-24 21:42:22 +02:00
[eric] agents: split per-run support methods into RunSupportMixin + p_-rename the private methods (convention-clean)
This commit is contained in:
ciregenz
@@ -65,6 +65,7 @@ from backend.apps.agents.tools.web import should_register_web_mcp
|
||||
from backend.apps.agents.manager.session.SessionLifecycleMixin import SessionLifecycleMixin
|
||||
from backend.apps.agents.manager.MessagingMixin import MessagingMixin
|
||||
from backend.apps.agents.manager.AgentLaunchMixin import AgentLaunchMixin
|
||||
from backend.apps.agents.manager.RunSupportMixin import RunSupportMixin
|
||||
from backend.apps.agents.manager.permissions import gate_hooks
|
||||
from backend.apps.agents.manager.session.workspace_git import _detect_git_identity, _ensure_cwd_git_repo
|
||||
from backend.apps.agents.manager.prompt.tool_catalog import (
|
||||
@@ -82,19 +83,13 @@ from backend.apps.agents.manager.session.history_compaction import (
|
||||
_get_branch_messages,
|
||||
)
|
||||
from backend.apps.agents.manager.prompt.prompt_context import resolve_mode
|
||||
from backend.apps.agents.manager.prompt.attachments import (
|
||||
_build_dir_tree,
|
||||
_build_prompt_content,
|
||||
_resolve_attachments,
|
||||
_resolve_context_paths,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")
|
||||
|
||||
|
||||
class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunSupportMixin):
|
||||
def __init__(self):
|
||||
self.sessions: dict[str, AgentSession] = {}
|
||||
self.tasks: dict[str, asyncio.Task] = {}
|
||||
@@ -103,83 +98,8 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# multi-second SDK teardown the cancel handler sits behind.
|
||||
self._live_partial: Dict[str, LivePartial] = {}
|
||||
|
||||
async def _build_mcp_servers(
|
||||
self,
|
||||
allowed_tools: list[str],
|
||||
active_mcps: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""Build the mcp_servers dict for ClaudeAgentOptions from installed MCP tools.
|
||||
|
||||
Filtering is two-stage:
|
||||
1. allowed_tools (mode/session permission), same as before.
|
||||
2. active_mcps (per-session activation gate), NEW. When this list is
|
||||
provided (non-None), only MCP servers whose sanitized name appears
|
||||
in it are forwarded to the SDK. Empty list means zero MCPs ship.
|
||||
None means legacy / non-gated path (used by sessions created
|
||||
before the gate existed, where active_mcps was implicit-all).
|
||||
|
||||
The activation gate is the dispatch-layer enforcement of the product
|
||||
invariant "all MCP actions only via ToolSearch": the model can only
|
||||
reach an MCP server's tools if the user has approved MCPActivate for
|
||||
that server, which appends to session.active_mcps. The model cannot
|
||||
bypass this by ignoring prompt instructions, the SDK simply receives
|
||||
no MCP definition for unactivated servers.
|
||||
|
||||
Servers whose every sub-tool is denied are skipped entirely.
|
||||
"""
|
||||
mcp_servers: dict = {}
|
||||
all_tools = load_all_tools()
|
||||
mcp_tools = [t for t in all_tools if t.mcp_config and t.enabled and t.auth_status in ("configured", "connected")]
|
||||
active_set = set(active_mcps) if active_mcps is not None else None
|
||||
logger.info(
|
||||
f"[MCP-DEBUG] Building MCP servers. {len(mcp_tools)} MCP tools found, "
|
||||
f"allowed_tools has {len(allowed_tools)} entries, "
|
||||
f"active_mcps={'<unset/all>' if active_set is None else sorted(active_set)}"
|
||||
)
|
||||
|
||||
for tool in mcp_tools:
|
||||
tool_ref = f"mcp:{tool.name}"
|
||||
if tool_ref not in allowed_tools and allowed_tools != get_all_tool_names():
|
||||
if not any(tool_ref == at for at in allowed_tools):
|
||||
logger.info(f"[MCP-DEBUG] SKIPPED {tool.name}: '{tool_ref}' not in allowed_tools")
|
||||
continue
|
||||
|
||||
server_name = _sanitize_server_name(tool.name)
|
||||
if active_set is not None and server_name not in active_set:
|
||||
logger.info(f"[MCP-DEBUG] GATED {server_name}: not in session.active_mcps, model must call MCPActivate first")
|
||||
continue
|
||||
|
||||
if is_fully_denied(tool):
|
||||
logger.info(f"[MCP-DEBUG] SKIPPED {tool.name}: fully denied")
|
||||
continue
|
||||
|
||||
if tool.auth_type == "oauth2" and tool.auth_status == "connected":
|
||||
if tool.name.lower() in ("discord", "github"):
|
||||
# Discord uses a shared bot token; GitHub OAuth-app tokens don't
|
||||
# expire and carry no refresh_token. Nothing to refresh either way.
|
||||
refreshed = True
|
||||
elif tool.name.lower() == "airtable":
|
||||
refreshed = await refresh_airtable_token(tool)
|
||||
elif tool.name.lower() == "hubspot":
|
||||
refreshed = await refresh_hubspot_token(tool)
|
||||
else:
|
||||
refreshed = await refresh_google_token(tool)
|
||||
logger.info(f"[MCP-DEBUG] {tool.name} token refresh: {'OK' if refreshed else 'FAILED'}")
|
||||
|
||||
config = derive_mcp_config(tool)
|
||||
if config:
|
||||
mcp_servers[server_name] = config
|
||||
env_keys = list(config.get("env", {}).keys())
|
||||
logger.info(f"[MCP-DEBUG] ADDED {server_name}: command={config.get('command')}, args={config.get('args')}, env_keys={env_keys}")
|
||||
else:
|
||||
logger.warning(f"[MCP-DEBUG] {tool.name}: derive_mcp_config returned None")
|
||||
|
||||
logger.info(f"[MCP-DEBUG] Final mcp_servers: {list(mcp_servers.keys())}")
|
||||
return mcp_servers
|
||||
|
||||
|
||||
def _build_dir_tree(self, root: str, max_depth: int = 4, prefix: str = "") -> list[str]:
|
||||
return _build_dir_tree(root, max_depth, prefix)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Compaction & token guard (Phase 2)
|
||||
@@ -195,33 +115,10 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# surfaces from the catch-all
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _maybe_compact(self, session: AgentSession, force: bool = False) -> bool:
|
||||
return context_budget.maybe_compact(session, force)
|
||||
|
||||
async def _emit_context_update(
|
||||
self,
|
||||
session_id: str,
|
||||
session: AgentSession,
|
||||
*,
|
||||
input_tokens: int | None = None,
|
||||
output_tokens: int | None = None,
|
||||
cache_read_tokens: int = 0,
|
||||
cache_read_pct: float = 0.0,
|
||||
) -> None:
|
||||
return await context_budget.emit_context_update(
|
||||
session_id, session,
|
||||
input_tokens=input_tokens, output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens, cache_read_pct=cache_read_pct,
|
||||
)
|
||||
|
||||
def _build_prompt_content(self, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, api_type: str = "anthropic", model: str = ""):
|
||||
return _build_prompt_content(prompt, images, context_paths, forced_tools, attached_skills, api_type, model)
|
||||
|
||||
def _resolve_attachments(self, context_paths: list | None, api_type: str, model: str) -> tuple[str, list[dict], list[str]]:
|
||||
return _resolve_attachments(context_paths, api_type, model)
|
||||
|
||||
def _resolve_context_paths(self, context_paths: list | None) -> str:
|
||||
return _resolve_context_paths(context_paths)
|
||||
|
||||
async def _run_agent_loop(self, session_id: str, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, fork_session: bool = False, selected_browser_ids: list[str] | None = None, selected_app_output_ids: list[str] | None = None, selected_setting_ids: list[str] | None = None):
|
||||
"""Run the Claude Agent SDK query loop for a session."""
|
||||
@@ -231,7 +128,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
|
||||
from backend.apps.agents.providers.registry import get_api_type as _get_api_type
|
||||
_api = _get_api_type(session.model)
|
||||
prompt_content = self._build_prompt_content(
|
||||
prompt_content = self.p_build_prompt_content(
|
||||
prompt, images, context_paths, forced_tools, attached_skills,
|
||||
api_type=_api, model=session.model,
|
||||
)
|
||||
@@ -299,7 +196,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# Reconcile active_mcps against currently-enabled tools (Phase 3).
|
||||
# If the user toggled a server off in the Tools page mid-session,
|
||||
# drop it from active_mcps automatically so the model isn't told
|
||||
# "X is active" while _build_mcp_servers silently filters it out.
|
||||
# "X is active" while p_build_mcp_servers silently filters it out.
|
||||
# Emit a context_status event so the model and UI both know.
|
||||
try:
|
||||
_enabled = {
|
||||
@@ -348,8 +245,8 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# Pass session.active_mcps as the activation filter. Empty list ⇒
|
||||
# no MCP tools shipped to the SDK; the model must MCPSearch and
|
||||
# MCPActivate first. The product invariant lives here at the
|
||||
# dispatch layer (see _build_mcp_servers docstring).
|
||||
mcp_servers = await self._build_mcp_servers(session.allowed_tools, session.active_mcps)
|
||||
# dispatch layer (see p_build_mcp_servers docstring).
|
||||
mcp_servers = await self.p_build_mcp_servers(session.allowed_tools, session.active_mcps)
|
||||
|
||||
_browser_delegation_tools = ["CreateBrowserAgent", "BrowserAgent", "BrowserAgents"]
|
||||
_browser_all_denied = all(
|
||||
@@ -410,7 +307,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# Always-on meta-MCP server. Exposes MCPList / MCPSearch /
|
||||
# MCPActivate so the model can discover and activate user MCPs at
|
||||
# runtime. The activation gate (active_mcps filter in
|
||||
# _build_mcp_servers above) ensures the model cannot reach any
|
||||
# p_build_mcp_servers above) ensures the model cannot reach any
|
||||
# other MCP server's tools without going through this layer first.
|
||||
mcp_meta_server_path = os.path.join(
|
||||
os.path.dirname(__file__), "mcp_meta_server.py"
|
||||
@@ -1039,14 +936,14 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
# programmatic summarization (no aux LLM call) so this adds
|
||||
# zero latency on the user's turn.
|
||||
try:
|
||||
if self._maybe_compact(session):
|
||||
if self.p_maybe_compact(session):
|
||||
new_input = _estimate_post_compact_input(session)
|
||||
await ws_manager.send_to_session(session_id, "agent:context_status", {
|
||||
"session_id": session_id,
|
||||
"reason": "compacted",
|
||||
"compacted_through_msg_id": session.compacted_through_msg_id,
|
||||
})
|
||||
await self._emit_context_update(
|
||||
await self.p_emit_context_update(
|
||||
session_id,
|
||||
session,
|
||||
input_tokens=new_input,
|
||||
@@ -1374,7 +1271,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
session.needs_fresh_session = True
|
||||
# Persist whatever streamed before the cancel (edit / branch
|
||||
# switch paths; the user-stop path already did this in stop_agent).
|
||||
await self._commit_partial_now(session)
|
||||
await self.p_commit_partial_now(session)
|
||||
turn.stream_text_msg_id = None
|
||||
turn.stream_text_accum = ""
|
||||
except Exception as e:
|
||||
@@ -1663,150 +1560,19 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin):
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to snapshot session {session_id}: {e}")
|
||||
|
||||
async def _stream_text(self, session_id: str, msg_id: str, text: str, delay: float = 0.03):
|
||||
"""Emit stream_start, word-by-word deltas, and stream_end for a text message."""
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_start", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"role": "assistant",
|
||||
})
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word if i == 0 else " " + word
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_delta", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"delta": chunk,
|
||||
})
|
||||
await asyncio.sleep(delay)
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
|
||||
async def _stream_tool_input(self, session_id: str, msg_id: str, tool_name: str, input_json: str, delay: float = 0.02):
|
||||
"""Emit stream_start, chunked deltas, and stream_end for a tool_call input."""
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_start", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"role": "tool_call",
|
||||
"tool_name": tool_name,
|
||||
})
|
||||
chunk_size = 12
|
||||
for i in range(0, len(input_json), chunk_size):
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_delta", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"delta": input_json[i:i + chunk_size],
|
||||
})
|
||||
await asyncio.sleep(delay)
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
|
||||
|
||||
|
||||
|
||||
async def _commit_partial_now(self, session) -> bool:
|
||||
"""Persist the in-flight streamed assistant text as a real message and
|
||||
push it to the client, idempotently. Lets a stop show the partial
|
||||
instantly instead of waiting out the SDK teardown the cancel handler
|
||||
sits behind. Returns True if it committed something."""
|
||||
live = self._live_partial.pop(session.id, None)
|
||||
if not live:
|
||||
return False
|
||||
text = live.text or ""
|
||||
msg_id = live.msg_id
|
||||
if not msg_id or not text.strip():
|
||||
return False
|
||||
if any(getattr(m, "id", None) == msg_id for m in session.messages):
|
||||
return False
|
||||
partial = Message(
|
||||
id=msg_id,
|
||||
role="assistant",
|
||||
content=text,
|
||||
branch_id=live.branch_id or session.active_branch_id,
|
||||
)
|
||||
upsert_message(session, partial)
|
||||
try:
|
||||
await ws_manager.send_to_session(session.id, "agent:message", {
|
||||
"session_id": session.id,
|
||||
"message": partial.model_dump(mode="json"),
|
||||
})
|
||||
await ws_manager.send_to_session(session.id, "agent:stream_end", {
|
||||
"session_id": session.id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
return True
|
||||
|
||||
async def _drain_task(self, task) -> None:
|
||||
"""Await a cancelled task's (possibly slow) teardown off the hot path."""
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
|
||||
async def generate_title(self, session_id: str, first_prompt: str) -> str:
|
||||
return await metadata.generate_title(self.sessions.get(session_id), session_id, first_prompt)
|
||||
|
||||
async def generate_turn_label(self, session_id: str, turn_id: str, user_prompt: str) -> None:
|
||||
return await metadata.generate_turn_label(self.sessions.get(session_id), session_id, turn_id, user_prompt)
|
||||
|
||||
async def warm_prompt_cache(self, session_id: str) -> None:
|
||||
"""Pre-warm Anthropic's prompt cache for a session by firing a
|
||||
max_tokens=1 dummy request through the same agent path. Anthropic
|
||||
processes the system+tools prefix and writes the cache; the next
|
||||
real user turn lands a cache hit instead of paying cold-start.
|
||||
|
||||
Skips silently if the session doesn't exist, isn't on Anthropic,
|
||||
or has no Anthropic credentials. Skips if a real request is
|
||||
already in flight on this session, Anthropic permits parallel
|
||||
requests but it just wastes the warm.
|
||||
"""
|
||||
session = self.sessions.get(session_id)
|
||||
if not session:
|
||||
return
|
||||
# If a real run is in flight, the cache will be warmed by it;
|
||||
# firing again is wasted tokens.
|
||||
existing = self.tasks.get(session_id)
|
||||
if existing and not existing.done():
|
||||
return
|
||||
|
||||
try:
|
||||
from backend.apps.agents.providers.registry import _find_builtin_model
|
||||
entry = _find_builtin_model(session.model)
|
||||
if not entry or entry.get("api") != "anthropic":
|
||||
return # other providers handle caching automatically
|
||||
|
||||
from backend.apps.settings.credentials import get_anthropic_client
|
||||
global_settings = load_settings()
|
||||
# Free lane rotates pool accounts per call, so a warm ping primes a cache
|
||||
# the next call won't hit, and worse it'd burn a metered run at idle (this
|
||||
# fires on dashboard mount, not a user query). Skip it on the free trial.
|
||||
if getattr(global_settings, "connection_mode", "own_key") == "free-trial":
|
||||
return
|
||||
client = get_anthropic_client(global_settings)
|
||||
|
||||
# Single ping with the same system + minimal user message.
|
||||
# max_tokens=1 keeps it cheap; we don't care about the output.
|
||||
await client.messages.create(
|
||||
model=entry.get("model_id", session.model),
|
||||
max_tokens=1,
|
||||
system="You are a helpful assistant. Reply with one character.",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
)
|
||||
logger.debug(f"Cache pre-warm fired for session {session_id}")
|
||||
except Exception as e:
|
||||
logger.debug(f"Cache pre-warm failed (non-fatal): {e}")
|
||||
|
||||
async def generate_group_meta(self, session_id: str, group_id: str, tool_calls: list[dict], results_summary: list[str] | None = None, is_refinement: bool = False) -> dict:
|
||||
return await metadata.generate_group_meta(self.sessions.get(session_id), session_id, group_id, tool_calls, results_summary, is_refinement)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Agent run entry points for AgentManager: launch a new top-level run, run the no-SDK mock
|
||||
fallback, and the staticmethod invoke_agent helper. Split into a mixin to keep the manager file
|
||||
under the size ceiling; self._run_agent_loop / self._stream_text / self.sessions resolve across
|
||||
under the size ceiling; self._run_agent_loop / self.p_stream_text / self.sessions resolve across
|
||||
the MRO exactly as before."""
|
||||
|
||||
import asyncio
|
||||
@@ -173,7 +173,7 @@ class AgentLaunchMixin:
|
||||
import json
|
||||
tool_input_content = {"tool": "Bash", "input": {"command": f"echo 'Processing: {prompt}'"}, "approved": decision.get("behavior") == "allow"}
|
||||
tool_msg_id = uuid4().hex
|
||||
await self._stream_tool_input(
|
||||
await self.p_stream_tool_input(
|
||||
session_id, tool_msg_id, "Bash",
|
||||
json.dumps(tool_input_content["input"], indent=2),
|
||||
)
|
||||
@@ -203,7 +203,7 @@ class AgentLaunchMixin:
|
||||
f"The agent was configured with:\n- Model: {session.model}\n- Mode: {session.mode}"
|
||||
)
|
||||
asst_msg_id = uuid4().hex
|
||||
await self._stream_text(session_id, asst_msg_id, asst_text)
|
||||
await self.p_stream_text(session_id, asst_msg_id, asst_text)
|
||||
|
||||
asst_msg = Message(id=asst_msg_id, role="assistant", content=asst_text, branch_id=session.active_branch_id)
|
||||
session.messages.append(asst_msg)
|
||||
|
||||
@@ -203,7 +203,7 @@ class MessagingMixin:
|
||||
# teardown, which can take several seconds; doing it here means the
|
||||
# streamed text stays put the instant Stop is pressed instead of
|
||||
# blinking out and reappearing once teardown finishes.
|
||||
await self._commit_partial_now(session)
|
||||
await self.p_commit_partial_now(session)
|
||||
await ws_manager.send_to_session(session_id, "agent:status", {
|
||||
"session_id": session_id,
|
||||
"status": "stopped",
|
||||
@@ -224,7 +224,7 @@ class MessagingMixin:
|
||||
task = self.tasks.pop(session_id, None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
asyncio.create_task(self._drain_task(task))
|
||||
asyncio.create_task(self.p_drain_task(task))
|
||||
|
||||
@typechecked
|
||||
def handle_approval(self, request_id: str, decision: Dict):
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Per-run support methods for AgentManager: build the gated MCP server set, warm the prompt
|
||||
cache, stream-emit helpers, commit/drain a stopped turn, context-update broadcast, and the aux
|
||||
metadata + prompt/attachment delegators. Split into a mixin to keep the manager file under the
|
||||
size ceiling; self.sessions / self.tasks / self._live_partial resolve across the MRO as before."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
from backend.apps.agents.core.models import AgentSession, Message
|
||||
from backend.apps.agents.core.ws_manager import ws_manager
|
||||
from backend.apps.settings.settings import load_settings
|
||||
from backend.apps.agents.manager import context_budget
|
||||
from backend.apps.agents.manager import metadata
|
||||
from backend.apps.agents.manager.streaming.upsert_message import upsert_message
|
||||
from backend.apps.agents.manager.prompt.tool_catalog import (
|
||||
get_all_tool_names,
|
||||
is_fully_denied,
|
||||
)
|
||||
from backend.apps.agents.manager.prompt.attachments import (
|
||||
_build_dir_tree as build_dir_tree,
|
||||
_build_prompt_content as build_prompt_content,
|
||||
_resolve_attachments as resolve_attachments,
|
||||
_resolve_context_paths as resolve_context_paths,
|
||||
)
|
||||
from backend.apps.tools_lib.tools_lib import (
|
||||
_load_all as load_all_tools,
|
||||
_sanitize_server_name as sanitize_server_name,
|
||||
derive_mcp_config,
|
||||
refresh_airtable_token,
|
||||
refresh_google_token,
|
||||
refresh_hubspot_token,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RunSupportMixin:
|
||||
@typechecked
|
||||
async def p_build_mcp_servers(
|
||||
self,
|
||||
allowed_tools: List[str],
|
||||
active_mcps: Optional[List[str]] = None,
|
||||
) -> Dict:
|
||||
"""Build the mcp_servers dict for ClaudeAgentOptions from installed MCP tools.
|
||||
|
||||
Filtering is two-stage:
|
||||
1. allowed_tools (mode/session permission), same as before.
|
||||
2. active_mcps (per-session activation gate), NEW. When this list is
|
||||
provided (non-None), only MCP servers whose sanitized name appears
|
||||
in it are forwarded to the SDK. Empty list means zero MCPs ship.
|
||||
None means legacy / non-gated path (used by sessions created
|
||||
before the gate existed, where active_mcps was implicit-all).
|
||||
|
||||
The activation gate is the dispatch-layer enforcement of the product
|
||||
invariant "all MCP actions only via ToolSearch": the model can only
|
||||
reach an MCP server's tools if the user has approved MCPActivate for
|
||||
that server, which appends to session.active_mcps. The model cannot
|
||||
bypass this by ignoring prompt instructions, the SDK simply receives
|
||||
no MCP definition for unactivated servers.
|
||||
|
||||
Servers whose every sub-tool is denied are skipped entirely.
|
||||
"""
|
||||
mcp_servers: dict = {}
|
||||
all_tools = load_all_tools()
|
||||
mcp_tools = [t for t in all_tools if t.mcp_config and t.enabled and t.auth_status in ("configured", "connected")]
|
||||
active_set = set(active_mcps) if active_mcps is not None else None
|
||||
logger.info(
|
||||
f"[MCP-DEBUG] Building MCP servers. {len(mcp_tools)} MCP tools found, "
|
||||
f"allowed_tools has {len(allowed_tools)} entries, "
|
||||
f"active_mcps={'<unset/all>' if active_set is None else sorted(active_set)}"
|
||||
)
|
||||
|
||||
for tool in mcp_tools:
|
||||
tool_ref = f"mcp:{tool.name}"
|
||||
if tool_ref not in allowed_tools and allowed_tools != get_all_tool_names():
|
||||
if not any(tool_ref == at for at in allowed_tools):
|
||||
logger.info(f"[MCP-DEBUG] SKIPPED {tool.name}: '{tool_ref}' not in allowed_tools")
|
||||
continue
|
||||
|
||||
server_name = sanitize_server_name(tool.name)
|
||||
if active_set is not None and server_name not in active_set:
|
||||
logger.info(f"[MCP-DEBUG] GATED {server_name}: not in session.active_mcps, model must call MCPActivate first")
|
||||
continue
|
||||
|
||||
if is_fully_denied(tool):
|
||||
logger.info(f"[MCP-DEBUG] SKIPPED {tool.name}: fully denied")
|
||||
continue
|
||||
|
||||
if tool.auth_type == "oauth2" and tool.auth_status == "connected":
|
||||
if tool.name.lower() in ("discord", "github"):
|
||||
# Discord uses a shared bot token; GitHub OAuth-app tokens don't
|
||||
# expire and carry no refresh_token. Nothing to refresh either way.
|
||||
refreshed = True
|
||||
elif tool.name.lower() == "airtable":
|
||||
refreshed = await refresh_airtable_token(tool)
|
||||
elif tool.name.lower() == "hubspot":
|
||||
refreshed = await refresh_hubspot_token(tool)
|
||||
else:
|
||||
refreshed = await refresh_google_token(tool)
|
||||
logger.info(f"[MCP-DEBUG] {tool.name} token refresh: {'OK' if refreshed else 'FAILED'}")
|
||||
|
||||
config = derive_mcp_config(tool)
|
||||
if config:
|
||||
mcp_servers[server_name] = config
|
||||
env_keys = list(config.get("env", {}).keys())
|
||||
logger.info(f"[MCP-DEBUG] ADDED {server_name}: command={config.get('command')}, args={config.get('args')}, env_keys={env_keys}")
|
||||
else:
|
||||
logger.warning(f"[MCP-DEBUG] {tool.name}: derive_mcp_config returned None")
|
||||
|
||||
logger.info(f"[MCP-DEBUG] Final mcp_servers: {list(mcp_servers.keys())}")
|
||||
return mcp_servers
|
||||
|
||||
@typechecked
|
||||
def p_build_dir_tree(self, root: str, max_depth: int = 4, prefix: str = "") -> List[str]:
|
||||
return build_dir_tree(root, max_depth, prefix)
|
||||
|
||||
@typechecked
|
||||
def p_maybe_compact(self, session: AgentSession, force: bool = False) -> bool:
|
||||
return context_budget.maybe_compact(session, force)
|
||||
|
||||
@typechecked
|
||||
async def p_emit_context_update(
|
||||
self,
|
||||
session_id: str,
|
||||
session: AgentSession,
|
||||
*,
|
||||
input_tokens: Optional[int] = None,
|
||||
output_tokens: Optional[int] = None,
|
||||
cache_read_tokens: int = 0,
|
||||
cache_read_pct: float = 0.0,
|
||||
) -> None:
|
||||
return await context_budget.emit_context_update(
|
||||
session_id, session,
|
||||
input_tokens=input_tokens, output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens, cache_read_pct=cache_read_pct,
|
||||
)
|
||||
|
||||
@typechecked
|
||||
def p_build_prompt_content(self, prompt: str, images: Optional[List] = None, context_paths: Optional[List] = None, forced_tools: Optional[List[str]] = None, attached_skills: Optional[List] = None, api_type: str = "anthropic", model: str = ""):
|
||||
return build_prompt_content(prompt, images, context_paths, forced_tools, attached_skills, api_type, model)
|
||||
|
||||
@typechecked
|
||||
def p_resolve_attachments(self, context_paths: Optional[List], api_type: str, model: str) -> Tuple[str, List[dict], List[str]]:
|
||||
return resolve_attachments(context_paths, api_type, model)
|
||||
|
||||
@typechecked
|
||||
def p_resolve_context_paths(self, context_paths: Optional[List]) -> str:
|
||||
return resolve_context_paths(context_paths)
|
||||
|
||||
@typechecked
|
||||
async def p_stream_text(self, session_id: str, msg_id: str, text: str, delay: float = 0.03):
|
||||
"""Emit stream_start, word-by-word deltas, and stream_end for a text message."""
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_start", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"role": "assistant",
|
||||
})
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word if i == 0 else " " + word
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_delta", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"delta": chunk,
|
||||
})
|
||||
await asyncio.sleep(delay)
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
|
||||
@typechecked
|
||||
async def p_stream_tool_input(self, session_id: str, msg_id: str, tool_name: str, input_json: str, delay: float = 0.02):
|
||||
"""Emit stream_start, chunked deltas, and stream_end for a tool_call input."""
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_start", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"role": "tool_call",
|
||||
"tool_name": tool_name,
|
||||
})
|
||||
chunk_size = 12
|
||||
for i in range(0, len(input_json), chunk_size):
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_delta", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
"delta": input_json[i:i + chunk_size],
|
||||
})
|
||||
await asyncio.sleep(delay)
|
||||
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
||||
"session_id": session_id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
|
||||
@typechecked
|
||||
async def p_commit_partial_now(self, session) -> bool:
|
||||
"""Persist the in-flight streamed assistant text as a real message and
|
||||
push it to the client, idempotently. Lets a stop show the partial
|
||||
instantly instead of waiting out the SDK teardown the cancel handler
|
||||
sits behind. Returns True if it committed something."""
|
||||
live = self._live_partial.pop(session.id, None)
|
||||
if not live:
|
||||
return False
|
||||
text = live.text or ""
|
||||
msg_id = live.msg_id
|
||||
if not msg_id or not text.strip():
|
||||
return False
|
||||
if any(getattr(m, "id", None) == msg_id for m in session.messages):
|
||||
return False
|
||||
partial = Message(
|
||||
id=msg_id,
|
||||
role="assistant",
|
||||
content=text,
|
||||
branch_id=live.branch_id or session.active_branch_id,
|
||||
)
|
||||
upsert_message(session, partial)
|
||||
try:
|
||||
await ws_manager.send_to_session(session.id, "agent:message", {
|
||||
"session_id": session.id,
|
||||
"message": partial.model_dump(mode="json"),
|
||||
})
|
||||
await ws_manager.send_to_session(session.id, "agent:stream_end", {
|
||||
"session_id": session.id,
|
||||
"message_id": msg_id,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
return True
|
||||
|
||||
@typechecked
|
||||
async def p_drain_task(self, task) -> None:
|
||||
"""Await a cancelled task's (possibly slow) teardown off the hot path."""
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
@typechecked
|
||||
async def generate_title(self, session_id: str, first_prompt: str) -> str:
|
||||
return await metadata.generate_title(self.sessions.get(session_id), session_id, first_prompt)
|
||||
|
||||
@typechecked
|
||||
async def generate_turn_label(self, session_id: str, turn_id: str, user_prompt: str) -> None:
|
||||
return await metadata.generate_turn_label(self.sessions.get(session_id), session_id, turn_id, user_prompt)
|
||||
|
||||
@typechecked
|
||||
async def warm_prompt_cache(self, session_id: str) -> None:
|
||||
"""Pre-warm Anthropic's prompt cache for a session by firing a
|
||||
max_tokens=1 dummy request through the same agent path. Anthropic
|
||||
processes the system+tools prefix and writes the cache; the next
|
||||
real user turn lands a cache hit instead of paying cold-start.
|
||||
|
||||
Skips silently if the session doesn't exist, isn't on Anthropic,
|
||||
or has no Anthropic credentials. Skips if a real request is
|
||||
already in flight on this session, Anthropic permits parallel
|
||||
requests but it just wastes the warm.
|
||||
"""
|
||||
session = self.sessions.get(session_id)
|
||||
if not session:
|
||||
return
|
||||
# If a real run is in flight, the cache will be warmed by it;
|
||||
# firing again is wasted tokens.
|
||||
existing = self.tasks.get(session_id)
|
||||
if existing and not existing.done():
|
||||
return
|
||||
|
||||
try:
|
||||
from backend.apps.agents.providers.registry import _find_builtin_model as find_builtin_model
|
||||
entry = find_builtin_model(session.model)
|
||||
if not entry or entry.get("api") != "anthropic":
|
||||
return # other providers handle caching automatically
|
||||
|
||||
from backend.apps.settings.credentials import get_anthropic_client
|
||||
global_settings = load_settings()
|
||||
# Free lane rotates pool accounts per call, so a warm ping primes a cache
|
||||
# the next call won't hit, and worse it'd burn a metered run at idle (this
|
||||
# fires on dashboard mount, not a user query). Skip it on the free trial.
|
||||
if getattr(global_settings, "connection_mode", "own_key") == "free-trial":
|
||||
return
|
||||
client = get_anthropic_client(global_settings)
|
||||
|
||||
# Single ping with the same system + minimal user message.
|
||||
# max_tokens=1 keeps it cheap; we don't care about the output.
|
||||
await client.messages.create(
|
||||
model=entry.get("model_id", session.model),
|
||||
max_tokens=1,
|
||||
system="You are a helpful assistant. Reply with one character.",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
)
|
||||
logger.debug(f"Cache pre-warm fired for session {session_id}")
|
||||
except Exception as e:
|
||||
logger.debug(f"Cache pre-warm failed (non-fatal): {e}")
|
||||
|
||||
@typechecked
|
||||
async def generate_group_meta(self, session_id: str, group_id: str, tool_calls: List[dict], results_summary: Optional[List[str]] = None, is_refinement: bool = False) -> Dict:
|
||||
return await metadata.generate_group_meta(self.sessions.get(session_id), session_id, group_id, tool_calls, results_summary, is_refinement)
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Formal proof (Z3 / SMT) of the MCP dispatch-gate security invariant.
|
||||
|
||||
The product rule "MCP tools are reachable only after MCPActivate" is enforced at
|
||||
dispatch in agent_manager._build_mcp_servers: for a gated session a server is
|
||||
dispatch in agent_manager.p_build_mcp_servers: for a gated session a server is
|
||||
forwarded to the model only if its sanitized name is in session.active_mcps.
|
||||
|
||||
tests/test_v2_invariants.py::test_mcp_gate_only_forwards_activated_servers
|
||||
|
||||
@@ -79,11 +79,11 @@ async def test_gate_blocks_when_active_mcps_empty():
|
||||
_fake_tool("Slack"),
|
||||
_fake_tool("Notion"),
|
||||
]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
# allowed_tools includes mcp:Gmail, but active_mcps is empty
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"],
|
||||
active_mcps=[],
|
||||
)
|
||||
@@ -99,10 +99,10 @@ async def test_gate_allows_only_activated_servers():
|
||||
_fake_tool("Slack"),
|
||||
_fake_tool("Notion"),
|
||||
]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"],
|
||||
active_mcps=["gmail"], # sanitized name of "Gmail"
|
||||
)
|
||||
@@ -117,10 +117,10 @@ async def test_gate_unset_active_mcps_legacy_allows_all():
|
||||
"""Pre-gate sessions use active_mcps=None → everything allowed (back-compat)."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool("Gmail"), _fake_tool("Slack")]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail", "mcp:Slack"],
|
||||
active_mcps=None, # legacy / unset
|
||||
)
|
||||
@@ -133,9 +133,9 @@ async def test_gate_disabled_tool_blocked_even_when_activated():
|
||||
"""Tool with enabled=False stays blocked even if in active_mcps."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool("Gmail", enabled=False)]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail"],
|
||||
active_mcps=["gmail"],
|
||||
)
|
||||
@@ -147,9 +147,9 @@ async def test_gate_unauthed_tool_blocked():
|
||||
"""Tool with auth_status='disconnected' stays blocked."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool("Gmail", auth_status="disconnected")]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail"],
|
||||
active_mcps=["gmail"],
|
||||
)
|
||||
@@ -161,10 +161,10 @@ async def test_gate_allowed_tools_filter_intersects_active_mcps():
|
||||
"""Activate gmail+slack but allowed_tools only has gmail → only gmail passes."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool("Gmail"), _fake_tool("Slack")]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail"], # mode-restricted
|
||||
active_mcps=["gmail", "slack"], # both activated
|
||||
)
|
||||
@@ -192,12 +192,12 @@ async def test_gate_stress_random_activations():
|
||||
# allowed_tools mirrors raw names of connected
|
||||
allowed = [f"mcp:{raw_names[i]}" for i in connected_idx]
|
||||
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_airtable_token", new=AsyncMock(return_value=True)), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_hubspot_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_airtable_token", new=AsyncMock(return_value=True)), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_hubspot_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=allowed,
|
||||
active_mcps=active,
|
||||
)
|
||||
@@ -642,16 +642,16 @@ async def test_mcp_gate_only_forwards_activated_servers():
|
||||
|
||||
# allowed_tools == get_all_tool_names() bypasses the (separate) permission
|
||||
# gate so we isolate the ACTIVATION gate. _sanitize_server_name -> identity.
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", side_effect=installed), \
|
||||
patch("backend.apps.agents.agent_manager.get_all_tool_names", return_value=["__ALL__"]), \
|
||||
patch("backend.apps.agents.agent_manager._sanitize_server_name", side_effect=lambda n: n), \
|
||||
patch("backend.apps.agents.agent_manager.is_fully_denied", return_value=False), \
|
||||
patch("backend.apps.agents.agent_manager.derive_mcp_config", side_effect=lambda t: {"command": "x"}):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", side_effect=installed), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.get_all_tool_names", return_value=["__ALL__"]), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.sanitize_server_name", side_effect=lambda n: n), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.is_fully_denied", return_value=False), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.derive_mcp_config", side_effect=lambda t: {"command": "x"}):
|
||||
allowed = ["__ALL__"]
|
||||
# Boundary 1: empty activation list -> zero servers, always.
|
||||
assert await mgr._build_mcp_servers(allowed, active_mcps=[]) == {}
|
||||
assert await mgr.p_build_mcp_servers(allowed, active_mcps=[]) == {}
|
||||
# Boundary 2: None (legacy) -> permission gate only, all forwarded.
|
||||
assert set((await mgr._build_mcp_servers(allowed, active_mcps=None)).keys()) == set(names)
|
||||
assert set((await mgr.p_build_mcp_servers(allowed, active_mcps=None)).keys()) == set(names)
|
||||
# Property: forwarded set is ALWAYS a subset of the activated set, and
|
||||
# equals exactly the activated-and-installed intersection.
|
||||
rng = random.Random(1234)
|
||||
@@ -660,7 +660,7 @@ async def test_mcp_gate_only_forwards_activated_servers():
|
||||
# throw in a bogus name the gate must never invent a server for
|
||||
if rng.random() < 0.3:
|
||||
active = active + ["ghost-not-installed"]
|
||||
forwarded = set((await mgr._build_mcp_servers(allowed, active_mcps=active)).keys())
|
||||
forwarded = set((await mgr.p_build_mcp_servers(allowed, active_mcps=active)).keys())
|
||||
assert forwarded <= set(active), f"leaked {forwarded - set(active)} for active={active}"
|
||||
assert forwarded == (set(active) & set(names)), f"mismatch for active={active}"
|
||||
|
||||
@@ -963,14 +963,14 @@ async def test_concurrent_gate_calls_isolated():
|
||||
"""Two concurrent _build_mcp_servers calls with different active_mcps must not cross-contaminate."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool("Gmail"), _fake_tool("Slack"), _fake_tool("Notion")]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
results = await asyncio.gather(
|
||||
mgr._build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["gmail"]),
|
||||
mgr._build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["slack"]),
|
||||
mgr._build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["notion"]),
|
||||
mgr._build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=[]),
|
||||
mgr.p_build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["gmail"]),
|
||||
mgr.p_build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["slack"]),
|
||||
mgr.p_build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=["notion"]),
|
||||
mgr.p_build_mcp_servers(allowed_tools=["mcp:Gmail", "mcp:Slack", "mcp:Notion"], active_mcps=[]),
|
||||
)
|
||||
gmail_only, slack_only, notion_only, empty = results
|
||||
assert set(gmail_only.keys()) == {"gmail"}
|
||||
@@ -1051,7 +1051,7 @@ async def test_context_update_emitter_refreshes_session_tokens(monkeypatch):
|
||||
s.framework_overhead_tokens = 42
|
||||
s.active_mcps = ["github"]
|
||||
|
||||
await AgentManager()._emit_context_update("x", s, input_tokens=250)
|
||||
await AgentManager().p_emit_context_update("x", s, input_tokens=250)
|
||||
|
||||
assert s.tokens == {"input": 250, "output": 7}
|
||||
assert sent == [(
|
||||
@@ -1184,9 +1184,9 @@ async def test_gate_handles_missing_refresh_token_gracefully():
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake = _fake_tool("MyApiTool", auth_status="configured")
|
||||
fake.auth_type = None # no oauth
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=[fake]):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=[fake]):
|
||||
mgr = AgentManager()
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:MyApiTool"],
|
||||
active_mcps=["myapitool"],
|
||||
)
|
||||
@@ -1384,7 +1384,7 @@ def test_resolve_attachments_handles_missing_path_gracefully():
|
||||
we emit a 'not found' refusal instead of crashing."""
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
mgr = AgentManager()
|
||||
text, native, refusals = mgr._resolve_attachments(
|
||||
text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": "/var/folders/nonexistent/definitely-gone.pdf", "type": "file"}],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -1402,7 +1402,7 @@ def test_resolve_attachments_handles_directory_path_not_file():
|
||||
tmpdir = tempfile.mkdtemp()
|
||||
open(os.path.join(tmpdir, "a.txt"), "w").write("hello")
|
||||
try:
|
||||
text, native, refusals = mgr._resolve_attachments(
|
||||
text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": tmpdir, "type": "directory"}],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -1431,7 +1431,7 @@ def test_resolve_attachments_mixed_kinds_total_size_guard():
|
||||
paths.append(fh.name)
|
||||
with tempfile.NamedTemporaryFile(suffix=".md", mode="w", delete=False) as fh:
|
||||
fh.write("# notes"); paths.append(fh.name)
|
||||
text, native, refusals = mgr._resolve_attachments(
|
||||
text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": p, "type": "file"} for p in paths],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -1491,7 +1491,7 @@ def test_sniff_recognises_macos_paths_with_spaces():
|
||||
try:
|
||||
with open(path, "wb") as f:
|
||||
f.write(b"%PDF-1.4\n")
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
assert native and native[0]["type"] == "document"
|
||||
@@ -1552,7 +1552,7 @@ def test_sniff_handles_windows_style_backslash_path_string():
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
mgr = AgentManager()
|
||||
# A path that doesn't exist (POSIX cannot interpret backslashes as separator)
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": r"C:\fake\path\nope.pdf", "type": "file"}],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -1591,7 +1591,7 @@ def test_resolve_attachments_classifies_renamed_binary_as_binary_not_pdf():
|
||||
fh.write(b"PK\x03\x04fake zip masquerading as pdf")
|
||||
path = fh.name
|
||||
try:
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
assert not native
|
||||
@@ -1664,7 +1664,7 @@ def test_anthropic_document_block_schema_matches_docs():
|
||||
fh.write(b"%PDF-1.4\n%canonical schema test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_t, native, _r = mgr._resolve_attachments(
|
||||
_t, native, _r = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
block = native[0]
|
||||
@@ -1909,7 +1909,7 @@ def test_resolve_attachments_openai_codex_refused_for_pdfs():
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="openai", model="gpt-5.3-codex",
|
||||
)
|
||||
assert not native
|
||||
@@ -1928,7 +1928,7 @@ def test_resolve_attachments_openai_codex_still_refuses_pdf():
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="openai", model="gpt-5.3-codex",
|
||||
)
|
||||
assert not native
|
||||
@@ -2029,7 +2029,7 @@ def test_resolve_attachments_anthropic_emits_native_document():
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
text, native, refusals = mgr._resolve_attachments(
|
||||
text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
assert native and native[0]["type"] == "document"
|
||||
@@ -2051,7 +2051,7 @@ def test_resolve_attachments_openai_accepts_pdf_via_bypass_translator():
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_text, native, refusals = mgr._resolve_attachments(
|
||||
_text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="openai", model="gpt-5.5",
|
||||
)
|
||||
assert native and native[0]["type"] == "document"
|
||||
@@ -2143,7 +2143,7 @@ def test_resolve_attachments_gemini_emits_native_document_after_translator_fix()
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_text, native, refusals = mgr._resolve_attachments(
|
||||
_text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="gemini", model="gemini-3.1-pro-api",
|
||||
)
|
||||
assert native and native[0]["type"] == "document"
|
||||
@@ -2162,7 +2162,7 @@ def test_resolve_attachments_text_file_inlined_not_native():
|
||||
fh.write("# hello\nworld")
|
||||
path = fh.name
|
||||
try:
|
||||
text, native, refusals = mgr._resolve_attachments(
|
||||
text, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="opus-4-7", model="opus-4-7",
|
||||
)
|
||||
assert not native
|
||||
@@ -2182,7 +2182,7 @@ def test_resolve_attachments_pdf_refused_when_too_large():
|
||||
fh.write(b"X" * (25 * 1024 * 1024))
|
||||
path = fh.name
|
||||
try:
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
assert not native
|
||||
@@ -2208,7 +2208,7 @@ def test_resolve_attachments_refuses_when_total_exceeds_request_cap():
|
||||
fh.write(b"%PDF-1.4\n")
|
||||
fh.write(b"X" * (8 * 1024 * 1024))
|
||||
paths.append(fh.name)
|
||||
_t, native, refusals = mgr._resolve_attachments(
|
||||
_t, native, refusals = mgr.p_resolve_attachments(
|
||||
[{"path": p, "type": "file"} for p in paths],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -2234,7 +2234,7 @@ def test_resolve_attachments_anthropic_marks_last_document_ephemeral_for_cache()
|
||||
with tempfile.NamedTemporaryFile(suffix=f"_{i}.pdf", delete=False) as fh:
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
paths.append(fh.name)
|
||||
_t, native, _r = mgr._resolve_attachments(
|
||||
_t, native, _r = mgr.p_resolve_attachments(
|
||||
[{"path": p, "type": "file"} for p in paths],
|
||||
api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
@@ -2258,11 +2258,11 @@ def test_resolve_attachments_anthropic_does_mark_ephemeral_but_only_anthropic():
|
||||
fh.write(b"%PDF-1.4\n%test\n")
|
||||
path = fh.name
|
||||
try:
|
||||
_t, ant_native, _r = mgr._resolve_attachments(
|
||||
_t, ant_native, _r = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="anthropic", model="opus-4-7",
|
||||
)
|
||||
assert ant_native and ant_native[0].get("cache_control") == {"type": "ephemeral"}
|
||||
_t, or_native, _r = mgr._resolve_attachments(
|
||||
_t, or_native, _r = mgr.p_resolve_attachments(
|
||||
[{"path": path, "type": "file"}], api_type="openrouter", model="openrouter/openai/gpt-5",
|
||||
)
|
||||
assert or_native and "cache_control" not in or_native[0]
|
||||
@@ -3026,14 +3026,14 @@ def test_view_builder_mode_has_default_folder():
|
||||
async def test_gate_100_sequential_calls_no_leak():
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
fake_tools = [_fake_tool(f"Server{i}") for i in range(10)]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
for i in range(100):
|
||||
n = i % 10
|
||||
active = [f"server{j}" for j in range(n)]
|
||||
allowed = [f"mcp:Server{j}" for j in range(10)]
|
||||
result = await mgr._build_mcp_servers(allowed_tools=allowed, active_mcps=active)
|
||||
result = await mgr.p_build_mcp_servers(allowed_tools=allowed, active_mcps=active)
|
||||
assert set(result.keys()) == set(active), \
|
||||
f"iteration {i}: expected {set(active)}, got {set(result.keys())}"
|
||||
|
||||
@@ -3141,13 +3141,13 @@ async def test_e2e_session_lifecycle_with_mcp_activation():
|
||||
from backend.apps.agents.agent_manager import AgentManager
|
||||
from backend.apps.agents.core.models import AgentSession
|
||||
fake_tools = [_fake_tool("Gmail"), _fake_tool("Slack")]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools", return_value=fake_tools), \
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
s = AgentSession(id="e2e", name="End-to-end", model="sonnet", mode="agent")
|
||||
|
||||
# Step 1: fresh, gate blocks everything
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail", "mcp:Slack"],
|
||||
active_mcps=s.active_mcps,
|
||||
)
|
||||
@@ -3161,7 +3161,7 @@ async def test_e2e_session_lifecycle_with_mcp_activation():
|
||||
s.pending_continuation = True
|
||||
|
||||
# Step 3: continuation turn, gate passes gmail
|
||||
result = await mgr._build_mcp_servers(
|
||||
result = await mgr.p_build_mcp_servers(
|
||||
allowed_tools=["mcp:Gmail", "mcp:Slack"],
|
||||
active_mcps=s.active_mcps,
|
||||
)
|
||||
@@ -3184,15 +3184,15 @@ async def test_e2e_50_random_activation_sequences():
|
||||
("Discord", "discord"), ("GitHub", "github"), ("Linear", "linear")]
|
||||
raw_names = [r for r, _ in server_pool]
|
||||
sanitized = [s for _, s in server_pool]
|
||||
with patch("backend.apps.agents.agent_manager.load_all_tools",
|
||||
with patch("backend.apps.agents.manager.RunSupportMixin.load_all_tools",
|
||||
return_value=[_fake_tool(r) for r in raw_names]), \
|
||||
patch("backend.apps.agents.agent_manager.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
patch("backend.apps.agents.manager.RunSupportMixin.refresh_google_token", new=AsyncMock(return_value=True)):
|
||||
mgr = AgentManager()
|
||||
for _ in range(50):
|
||||
n = random.randint(0, len(sanitized))
|
||||
active = random.sample(sanitized, n)
|
||||
allowed = [f"mcp:{r}" for r in raw_names]
|
||||
result = await mgr._build_mcp_servers(allowed, active)
|
||||
result = await mgr.p_build_mcp_servers(allowed, active)
|
||||
keys = set(result.keys())
|
||||
assert keys == set(active), f"mismatch: active={active} keys={keys}"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user