feat: Fix A2A push-config race in eval client (flaky test_set_push_config_round_trips)

CHANGELOG.md

@@ -22,6 +22,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   The original six curated keys ride along unchanged, so older plugin-kits keep working.
 
 ### Fixed
+- **The eval client no longer races push-config registration (#2230).** A task started
+  over ``SendStreamingMessage`` is live slightly before its store write lands, so an
+  immediate ``CreateTaskPushNotificationConfig`` could bounce with TASK_NOT_FOUND —
+  the source of the flaky ``test_set_push_config_round_trips`` (and its event-loop
+  teardown noise). The push-config helpers now retry exactly that error with a short
+  bounded backoff (5 attempts, ~0.75s worst case); any other JSON-RPC error, or a
+  genuinely unknown task id, still fails immediately/after the bound.
 - **Six status tones now actually theme (#2224).** Chat notes, the keybindings conflict
   state, and the knowledge delete-armed state referenced `--pl-color-info/warning/danger`
   — names the design package never defines — so their hex fallbacks rendered permanently:

evals/client.py

@@ -70,6 +70,13 @@ def _resolve_auth_env() -> tuple[str, str]:
     return bearer, api_key
 
 
+# Bounded backoff for the push-config helpers' TASK_NOT_FOUND retry
+# (``AgentClient._rpc_until_task_visible``). Delays double per attempt:
+# 0.05 → 0.1 → 0.2 → 0.4, ~0.75s worst case before the error propagates.
+_TASK_VISIBLE_ATTEMPTS = 5
+_TASK_VISIBLE_BACKOFF_S = 0.05
+
+
 class AgentClient:
     """Thin A2A client tied to one agent instance."""
 
@@ -416,6 +423,26 @@ async def _rpc(self, method: str, params: dict, *, timeout_s: int = 10) -> dict:
             raise A2ARpcError(method, resp["error"])
         return resp.get("result") or {}
 
+    async def _rpc_until_task_visible(self, method: str, params: dict, *, timeout_s: int = 10) -> dict:
+        """``_rpc``, absorbing the task-registration race (#2230).
+
+        A task started via ``SendStreamingMessage`` is live (streaming frames,
+        active-task registry) slightly before its store write lands, so an
+        immediate task-scoped call can get ``TASK_NOT_FOUND`` for a task that
+        provably exists. Retry exactly that error with a short doubling
+        backoff, bounded at ``_TASK_VISIBLE_ATTEMPTS`` attempts — a genuinely
+        wrong task id still fails, just ~0.75s later. Any other JSON-RPC error
+        propagates immediately.
+        """
+        for attempt in range(_TASK_VISIBLE_ATTEMPTS):
+            try:
+                return await self._rpc(method, params, timeout_s=timeout_s)
+            except A2ARpcError as e:
+                if not e.is_task_not_found or attempt == _TASK_VISIBLE_ATTEMPTS - 1:
+                    raise
+            await asyncio.sleep(_TASK_VISIBLE_BACKOFF_S * 2**attempt)
+        raise AssertionError("unreachable: the last attempt returns or raises")
+
     # ── push notification config ────────────────────────────────────────────
 
     async def set_push_config(self, task_id: str, url: str, token: str | None = None) -> dict:
@@ -430,40 +457,51 @@ async def set_push_config(self, task_id: str, url: str, token: str | None = None
         The task must already exist server-side (the handler looks it up and
         raises TaskNotFound otherwise), so call this on an in-flight task —
         or pass the config inline on ``SendMessage`` via
-        ``configuration.taskPushNotificationConfig``.
+        ``configuration.taskPushNotificationConfig``. A just-started task may
+        not have hit the store yet, so TASK_NOT_FOUND is retried briefly
+        (``_rpc_until_task_visible``) instead of failing the first call.
 
         When ``token`` is set the agent echoes it back as the
         ``X-A2A-Notification-Token`` header on every delivery.
         """
         params: dict = {"taskId": task_id, "url": url}
         if token:
             params["token"] = token
-        return await self._rpc("CreateTaskPushNotificationConfig", params)
+        return await self._rpc_until_task_visible("CreateTaskPushNotificationConfig", params)
 
     async def get_push_config(self, task_id: str, config_id: str | None = None) -> dict:
         """``GetTaskPushNotificationConfig`` → the stored config.
 
         ``config_id`` defaults to the task id — that is what the store assigns
         when a config is created without an explicit ``id``."""
-        return await self._rpc(
+        return await self._rpc_until_task_visible(
             "GetTaskPushNotificationConfig",
             {"taskId": task_id, "id": config_id or task_id},
         )
 
     async def list_push_configs(self, task_id: str) -> list[dict]:
         """``ListTaskPushNotificationConfigs`` → the task's configs."""
-        res = await self._rpc("ListTaskPushNotificationConfigs", {"taskId": task_id})
+        res = await self._rpc_until_task_visible("ListTaskPushNotificationConfigs", {"taskId": task_id})
         return res.get("configs") or []
 
 
 class A2ARpcError(RuntimeError):
     """A JSON-RPC error frame returned by the agent."""
 
+    #: The A2A spec code for ``TaskNotFoundError`` (a2a-sdk ``utils/errors.py``).
+    TASK_NOT_FOUND_CODE = -32001
+
     def __init__(self, method: str, error: dict):
         self.method = method
         self.error = error
         super().__init__(f"{method}: {error}")
 
+    @property
+    def is_task_not_found(self) -> bool:
+        """True when the frame is the server's TaskNotFound (``-32001``) —
+        the one error ``_rpc_until_task_visible`` treats as transient."""
+        return self.error.get("code") == self.TASK_NOT_FOUND_CODE
+
 
 def _norm_state(state: str) -> str:
     """``TASK_STATE_COMPLETED`` → ``completed``; pass through already-lowercased

tests/test_eval_client_subscribe_push.py

@@ -20,6 +20,12 @@
   ``ParseDict``s params strictly, sending one is a ``-32602``
   (``test_v0_3_nested_push_config_shape_is_rejected``).
 
+The ``_scripted_push_app`` tests pin the client's TASK_NOT_FOUND retry
+(``_rpc_until_task_visible``): a just-started task is live in the active
+registry before its store write lands, so the first push-config call can race
+it (#2230). The client retries exactly that error with a bounded backoff; the
+scripted server makes the race deterministic instead of timing-dependent.
+
 The last test is a genuine end-to-end delivery: a real ``BasePushNotificationSender``
 POSTs over real TCP to ``evals/webhook.py``'s listener, so ``capture.received``
 is proof the agent actually called the webhook.
@@ -41,7 +47,7 @@
     InMemoryTaskStore,
 )
 from a2a.types import AgentSkill
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 
 import protolabs_a2a as pa
 import evals.client as ec
@@ -215,6 +221,92 @@ async def test_set_push_config_round_trips(monkeypatch):
         await asyncio.wait_for(streaming, 10)
 
 
+def _scripted_push_app(fail_first_n: int) -> tuple[FastAPI, list[str]]:
+    """A scripted ``/a2a`` endpoint for pinning the client's retry behavior.
+
+    Answers the first ``fail_first_n`` calls with the exact TaskNotFound frame
+    a2a-sdk emits (code ``-32001``) and succeeds after — the deterministic
+    stand-in for the registration race, where a just-started task is live in
+    the active registry before its store write lands (#2230). Returns
+    ``(app, calls)``; ``calls`` logs each JSON-RPC method received.
+    """
+    calls: list[str] = []
+    app = FastAPI()
+
+    @app.post("/a2a")
+    async def a2a(request: Request) -> dict:
+        body = await request.json()
+        calls.append(body["method"])
+        if len(calls) <= fail_first_n:
+            return {
+                "jsonrpc": "2.0",
+                "id": body["id"],
+                "error": {"code": -32001, "message": "Task not found"},
+            }
+        return {
+            "jsonrpc": "2.0",
+            "id": body["id"],
+            "result": {"taskId": body["params"]["taskId"], "url": body["params"].get("url", "")},
+        }
+
+    return app, calls
+
+
+@pytest.mark.asyncio
+async def test_set_push_config_retries_once_when_task_is_not_yet_visible(monkeypatch):
+    """Deterministic reproduction of the #2230 flake: the server says
+    TASK_NOT_FOUND once (task not in the store yet), then succeeds. The client
+    absorbs it — exactly one retry, then the created config comes back."""
+    app, calls = _scripted_push_app(fail_first_n=1)
+    client = _route(monkeypatch, app)
+
+    created = await client.set_push_config("task-1", "https://example.test/hook")
+
+    assert created["taskId"] == "task-1"
+    assert calls == ["CreateTaskPushNotificationConfig"] * 2  # one retry, no more
+
+
+@pytest.mark.asyncio
+async def test_set_push_config_gives_up_after_bounded_retries(monkeypatch):
+    """A genuinely unknown task still fails: TASK_NOT_FOUND on every attempt
+    exhausts the bounded backoff and the error propagates."""
+    app, calls = _scripted_push_app(fail_first_n=10_000)  # never succeeds
+    client = _route(monkeypatch, app)
+    monkeypatch.setattr(ec, "_TASK_VISIBLE_BACKOFF_S", 0.0)  # keep the exhaustion instant
+
+    with pytest.raises(ec.A2ARpcError) as excinfo:
+        await client.set_push_config("no-such-task", "https://example.test/hook")
+
+    assert excinfo.value.is_task_not_found
+    assert len(calls) == ec._TASK_VISIBLE_ATTEMPTS  # bounded, not infinite
+
+
+@pytest.mark.asyncio
+async def test_set_push_config_does_not_retry_other_rpc_errors(monkeypatch):
+    """Only TaskNotFound is transient — a ``-32602`` (wrong param shape, e.g.
+    the 0.3 nested config) raises on the first call, no retry."""
+    calls: list[str] = []
+    app = FastAPI()
+
+    @app.post("/a2a")
+    async def a2a(request: Request) -> dict:
+        body = await request.json()
+        calls.append(body["method"])
+        return {
+            "jsonrpc": "2.0",
+            "id": body["id"],
+            "error": {"code": -32602, "message": "Invalid parameters"},
+        }
+
+    client = _route(monkeypatch, app)
+
+    with pytest.raises(ec.A2ARpcError) as excinfo:
+        await client.set_push_config("task-1", "https://example.test/hook")
+
+    assert not excinfo.value.is_task_not_found
+    assert len(calls) == 1
+
+
 @pytest.mark.asyncio
 async def test_v0_3_nested_push_config_shape_is_rejected(monkeypatch):
     """Pins why ``set_push_config`` sends flat params: the 0.3 wrapper