openswarm/backend/apps/agents/manager/run/run_options_helpers.py

"""Self-contained sub-steps of the per-turn options build, pulled out of run_agent_loop's options
assembly so each file stays under the ceiling. Free functions taking the manager (for maybe_compact /
emit_context_update); pure relocation."""

import logging
from typing import Dict, List, Optional
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.agents.manager.session.history_compaction import estimate_post_compact_input, wrap_platform_note

logger = logging.getLogger(__name__)

# Always SDK-blocked regardless of permissions: claude.ai partner MCPs bypass our MCPActivate gate, and the CLI's built-in sub-agent tool (Task on 2.1.122, Agent on older builds) is replaced by our SpawnAgent MCP.
HARD_BLOCKED_TOOLS: List[str] = ["mcp__claude_ai_*", "Agent", "Task"]


@typechecked
def merge_hard_blocked_tools(effective_disallowed: List[str]) -> List[str]:
    """The SDK deny list = the computed per-turn denies PLUS the unconditional hard blocks. A plain assignment here once silently discarded effective_disallowed (Cron*/Skill/web-swap/per-tool MCP denies), leaving the runtime gate as the only wall; merge, never overwrite."""
    return effective_disallowed + [t for t in HARD_BLOCKED_TOOLS if t not in effective_disallowed]


# `manager` is the AgentManager; it isn't annotated because typing it would import agent_manager back into a module agent_manager already imports (a cycle). Same reason self is never annotated.
@typechecked
async def pre_send_context_guard(manager, session: AgentSession, session_id: str) -> None:
    # The second trigger (hermes lift): the threshold below fires at a percentage of the window,
    # which on a 1M lane almost never arrives, so history rides untouched to the cliff. This one
    # fires on COST and pays for its own cache miss.
    try:
        from backend.apps.agents.manager.session.proactive_prune import (
            arm_proactive_prune,
            estimate_aged_rebuild_tokens,
            should_proactively_prune,
        )
        if should_proactively_prune(session):
            arm_proactive_prune(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 manager.emit_context_update(
                session_id, session,
                input_tokens=estimate_aged_rebuild_tokens(session),
                output_tokens=session.tokens.get("output", 0),
            )
    except Exception:
        logger.exception("proactive prune failed; proceeding without it")

    try:
        if manager.maybe_compact(session):
            # A mark alone never applies on the resume path (the CLI replays its own untrimmed transcript), so pay for the rebuild too: next turn drops the SDK convo and rebuilds with the cutoff + distilled summary. One respawn per compaction epoch is the price of never reaching the wall.
            session.needs_fresh_session = True
            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 manager.emit_context_update(
                session_id,
                session,
                input_tokens=new_input,
                output_tokens=session.tokens.get("output", 0),
            )
    except Exception:
        logger.exception("compaction failed; proceeding without it")

    # Pre-send hard guard (Phase 2). After compaction, if the session is still over context_soft_cap_pct of the window, LRU-trim oldest active_mcps. Stops the 429 from ever firing on predictable overflow paths.
    try:
        # Use the most recent measurement (the prior turn's input_tokens) as the estimate. Conservative because the current turn's user prompt + any new history adds on top, but the first turn of a fresh session has tokens=0 so we only act once we've seen real numbers.
        p_est_tokens = session.tokens.get("input", 0)
        p_hard_cap = int(session.context_window * session.context_soft_cap_pct)
        if p_est_tokens >= p_hard_cap:
            trimmed: List[str] = []
            while p_est_tokens >= p_hard_cap and len(session.active_mcps) > 1:
                # Keep at least one MCP active so the model can finish whatever it was doing; trim from oldest which is FIFO order in the list.
                trimmed.append(f"mcp:{session.active_mcps.pop(0)}")
                p_est_tokens -= 8_000  # rough per-MCP schema cost
            if trimmed:
                await ws_manager.send_to_session(session_id, "agent:context_status", {
                    "session_id": session_id,
                    "reason": "trimmed",
                    "trimmed": trimmed,
                    "estimate_after": p_est_tokens,
                })
                # Surface a visible system breadcrumb in the chat so the user (and the model on the next turn) know which MCPs got dropped. Without this, the model may keep trying to call a now-missing tool and the user has no idea why.
                try:
                    p_names = ", ".join(t.replace("mcp:", "") for t in trimmed)
                    p_trim_msg = Message(
                        role="system",
                        content=wrap_platform_note(
                            f"Trimmed {len(trimmed)} app{'s' if len(trimmed) != 1 else ''} from this session to fit "
                            f"the model's context: {p_names}. Re-activate via MCPSearch + MCPActivate "
                            "if you still need them."
                        ),
                        branch_id=session.active_branch_id,
                    )
                    session.messages.append(p_trim_msg)
                    await ws_manager.send_to_session(session_id, "agent:message", {
                        "session_id": session_id,
                        "message": p_trim_msg.model_dump(mode="json"),
                    })
                except Exception:
                    logger.exception("failed to emit MCP-trimmed breadcrumb")
                # Trimming changes mcp_servers / outputs context → rebuild options. The cheapest correct path is to flag for fork on next turn via needs_fork and let the existing fork path handle it.
                session.needs_fork = True
    except Exception:
        logger.exception("pre-send token guard failed; proceeding")


@typechecked
def set_framework_overhead(session: AgentSession, composed_prompt: Optional[str]) -> None:
    """Per-turn estimate of framework overhead (subtracted from displayed input). Conservative on
    purpose so honest over-shows beat lies: 16K Claude Code preset, 12K base+deferred tools, ~3K/MCP
    (real defs span 1-10K; 3K median keeps the meter honest), char/4 of the composed prompt."""
    p_PRESET_OVERHEAD = 16_000
    p_TOOL_DEFS_OVERHEAD = 12_000
    p_PER_MCP_OVERHEAD = 3_000
    p_composed_tokens = len(composed_prompt or "") // 4
    p_mcp_tokens = len(session.active_mcps) * p_PER_MCP_OVERHEAD
    session.framework_overhead_tokens = (
        p_PRESET_OVERHEAD + p_TOOL_DEFS_OVERHEAD + p_composed_tokens + p_mcp_tokens
    )


@typechecked
def register_web_mcp_server(mcp_servers: Dict, p_m: str, browser_ok: bool = False, rich_ui_ok: bool = False) -> None:
    """Add the DDG-backed web tools when the primary has no reliable native web path. They ride the
    combined openswarm-core process now (ENG-208): flip the module flag and env on the entry that
    register_builtin_mcp_servers already made, instead of spawning an eleventh interpreter."""
    # Tell the MCP which primary the session is using so it can route to that provider's native search tool.
    if p_m.startswith(("gc/", "gemini/", "ag/")):
        p_primary_hint = "gemini"
    elif p_m.startswith("cx/"):
        p_primary_hint = "openai"
    else:
        p_primary_hint = ""
    core = mcp_servers.get("openswarm-core")
    if core is not None:
        env = core["env"]
        modules = [m for m in env.get("OSW_MCP_MODULES", "").split(",") if m]
        if "web" not in modules:
            modules.append("web")
        env["OSW_MCP_MODULES"] = ",".join(modules)
        env["OPENSWARM_PRIMARY_API"] = p_primary_hint
        env["OPENSWARM_BROWSER_OK"] = "1" if browser_ok else "0"
        env["OPENSWARM_RICH_UI_OK"] = "1" if rich_ui_ok else "0"
        return
    # Belt for a caller that skipped register_builtin (none today): web tools still arrive, alone.
    import os
    import sys
    import backend.apps.agents as p_agents_pkg
    combined_path = os.path.join(os.path.dirname(p_agents_pkg.__file__), "combined_meta_mcp_server.py")
    from backend.auth import get_auth_token as p_get_auth_token3
    mcp_servers["openswarm-core"] = {
        "command": sys.executable,
        "args": [combined_path],
        "env": {
            "OSW_MCP_MODULES": "web",
            "OPENSWARM_PORT": os.environ.get("OPENSWARM_PORT", "8324"),
            "OPENSWARM_AUTH_TOKEN": p_get_auth_token3(),
            "OPENSWARM_PRIMARY_API": p_primary_hint,
            "OPENSWARM_BROWSER_OK": "1" if browser_ok else "0",
            "OPENSWARM_RICH_UI_OK": "1" if rich_ui_ok else "0",
        },
        "type": "stdio",
    }
    logger.info(
        f"[MCP-DEBUG] Primary {p_m} has no reliable native web search, "
        f"registering openswarm-web (DDG search + trafilatura fetch, free)"
    )


@typechecked
def append_web_tools_hint(composed_prompt: Optional[str], need_web_mcp: bool, effective_allowed: List[str]) -> str:
    """Append a <web_tools> block naming the MCP-backed WebSearch/WebFetch when the deferred bare
    WebSearch tool isn't usable on this session, so smaller models don't thrash on ToolSearch."""
    p_web_tools_available = need_web_mcp and (
        "mcp__openswarm-core__WebSearch" in effective_allowed
        or "mcp__openswarm-core__WebFetch" in effective_allowed
    )
    if not p_web_tools_available:
        return composed_prompt
    p_hint_lines = ["<web_tools>"]
    p_hint_lines.append(
        "This session does NOT have the built-in `WebSearch` / "
        "`WebFetch` tools (they delegate to Anthropic Haiku, which "
        "isn't reachable on this primary). Use the MCP-backed "
        "equivalents instead, call them DIRECTLY, no ToolSearch "
        "step needed:"
    )
    if "mcp__openswarm-core__WebSearch" in effective_allowed:
        p_hint_lines.append(
            "- `mcp__openswarm-core__WebSearch(query: str, "
            "num_results?: int)`, DuckDuckGo search."
        )
    if "mcp__openswarm-core__WebFetch" in effective_allowed:
        p_hint_lines.append(
            "- `mcp__openswarm-core__WebFetch(url: str, prompt?: "
            "str)`, fetch a URL and return readable text."
        )
    p_hint_lines.append(
        "Do not call `ToolSearch(select:WebSearch)`, bare "
        "`WebSearch` is unavailable on this session and that path "
        "will return empty matches."
    )
    p_hint_lines.append("</web_tools>")
    p_web_hint = "\n".join(p_hint_lines)
    return f"{composed_prompt}\n\n{p_web_hint}" if composed_prompt else p_web_hint


@typechecked
def inject_thinking_options(options_kwargs: Dict, session: AgentSession, prompt: str, resolved_model: str, api_type: str) -> None:
    """Map the session's thinking_level onto the SDK options (anthropic thinking/effort, openai/codex
    reasoning_effort), with the short-prompt + gc/gemini-3 force-off overrides. Best-effort."""
    try:
        level = getattr(session, "thinking_level", "auto") or "auto"
        # Trivially short prompts ("hi", "thanks") don't benefit from 5-30s of hidden reasoning, but
        # this flip rides the BOOT fingerprint: a short first message drifted it and threw away the
        # pre-warmed CLI, measured as a second spawn on 5 of 5 short-prompt sessions. A wasted 0.9s
        # respawn costs more than the reasoning it saves, so the flip only applies once a session is
        # already running on a live client (later turns), never on the first message.
        p_prompt_len = len((prompt or "").strip())
        p_first_turn = not getattr(session, "sdk_session_id", None)
        if 0 < p_prompt_len < 50 and level != "off" and not p_first_turn:
            level = "off"
        # gc/gemini-3* without Antigravity 400s every multi-step turn on thoughtSignature continuity.
        if (
            isinstance(resolved_model, str)
            and resolved_model.startswith("gc/gemini-3")
            and level != "off"
        ):
            logger.info(
                "Forcing thinking_level=off for %s (gc/ thoughtSignature isn't roundtrippable; connect Antigravity for reasoning).",
                resolved_model,
            )
            level = "off"
        if api_type == "anthropic":
            if level == "off":
                options_kwargs["thinking"] = {"type": "disabled"}
            elif level in ("low", "medium", "high"):
                options_kwargs["effort"] = level
        elif api_type in ("openai", "codex"):
            # GPT-5 + Codex take reasoning_effort; 9Router carries the Anthropic-shaped `effort`.
            if level in ("low", "medium", "high"):
                options_kwargs["effort"] = level
    except Exception as e:
        logger.debug(f"thinking_level param injection skipped: {e}")