diff --git a/coding-agent/docs/adr/0015-enforce-owner-routed-http-event-streams.md b/coding-agent/docs/adr/0015-enforce-owner-routed-http-event-streams.md new file mode 100644 index 000000000..3f43a8f61 --- /dev/null +++ b/coding-agent/docs/adr/0015-enforce-owner-routed-http-event-streams.md @@ -0,0 +1,50 @@ +# ADR-0015: Enforce owner-routed HTTP event streams for Phase 2 sticky routing + +**Status**: Proposed +**Date**: 2026-04-21 + +## Context + +ADR-0013 defines Phase 2 multi-instance ownership as an owner/lease/fencing layer on top of PostgreSQL-backed session durability. That ADR already says runtime execution, approval handling, and event streaming are owner-routed operations, and it explicitly allows sticky routing as the first event-routing strategy. + +The current HTTP implementation does not yet enforce that rule for live SSE streams. `GET /sessions/{session_id}/events` only checks that the session exists, then registers an in-memory event queue and starts a keepalive loop. In a multi-instance deployment, a non-owner instance can still attach a live event stream for a session it does not own, and an already-open stream on the old owner is not forced to stop when ownership changes. + +That gap matters because Phase 2 is already using owner checks to fence runtime-sensitive mutations (`run_agent`, checkpoint restore, close, approval submission). Leaving `/events` unfenced means the repository still allows a stale instance to appear live even after the owner has changed. This does not violate at-most-once execution directly, but it does violate the owner-routed contract that ADR-0013 and the PG phase 2 checklist already claim. + +## Decision + +Adopt owner-routed sticky routing as the first concrete Phase 2 policy for HTTP event streams. + +Specifically: + +- `GET /sessions/{session_id}/events` must reject non-owner instances before registering a live event queue. +- The rejection contract for a non-owner or stale owner is `409 Conflict` with the existing `SessionOwnershipConflictError` detail, matching the current Phase 2 HTTP conflict surface. +- Once a stream is open, the keepalive loop must re-check ownership before emitting keepalive pings. If ownership is lost, the generator must terminate and allow queue cleanup to run. +- The first cut remains sticky-routing only. It does not add brokered event routing, redirect payloads, or cross-instance queue transfer. +- The owner check should be exposed through public `SessionManager` methods for event-stream authorization rather than calling `_assert_owner` directly from HTTP handlers. + +This ADR intentionally keeps approval metadata persistence out of scope. It closes the narrowest, clearest routing gap first: only the current owner may serve the live SSE stream for a session. + +## Alternatives Rejected + +- Keep `/events` unfenced and rely on load-balancer stickiness alone — rejected because it leaves stale instances free to attach live streams even when PostgreSQL ownership already says they are not authoritative. +- Add brokered event routing now — rejected because it expands Phase 2 into cross-instance fan-out before the repository has fully closed the simpler owner-routed sticky-routing contract. +- Return a redirect payload instead of `409 Conflict` — rejected for the first cut because the HTTP API already uses `409` for Phase 2 ownership conflicts and there is no canonical owner-address discovery mechanism yet. +- Continue serving already-open streams after owner loss until disconnect — rejected because it weakens the owner-routed contract and lets stale instances keep presenting themselves as authoritative. + +## Acceptance Criteria + +- [ ] `test_get_events_rejects_stale_owner_before_stream_registration` +- [ ] `test_get_events_stops_stream_after_owner_change` +- [ ] `test_get_events_keeps_stream_alive_for_current_owner` +- [ ] `test_get_events_returns_404_before_owner_check_for_missing_session` +- [ ] `uv run pytest tests/ui/test_http_server_failover.py -v` +- [ ] `uv run pytest tests/ui/test_http_server.py tests/ui/test_http_server_failover.py tests/ui/test_session_manager_owner_checks.py -k "events or owner or failover" -v` + +## References + +- `docs/adr/0013-define-phase2-multi-instance-session-ownership-for-pg-http-sessions.md` +- `docs/pg-phase2-multi-instance-checklist.md` +- `src/coding_agent/ui/http_server.py` +- `src/coding_agent/ui/session_manager.py` +- `tests/ui/test_http_server.py` diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 73e86eec7..d7f1da027 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -652,12 +652,23 @@ async def get_events( api_key: str | None = Depends(verify_api_key), ) -> EventSourceResponse: """Persistent SSE event stream (fan-out supported).""" - if not await session_manager.has_session_async(session_id): - raise HTTPException(status_code=404, detail="Session not found") + try: + await session_manager.get_session_async(session_id) + except KeyError as exc: + raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc + + try: + await session_manager.authorize_event_stream(session_id) + except SessionOwnershipConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc - session = await session_manager.get_session_async(session_id) queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=100) - await session_manager.add_event_queue_async(session_id, queue) + try: + await session_manager.register_owned_event_queue_async(session_id, queue) + except KeyError as exc: + raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc + except SessionOwnershipConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc async def event_generator() -> AsyncIterator[dict[str, str]]: """Generate events from queue.""" @@ -665,12 +676,20 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: while True: try: event = await asyncio.wait_for(queue.get(), timeout=30.0) + try: + await session_manager.verify_event_stream_ownership(session_id) + except SessionOwnershipConflictError: + break yield event if event.get("event") == "SessionClosed": break except asyncio.TimeoutError: if not await session_manager.has_session_async(session_id): break + try: + await session_manager.verify_event_stream_ownership(session_id) + except SessionOwnershipConflictError: + break try: if not await session_manager.has_event_queue_async( session_id, queue diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 2ac09107f..5be543a8e 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -516,6 +516,12 @@ async def _assert_owner(self, session_id: str) -> None: ): raise SessionOwnershipConflictError("stale owner or fencing token rejected") + async def authorize_event_stream(self, session_id: str) -> None: + await self._assert_owner(session_id) + + async def verify_event_stream_ownership(self, session_id: str) -> None: + await self._assert_owner(session_id) + async def _run_store_io(self, func: Callable[..., T], /, *args: object) -> T: async with self._store_io_guard: loop = asyncio.get_running_loop() @@ -571,6 +577,22 @@ async def add_event_queue_async( session = await self.get_session_async(session_id) session.event_queues.append(queue) + async def register_owned_event_queue_async( + self, + session_id: str, + queue: asyncio.Queue[dict[str, Any]], + ) -> None: + async with self._lock: + await self._assert_owner(session_id) + session = await self.get_session_async(session_id) + session.event_queues.append(queue) + try: + await self._assert_owner(session_id) + except (Exception, asyncio.CancelledError): + if queue in session.event_queues: + session.event_queues.remove(queue) + raise + async def remove_event_queue_async( self, session_id: str, @@ -591,7 +613,7 @@ async def _close_resource_async(self, resource: object) -> None: if isawaitable(close_result): await close_result - async def remove_session_async(self, session_id: str) -> None: + async def _remove_session_async_no_lock(self, session_id: str) -> None: session = await self.get_session_async(session_id) await self._close_runtime(session) self._session_cache.pop(session_id, None) @@ -599,6 +621,10 @@ async def remove_session_async(self, session_id: str) -> None: self._approval_stores.pop(session_id, None) self._session_turn_locks.pop(session_id, None) + async def remove_session_async(self, session_id: str) -> None: + async with self._lock: + await self._remove_session_async_no_lock(session_id) + async def _restore_tape(self, tape_id: str | None) -> Tape | None: if tape_id is None: return None @@ -983,7 +1009,7 @@ async def close_session(self, session_id: str) -> None: f"Session task for {session_id} did not stop after cancellation" ) - await self.remove_session_async(session_id) + await self._remove_session_async_no_lock(session_id) logger.info(f"Closed session: {session_id}") diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index ddb0fd976..9fd12320c 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -866,7 +866,7 @@ async def test_event_generator_uses_public_session_apis_for_keepalive_exit( has_session_calls: list[str] = [] get_session_calls: list[str] = [] - has_session_results = iter([True, True, False]) + has_session_results = iter([True, False]) class FakeEventSourceResponse: def __init__(self, body_iterator): @@ -920,7 +920,7 @@ async def fake_wait_for(awaitable, timeout): with pytest.raises(StopAsyncIteration): await anext(event_generator) - assert has_session_calls == [session_id, session_id, session_id] + assert has_session_calls == [session_id, session_id] assert len(get_session_calls) >= 2 assert all( call_session_id == session_id for call_session_id in get_session_calls diff --git a/coding-agent/tests/ui/test_http_server_failover.py b/coding-agent/tests/ui/test_http_server_failover.py new file mode 100644 index 000000000..995e57737 --- /dev/null +++ b/coding-agent/tests/ui/test_http_server_failover.py @@ -0,0 +1,590 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from collections.abc import AsyncIterator +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta +from typing import Callable +from typing import Any +from typing import cast + +import pytest +from httpx import ASGITransport, AsyncClient +from fastapi import HTTPException +from starlette.requests import Request + +from coding_agent.ui.http_server import ( + _broadcast_event, + app, + get_events, + limiter, + session_manager, +) +from coding_agent.ui.session_owner_store import SessionOwnerRecord +from coding_agent.ui.session_store import InMemorySessionStore + + +class FakeOwnerStore: + def __init__(self) -> None: + self._owners: dict[str, SessionOwnerRecord] = {} + + async def acquire( + self, + session_id: str, + owner_id: str, + lease_seconds: float = 30.0, + fencing_token: int = 1, + ) -> bool: + self._owners[session_id] = SessionOwnerRecord( + owner_id=owner_id, + lease_expires_at=datetime.now(UTC) + timedelta(seconds=lease_seconds), + fencing_token=fencing_token, + ) + return True + + async def get_owner(self, session_id: str) -> SessionOwnerRecord | None: + return self._owners.get(session_id) + + async def renew( + self, + session_id: str, + owner_id: str, + lease_seconds: float = 30.0, + new_fencing_token: int = 2, + current_fencing_token: int = 1, + ) -> bool: + owner = self._owners.get(session_id) + if owner is None: + return False + if owner.owner_id != owner_id or owner.fencing_token != current_fencing_token: + return False + self._owners[session_id] = SessionOwnerRecord( + owner_id=owner_id, + lease_expires_at=datetime.now(UTC) + timedelta(seconds=lease_seconds), + fencing_token=new_fencing_token, + ) + return True + + async def release( + self, + session_id: str, + owner_id: str, + fencing_token: int, + ) -> bool: + owner = self._owners.get(session_id) + if owner is None: + return False + if owner.owner_id != owner_id or owner.fencing_token != fencing_token: + return False + del self._owners[session_id] + return True + + +@pytest.fixture(autouse=True) +async def clear_sessions(): + session_manager.clear_sessions() + limiter.reset() + yield + session_manager.clear_sessions() + limiter.reset() + + +@pytest.fixture +async def client(): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + yield ac + + +@pytest.fixture +def owner_store(monkeypatch: pytest.MonkeyPatch) -> Iterator[FakeOwnerStore]: + fake_owner_store = FakeOwnerStore() + original_owner_store = session_manager._owner_store + original_owner_id = session_manager._owner_id + original_fencing_token = session_manager._fencing_token + + monkeypatch.setattr(session_manager, "_store", InMemorySessionStore()) + monkeypatch.setattr(session_manager, "_session_cache", {}) + monkeypatch.setattr(session_manager, "_owner_store", fake_owner_store) + monkeypatch.setattr(session_manager, "_owner_id", "owner-a") + monkeypatch.setattr(session_manager, "_fencing_token", 1) + + yield fake_owner_store + + monkeypatch.setattr(session_manager, "_owner_store", original_owner_store) + monkeypatch.setattr(session_manager, "_owner_id", original_owner_id) + monkeypatch.setattr(session_manager, "_fencing_token", original_fencing_token) + + +def _events_request(session_id: str) -> Request: + return Request( + { + "type": "http", + "method": "GET", + "path": f"/sessions/{session_id}/events", + "headers": [], + } + ) + + +class FakeEventSourceResponse: + def __init__(self, body_iterator: AsyncIterator[dict[str, str]]): + self.body_iterator = body_iterator + + +@pytest.mark.asyncio +async def test_clear_sessions_fixture_resets_limiter_in_teardown( + monkeypatch: pytest.MonkeyPatch, +): + session_clear_calls = 0 + limiter_reset_calls = 0 + + def fake_clear_sessions() -> None: + nonlocal session_clear_calls + session_clear_calls += 1 + + def fake_limiter_reset() -> None: + nonlocal limiter_reset_calls + limiter_reset_calls += 1 + + monkeypatch.setattr(session_manager, "clear_sessions", fake_clear_sessions) + monkeypatch.setattr(limiter, "reset", fake_limiter_reset) + + fixture_factory = cast( + "Callable[[], AsyncGenerator[None, None]]", + getattr(clear_sessions, "__wrapped__"), + ) + fixture_gen = fixture_factory() + await anext(fixture_gen) + + assert session_clear_calls == 1 + assert limiter_reset_calls == 1 + + with pytest.raises(StopAsyncIteration): + await anext(fixture_gen) + + assert session_clear_calls == 2 + assert limiter_reset_calls == 2 + + +@pytest.mark.asyncio +async def test_get_events_returns_404_before_owner_check_for_missing_session( + owner_store: FakeOwnerStore, +): + del owner_store + with pytest.raises(HTTPException) as exc_info: + await get_events(_events_request("missing-session"), "missing-session", None) + + response = exc_info.value + assert getattr(response, "status_code", None) == 404 + assert getattr(response, "detail", None) == "Session not found: missing-session" + + +@pytest.mark.asyncio +async def test_get_events_rejects_stale_owner_before_stream_registration( + client: AsyncClient, + owner_store: FakeOwnerStore, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_events(_events_request(session_id), session_id, None) + + response = exc_info.value + assert getattr(response, "status_code", None) == 409 + assert getattr(response, "detail", None) == "stale owner or fencing token rejected" + + +@pytest.mark.asyncio +async def test_get_events_stops_stream_after_owner_change( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + monkeypatch.setattr( + "coding_agent.ui.http_server.EventSourceResponse", FakeEventSourceResponse + ) + + real_wait_for = asyncio.wait_for + timeout_count = 0 + + async def fake_wait_for(awaitable, timeout): + nonlocal timeout_count + if timeout == 30.0: + awaitable.close() + timeout_count += 1 + if timeout_count == 1: + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + raise asyncio.TimeoutError + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr("coding_agent.ui.http_server.asyncio.wait_for", fake_wait_for) + + response = await get_events(_events_request(session_id), session_id, None) + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) + + with pytest.raises(StopAsyncIteration): + await anext(event_generator) + + +@pytest.mark.asyncio +async def test_get_events_drops_queued_event_after_owner_change( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + monkeypatch.setattr( + "coding_agent.ui.http_server.EventSourceResponse", FakeEventSourceResponse + ) + + original_verify = session_manager.verify_event_stream_ownership + verify_calls = 0 + + async def fake_verify_event_stream_ownership(current_session_id: str) -> None: + nonlocal verify_calls + assert current_session_id == session_id + verify_calls += 1 + if verify_calls == 1: + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + await original_verify(current_session_id) + + monkeypatch.setattr( + session_manager, + "verify_event_stream_ownership", + fake_verify_event_stream_ownership, + ) + + response = await get_events(_events_request(session_id), session_id, None) + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) + session = session_manager.get_session(session_id) + await session.event_queues[0].put({"event": "message", "data": "stale"}) + + with pytest.raises(StopAsyncIteration): + await anext(event_generator) + + +@pytest.mark.asyncio +async def test_get_events_returns_404_when_session_disappears_during_queue_registration( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_authorize = session_manager.authorize_event_stream + + async def fake_authorize_event_stream(current_session_id: str) -> None: + assert current_session_id == session_id + await original_authorize(current_session_id) + session_manager.clear_sessions() + + monkeypatch.setattr( + session_manager, + "authorize_event_stream", + fake_authorize_event_stream, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_events(_events_request(session_id), session_id, None) + + response = exc_info.value + assert response.status_code == 404 + assert response.detail == f"Session not found: {session_id}" + + +@pytest.mark.asyncio +async def test_register_owned_event_queue_cleans_up_queue_on_unexpected_error( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_assert_owner = session_manager._assert_owner + assert_owner_calls = 0 + queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=100) + + async def fake_assert_owner(current_session_id: str) -> None: + nonlocal assert_owner_calls + assert current_session_id == session_id + assert_owner_calls += 1 + if assert_owner_calls == 2: + raise RuntimeError("owner store decode failed") + await original_assert_owner(current_session_id) + + monkeypatch.setattr( + session_manager, + "_assert_owner", + fake_assert_owner, + ) + + with pytest.raises(RuntimeError, match="owner store decode failed"): + await session_manager.register_owned_event_queue_async(session_id, queue) + + assert queue not in session_manager.get_session(session_id).event_queues + + +@pytest.mark.asyncio +async def test_register_owned_event_queue_cleans_up_queue_on_cancellation( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_assert_owner = session_manager._assert_owner + assert_owner_calls = 0 + queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=100) + + async def fake_assert_owner(current_session_id: str) -> None: + nonlocal assert_owner_calls + assert current_session_id == session_id + assert_owner_calls += 1 + if assert_owner_calls == 2: + raise asyncio.CancelledError + await original_assert_owner(current_session_id) + + monkeypatch.setattr( + session_manager, + "_assert_owner", + fake_assert_owner, + ) + + with pytest.raises(asyncio.CancelledError): + await session_manager.register_owned_event_queue_async(session_id, queue) + + assert queue not in session_manager.get_session(session_id).event_queues + + +@pytest.mark.asyncio +async def test_get_events_rejects_owner_change_after_queue_registration( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_assert_owner = session_manager._assert_owner + assert_owner_calls = 0 + + async def fake_assert_owner(current_session_id: str) -> None: + nonlocal assert_owner_calls + assert current_session_id == session_id + assert_owner_calls += 1 + if assert_owner_calls == 3: + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + await original_assert_owner(current_session_id) + + monkeypatch.setattr( + session_manager, + "_assert_owner", + fake_assert_owner, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_events(_events_request(session_id), session_id, None) + + response = exc_info.value + assert response.status_code == 409 + assert response.detail == "stale owner or fencing token rejected" + assert session_manager.get_session(session_id).event_queues == [] + + +@pytest.mark.asyncio +async def test_get_events_rejects_broadcast_during_append_to_recheck_window( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_assert_owner = session_manager._assert_owner + assert_owner_calls = 0 + + async def fake_assert_owner(current_session_id: str) -> None: + nonlocal assert_owner_calls + assert current_session_id == session_id + assert_owner_calls += 1 + if assert_owner_calls == 3: + session = session_manager.get_session(session_id) + await _broadcast_event(session, {"event": "message", "data": "stale"}) + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + await original_assert_owner(current_session_id) + + monkeypatch.setattr( + session_manager, + "_assert_owner", + fake_assert_owner, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_events(_events_request(session_id), session_id, None) + + response = exc_info.value + assert response.status_code == 409 + assert response.detail == "stale owner or fencing token rejected" + assert session_manager.get_session(session_id).event_queues == [] + + +@pytest.mark.asyncio +async def test_get_events_rejects_close_during_append_to_recheck_window( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + original_close_runtime = session_manager._close_runtime + close_started = asyncio.Event() + allow_close_to_finish = asyncio.Event() + + async def fake_close_runtime(session) -> None: + close_started.set() + await allow_close_to_finish.wait() + await original_close_runtime(session) + + monkeypatch.setattr( + session_manager, + "_close_runtime", + fake_close_runtime, + ) + + remove_task = asyncio.create_task(session_manager.remove_session_async(session_id)) + await close_started.wait() + + get_events_task = asyncio.create_task( + get_events(_events_request(session_id), session_id, None) + ) + + allow_close_to_finish.set() + await remove_task + + with pytest.raises(HTTPException) as exc_info: + await get_events_task + + response = exc_info.value + assert response.status_code == 404 + assert response.detail == f"Session not found: {session_id}" + + +@pytest.mark.asyncio +async def test_get_events_keeps_stream_alive_for_current_owner( + client: AsyncClient, + owner_store: FakeOwnerStore, + monkeypatch: pytest.MonkeyPatch, +): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=1, + ) + + monkeypatch.setattr( + "coding_agent.ui.http_server.EventSourceResponse", FakeEventSourceResponse + ) + + real_wait_for = asyncio.wait_for + timeout_count = 0 + + async def fake_wait_for(awaitable, timeout): + nonlocal timeout_count + if timeout == 30.0: + awaitable.close() + timeout_count += 1 + if timeout_count == 1: + raise asyncio.TimeoutError + owner_store._owners[session_id] = SessionOwnerRecord( + owner_id="owner-b", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + fencing_token=2, + ) + raise asyncio.TimeoutError + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr("coding_agent.ui.http_server.asyncio.wait_for", fake_wait_for) + + response = await get_events(_events_request(session_id), session_id, None) + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) + + event = await anext(event_generator) + assert event == {"event": "ping", "data": ""} + + with pytest.raises(StopAsyncIteration): + await anext(event_generator)