fix(a2a): give up on an agent whose sandbox died, and cancel a prompt that timed out

packages/agentenv-protocol/README.md

@@ -301,6 +301,11 @@ A successful `TaskResult` must contain at least one text, file, or data part;
 the framework rejects empty successes rather than emitting an ungradeable task.
 The SDK does not retry tasks.
 
+A client that gives up on a task sends `tasks/cancel`. The framework marks the
+task canceled and cancels `run()`, which gets `asyncio.CancelledError` at its
+next `await`. A process `run()` started keeps running unless `run()` stops it,
+so kill it before re-raising.
+
 `enable(..., description="...")` is reserved for declarations carrying
 configuration or metadata. It preserves the agent-specific extension prose
 published in the Agent Card. The versioned SDK definition provides a generic

packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py

@@ -1226,12 +1226,12 @@ def _validate_status_event(cls, item: Any, request: TaskRequest[Any]) -> None:
                 "streaming TaskStatusUpdateEvent must be non-terminal and final=False"
             )
 
-    async def cancel(self, _context: Any, _event_queue: Any) -> None:
-        """Satisfy the upstream executor interface without offering cancellation."""
-        from a2a.types import UnsupportedOperationError
-        from a2a.utils.errors import ServerError
+    async def cancel(self, context: Any, event_queue: Any) -> None:
+        """Mark the task canceled. The SDK then cancels the coroutine running it, so ``agent.run`` gets
+        ``CancelledError`` at its next ``await`` and should stop whatever it started, such as a CLI process."""
+        from a2a.server.tasks import TaskUpdater
 
-        raise ServerError(error=UnsupportedOperationError())
+        await TaskUpdater(event_queue, context.task_id, context.context_id).cancel()
 
 
 def _from_a2a_part(part: Any) -> Any:

packages/agentenv-protocol/tests/test_a2a_agent.py

@@ -2908,3 +2908,76 @@ async def shadow(self):
 
     with pytest.raises(ValueError, match=rf"conflicts with {owner} route"):
         Agent().create_app()
+
+
+def _send_in_context(client: TestClient, text: str, message_id: str) -> str:
+    response = client.post(
+        "/a2a",
+        json={
+            "jsonrpc": "2.0",
+            "id": message_id,
+            "method": "message/send",
+            "params": {
+                "message": {
+                    "kind": "message",
+                    "messageId": message_id,
+                    "role": "user",
+                    "contextId": "context-1",
+                    "parts": [{"kind": "text", "text": text}],
+                },
+                "configuration": {"blocking": False},
+            },
+        },
+    )
+    return response.json()["result"]["id"]
+
+
+def _task_rpc(client: TestClient, method: str, task_id: str) -> dict[str, Any]:
+    return client.post(
+        "/a2a",
+        json={"jsonrpc": "2.0", "id": method, "method": method, "params": {"id": task_id}},
+    ).json()
+
+
+def _within(seconds: float, condition) -> bool:
+    deadline = time.monotonic() + seconds
+    while not condition():
+        if time.monotonic() > deadline:
+            return False
+        time.sleep(0.01)
+    return True
+
+
+def test_cancel_stops_a_running_task_and_frees_its_context() -> None:
+    ran: list[str] = []
+    cancelled: list[str] = []
+
+    @a2a_agent(identity=AgentIdentity(name="slow", description="test", version="1"))
+    class Agent(AgentEnvAgent):
+        async def run(self, request: TaskRequest) -> TaskResult:
+            text = request.parts[0].text
+            ran.append(text)
+            if text == "hang":
+                try:
+                    await asyncio.Event().wait()
+                except asyncio.CancelledError:
+                    cancelled.append(text)
+                    raise
+            return TaskResult.text(f"did {text}")
+
+    def state(client: TestClient, task_id: str) -> str:
+        return _task_rpc(client, "tasks/get", task_id)["result"]["status"]["state"]
+
+    with TestClient(Agent().create_app()) as client:
+        hung = _send_in_context(client, "hang", "message-1")
+        assert _within(5, lambda: ran == ["hang"])
+
+        canceled = _task_rpc(client, "tasks/cancel", hung)
+        assert canceled["result"]["status"]["state"] == "canceled"
+
+        following = _send_in_context(client, "next", "message-2")
+        assert _within(5, lambda: state(client, following) == "completed")
+        assert state(client, hung) == "canceled"
+
+    assert cancelled == ["hang"]
+    assert ran == ["hang", "next"]

src/agent_env/a2a_agent/protocol.py

@@ -18,6 +18,19 @@
     TaskState.completed, TaskState.failed, TaskState.canceled, TaskState.rejected,
 })
 
+# How long an agent on a sandbox may go without answering a poll before ``poll_a2a_task`` gives up on it, over at
+# least _UNREACHABLE_MIN_POLLS polls in a row: a sandbox that died answers none, a tunnel blip drops one or two.
+UNREACHABLE_AFTER_SECONDS = 60
+_UNREACHABLE_MIN_POLLS = 3
+# What a proxy in front of the agent answers when it can't reach it; any other status is the agent answering.
+_GATEWAY_STATUSES = frozenset({502, 503, 504})
+_CANCEL_TIMEOUT_SECONDS = 10
+
+
+class AgentUnreachableError(TimeoutError):
+    """The agent stopped answering, as one whose sandbox died does. A ``TimeoutError``, so whatever handles an
+    agent running out of time handles this the same way, only sooner."""
+
 
 _ERROR_CODE = re.compile(r"[a-z][a-z0-9_]{0,63}")
 
@@ -134,10 +147,18 @@ async def poll_a2a_task(
     task_id: str,
     timeout_seconds: int,
     poll_interval_seconds: int = 10,
+    *,
+    sandbox_id: Optional[str] = None,
 ) -> dict:
-    """POST /a2a tasks/get until status.state is 'completed' or 'failed'. Returns the result dict."""
+    """POST /a2a tasks/get until the task reaches a terminal state, and return it.
+
+    Out of time, it raises ``TimeoutError``. ``sandbox_id`` names the sandbox the agent runs on: then it raises
+    ``AgentUnreachableError``, naming it, once the agent has gone UNREACHABLE_AFTER_SECONDS without answering,
+    rather than waiting out the timeout. Either way it first asks the agent to cancel the task, so an agent still
+    working on it stops."""
     deadline = time.monotonic() + timeout_seconds
     consecutive_failures = 0
+    answered_at, unanswered, last_failure = time.monotonic(), 0, ""
     while time.monotonic() < deadline:
         backoff = min(poll_interval_seconds + consecutive_failures * 5, 60)
         await asyncio.sleep(backoff)
@@ -149,30 +170,60 @@ async def poll_a2a_task(
                 }, timeout=30)
             resp.raise_for_status()
             data = resp.json()
-            if "error" in data:
-                consecutive_failures += 1
-                logger.warning(f"A2A poll returned error (consec={consecutive_failures}, will retry): {data['error']}")
-                continue
-            result = data["result"]
         except httpx.HTTPStatusError as e:
             if 400 <= e.response.status_code < 500:
                 raise
             consecutive_failures += 1
             logger.warning(f"A2A poll got {e.response.status_code} (consec={consecutive_failures}, will retry)")
-            continue
+            if e.response.status_code in _GATEWAY_STATUSES:
+                unanswered, last_failure = unanswered + 1, f"HTTP {e.response.status_code}"
+            else:
+                answered_at, unanswered = time.monotonic(), 0
         except httpx.HTTPError as e:
             consecutive_failures += 1
             logger.warning(f"A2A poll failed (consec={consecutive_failures}, will retry): {type(e).__name__}: {e}")
+            unanswered, last_failure = unanswered + 1, f"{type(e).__name__}: {e}" if str(e) else type(e).__name__
+        else:
+            answered_at, unanswered = time.monotonic(), 0
+            if "error" in data:
+                consecutive_failures += 1
+                logger.warning(f"A2A poll returned error (consec={consecutive_failures}, will retry): {data['error']}")
+                continue
+            consecutive_failures = 0
+            if data["result"]["status"]["state"] in _TERMINAL_TASK_STATES:
+                return data["result"]
             continue
+        silent = time.monotonic() - answered_at
+        if sandbox_id is not None and unanswered >= _UNREACHABLE_MIN_POLLS and silent >= UNREACHABLE_AFTER_SECONDS:
+            await cancel_a2a_task(a2a_url, task_id)
+            raise AgentUnreachableError(
+                f"The agent on sandbox {sandbox_id} stopped answering: {unanswered} polls in a row over "
+                f"{silent:.0f}s got no reply (last: {last_failure})")
 
-        consecutive_failures = 0
-        state = result["status"]["state"]
-        if state in _TERMINAL_TASK_STATES:
-            return result
-
+    await cancel_a2a_task(a2a_url, task_id)
     raise TimeoutError(f"A2A task {task_id} did not complete within {timeout_seconds}s")
 
 
+async def cancel_a2a_task(a2a_url: str, task_id: str) -> None:
+    """POST /a2a tasks/cancel, so the agent stops working on ``task_id``. Best effort: an agent that can't be
+    reached, or won't cancel, is only logged."""
+    try:
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(f"{a2a_url}/a2a", json={
+                "jsonrpc": "2.0", "id": "cancel", "method": "tasks/cancel",
+                "params": {"id": task_id},
+            }, timeout=_CANCEL_TIMEOUT_SECONDS)
+        resp.raise_for_status()
+        error = resp.json().get("error")
+    except (httpx.HTTPError, ValueError) as e:
+        logger.warning(f"Couldn't cancel A2A task {task_id}: {type(e).__name__}: {e}")
+        return
+    if error:
+        logger.warning(f"The agent didn't cancel A2A task {task_id}: {error}")
+    else:
+        logger.info(f"Canceled A2A task {task_id}")
+
+
 @dataclass(frozen=True)
 class TerminalResponse:
     """What's pulled from an A2A terminal status message: the text reply, the typed

src/agent_env/task_step/task_steps/prompt_agent.py

@@ -508,6 +508,7 @@ async def _execute_conversation(
                 )
                 result = await protocol.poll_a2a_task(
                     target_url, sent_task_id, self.timeout_seconds, self.poll_interval_seconds,
+                    sandbox_id=agent.sandbox_id,
                 )
             target_state = result["status"]["state"]
             status_msg = (result.get("status") or {}).get("message") or {}
@@ -579,9 +580,10 @@ async def _execute_conversation(
                 user_result = await protocol.poll_a2a_task(
                     user_url, sent_user_task_id,
                     self.user_agent_timeout_seconds, self.poll_interval_seconds,
+                    sandbox_id=user_sim.sandbox_id if is_user_sim else None,
                 )
-            except TimeoutError:
-                logger.warning(f"user_a2a_url timeout for conversation {conversation_id}; marking abandoned")
+            except TimeoutError as e:
+                logger.warning(f"user_a2a_url timeout for conversation {conversation_id} ({e}); marking abandoned")
                 conversation_store.mark_closed(conversation_id)
                 break
 

src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py

@@ -1022,6 +1022,7 @@ async def _invoke_judge_a2a(
             task_id,
             timeout_seconds=self.judge_timeout_seconds,
             poll_interval_seconds=self.DEFAULT_POLL_INTERVAL_SECONDS,
+            sandbox_id=getattr(judge_agent, "sandbox_id", None),
         )
         state = result["status"]["state"]
         status_msg = (result.get("status") or {}).get("message") or {}