Merge pull request #661 from davidhadas/wrap_mcp_init

fix mcp never failing till timeout bug

Author: pdettori

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():

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,
         )

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")

a2a/git_issue_agent/git_issue_agent/main.py

@@ -2,18 +2,24 @@
 import logging
 import re
 import sys
+import threading
 
 from crewai_tools.adapters.tool_collection import ToolCollection
 
 from git_issue_agent.agents import GitAgents
 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