fix: synchronize gateway concurrency tests causally

Wait for browser-controller protocol readiness, drive live-owner receipts from durable admission, and begin cleanup observation only when teardown starts. This removes CI scheduling races without extending production deadlines.

tests/gateway/test_browser_control_api.py

@@ -91,6 +91,33 @@ def _registration_body(**overrides):
     return payload
 
 
+async def _wait_for_controller(ws, nonce: str) -> None:
+    """Complete a protocol round trip before dispatching through the broker."""
+    await ws.send_json(
+        {
+            "method": "browser.controller.heartbeat",
+            "params": {"nonce": nonce},
+        }
+    )
+    assert await ws.receive_json() == {
+        "method": "browser.controller.heartbeat",
+        "params": {"nonce": nonce, "ok": True},
+    }
+
+
+async def _receive_command(ws, pending: asyncio.Task) -> dict:
+    """Wait for either the controller frame or an early dispatch failure."""
+    receive = asyncio.create_task(ws.receive_json())
+    done, _ = await asyncio.wait((receive, pending), return_when=asyncio.FIRST_COMPLETED)
+    if receive in done:
+        return receive.result()
+    receive.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await receive
+    await pending
+    raise AssertionError("dispatch completed without sending a controller command")
+
+
 @pytest.mark.asyncio
 async def test_registration_grants_only_the_exact_real_action_allowlist(monkeypatch):
     adapter = _adapter()
@@ -404,17 +431,7 @@ async def test_local_api_ticket_ws_noop_round_trip_filters_spoofed_identity_and_
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(registration["ticket"])],
         )
-        await ws.send_json(
-            {
-                "method": "browser.controller.heartbeat",
-                "params": {"nonce": "heartbeat-api-fixture"},
-            }
-        )
-        heartbeat = await ws.receive_json(timeout=2.0)
-        assert heartbeat == {
-            "method": "browser.controller.heartbeat",
-            "params": {"nonce": "heartbeat-api-fixture", "ok": True},
-        }
+        await _wait_for_controller(ws, "heartbeat-api-fixture")
         scope = ControllerScope(
             principal_id=registration["scope"]["principal_id"],
             profile_id=registration["scope"]["profile_id"],
@@ -434,7 +451,7 @@ async def test_local_api_ticket_ws_noop_round_trip_filters_spoofed_identity_and_
                 tool_call_id="tool-call-fixture",
             )
         )
-        command = await ws.receive_json(timeout=2.0)
+        command = await _receive_command(ws, pending)
         assert command["method"] == "browser.controller.command"
         assert command["params"]["action"] == "controller.noop"
         await ws.send_json(
@@ -447,7 +464,7 @@ async def test_local_api_ticket_ws_noop_round_trip_filters_spoofed_identity_and_
                 },
             }
         )
-        assert await asyncio.wait_for(pending, timeout=2.0) == {"echo": "local-api"}
+        assert await pending == {"echo": "local-api"}
 
         rejected = asyncio.create_task(
             asyncio.to_thread(
@@ -458,7 +475,7 @@ async def test_local_api_ticket_ws_noop_round_trip_filters_spoofed_identity_and_
                 tool_call_id="tool-call-rejected",
             )
         )
-        rejected_command = await ws.receive_json(timeout=2.0)
+        rejected_command = await _receive_command(ws, rejected)
         await ws.send_json(
             {
                 "method": "browser.controller.result",
@@ -470,7 +487,7 @@ async def test_local_api_ticket_ws_noop_round_trip_filters_spoofed_identity_and_
             }
         )
         with pytest.raises(ControllerRejected, match="controller_rejected"):
-            await asyncio.wait_for(rejected, timeout=2.0)
+            await rejected
         await ws.close()
 
         with pytest.raises(WSServerHandshakeError) as replay:
@@ -497,6 +514,7 @@ async def test_real_browser_action_routes_through_controller_without_legacy_fall
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(registration["ticket"])],
         )
+        await _wait_for_controller(ws, "real-action-ready")
 
         legacy_calls = []
         pending = asyncio.create_task(
@@ -513,7 +531,7 @@ async def test_real_browser_action_routes_through_controller_without_legacy_fall
                 tool_call_id="tool-call-real-action",
             )
         )
-        command = await ws.receive_json(timeout=2.0)
+        command = await _receive_command(ws, pending)
         assert command["method"] == "browser.controller.command"
         assert command["params"]["action"] == "browser_snapshot"
         assert command["params"]["arguments"] == {"include": "accessibility"}
@@ -532,7 +550,7 @@ async def test_real_browser_action_routes_through_controller_without_legacy_fall
             }
         )
 
-        assert await asyncio.wait_for(pending, timeout=2.0) == (
+        assert await pending == (
             '{"title": "Example Domain", "url": "https://example.test/", "refs": []}'
         )
         assert legacy_calls == []
@@ -554,6 +572,7 @@ async def test_local_api_same_identity_reconnect_completes_command_started_on_ol
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(first["ticket"])],
         )
+        await _wait_for_controller(first_ws, "first-reconnect-generation-ready")
 
         pending = asyncio.create_task(
             asyncio.to_thread(
@@ -569,7 +588,7 @@ async def test_local_api_same_identity_reconnect_completes_command_started_on_ol
                 tool_call_id="tool-call-reconnect",
             )
         )
-        command = await first_ws.receive_json(timeout=10.0)
+        command = await _receive_command(first_ws, pending)
         await first_ws.close()
         await asyncio.sleep(0)
         assert not pending.done()
@@ -584,6 +603,7 @@ async def test_local_api_same_identity_reconnect_completes_command_started_on_ol
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(second["ticket"])],
         )
+        await _wait_for_controller(second_ws, "second-reconnect-generation-ready")
         await second_ws.send_json(
             {
                 "method": "browser.controller.result",
@@ -594,7 +614,7 @@ async def test_local_api_same_identity_reconnect_completes_command_started_on_ol
                 },
             }
         )
-        assert await asyncio.wait_for(pending, timeout=10.0) == '{"reconnected": true}'
+        assert await pending == '{"reconnected": true}'
         await second_ws.close()
 
 
@@ -613,6 +633,7 @@ async def test_local_api_explicit_detach_is_hard_and_stale_socket_cannot_detach_
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(first["ticket"])],
         )
+        await _wait_for_controller(first_ws, "stale-detach-generation-ready")
         second_response = await client.post(
             "/v1/browser-control/register",
             json=_registration_body(capabilities=["controller.noop"]),
@@ -623,12 +644,11 @@ async def test_local_api_explicit_detach_is_hard_and_stale_socket_cannot_detach_
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(second["ticket"])],
         )
+        await _wait_for_controller(second_ws, "detach-generation-ready")
 
         await first_ws.send_json(
             {"method": "browser.controller.detach", "params": {}}
         )
-        with pytest.raises(asyncio.TimeoutError):
-            await first_ws.receive_json(timeout=0.05)
 
         pending = asyncio.create_task(
             asyncio.to_thread(
@@ -646,17 +666,17 @@ async def test_local_api_explicit_detach_is_hard_and_stale_socket_cannot_detach_
                 tool_call_id="tool-call-explicit-detach",
             )
         )
-        command = await second_ws.receive_json(timeout=2.0)
+        command = await _receive_command(second_ws, pending)
         await second_ws.send_json(
             {"method": "browser.controller.detach", "params": {}}
         )
-        detached = await second_ws.receive_json(timeout=2.0)
+        detached = await second_ws.receive_json()
         assert detached == {
             "method": "browser.controller.detach",
             "params": {"ok": True},
         }
         with pytest.raises(ControllerCancelled):
-            await asyncio.wait_for(pending, timeout=2.0)
+            await pending
         assert command["method"] == "browser.controller.command"
         await first_ws.close()
         await second_ws.close()
@@ -685,6 +705,7 @@ async def test_remote_api_uses_the_same_authenticated_noop_round_trip(monkeypatc
             "/v1/browser-control/ws",
             protocols=[CONTROL_PROTOCOL, _ticket_protocol(registration["ticket"])],
         )
+        await _wait_for_controller(ws, "remote-ready")
         scope = ControllerScope(
             principal_id=registration["scope"]["principal_id"],
             profile_id=registration["scope"]["profile_id"],
@@ -703,7 +724,7 @@ async def test_remote_api_uses_the_same_authenticated_noop_round_trip(monkeypatc
                 tool_call_id="tool-call-remote",
             )
         )
-        command = await ws.receive_json(timeout=2.0)
+        command = await _receive_command(ws, pending)
         await ws.send_json(
             {
                 "method": "browser.controller.result",
@@ -714,7 +735,5 @@ async def test_remote_api_uses_the_same_authenticated_noop_round_trip(monkeypatc
                 },
             }
         )
-        assert await asyncio.wait_for(pending, timeout=2.0) == {
-            "family": "remote-api"
-        }
+        assert await pending == {"family": "remote-api"}
         await ws.close()

tests/gateway/test_compress_command.py

@@ -483,16 +483,10 @@ async def test_compress_command_cleanup_does_not_block_event_loop():
     ]
     runner = _make_runner(history)
 
-    close_started = threading.Event()
     release_close = threading.Event()
 
-    def slow_close():
-        close_started.set()
-        release_close.wait(timeout=5)
-
     agent_instance = MagicMock()
     agent_instance.shutdown_memory_provider = MagicMock()
-    agent_instance.close = slow_close
     agent_instance._cached_system_prompt = ""
     agent_instance.tools = None
     agent_instance.context_compressor.has_content_to_compress.return_value = True
@@ -516,37 +510,46 @@ async def test_compress_command_cleanup_does_not_block_event_loop():
             ticks["n"] += 1
             await asyncio.sleep(0.005)
 
-    def _observer():
-        # threading.Event wait does not need the event loop. Sample ticks
-        # while close() is still held so an on-loop teardown is visible.
-        if not close_started.wait(timeout=5):
-            observed["error"] = "close() never started"
-            release_close.set()
-            return
+    def slow_close():
+        observed["close_started"] = True
         baseline = ticks["n"]
-        time.sleep(0.12)
-        observed["ticks_during_block"] = ticks["n"] - baseline
-        release_close.set()
 
+        def _observer():
+            # Start observation from the cleanup call itself. This excludes all
+            # unrelated setup/import time before teardown begins.
+            time.sleep(0.12)
+            observed["ticks_during_block"] = ticks["n"] - baseline
+            release_close.set()
+
+        observer = threading.Thread(target=_observer, name="compress-cleanup-observer", daemon=True)
+        observer.start()
+        release_close.wait()
+        observer.join()
+
+    agent_instance.close = slow_close
     hb = asyncio.create_task(_heartbeat())
-    observer = threading.Thread(target=_observer, name="compress-cleanup-observer", daemon=True)
-    observer.start()
 
     with (
-        patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "***"}),
-        patch("gateway.run._resolve_gateway_model", return_value="test-model"),
-        patch("run_agent.AIAgent", return_value=agent_instance),
+        patch.object(
+            runner,
+            "_resolve_session_agent_runtime",
+            return_value=("test-model", {"api_key": "***"}),
+        ),
+        patch.object(
+            runner,
+            "_build_manual_compression_agent",
+            AsyncMock(return_value=agent_instance),
+        ),
         patch("agent.model_metadata.estimate_request_tokens_rough", return_value=100),
     ):
         result = await runner._handle_compress_command(_make_event())
 
-    observer.join(timeout=5)
     stop.set()
     await hb
     runner._shutdown_executor()
 
     assert "Compressed:" in result
-    assert "error" not in observed, observed.get("error")
+    assert observed.get("close_started") is True
     assert observed.get("ticks_during_block", 0) >= 5, (
         "event loop was blocked during manual /compress cleanup: only "
         f"{observed.get('ticks_during_block')} ticks while agent.close() was running"

tests/gateway/test_peer_dm_into_open_bot_chat.py

@@ -10,8 +10,8 @@ from __future__ import annotations
 
 import asyncio
 import json
+import queue
 import threading
-import time
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch
 
@@ -34,20 +34,52 @@ def _app(adapter):
     return app
 
 
-def _owner_answers(home, reply):
-    """What the Desktop's live session does: claim the delivery, run it as its next turn, settle it."""
+def _owner_settles(monkeypatch, home, *, status, reply="", error="", reason=""):
+    """Start an owner that claims the exact delivery after durable admission."""
+    admitted: queue.Queue[str] = queue.Queue()
+    original_deliver = mailbox.deliver_to_live_owner
+
+    def _deliver(*args, **kwargs):
+        record = original_deliver(*args, **kwargs)
+        admitted.put(record["delivery_id"])
+        return record
+
+    monkeypatch.setattr(mailbox, "deliver_to_live_owner", _deliver)
+    ready = threading.Event()
+    errors = []
+
     def _run():
-        owner = mailbox.find_canonical_live_owner(home)
-        for _ in range(200):
+        try:
+            owner = mailbox.find_canonical_live_owner(home)
+            ready.set()
+            delivery_id = admitted.get()
             claimed = mailbox.claim_pending_delivery(home, owner)
-            if claimed is not None:
-                mailbox.complete_delivery(home, claimed["delivery_id"], status="settled", reply=reply)
-                return
-            time.sleep(0.02)
+            assert claimed is not None and claimed["delivery_id"] == delivery_id
+            mailbox.complete_delivery(
+                home, delivery_id, status=status, reply=reply, error=error, reason=reason
+            )
+        except BaseException as exc:
+            errors.append(exc)
+            ready.set()
 
     thread = threading.Thread(target=_run, daemon=True)
     thread.start()
-    return thread
+    ready.wait()
+    if errors:
+        raise errors[0]
+    return thread, errors
+
+
+def _owner_answers(monkeypatch, home, reply):
+    """What the Desktop's live session does: claim the delivery, run it as its next turn, settle it."""
+    return _owner_settles(monkeypatch, home, status="settled", reply=reply)
+
+
+def _join_owner(owner) -> None:
+    thread, errors = owner
+    thread.join()
+    if errors:
+        raise errors[0]
 
 
 @pytest.mark.asyncio
@@ -81,7 +113,7 @@ async def test_a_peer_turn_into_an_open_bot_chat_is_answered_by_its_live_owner(
             session_id="bot-chat", surface="desktop", config={}, registry_home=home, track_liveness=True,
             metadata={"live_session_id": "live-1", "bot_live_delivery_consumer": True})
         assert lease is not None and refusal is None
-    owner = _owner_answers(home, "pong") if owner_replies and target == "bot-chat" else None
+    owner = _owner_answers(monkeypatch, home, "pong") if owner_replies and target == "bot-chat" else None
     adapter = APIServerAdapter(PlatformConfig(enabled=True))
     adapter._session_db = db
     try:
@@ -90,7 +122,7 @@ async def test_a_peer_turn_into_an_open_bot_chat_is_answered_by_its_live_owner(
                 resp = await cli.post(f"/api/sessions/{target}/chat", json={"message": "ping", "author": AUTHOR})
                 body = await resp.json()
         if owner is not None:
-            owner.join(5)
+            _join_owner(owner)
         assert resp.status == status, body
         assert run.called is turn_ran_here
         admitted = sorted((home / "runtime" / "bot_live_delivery").glob("*.json"))
@@ -139,7 +171,7 @@ async def test_a_streamed_peer_turn_into_an_open_bot_chat_is_answered_by_its_liv
         session_id="bot-chat", surface="desktop", config={}, registry_home=home, track_liveness=True,
         metadata={"live_session_id": "live-1", "bot_live_delivery_consumer": True})
     assert lease is not None and refusal is None
-    owner = _owner_answers(home, "pong") if owner_replies else None
+    owner = _owner_answers(monkeypatch, home, "pong") if owner_replies else None
     adapter = APIServerAdapter(PlatformConfig(enabled=True))
     adapter._session_db = db
     app = web.Application()
@@ -151,7 +183,7 @@ async def test_a_streamed_peer_turn_into_an_open_bot_chat_is_answered_by_its_liv
                 assert resp.status == 200 and resp.content_type == "text/event-stream"
                 events = _sse_events(await resp.text())
         if owner is not None:
-            owner.join(5)
+            _join_owner(owner)
         assert not run.called
         [record] = [json.loads(p.read_text()) for p in (home / "runtime" / "bot_live_delivery").glob("*.json")]
         assert (record["message"], record["author"]) == ("ping", AUTHOR)
@@ -176,24 +208,6 @@ def _runs_app(adapter):
     return app
 
 
-def _owner_settles(home, *, status, reply="", error="", reason="", after=0.05):
-    """The Desktop's live session: claim the delivery, run it, write the receipt."""
-    def _run():
-        owner = mailbox.find_canonical_live_owner(home)
-        for _ in range(200):
-            claimed = mailbox.claim_pending_delivery(home, owner)
-            if claimed is not None:
-                time.sleep(after)
-                mailbox.complete_delivery(home, claimed["delivery_id"], status=status, reply=reply,
-                                          error=error, reason=reason)
-                return
-            time.sleep(0.02)
-
-    thread = threading.Thread(target=_run, daemon=True)
-    thread.start()
-    return thread
-
-
 async def _poll_terminal(cli, run_id, *, until=("completed", "failed", "cancelled"), tries=100):
     status = {}
     for _ in range(tries):
@@ -232,7 +246,13 @@ async def test_a_peer_run_into_an_open_bot_chat_is_driven_by_its_owners_receipt(
         session_id="bot-chat", surface="desktop", config={}, registry_home=home, track_liveness=True,
         metadata={"live_session_id": "live-1", "bot_live_delivery_consumer": True})
     assert lease is not None and refusal is None
-    owner = _owner_settles(home, status=receipt[0], reply=receipt[1], error=receipt[2], reason=receipt[3]) if receipt else None
+    owner = (
+        _owner_settles(
+            monkeypatch, home, status=receipt[0], reply=receipt[1], error=receipt[2], reason=receipt[3]
+        )
+        if receipt
+        else None
+    )
     adapter = APIServerAdapter(PlatformConfig(enabled=True))
     adapter._session_db = db
     ran_here = []
@@ -256,7 +276,7 @@ async def test_a_peer_run_into_an_open_bot_chat_is_driven_by_its_owners_receipt(
                     assert (await cli.post(f"/v1/runs/{run_id}/stop")).status == 200
                 status = await _poll_terminal(cli, run_id)
         if owner is not None:
-            owner.join(5)
+            _join_owner(owner)
         assert status["status"] == expected[0], status
         assert bool(ran_here) is turn_ran_here
         admitted = sorted((home / "runtime" / "bot_live_delivery").glob("*.json"))