From 9459c410a65a7e68d2f82dbbc11c514e3f52d400 Mon Sep 17 00:00:00 2001 From: David Hadas Date: Sun, 5 Jul 2026 10:10:18 +0300 Subject: [PATCH] fix mcp never failing till timeout bug Signed-off-by: David Hadas --- a2a/git_issue_agent/a2a_agent.py | 72 ++++++- a2a/git_issue_agent/git_issue_agent/agents.py | 11 +- a2a/git_issue_agent/git_issue_agent/config.py | 15 ++ a2a/git_issue_agent/git_issue_agent/main.py | 57 ++++- .../git_issue_agent/mcp_connect.py | 197 ++++++++++++++++++ .../git_issue_agent/prompts.py | 2 + .../git_issue_agent/tool_limits.py | 157 ++++++++++++++ 7 files changed, 499 insertions(+), 12 deletions(-) create mode 100644 a2a/git_issue_agent/git_issue_agent/mcp_connect.py create mode 100644 a2a/git_issue_agent/git_issue_agent/tool_limits.py diff --git a/a2a/git_issue_agent/a2a_agent.py b/a2a/git_issue_agent/a2a_agent.py index 1cbb2761..6c394d31 100644 --- a/a2a/git_issue_agent/a2a_agent.py +++ b/a2a/git_issue_agent/a2a_agent.py @@ -2,13 +2,14 @@ Module for A2A Agent. """ +import asyncio import logging import os import sys +import threading import traceback import uvicorn -from crewai_tools import MCPServerAdapter from crewai_tools.adapters.tool_collection import ToolCollection @@ -32,7 +33,8 @@ from git_issue_agent.config import settings, Settings from git_issue_agent.event import Event -from git_issue_agent.main import GitIssueAgent +from git_issue_agent.main import GitIssueAgent, TaskCancelled +from git_issue_agent.mcp_connect import mcp_tools_session logger = logging.getLogger(__name__) logging.basicConfig(level=settings.LOG_LEVEL, stream=sys.stdout, format="%(levelname)s: %(message)s") @@ -115,11 +117,26 @@ class GithubExecutor(AgentExecutor): A class to handle research execution for A2A Agent. """ - async def _run_agent(self, messages: dict, settings: Settings, event_emitter: Event, toolkit: ToolCollection): + def __init__(self): + # Per-request cooperative-cancel flags, keyed by task id. CrewAI runs its + # crew synchronously in a worker thread (kickoff_async -> asyncio.to_thread), + # so cancelling the awaiting coroutine cannot stop the work; instead the + # crew checks this Event between steps and stops itself. + self._cancel_events: dict[str, threading.Event] = {} + + async def _run_agent( + self, + messages: dict, + settings: Settings, + event_emitter: Event, + toolkit: ToolCollection, + cancel_event: threading.Event, + ): git_issue_agent = GitIssueAgent( config=settings, eventer=event_emitter, mcp_toolkit=toolkit, + cancel_event=cancel_event, ) result = await git_issue_agent.execute(messages) await event_emitter.emit_event(result, True) @@ -163,6 +180,11 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): } ) + # Register a cooperative-cancel flag for this task so cancel()/disconnect + # can stop the (thread-bound) crew run. + cancel_event = threading.Event() + self._cancel_events[task.id] = cancel_event + # Hook up MCP tools try: if settings.MCP_URL: @@ -173,7 +195,13 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): "transport": "streamable-http", "headers": headers, } - with MCPServerAdapter(server_params, connect_timeout=settings.MCP_TIMEOUT) as mcp_tools: + # mcp_tools_session fails fast if the connection errors out, while + # still allowing up to MCP_TIMEOUT for a slow (e.g. OAuth) connect. + async with mcp_tools_session( + server_params, + connect_timeout=settings.MCP_TIMEOUT, + poll_interval=settings.MCP_POLL_INTERVAL, + ) as mcp_tools: # Keep only search and list issue-related tools. issue_tools = [ tool @@ -187,21 +215,47 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): "No issue-related tools found from the GitHub MCP server. " "Ensure your PAT scopes allow issue access and the server is reachable." ) - await self._run_agent(messages, settings, event_emitter, issue_tools) + # Tool output is bounded inside GitIssueAgent (wrap_tool_output). + await self._run_agent(messages, settings, event_emitter, issue_tools, cancel_event) else: - await self._run_agent(messages, settings, event_emitter, None) - + await self._run_agent(messages, settings, event_emitter, None, cancel_event) + + except (asyncio.CancelledError, TaskCancelled): + # A2A raises CancelledError in this task on cancel/disconnect; the crew + # may also surface TaskCancelled after observing the flag. Signal the + # worker thread to stop and report the task as cancelled. + cancel_event.set() + logging.info("Task %s cancelled; stopping crew run", task.id) + try: + await task_updater.cancel() + except Exception: # noqa: BLE001 - best-effort status on an already-torn-down queue + pass + raise except Exception as e: traceback.print_exc() await event_emitter.emit_event( f"I'm sorry I was unable to fulfill your request. I encountered the following exception: {str(e)}", True ) + finally: + self._cancel_events.pop(task.id, None) async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: """ - Not implemented + Signal the running crew for this task to stop, and report cancellation. + + The crew runs in a worker thread, so we flip its cooperative-cancel flag; + it stops at the next step boundary rather than mid-LLM-call. """ - raise Exception("cancel not supported") + task = context.current_task + task_id = task.id if task else None + cancel_event = self._cancel_events.get(task_id) if task_id else None + if cancel_event is not None: + cancel_event.set() + logging.info("Cancellation requested for task %s", task_id) + + if task_id: + task_updater = TaskUpdater(event_queue, task_id, task.context_id) + await task_updater.cancel() def run(): diff --git a/a2a/git_issue_agent/git_issue_agent/agents.py b/a2a/git_issue_agent/git_issue_agent/agents.py index 33ecd5a1..c9506c30 100644 --- a/a2a/git_issue_agent/git_issue_agent/agents.py +++ b/a2a/git_issue_agent/git_issue_agent/agents.py @@ -5,8 +5,11 @@ class GitAgents: - def __init__(self, config: Settings, issue_tools): + def __init__(self, config: Settings, issue_tools, step_callback=None): self.llm = CrewLLM(config) + # step_callback is invoked by CrewAI after each agent step; we use it to + # cooperatively stop a run on cancellation. + self.step_callback = step_callback ################### # Pre-requisite validator @@ -33,6 +36,7 @@ def __init__(self, config: Settings, issue_tools): tasks=[self.prereq_identifier_task], process=Process.sequential, verbose=True, + step_callback=step_callback, ) ################### @@ -49,6 +53,10 @@ def __init__(self, config: Settings, issue_tools): verbose=True, llm=self.llm.llm, inject_date=True, + # max_iter is the backstop for the context-overflow fallback: even if a + # tool result slips past the size caps (tool_limits.py) and triggers + # CrewAI's summarize-and-retry loop, the run terminates after this many + # iterations with a partial answer rather than looping unbounded. max_iter=6, max_retry_limit=3, respect_context_window=True, @@ -75,4 +83,5 @@ def __init__(self, config: Settings, issue_tools): tasks=[self.issue_query_task], process=Process.sequential, verbose=True, + step_callback=step_callback, ) diff --git a/a2a/git_issue_agent/git_issue_agent/config.py b/a2a/git_issue_agent/git_issue_agent/config.py index 98d7a7bc..a2e4c340 100644 --- a/a2a/git_issue_agent/git_issue_agent/config.py +++ b/a2a/git_issue_agent/git_issue_agent/config.py @@ -30,6 +30,21 @@ class Settings(BaseSettings): os.getenv("MCP_URL", "https://api.githubcopilot.com/mcp/"), description="Endpoint for an option MCP server" ) MCP_TIMEOUT: int = Field(os.getenv("MCP_TIMEOUT", 600), description="Timeout in seconds for MCP server connection") + MCP_POLL_INTERVAL: float = Field( + os.getenv("MCP_POLL_INTERVAL", 1.0), + description="How often (seconds) to poll the MCP connection state so failures surface fast", + gt=0, + ) + MAX_ISSUES: int = Field( + os.getenv("MAX_ISSUES", 30), + description="Max number of items kept from an MCP issue result set before it reaches the LLM", + gt=0, + ) + MAX_TOOL_CHARS: int = Field( + os.getenv("MAX_TOOL_CHARS", 24000), + description="Hard character budget for any single tool observation fed to the LLM", + gt=0, + ) # auth variables for token validation ISSUER: Optional[str] = Field(os.getenv("ISSUER", None), description="The issuer for incoming JWT tokens") diff --git a/a2a/git_issue_agent/git_issue_agent/main.py b/a2a/git_issue_agent/git_issue_agent/main.py index bcd3966d..b32ca0a3 100644 --- a/a2a/git_issue_agent/git_issue_agent/main.py +++ b/a2a/git_issue_agent/git_issue_agent/main.py @@ -2,6 +2,7 @@ import logging import re import sys +import threading from crewai_tools.adapters.tool_collection import ToolCollection @@ -9,11 +10,16 @@ from git_issue_agent.config import Settings, settings from git_issue_agent.data_types import IssueSearchInfo from git_issue_agent.event import Event +from git_issue_agent.tool_limits import wrap_tool_output logger = logging.getLogger(__name__) logging.basicConfig(level=settings.LOG_LEVEL, stream=sys.stdout, format="%(levelname)s: %(message)s") +class TaskCancelled(Exception): + """Raised to unwind a CrewAI run once cancellation has been requested.""" + + def _parse_prereq_from_raw(raw: str) -> IssueSearchInfo: """Parse IssueSearchInfo from raw LLM text when instructor/pydantic parsing fails. @@ -41,11 +47,38 @@ def __init__( eventer: Event = None, mcp_toolkit: ToolCollection = None, logger=None, + cancel_event: threading.Event = None, ): - self.agents = GitAgents(config, mcp_toolkit) + self.cancel_event = cancel_event or threading.Event() + self._truncation_notes: list[str] = [] + # Bound each MCP tool's output so a large result set can't overflow the + # model context window (which triggers CrewAI's slow summarize fallback). + if mcp_toolkit: + mcp_toolkit = [ + wrap_tool_output( + tool, + max_items=config.MAX_ISSUES, + max_chars=config.MAX_TOOL_CHARS, + on_truncate=self.add_truncation_note, + ) + for tool in mcp_toolkit + ] + # step_callback runs on the CrewAI worker thread after each agent step; + # raising here propagates out of the crew loop and stops it promptly. + self.agents = GitAgents(config, mcp_toolkit, step_callback=self._on_step) self.eventer = eventer self.logger = logger or logging.getLogger(__name__) + def _on_step(self, _step) -> None: + """CrewAI step callback: abort the run if cancellation was requested.""" + if self.cancel_event.is_set(): + raise TaskCancelled() + + def _check_cancelled(self) -> None: + """Abort between crews (e.g. prereq -> main) if cancellation was requested.""" + if self.cancel_event.is_set(): + raise TaskCancelled() + async def _send_event(self, message: str, final: bool = False): logger.info(message) if self.eventer: @@ -82,10 +115,26 @@ async def _get_prereq_output(self, query: str) -> IssueSearchInfo: raw = self.agents.prereq_identifier_task.output.raw self.logger.info(f"Prereq raw output: {raw}") return _parse_prereq_from_raw(raw) + except TaskCancelled: + # Never swallow cancellation as a "prereq failed" fallback. + raise except Exception as e: self.logger.warning(f"Prereq crew failed: {e}") return IssueSearchInfo() + def add_truncation_note(self, tool_name: str, kept: int, total: int) -> None: + """Record that a tool result was trimmed, for a user-facing note. + + Used as the ``on_truncate`` callback for wrap_tool_output. ``kept == -1`` + marks a raw character-budget cut where item counts are unknown. + """ + if kept == -1: + note = "Some tool output was truncated to fit the model context window." + else: + note = f"Results were limited to the first {kept} of {total} items." + if note not in self._truncation_notes: + self._truncation_notes.append(note) + async def execute(self, user_input): query = self.extract_user_input(user_input) await self._send_event("šŸ§ Evaluating requirements...") @@ -98,6 +147,7 @@ async def execute(self, user_input): if not repo_id_task_output.owner: return "When supplying a repository name, you must also provide an owner of the repo." + self._check_cancelled() await self._send_event("šŸ”Ž Searching for issues...") await self.agents.crew.kickoff_async( inputs={ @@ -107,4 +157,7 @@ async def execute(self, user_input): "issues": repo_id_task_output.issue_numbers, } ) - return self.agents.issue_query_task.output.raw + answer = self.agents.issue_query_task.output.raw + if self._truncation_notes: + answer = f"{answer}\n\n_Note: {' '.join(self._truncation_notes)}_" + return answer diff --git a/a2a/git_issue_agent/git_issue_agent/mcp_connect.py b/a2a/git_issue_agent/git_issue_agent/mcp_connect.py new file mode 100644 index 00000000..98c0e809 --- /dev/null +++ b/a2a/git_issue_agent/git_issue_agent/mcp_connect.py @@ -0,0 +1,197 @@ +"""Reusable helper for connecting to an MCP server without hanging on failure. + +Background +---------- +``crewai_tools.MCPServerAdapter`` connects via ``mcpadapt``. Under the hood +``mcpadapt`` runs the connection inside a *daemon thread* and the synchronous +entry point blocks the caller on a single ``threading.Event`` wait: + + if not self.ready.wait(timeout=connect_timeout): + raise TimeoutError(...) + +If the connection *fails* for any reason (auth refused, HTTP 401, connection +refused, DNS error, ...), the exception is raised inside the daemon thread and +is **never re-raised on the calling thread**. ``ready`` is simply never set, so +the caller blocks for the full ``connect_timeout`` (600s in this agent) before +getting a generic ``TimeoutError``. Because this happens synchronously inside an +``async def`` handler, it also freezes the event loop and starves other tasks. + +This module fixes that by owning the wait itself. It starts the connection and +then polls two signals on the underlying ``MCPAdapt`` object: + +* ``ready`` (a ``threading.Event``) -- set once the connection succeeds. +* ``thread`` (a daemon ``threading.Thread``) -- dies when the connect coroutine + raises. + +That yields three clean states: + +============================ ============= ========== ==================== +state thread alive? ready set? action +============================ ============= ========== ==================== +connecting / OAuth login yes no keep waiting +connect failed **no** no **fail immediately** +connected yes yes proceed with tools +============================ ============= ========== ==================== + +So a user taking minutes to complete an interactive OAuth login is still +supported (thread stays alive -> we wait, up to ``connect_timeout``), while a +hard failure aborts within one poll interval instead of after ``connect_timeout``. + +The helper is written generically -- there is nothing GitHub-specific here -- so +any agent in this repo can copy it and use ``mcp_tools_session``. + +Version note +------------ +The fast-fail path pokes at ``mcpadapt`` internals (``MCPAdapt.ready`` / +``MCPAdapt.thread`` / ``MCPAdapt.tools()``), verified against ``crewai-tools`` +1.6.1 (mcpadapt bundled therein). All internal access is guarded with +``getattr``; if the attributes are missing (library changed), we fall back to +the original blocking ``MCPServerAdapter`` behavior so we never crash -- we only +lose the fast-fail optimization, with a warning logged. +""" + +import asyncio +import logging +from contextlib import asynccontextmanager + +from crewai_tools import MCPServerAdapter +from crewai_tools.adapters.tool_collection import ToolCollection + +logger = logging.getLogger(__name__) + + +class McpConnectError(RuntimeError): + """Raised when an MCP connection fails or times out. + + The real underlying cause is chained via ``raise ... from`` when available. + """ + + +@asynccontextmanager +async def mcp_tools_session( + server_params: dict, + *, + connect_timeout: int, + poll_interval: float = 1.0, +): + """Connect to an MCP server, yielding its tools for the duration of the block. + + Fails fast on connection errors while capping a genuine never-returns hang at + ``connect_timeout`` seconds. The MCP session stays open for the body of the + ``async with`` and is torn down on exit (success or failure). + + Args: + server_params: Passed straight to ``MCPServerAdapter`` (e.g. ``url`` / + ``transport`` / ``headers``). + connect_timeout: Last-resort ceiling, in seconds, for a connection that + neither succeeds nor fails (e.g. a stuck interactive OAuth flow). + poll_interval: How often, in seconds, to check the connection state. + Small values surface failures almost instantly at negligible cost. + + Yields: + A ``ToolCollection`` of the adapted MCP tools. + + Raises: + McpConnectError: on connection failure or on hitting ``connect_timeout``. + """ + adapter = _build_adapter(server_params, connect_timeout) + + # Fallback: internals not as expected -> use the library's own blocking + # context manager. Correct, just without the fast-fail optimization. + if adapter is None: + logger.warning( + "mcpadapt internals not recognized; falling back to blocking connect " + "(a failed connection may take up to %ss to surface)", + connect_timeout, + ) + blocking = MCPServerAdapter(server_params, connect_timeout=connect_timeout) + try: + with blocking as tools: + yield tools + finally: + pass # `with` handles teardown + return + + server_adapter, mcp_adapter = adapter + try: + tools = await _connect_with_fast_fail(mcp_adapter, connect_timeout, poll_interval) + # Wrap the raw adapted tools the same way MCPServerAdapter.tools does. + yield ToolCollection(tools) + finally: + # Cancel the keep-alive task, join the daemon thread, close the loop. + try: + server_adapter.stop() + except Exception as exc: # noqa: BLE001 - teardown must not mask the real error + logger.warning("Error during MCP adapter teardown: %s", exc) + + +def _build_adapter(server_params: dict, connect_timeout: int): + """Construct an MCPServerAdapter and start its connection WITHOUT blocking. + + Returns ``(server_adapter, mcp_adapter)`` on success, or ``None`` if the + library internals we depend on are not present (caller should fall back). + + We replicate ``mcpadapt``'s own start sequence -- ``thread.start()`` -- rather + than calling the library's ``start()``, which would block on ``ready.wait``. + """ + try: + # Build the object graph without triggering the blocking start(). We + # can't call MCPServerAdapter(...) directly because its __init__ calls + # start(); instead build the underlying MCPAdapt ourselves, mirroring + # crewai_tools' construction, and attach it to a bare MCPServerAdapter. + from mcpadapt.core import MCPAdapt + from mcpadapt.crewai_adapter import CrewAIAdapter + + server_adapter = MCPServerAdapter.__new__(MCPServerAdapter) + server_adapter._adapter = None + server_adapter._tools = None + server_adapter._tool_names = None + server_adapter._serverparams = server_params + + mcp_adapter = MCPAdapt(server_params, CrewAIAdapter(), connect_timeout) + server_adapter._adapter = mcp_adapter + + # Sanity-check the internals we will poll before committing to this path. + thread = getattr(mcp_adapter, "thread", None) + ready = getattr(mcp_adapter, "ready", None) + if thread is None or ready is None or not hasattr(mcp_adapter, "tools"): + return None + + thread.start() # kick off the connection; do NOT call start()/ __enter__ + return server_adapter, mcp_adapter + except Exception as exc: # noqa: BLE001 - any surprise -> fall back to blocking + logger.warning("Could not set up fast-fail MCP connect (%s); will fall back", exc) + return None + + +async def _connect_with_fast_fail(mcp_adapter, connect_timeout: int, poll_interval: float): + """Poll the connection until ready, dead, or timed out. + + Returns the adapted tool list on success. Raises ``McpConnectError`` on a + dead connect thread (fast path) or on exceeding ``connect_timeout``. + """ + ready = mcp_adapter.ready + thread = mcp_adapter.thread + + waited = 0.0 + while True: + if ready.is_set(): + break + if not thread.is_alive(): + # The connect coroutine raised and killed its thread. mcpadapt + # discards the exception, so we surface a clear error of our own. + raise McpConnectError( + "MCP connection failed (the connection attempt terminated before " + "becoming ready). Check the MCP URL, network reachability, and auth." + ) + if waited >= connect_timeout: + raise McpConnectError(f"MCP connection did not become ready within {connect_timeout}s.") + await asyncio.sleep(poll_interval) + waited += poll_interval + + # Connection is ready; produce the adapted tools. tools() drives the loop in + # the worker thread, so failures here are real and should also surface fast. + try: + return await asyncio.to_thread(mcp_adapter.tools) + except Exception as exc: # noqa: BLE001 - surface the real cause + raise McpConnectError(f"Failed to list MCP tools: {exc}") from exc diff --git a/a2a/git_issue_agent/git_issue_agent/prompts.py b/a2a/git_issue_agent/git_issue_agent/prompts.py index aa3c0239..5994b914 100644 --- a/a2a/git_issue_agent/git_issue_agent/prompts.py +++ b/a2a/git_issue_agent/git_issue_agent/prompts.py @@ -33,6 +33,7 @@ - Use only one tool per call. - Always copy owner/organization names, repository names, and issue numbers exactly as provided by the user or upstream context. Do not truncate, split, normalize, or otherwise modify these identifiers. Example: for "github/github-mcp-server" you must pass `"owner": "github", "repo": "github-mcp-server"`. - When a tool call succeeds, do not immediately call another tool just to reformat or summarize the same results. Instead, work with the observation you have unless additional data is explicitly required. +- When the user asks for a specific count or a narrow scope (e.g. "top 5", "open bugs"), pass an explicit limit/filter (such as `per_page`, `state`, or `labels`) in the Action Input rather than fetching everything. Ask only for what you need. CORRECT EXAMPLE Thought: The user provided owner and repo; list_issues fits. @@ -49,6 +50,7 @@ - Summarize or aggregate long lists instead of echoing raw JSON. Provide counts or grouped highlights when appropriate. - Clearly cite or reference the tool results. - If a tool failed or inputs were missing, say so explicitly. Don't attempt to guess the answer. +- Tool results may be capped for size, so the items you see can be fewer than the total that exist. Use a total field (e.g. `total_count`) as the true total when present, and phrase answers as "showing the top N of M" rather than asserting that the number of returned items is the total. - After the Observation is provided by the system, output: Thought: I now know the final answer diff --git a/a2a/git_issue_agent/git_issue_agent/tool_limits.py b/a2a/git_issue_agent/git_issue_agent/tool_limits.py new file mode 100644 index 00000000..0684feb5 --- /dev/null +++ b/a2a/git_issue_agent/git_issue_agent/tool_limits.py @@ -0,0 +1,157 @@ +"""Bound the size of MCP tool output before it reaches the LLM. + +Why +--- +A broad issue query can return a very large payload (e.g. 81 issues for a "top 5" +request). Feeding that raw into the prompt overflows the model context window, +which then triggers CrewAI's expensive summarize-and-retry fallback. The reliable +fix is to cap what a tool returns *before* it is appended to the prompt, since a +prompt instruction alone cannot guarantee the model asks for less. + +This module wraps a CrewAI ``BaseTool`` so that, transparently to CrewAI: + +* an explicit fetch bound (``per_page``) is injected into the tool arguments when + the tool's schema supports it and the caller did not already ask for less; and +* the returned observation is sliced to at most ``max_items`` list entries and + then hard-capped to ``max_chars`` characters. + +It is written generically -- it keys off the common ``{total_count, items:[...]}`` +/ ``results`` MCP result shape and falls back to plain character truncation for +anything else -- so any agent can reuse it. There is nothing GitHub-specific here. +""" + +import json +import logging +from typing import Any, Callable, Optional + +logger = logging.getLogger(__name__) + +# Common keys used by MCP servers to hold a result list and its true total. +_LIST_KEYS = ("items", "results", "issues", "data") +_TOTAL_KEYS = ("total_count", "total", "count") + +# Common parameter names a tool schema may expose to bound page size. +_LIMIT_PARAM_CANDIDATES = ("per_page", "perPage", "limit", "page_size", "pageSize") + +_TRUNCATION_MARKER = "\nā€¦ [output truncated to fit the model context window]" + + +def wrap_tool_output( + tool, + *, + max_items: int, + max_chars: int, + on_truncate: Optional[Callable[[str, int, int], None]] = None, +): + """Return ``tool`` with its ``_run`` wrapped to bound argument and result size. + + Args: + tool: A CrewAI ``BaseTool`` (both ``run()`` and ``to_structured_tool()`` + route through ``_run``, so wrapping ``_run`` covers every call path). + max_items: Max list entries kept from a structured result. + max_chars: Hard character budget for the returned observation string. + on_truncate: Optional callback ``(tool_name, kept, total)`` invoked when a + result is trimmed, so the caller can surface a user-facing note. + + Returns: + The same tool instance, mutated in place. Returns it unchanged (with a + warning) if the expected ``_run`` attribute is absent. + """ + original_run = getattr(tool, "_run", None) + if not callable(original_run): + logger.warning("Tool %r has no callable _run; leaving unbounded", getattr(tool, "name", tool)) + return tool + + limit_param = _find_limit_param(tool, max_items) + + def bounded_run(*args: Any, **kwargs: Any) -> Any: + # 1. Inject a fetch bound when the schema supports it and the model did + # not already request an equal-or-smaller page size. + if limit_param is not None: + existing = kwargs.get(limit_param) + if not isinstance(existing, int) or existing > max_items: + kwargs[limit_param] = max_items + + result = original_run(*args, **kwargs) + + # 2. Cap the observation before it becomes a prompt message. + return _cap_result( + result, + tool_name=getattr(tool, "name", "tool"), + max_items=max_items, + max_chars=max_chars, + on_truncate=on_truncate, + ) + + # BaseTool is a pydantic model; bypass validation to attach the bound method. + object.__setattr__(tool, "_run", bounded_run) + return tool + + +def _find_limit_param(tool, max_items: int) -> Optional[str]: + """Return a page-size parameter name from the tool's schema, if any.""" + schema = getattr(tool, "args_schema", None) + fields = getattr(schema, "model_fields", None) if schema is not None else None + if not fields: + return None + for candidate in _LIMIT_PARAM_CANDIDATES: + if candidate in fields: + return candidate + return None + + +def _cap_result( + result: Any, + *, + tool_name: str, + max_items: int, + max_chars: int, + on_truncate: Optional[Callable[[str, int, int], None]], +) -> Any: + """Slice a structured result to max_items, then hard-cap to max_chars.""" + text = result if isinstance(result, str) else str(result) + + # Try to interpret as JSON so we can slice the list intelligently. + parsed = None + try: + parsed = json.loads(text) + except (ValueError, TypeError): + parsed = None + + if isinstance(parsed, dict): + list_key = next((k for k in _LIST_KEYS if isinstance(parsed.get(k), list)), None) + if list_key is not None: + items = parsed[list_key] + total = _read_total(parsed, default=len(items)) + if len(items) > max_items: + parsed[list_key] = items[:max_items] + _notify(on_truncate, tool_name, max_items, total) + text = json.dumps(parsed) + elif isinstance(parsed, list) and len(parsed) > max_items: + total = len(parsed) + text = json.dumps(parsed[:max_items]) + _notify(on_truncate, tool_name, max_items, total) + + # Final hard character budget, regardless of structure. + if len(text) > max_chars: + _notify(on_truncate, tool_name, -1, -1) + text = text[:max_chars] + _TRUNCATION_MARKER + + return text + + +def _read_total(parsed: dict, default: int) -> int: + for key in _TOTAL_KEYS: + value = parsed.get(key) + if isinstance(value, int): + return value + return default + + +def _notify(cb: Optional[Callable[[str, int, int], None]], tool_name: str, kept: int, total: int) -> None: + if cb is None: + return + try: + cb(tool_name, kept, total) + except Exception as exc: # noqa: BLE001 - a note callback must never break tool execution + logger.warning("Truncation callback failed: %s", exc)