Repository navigation
fix(ui): enforce owner-routed event streams - #81
Conversation
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughEnforces Phase 2 owner-routed policy for HTTP SSE session event streams: session lookup and authorization occur before event-queue registration, ownership is re-verified on keepalive and per-event, ownership loss returns/terminates with 409 Conflict, SessionManager exposes public authorize/verify/register APIs, and tests covering failover and staleness were added. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant HTTPServer
participant SessionMgr
participant OwnerStore
participant StreamGen
Client->>HTTPServer: GET /sessions/{id}/events
HTTPServer->>SessionMgr: get_session_async(session_id)
alt session missing
SessionMgr-->>HTTPServer: raise KeyError (mapped to 404)
HTTPServer-->>Client: 404 Session not found
else session exists
HTTPServer->>SessionMgr: authorize_event_stream(session_id)
SessionMgr->>OwnerStore: get_owner / validate fencing token
OwnerStore-->>SessionMgr: owner info or stale
SessionMgr-->>HTTPServer: ok or raise SessionOwnershipConflictError (409)
alt Owner Valid
HTTPServer->>SessionMgr: register_owned_event_queue_async(session_id, queue)
SessionMgr-->>HTTPServer: queue registered or raise SessionOwnershipConflictError (409)
HTTPServer-->>Client: 200 SSE stream established
loop keepalive / per-event
StreamGen->>SessionMgr: verify_event_stream_ownership(session_id)
SessionMgr->>OwnerStore: get_owner / validate
OwnerStore-->>SessionMgr: owner valid or changed
alt still owner
StreamGen-->>Client: event or ping
else owner changed/stale
SessionMgr-->>StreamGen: raise SessionOwnershipConflictError
StreamGen->>StreamGen: terminate generator and cleanup queue
StreamGen-->>Client: stream closed
end
end
else Owner Invalid/Stale
HTTPServer-->>Client: 409 Conflict (SessionOwnershipConflictError)
end
end
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly Related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Pull request overview
This PR tightens Phase 2 multi-instance semantics for the UI’s SSE event stream by enforcing session ownership (“owner-routed” behavior) on GET /sessions/{session_id}/events, and documents the decision in a new ADR.
Changes:
- Add
SessionManager.authorize_event_stream()andSessionManager.verify_event_stream_ownership()as public wrappers around the existing owner fencing logic. - Enforce ownership before SSE queue registration and re-validate ownership on keepalive timeouts to terminate stale streams.
- Add failover-focused SSE tests and an ADR describing the Phase 2 sticky-routing policy for event streams.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
coding-agent/src/coding_agent/ui/http_server.py |
Adds owner checks to the SSE /events handler before stream registration and during keepalive revalidation. |
coding-agent/src/coding_agent/ui/session_manager.py |
Introduces public event-stream ownership authorization/revalidation methods built on existing _assert_owner. |
coding-agent/tests/ui/test_http_server_failover.py |
Adds coverage ensuring missing sessions 404 before owner checks, stale owners are rejected, and streams stop when ownership changes. |
coding-agent/docs/adr/0015-enforce-owner-routed-http-event-streams.md |
Documents the Phase 2 decision/acceptance criteria for owner-routed SSE with sticky routing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from coding_agent.ui.session_owner_store import ( | ||
| SessionOwnerRecord, | ||
| SessionOwnershipConflictError, | ||
| ) |
There was a problem hiding this comment.
SessionOwnershipConflictError is imported but never used in this test module. Please remove the unused import to keep the test file clean (and avoid failing future linting if introduced).
| from coding_agent.ui.session_owner_store import ( | |
| SessionOwnerRecord, | |
| SessionOwnershipConflictError, | |
| ) | |
| from coding_agent.ui.session_owner_store import SessionOwnerRecord |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
coding-agent/src/coding_agent/ui/http_server.py (1)
667-682:⚠️ Potential issue | 🟠 MajorRe-check ownership before yielding queued events, not only on keepalive.
This only fences the stream after a 30s idle timeout. If the queue keeps receiving events, the timeout path never runs, so a stale instance can continue serving an already-open SSE stream after ownership has moved. The same gap exists if ownership flips after the preflight check but before queue registration.
Suggested fix
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🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@coding-agent/src/coding_agent/ui/http_server.py` around lines 667 - 682, event_generator currently only verifies ownership on timeout, so if events keep arriving a stale instance can yield them; modify event_generator to re-check session ownership immediately after awaiting queue.get() and before yielding each event by calling session_manager.verify_event_stream_ownership(session_id) (and/or session_manager.has_session_async(session_id) as appropriate) and handle SessionOwnershipConflictError by breaking the loop (and not yielding the event); ensure the SessionClosed behavior still works and preserve the existing timeout-based checks.
🧹 Nitpick comments (1)
coding-agent/tests/ui/test_http_server_failover.py (1)
78-94: Drop the dead cleanup loops or stop swallowing teardown failures.
session_manager.clear_sessions()already empties the store before eachlist_sessions()call here, so these loops are currently no-ops. Keeping them behindexcept Exception: passmeans a future cleanup regression will be silently masked.Suggested simplification
`@pytest.fixture`(autouse=True) async def clear_sessions(): session_manager.clear_sessions() limiter.reset() - for session_id in list(session_manager.list_sessions()): - try: - await session_manager.close_session(session_id) - except Exception: - pass yield session_manager.clear_sessions() - for session_id in list(session_manager.list_sessions()): - try: - await session_manager.close_session(session_id) - except Exception: - pass🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@coding-agent/tests/ui/test_http_server_failover.py` around lines 78 - 94, The fixture clear_sessions has redundant loops that iterate session_manager.list_sessions() immediately after calling session_manager.clear_sessions() and also swallows all teardown errors; remove the dead loops and/or stop masking failures: delete both for-blocks that call session_manager.list_sessions() and await session_manager.close_session(session_id), leaving session_manager.clear_sessions() and limiter.reset() in setup and session_manager.clear_sessions() in teardown, or if explicit close is required keep the loops but replace the bare except Exception: pass with proper error handling (log and re-raise or assert) so failures from session_manager.close_session(...) are not silently ignored; refer to clear_sessions, session_manager.clear_sessions(), limiter.reset(), session_manager.list_sessions(), and session_manager.close_session() when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@coding-agent/src/coding_agent/ui/http_server.py`:
- Around line 667-682: event_generator currently only verifies ownership on
timeout, so if events keep arriving a stale instance can yield them; modify
event_generator to re-check session ownership immediately after awaiting
queue.get() and before yielding each event by calling
session_manager.verify_event_stream_ownership(session_id) (and/or
session_manager.has_session_async(session_id) as appropriate) and handle
SessionOwnershipConflictError by breaking the loop (and not yielding the event);
ensure the SessionClosed behavior still works and preserve the existing
timeout-based checks.
---
Nitpick comments:
In `@coding-agent/tests/ui/test_http_server_failover.py`:
- Around line 78-94: The fixture clear_sessions has redundant loops that iterate
session_manager.list_sessions() immediately after calling
session_manager.clear_sessions() and also swallows all teardown errors; remove
the dead loops and/or stop masking failures: delete both for-blocks that call
session_manager.list_sessions() and await
session_manager.close_session(session_id), leaving
session_manager.clear_sessions() and limiter.reset() in setup and
session_manager.clear_sessions() in teardown, or if explicit close is required
keep the loops but replace the bare except Exception: pass with proper error
handling (log and re-raise or assert) so failures from
session_manager.close_session(...) are not silently ignored; refer to
clear_sessions, session_manager.clear_sessions(), limiter.reset(),
session_manager.list_sessions(), and session_manager.close_session() when making
the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 94a4d923-308e-412f-b17e-7ea7dd6111d9
📒 Files selected for processing (4)
coding-agent/docs/adr/0015-enforce-owner-routed-http-event-streams.mdcoding-agent/src/coding_agent/ui/http_server.pycoding-agent/src/coding_agent/ui/session_manager.pycoding-agent/tests/ui/test_http_server_failover.py
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if not await session_manager.has_session_async(session_id): | ||
| raise HTTPException(status_code=404, detail="Session not found") | ||
|
|
||
| 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) |
There was a problem hiding this comment.
get_events() does an existence check via has_session_async() and then calls get_session_async(). For uncached sessions this results in two store loads and still leaves a TOCTOU window where the session can disappear between calls, causing an uncaught KeyError (500) from get_session_async(). Consider using a single get_session_async() in a try/except KeyError to return 404, then run authorize_event_stream() afterward (and proceed to queue registration).
| if not await session_manager.has_session_async(session_id): | |
| raise HTTPException(status_code=404, detail="Session not found") | |
| 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.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 | |
| queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=100) | |
| try: | |
| await session_manager.add_event_queue_async(session_id, queue) | |
| except KeyError as exc: | |
| raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc |
| 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 |
There was a problem hiding this comment.
verify_event_stream_ownership() is called for every queued event (and again on each keepalive timeout). In PG-backed deployments this likely means a database lookup per SSE message via SessionOwnerStore.get_owner(), which can become a throughput bottleneck for active streams. Consider verifying ownership only on keepalive timeouts (or with a time-based debounce / cached-lease approach) to bound the verification IO rate while still terminating streams promptly on owner loss.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if timeout_count == 1: | ||
| raise asyncio.TimeoutError | ||
| raise StopAsyncIteration |
There was a problem hiding this comment.
fake_wait_for raises StopAsyncIteration to try to terminate the SSE generator. Raising StopAsyncIteration from inside an awaited call in an async generator can be converted into a RuntimeError ("async generator raised StopAsyncIteration") if that code path is ever exercised, making the test brittle. Prefer driving generator termination via the existing control paths (e.g., raise asyncio.TimeoutError and have has_session_async return False / or force an ownership conflict) rather than raising StopAsyncIteration from the mocked wait_for.
| if timeout_count == 1: | |
| raise asyncio.TimeoutError | |
| raise StopAsyncIteration | |
| if timeout_count == 2: | |
| 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 |
| await session_manager.add_event_queue_async(session_id, queue) | ||
| except KeyError as exc: | ||
| raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc | ||
|
|
There was a problem hiding this comment.
There is a small race window between authorize_event_stream() and add_event_queue_async(): ownership can change after the authorization check but before the queue is registered, allowing a non-owner instance to register an in-memory queue (even if it will be dropped on the next keepalive/ownership check). If the intent is to guarantee “reject before stream registration”, consider coupling these operations (e.g., a single SessionManager method that asserts ownership and registers the queue under a shared lock / critical section, or re-check ownership immediately after registration and remove the queue on failure).
| try: | |
| await session_manager.verify_event_stream_ownership(session_id) | |
| except SessionOwnershipConflictError as exc: | |
| await _cleanup_event_queue_on_disconnect(session_id, queue) | |
| raise HTTPException(status_code=409, detail=str(exc)) from exc |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@coding-agent/src/coding_agent/ui/session_manager.py`:
- Around line 580-593: register_owned_event_queue_async mutates
session.event_queues without using the same session-level lock used by
close_session(), allowing a race where close_session() deletes the session after
get_session_async() returns; fix by acquiring the same lock (self._lock) around
the check-and-append sequence: inside an "async with self._lock" block re-verify
ownership (call _assert_owner or check session presence/owner), fetch or
validate the session (get_session_async or direct lookup), append the queue to
session.event_queues, and only release the lock after append; keep the existing
except SessionOwnershipConflictError cleanup logic (remove queue if appended)
but ensure the append and removal are performed under the lock so a concurrent
close cannot leave a dangling registration; reference functions/members:
register_owned_event_queue_async, _assert_owner, get_session_async,
close_session, self._lock, session.event_queues.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 418ec431-2704-48f5-b2bb-bfeddc5717cb
📒 Files selected for processing (3)
coding-agent/src/coding_agent/ui/http_server.pycoding-agent/src/coding_agent/ui/session_manager.pycoding-agent/tests/ui/test_http_server_failover.py
🚧 Files skipped from review as they are similar to previous changes (1)
- coding-agent/tests/ui/test_http_server_failover.py
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| session.event_queues.append(queue) | ||
| try: | ||
| await self._assert_owner(session_id) | ||
| except SessionOwnershipConflictError: |
There was a problem hiding this comment.
register_owned_event_queue_async() only removes the newly-appended queue when _assert_owner() raises SessionOwnershipConflictError. If _assert_owner() (or the owner backend) raises any other exception after the append (e.g., backend/codec errors), the queue will remain registered on the session, which can leak memory and cause future broadcasts to enqueue into an orphaned queue. Consider ensuring the queue is removed on any exception raised after session.event_queues.append(queue) (not just ownership conflicts), and then re-raise the original error.
| except SessionOwnershipConflictError: | |
| except BaseException: |
| session_manager.clear_sessions() | ||
| limiter.reset() | ||
| yield | ||
| session_manager.clear_sessions() |
There was a problem hiding this comment.
The clear_sessions autouse fixture resets the global limiter before each test, but not in teardown. Because limiter is a module-level singleton shared across the whole test suite, leftover rate-limit state from the final test in this module can leak into subsequent tests and cause order-dependent 429s. Consider also calling limiter.reset() after yield (matching the pattern in tests/ui/test_http_server.py).
| session_manager.clear_sessions() | |
| session_manager.clear_sessions() | |
| limiter.reset() |
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Summary
GET /sessions/{session_id}/eventsbefore stream registration and during keepalive revalidation.Verification
uv run pytest tests/ui/test_http_server_failover.py tests/ui/test_http_server.py::TestEventsFanOut::test_event_queues_registered tests/ui/test_http_server.py::TestEventsFanOut::test_multiple_queues_in_session tests/ui/test_http_server.py::TestEventsFanOut::test_event_queue_cleanup_is_shielded_on_disconnect tests/ui/test_http_server.py::TestEventsFanOut::test_event_queue_cleanup_ignores_missing_session tests/ui/test_http_server.py::TestEventsFanOut::test_event_generator_uses_public_session_apis_for_keepalive_exit tests/ui/test_http_server.py::TestEventsFanOut::test_event_generator_exits_cleanly_when_session_disappears_during_keepalive tests/ui/test_session_manager_owner_checks.py::test_run_agent_rejects_non_owner_instance tests/ui/test_session_manager_owner_checks.py::test_restore_checkpoint_rejects_stale_owner tests/ui/test_session_manager_owner_checks.py::test_close_session_rejects_stale_owner tests/ui/test_session_manager_owner_checks.py::test_ensure_session_runtime_rejects_stale_owner tests/ui/test_session_manager_owner_checks.py::test_submit_approval_rejects_stale_owner -vSummary by CodeRabbit
New Features
Documentation
Tests