From b39386e40dcf596bbc31ae73d49242e479235953 Mon Sep 17 00:00:00 2001 From: Tomu Hirata Date: Fri, 31 Jul 2026 18:44:37 +0900 Subject: [PATCH 1/3] fix(sessions): stream session events over a WebSocket to dodge HTTP connection exhaustion Each open tab held its session event stream on a fetch-based SSE GET (`/v1/sessions/{id}/stream`), which counts against the browser's ~6 HTTP/1.1 connections-per-origin cap shared across all tabs. With 5-6 tabs open, the SSE streams filled every slot and unrelated navigation / API requests queued behind them, so the UI appeared hung. Add a dedicated per-conversation event WebSocket (`WS /v1/sessions/{id}/stream/ws`) that carries the same `ServerStreamEvent` payloads, moving the stream onto the browser's separate, effectively unbounded WebSocket pool. The web client defaults to it; the SSE route stays as a fallback (`VITE_EVENT_STREAM_TRANSPORT=sse`). Server: factor the SSE generator into a transport-agnostic `_iter_session_events` core (presence + subscribe + validation) so the SSE wire output is byte-identical and the WS route shares the same snapshot-on-connect and presence semantics. Client: split the pump into `pumpParsedEvents` (transport-agnostic) plus `pumpStreamEvents` (SSE wrapper); add the `sessionEventSocket` transport; `startStreamPump` selects transport via `useEventStreamWebSocket()`. Signed-off-by: Tomu Hirata --- omnigent/server/routes/_sessions/helpers.py | 141 ++++++--- .../server/routes/sessions/routes_events.py | 299 +++++++++++++----- tests/server/routes/test_session_stream_ws.py | 185 +++++++++++ tests/server/test_stream_events.py | 61 ++++ web/src/lib/sessionEventSocket.test.ts | 130 ++++++++ web/src/lib/sessionEventSocket.ts | 173 ++++++++++ web/src/store/chatStore.test.ts | 57 ++++ web/src/store/chatStore.ts | 133 +++++++- web/src/test-setup.ts | 9 + 9 files changed, 1069 insertions(+), 119 deletions(-) create mode 100644 tests/server/routes/test_session_stream_ws.py create mode 100644 web/src/lib/sessionEventSocket.test.ts create mode 100644 web/src/lib/sessionEventSocket.ts diff --git a/omnigent/server/routes/_sessions/helpers.py b/omnigent/server/routes/_sessions/helpers.py index 46e801f72f..6e6c4b977e 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/omnigent/server/routes/_sessions/helpers.py @@ -7193,6 +7193,97 @@ async def _evaluate_output_policy( } +async def _iter_session_events( + is_disconnected: Callable[[], Awaitable[bool]], + session_id: str, + on_subscribed: Callable[[], Awaitable[Iterable[dict[str, Any]]]] | None = None, + viewer_user_id: str | None = None, + viewer_idle: bool = False, + presence_root_id: str | None = None, +) -> AsyncIterator[tuple[str, dict[str, Any]]]: + """ + Yield validated ``(event_type, payload)`` pairs from the live stream. + + Transport-agnostic core shared by the SSE route (via + :func:`_stream_live_events`, which SSE-formats each pair) and the + per-conversation event WebSocket (which ``json.dumps`` each pair). + Owns presence registration, the :func:`session_stream.subscribe` + live-tail with snapshot hooks, and ``ServerStreamEvent`` validation + at the wire boundary. + + Unlike the SSE wrapper this yields NO ``[DONE]`` sentinel — the + end-of-iteration IS the terminal signal, and each transport encodes + it in its own idiom (SSE ``data: [DONE]``, WS a normal close). A + subscriber-queue overflow ends the iteration WITHOUT any terminal + marker so both transports treat it as a dropped stream and the client + reconnects through the persisted snapshot. + + :param is_disconnected: Async predicate polled on each event to + detect client disconnect; the SSE route passes + ``request.is_disconnected`` and the WS route a receive-loop + flag. Kept as a callable so this core depends on neither the + FastAPI ``Request`` nor the ``WebSocket`` type. + :param session_id: Session/conversation identifier whose stream to + subscribe to, e.g. ``"conv_abc123"``. + :param on_subscribed: Optional snapshot-on-connect hook forwarded to + :func:`session_stream.subscribe`; see :func:`_stream_live_events`. + :param viewer_user_id: Presence identity for this stream's lifetime, + or ``None`` to skip presence tracking. + :param viewer_idle: Connect-time idle flag; ignored when + *viewer_user_id* is ``None``. + :param presence_root_id: Root conversation of the session tree; + required when *viewer_user_id* is set. + :returns: An async iterator of ``(event_type, validated_payload)``. + :raises ValueError: If *viewer_user_id* is set without + *presence_root_id*. + :raises session_stream.SubscriberOverflowError: propagated so the + transport wrappers can decide how to end the stream. + """ + presence_token: str | None = None + if viewer_user_id is not None: + if presence_root_id is None: + raise ValueError("presence_root_id is required when viewer_user_id is set") + presence_token = presence.connect( + presence_root_id, session_id, viewer_user_id, viewer_idle + ) + try: + # ``aclosing`` propagates outer ``aclose`` into ``subscribe``; + # a bare ``async for`` would leave the subscriber slot until GC. + async with contextlib.aclosing( + session_stream.subscribe( + session_id, + heartbeat_interval_s=_SESSION_STREAM_HEARTBEAT_INTERVAL_S, + ready_event={"type": "session.heartbeat"}, + # In-flight text replay must be captured synchronously at slot + # registration (before ``ready_event`` suspends), not in the + # async ``on_subscribed`` hook, or window deltas double-render. + # Resource state stays in ``on_subscribed`` — it needs + # awaits and is not dedup-sensitive. + pre_ready_snapshot=lambda: inflight_text.snapshot_for(session_id), + on_subscribed=on_subscribed, + ) + ) as live_events: + async for event in live_events: + if await is_disconnected(): + break + event_type = event.get("type") + if not isinstance(event_type, str): + raise ValueError( + f"session stream event missing string ``type`` field: {event!r}", + ) + validated = _SERVER_STREAM_EVENT_ADAPTER.validate_python(event) + yield event_type, validated.model_dump() + finally: + # The non-None checks besides presence_token's are type + # narrowing only: a minted token implies both were set above. + if ( + presence_token is not None + and viewer_user_id is not None + and presence_root_id is not None + ): + presence.disconnect(presence_root_id, viewer_user_id, presence_token) + + async def _stream_live_events( request: Request, session_id: str, @@ -7275,40 +7366,22 @@ async def _stream_live_events( # fans out to ALREADY-subscribed co-viewers, while this stream # learns the full list (self included) from the snapshot-on-connect # presence event — full-state events make that ordering race benign. - presence_token: str | None = None - if viewer_user_id is not None: - if presence_root_id is None: - raise ValueError("presence_root_id is required when viewer_user_id is set") - presence_token = presence.connect( - presence_root_id, session_id, viewer_user_id, viewer_idle - ) + # The shared core (:func:`_iter_session_events`) owns presence, + # subscribe, and validation; this wrapper only SSE-formats each pair + # and appends the ``[DONE]`` sentinel on clean completion. try: - # ``aclosing`` propagates outer ``aclose`` into ``subscribe``; - # a bare ``async for`` would leave the subscriber slot until GC. async with contextlib.aclosing( - session_stream.subscribe( + _iter_session_events( + request.is_disconnected, session_id, - heartbeat_interval_s=_SESSION_STREAM_HEARTBEAT_INTERVAL_S, - ready_event={"type": "session.heartbeat"}, - # In-flight text replay must be captured synchronously at slot - # registration (before ``ready_event`` suspends), not in the - # async ``on_subscribed`` hook, or window deltas double-render. - # Resource state stays in ``on_subscribed`` — it needs - # awaits and is not dedup-sensitive. - pre_ready_snapshot=lambda: inflight_text.snapshot_for(session_id), on_subscribed=on_subscribed, + viewer_user_id=viewer_user_id, + viewer_idle=viewer_idle, + presence_root_id=presence_root_id, ) - ) as live_events: - async for event in live_events: - if await request.is_disconnected(): - break - event_type = event.get("type") - if not isinstance(event_type, str): - raise ValueError( - f"session stream event missing string ``type`` field: {event!r}", - ) - validated = _SERVER_STREAM_EVENT_ADAPTER.validate_python(event) - yield _format_sse(event_type, validated.model_dump()) + ) as events: + async for event_type, payload in events: + yield _format_sse(event_type, payload) except session_stream.SubscriberOverflowError: _logger.warning( "session stream subscriber overflowed for %s; closing for snapshot reconnect", @@ -7318,15 +7391,6 @@ async def _stream_live_events( # Normal completion only — never yield from ``finally`` (aclose / # GeneratorExit would raise ``async generator ignored GeneratorExit``). yield "data: [DONE]\n\n" - finally: - # The non-None checks besides presence_token's are type - # narrowing only: a minted token implies both were set above. - if ( - presence_token is not None - and viewer_user_id is not None - and presence_root_id is not None - ): - presence.disconnect(presence_root_id, viewer_user_id, presence_token) def _validate_terminal_launch_args(value: list[str] | None) -> list[str] | None: @@ -9236,6 +9300,7 @@ async def _load_model_options_from_host(session_id: str, host_id: str) -> None: "_invalidate_runner_backed_snapshot_state", "_is_codex_native_subagent", "_is_kiro_native_session", + "_iter_session_events", "_last_task_error_from_labels", "_latest_assistant_text_from_store", "_latest_message_preview", diff --git a/omnigent/server/routes/sessions/routes_events.py b/omnigent/server/routes/sessions/routes_events.py index a3507f9b6e..e896a1c42d 100644 --- a/omnigent/server/routes/sessions/routes_events.py +++ b/omnigent/server/routes/sessions/routes_events.py @@ -3,6 +3,8 @@ from __future__ import annotations import asyncio +import contextlib +import json import secrets from collections.abc import Callable from typing import Any, Literal, cast @@ -11,6 +13,10 @@ from fastapi import ( APIRouter, Request, + WebSocket, + WebSocketDisconnect, + WebSocketException, + status, ) from fastapi.responses import StreamingResponse @@ -1581,84 +1587,7 @@ async def stream_session( ) async def _resource_snapshot() -> list[dict[str, Any]]: - """Gather current resource state to emit as snapshot-on-connect. - - Best-effort: every runner-touching gather is time-boxed and - guarded so a slow/unavailable runner never blocks the live - tail. Terminals arrive as ``session.resource.created`` (the - same shape the web's live handler already consumes); child - sessions as ``session.child_session.updated``; changed files - as a single invalidate that triggers a client refetch. - - The in-flight assistant-text replay is NOT read here: it is - dedup-sensitive and must be captured synchronously at slot - registration via ``subscribe``'s ``pre_ready_snapshot`` hook, - before ``ready_event`` suspends. The resource - gathers below need awaits and are not dedup-sensitive, so they - stay in this async hook. - """ - events: list[dict[str, Any]] = [] - try: - page = await asyncio.to_thread( - conversation_store.list_conversations, - limit=100, - kind="sub_agent", - parent_conversation_id=session_id, - order="desc", - sort_by="created_at", - ) - summaries = await _child_session_summaries_from_conversations( - page.data, - session_id, - conversation_store, - ) - for summary in summaries: - events.append( - { - "type": "session.child_session.updated", - "conversation_id": session_id, - "child_session_id": summary.id, - "child": summary.model_dump(mode="json"), - } - ) - except Exception: - _logger.debug("snapshot: child sessions failed for %s", session_id, exc_info=True) - if runner_client is not None: - try: - resp = await asyncio.wait_for( - # order=asc: the web cache appends each replayed - # ``created`` event, so the replay must arrive in - # creation order or the session's own terminal (always - # created first) lands behind later agent-launched - # ones. limit=1000 (the runner endpoint max) keeps the - # oldest-first window from dropping the newest - # terminals past the default page of 20. - runner_client.get( - f"/v1/sessions/{session_id}/resources/terminals", - params={"order": "asc", "limit": "1000"}, - ), - timeout=_SNAPSHOT_RUNNER_TIMEOUT_S, - ) - if resp.status_code == 200: - for item in resp.json().get("data", []): - events.append({"type": "session.resource.created", "resource": item}) - except Exception: - _logger.debug("snapshot: terminals failed for %s", session_id, exc_info=True) - # Tell the client to (re)fetch the changed-files list rather - # than fetching it here (avoids a second runner round-trip). - events.append( - { - "type": "session.changed_files.invalidated", - "session_id": session_id, - "environment_id": "default", - } - ) - # Current viewer list (full state, includes this stream's own - # registration) so a joiner never waits for the next presence - # edge to learn who's here. Scoped to the session tree's root - # so a sub-agent page sees viewers of every agent in the tree. - events.append(presence.snapshot(conv.root_conversation_id, session_id)) - return events + return await _build_resource_snapshot(session_id, conv, runner_client) return StreamingResponse( _stream_live_events( @@ -1691,6 +1620,220 @@ async def _resource_snapshot() -> list[dict[str, Any]]: }, ) + async def _build_resource_snapshot( + session_id: str, + conv: Any, + runner_client: httpx.AsyncClient | None, + ) -> list[dict[str, Any]]: + """Gather current resource state to emit as snapshot-on-connect. + + Shared by the SSE route (``GET /sessions/{id}/stream``) and the + event WebSocket (``WS /sessions/{id}/stream/ws``). + + Best-effort: every runner-touching gather is time-boxed and + guarded so a slow/unavailable runner never blocks the live + tail. Terminals arrive as ``session.resource.created`` (the + same shape the web's live handler already consumes); child + sessions as ``session.child_session.updated``; changed files + as a single invalidate that triggers a client refetch. + + The in-flight assistant-text replay is NOT read here: it is + dedup-sensitive and must be captured synchronously at slot + registration via ``subscribe``'s ``pre_ready_snapshot`` hook, + before ``ready_event`` suspends. The resource + gathers below need awaits and are not dedup-sensitive, so they + stay in this async hook. + """ + events: list[dict[str, Any]] = [] + try: + page = await asyncio.to_thread( + conversation_store.list_conversations, + limit=100, + kind="sub_agent", + parent_conversation_id=session_id, + order="desc", + sort_by="created_at", + ) + summaries = await _child_session_summaries_from_conversations( + page.data, + session_id, + conversation_store, + ) + for summary in summaries: + events.append( + { + "type": "session.child_session.updated", + "conversation_id": session_id, + "child_session_id": summary.id, + "child": summary.model_dump(mode="json"), + } + ) + except Exception: + _logger.debug("snapshot: child sessions failed for %s", session_id, exc_info=True) + if runner_client is not None: + try: + resp = await asyncio.wait_for( + # order=asc: the web cache appends each replayed + # ``created`` event, so the replay must arrive in + # creation order or the session's own terminal (always + # created first) lands behind later agent-launched + # ones. limit=1000 (the runner endpoint max) keeps the + # oldest-first window from dropping the newest + # terminals past the default page of 20. + runner_client.get( + f"/v1/sessions/{session_id}/resources/terminals", + params={"order": "asc", "limit": "1000"}, + ), + timeout=_SNAPSHOT_RUNNER_TIMEOUT_S, + ) + if resp.status_code == 200: + for item in resp.json().get("data", []): + events.append({"type": "session.resource.created", "resource": item}) + except Exception: + _logger.debug("snapshot: terminals failed for %s", session_id, exc_info=True) + # Tell the client to (re)fetch the changed-files list rather + # than fetching it here (avoids a second runner round-trip). + events.append( + { + "type": "session.changed_files.invalidated", + "session_id": session_id, + "environment_id": "default", + } + ) + # Current viewer list (full state, includes this stream's own + # registration) so a joiner never waits for the next presence + # edge to learn who's here. Scoped to the session tree's root + # so a sub-agent page sees viewers of every agent in the tree. + events.append(presence.snapshot(conv.root_conversation_id, session_id)) + return events + + # ── WS /sessions/{session_id}/stream/ws ──────────────────────── + + @router.websocket("/sessions/{session_id}/stream/ws") + async def stream_session_ws( + websocket: WebSocket, + session_id: str, + idle: bool = False, + ) -> None: + """ + Per-conversation event stream over a WebSocket. + + Functional twin of ``GET /sessions/{session_id}/stream`` — same + live-tail contract, same snapshot-on-connect, same presence + semantics — carried over a WebSocket so it rides the browser's + separate (effectively unbounded) WS connection pool instead of + the ~6-per-origin HTTP/1.1 pool the SSE stream competes in. With + several tabs open the SSE streams alone can exhaust that HTTP + pool and stall every other request; this endpoint is the web + client's default, with the SSE route kept as a fallback. + + Protocol: the server pushes one JSON text frame per event + (``{"type": ..., ...}`` — the same ``ServerStreamEvent`` payloads + the SSE route emits). The client sends nothing after the + handshake; a normal close is the analog of the SSE ``[DONE]`` + sentinel. Like the SSE stream this does NOT replay history — + clients reconcile via ``GET /v1/sessions/{id}`` and dedupe by + item id. + + :param websocket: The incoming FastAPI :class:`WebSocket`. + :param session_id: Session/conversation identifier, + e.g. ``"conv_abc123"``. + :param idle: Presence idle flag from the query string, matching + the SSE route's ``idle`` param. An idle flip mid-view arrives + as a reconnect carrying the new value. + """ + user_id = auth_provider.get_user_id(websocket) if auth_provider is not None else None + # An unauthenticated socket must not probe session ids; reject the + # handshake before accept (mirrors the session-updates gate). + if permission_store is not None and user_id is None: + raise WebSocketException( + code=status.WS_1008_POLICY_VIOLATION, + reason="authentication required", + ) + # Access failures surface as HTTP-shaped OmnigentError (403/404); a + # pre-accept WebSocket can only reject via WebSocketException, so map + # both to a policy-violation close (the same code the auth gate uses). + try: + access = await _require_access_and_level( + user_id, session_id, LEVEL_READ, permission_store, conversation_store + ) + except OmnigentError as exc: + raise WebSocketException( + code=status.WS_1008_POLICY_VIOLATION, + reason=str(exc.code), + ) from exc + conv = access.conversation + if conv is None: + conv = await asyncio.to_thread(conversation_store.get_conversation, session_id) + if conv is None: + raise WebSocketException( + code=status.WS_1008_POLICY_VIOLATION, + reason="session not found", + ) + runner_client = await _get_runner_client(session_id, runner_router) + await _ensure_runner_relay_ready( + session_id, + conv.runner_id, + runner_client, + conversation_store, + ) + await websocket.accept() + + # A concurrent receive loop is the WS analog of the SSE route's + # ``request.is_disconnected()`` poll: a browser tab close / navigation + # surfaces as ``WebSocketDisconnect`` here, flipping the flag so the + # emit loop stops at its next event instead of blocking on a dead + # socket until the next heartbeat write fails. + disconnected = asyncio.Event() + + async def _watch_disconnect() -> None: + try: + while True: + await websocket.receive_text() + except WebSocketDisconnect: + disconnected.set() + + async def _resource_snapshot() -> list[dict[str, Any]]: + return await _build_resource_snapshot(session_id, conv, runner_client) + + async def _is_disconnected() -> bool: + return disconnected.is_set() + + recv_task = asyncio.create_task(_watch_disconnect(), name="session-stream-ws-recv") + try: + async with contextlib.aclosing( + _iter_session_events( + _is_disconnected, + session_id, + on_subscribed=_resource_snapshot, + viewer_user_id=_attribution_user(user_id), + viewer_idle=idle, + presence_root_id=conv.root_conversation_id, + ) + ) as events: + async for _event_type, payload in events: + if disconnected.is_set(): + break + await websocket.send_text(json.dumps(payload)) + except session_stream.SubscriberOverflowError: + # A dropped-transport signal: end WITHOUT a clean [DONE]-equivalent + # so the client reconnects and reconciles from the snapshot. Close + # with a going-away code rather than a normal 1000. + _logger.warning( + "session stream ws subscriber overflowed for %s; closing for reconnect", + session_id, + ) + with contextlib.suppress(RuntimeError): + await websocket.close(code=status.WS_1001_GOING_AWAY) + except WebSocketDisconnect: + pass + finally: + recv_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await recv_task + with contextlib.suppress(RuntimeError): + await websocket.close() + # ── DELETE /sessions/{session_id} ────────────────────────────── @router.delete( diff --git a/tests/server/routes/test_session_stream_ws.py b/tests/server/routes/test_session_stream_ws.py new file mode 100644 index 0000000000..adea55231c --- /dev/null +++ b/tests/server/routes/test_session_stream_ws.py @@ -0,0 +1,185 @@ +"""Tests for the ``WS /v1/sessions/{id}/stream/ws`` event stream. + +The functional twin of the SSE ``GET /v1/sessions/{id}/stream`` route, +carried over a WebSocket so it rides the browser's separate connection +pool instead of the ~6-per-origin HTTP/1.1 pool the SSE stream competes +in. These tests drive the real route (no store/auth mocks) against +file-backed SQLite stores: they assert the subscription acknowledgment +heartbeat, live fan-out of a published event, the resource +snapshot-on-connect, and the unauthenticated-handshake reject. + +The wire protocol is one JSON text frame per ``ServerStreamEvent``; a +normal close is the analog of the SSE ``[DONE]`` sentinel. +""" + +from __future__ import annotations + +import json + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from omnigent.runtime import session_stream +from omnigent.server.auth import LEVEL_OWNER, UnifiedAuthProvider +from omnigent.server.routes.sessions import create_sessions_router +from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore +from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore +from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore + +ALICE = "alice@example.com" +BOB = "bob@example.com" + + +class _NoIdentityAuthProvider: + """Auth provider whose handshake yields no identity (see the updates-ws + test) — exercises the reject-when-unauthenticated gate deterministically.""" + + def get_user_id(self, request: object) -> None: + """Always return ``None`` (no authenticated identity).""" + del request + return + + +@pytest.fixture +def stores( + db_uri: str, +) -> tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore]: + """Real file-backed stores so writes from the test thread are visible to + the WS handler thread.""" + return ( + SqlAlchemyConversationStore(db_uri), + SqlAlchemyAgentStore(db_uri), + SqlAlchemyPermissionStore(db_uri), + ) + + +@pytest.fixture +def app( + stores: tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore], +) -> FastAPI: + """Minimal app mounting the sessions router with header auth and a real + permission store — the surface the event WebSocket exercises.""" + conversation_store, agent_store, permission_store = stores + app = FastAPI() + app.include_router( + create_sessions_router( + conversation_store=conversation_store, + agent_store=agent_store, + auth_provider=UnifiedAuthProvider(source="header"), + permission_store=permission_store, + ), + prefix="/v1", + ) + return app + + +def _seed_session( + stores: tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore], + *, + owner: str, + title: str, +) -> str: + """Create a session-shaped conversation owned by ``owner`` and return its id.""" + conversation_store, agent_store, permission_store = stores + if agent_store.get("087b7cb7ac30abf4debfaa578d052ec6") is None: + agent_store.create( + agent_id="087b7cb7ac30abf4debfaa578d052ec6", + name="test-agent", + bundle_location="087b7cb7ac30abf4debfaa578d052ec6/bundle", + ) + conv = conversation_store.create_conversation( + title=title, agent_id="087b7cb7ac30abf4debfaa578d052ec6" + ) + permission_store.ensure_user(owner) + permission_store.grant(owner, conv.id, LEVEL_OWNER) + return conv.id + + +def _recv_until(ws: object, wanted: set[str], *, max_frames: int = 50) -> dict[str, object]: + """Read frames until one whose ``type`` is in ``wanted`` arrives, skipping + heartbeats/snapshot events the test isn't awaiting.""" + for _ in range(max_frames): + frame = json.loads(ws.receive_text()) # type: ignore[attr-defined] + if frame.get("type") in wanted: + return frame + raise AssertionError(f"no frame in {wanted} after {max_frames} frames") + + +def test_stream_ws_acks_with_heartbeat(app: FastAPI, stores) -> None: + """Connecting sends a ``session.heartbeat`` ack once the live-tail slot is + registered — the WS analog of the SSE ready-event, so the client can wait + for a concrete subscription before posting a fast one-shot turn.""" + sid = _seed_session(stores, owner=ALICE, title="s") + with TestClient(app).websocket_connect( + f"/v1/sessions/{sid}/stream/ws", headers={"X-Forwarded-Email": ALICE} + ) as ws: + first = json.loads(ws.receive_text()) + assert first["type"] == "session.heartbeat" + + +def test_stream_ws_fans_out_published_event(app: FastAPI, stores) -> None: + """An event published onto the session after connect is delivered live as + one JSON frame — the core fan-out the SSE route also relies on.""" + sid = _seed_session(stores, owner=ALICE, title="s") + with TestClient(app).websocket_connect( + f"/v1/sessions/{sid}/stream/ws", headers={"X-Forwarded-Email": ALICE} + ) as ws: + # Wait for the subscription ack so the subscriber slot is registered + # before we publish (the broker has no buffer / no replay). + assert json.loads(ws.receive_text())["type"] == "session.heartbeat" + session_stream.publish( + sid, + {"type": "response.output_text.delta", "delta": "hello", "item_id": "m1"}, + ) + evt = _recv_until(ws, {"response.output_text.delta"}) + assert evt["delta"] == "hello" + + +def test_stream_ws_snapshot_includes_presence(app: FastAPI, stores) -> None: + """The snapshot-on-connect carries the same resource events as the SSE + route — at minimum the changed-files invalidate and a presence frame — so + a fresh WS client hydrates without a separate poll.""" + sid = _seed_session(stores, owner=ALICE, title="s") + with TestClient(app).websocket_connect( + f"/v1/sessions/{sid}/stream/ws", headers={"X-Forwarded-Email": ALICE} + ) as ws: + invalidate = _recv_until(ws, {"session.changed_files.invalidated"}) + assert invalidate["session_id"] == sid + presence = _recv_until(ws, {"session.presence"}) + # The connecting viewer registered itself, so the snapshot presence + # frame is scoped to this session and non-empty. + assert presence["conversation_id"] == sid + + +def test_stream_ws_rejects_unauthenticated(stores) -> None: + """With permissions enabled, a handshake with no identity is closed at the + handshake (1008 policy violation) before any session data is read.""" + conversation_store, agent_store, permission_store = stores + sid = _seed_session(stores, owner=ALICE, title="s") + app = FastAPI() + app.include_router( + create_sessions_router( + conversation_store=conversation_store, + agent_store=agent_store, + auth_provider=_NoIdentityAuthProvider(), # type: ignore[arg-type] + permission_store=permission_store, + ), + prefix="/v1", + ) + with pytest.raises(WebSocketDisconnect) as exc_info: + with TestClient(app).websocket_connect(f"/v1/sessions/{sid}/stream/ws"): + pass + assert exc_info.value.code == 1008 + + +def test_stream_ws_rejects_unauthorized_user(app: FastAPI, stores) -> None: + """A user without access to the session is denied — the id is never + trusted from the client for authorization.""" + sid = _seed_session(stores, owner=ALICE, title="s") + with pytest.raises(WebSocketDisconnect): + with TestClient(app).websocket_connect( + f"/v1/sessions/{sid}/stream/ws", headers={"X-Forwarded-Email": BOB} + ): + pass diff --git a/tests/server/test_stream_events.py b/tests/server/test_stream_events.py index 8396bddbc1..d282ef70a2 100644 --- a/tests/server/test_stream_events.py +++ b/tests/server/test_stream_events.py @@ -494,6 +494,67 @@ async def is_disconnected(self) -> bool: assert all("[DONE]" not in frame for frame in frames) +@pytest.mark.asyncio +async def test_iter_session_events_yields_validated_pairs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The transport-agnostic core yields ``(type, validated_payload)`` pairs. + + This is the shared source under both the SSE route (which SSE-formats + each pair) and the event WebSocket (which ``json.dumps`` each payload), + so it must emit the validated dict — not the raw publish payload — and + NO ``[DONE]`` sentinel (each transport encodes its own terminal signal). + """ + from omnigent.runtime import session_stream + from omnigent.server.routes.sessions import _iter_session_events + + async def fake_subscribe(*_args: Any, **_kwargs: Any): + yield {"type": "session.heartbeat"} + yield {"type": "response.output_text.delta", "delta": "hi", "item_id": "m1"} + + monkeypatch.setattr(session_stream, "subscribe", fake_subscribe) + + async def _never_disconnected() -> bool: + return False + + pairs = [pair async for pair in _iter_session_events(_never_disconnected, "conv_ws")] + + types = [t for t, _ in pairs] + assert types == ["session.heartbeat", "response.output_text.delta"] + # Payloads are the validated model dump, and no [DONE]-equivalent leaks + # into the pair stream (the WS route closes the socket instead). + assert pairs[1][1]["delta"] == "hi" + assert all(payload.get("type") != "[DONE]" for _, payload in pairs) + + +@pytest.mark.asyncio +async def test_iter_session_events_stops_on_disconnect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A disconnect flip ends the core at its next event (WS receive-loop path).""" + from omnigent.runtime import session_stream + from omnigent.server.routes.sessions import _iter_session_events + + async def fake_subscribe(*_args: Any, **_kwargs: Any): + yield {"type": "session.heartbeat"} + yield {"type": "session.heartbeat"} + + monkeypatch.setattr(session_stream, "subscribe", fake_subscribe) + + disconnected = False + + async def _is_disconnected() -> bool: + return disconnected + + events = _iter_session_events(_is_disconnected, "conv_ws") + first = await events.__anext__() + assert first[0] == "session.heartbeat" + disconnected = True + # The next poll sees the disconnect and ends iteration. + with pytest.raises(StopAsyncIteration): + await events.__anext__() + + def test_publish_session_status_helper_uses_waiting_literal() -> None: """``workflow._publish_session_status`` publishes a typed waiting event. diff --git a/web/src/lib/sessionEventSocket.test.ts b/web/src/lib/sessionEventSocket.test.ts new file mode 100644 index 0000000000..9a9bc59a64 --- /dev/null +++ b/web/src/lib/sessionEventSocket.test.ts @@ -0,0 +1,130 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { streamSessionEventsWs, type WsStreamResult } from "./sessionEventSocket"; + +// Minimal stand-in for the browser WebSocket, mirroring sessionUpdatesSocket's +// test double: a real socket can't open in jsdom, and we need deterministic +// control over when frames / close arrive relative to the consumer's `await`. +class FakeWebSocket { + static readonly CONNECTING = 0; + static readonly OPEN = 1; + static readonly CLOSING = 2; + static readonly CLOSED = 3; + static instances: FakeWebSocket[] = []; + + readyState = FakeWebSocket.OPEN; + onopen: (() => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: (() => void) | null = null; + onclose: ((event: { code: number }) => void) | null = null; + closeCount = 0; + readonly url: string; + + constructor(url: string) { + this.url = url; + FakeWebSocket.instances.push(this); + } + + close(): void { + this.closeCount += 1; + this.readyState = FakeWebSocket.CLOSED; + // The real socket fires onclose asynchronously; emulate that so a caller + // that closes then awaits still observes the terminal frame. + this.onclose?.({ code: 1000 }); + } + + /** Test helper: deliver one server text frame. */ + emit(payload: unknown): void { + this.onmessage?.({ data: JSON.stringify(payload) } as MessageEvent); + } + + /** Test helper: server-initiated close with a given code. */ + serverClose(code: number): void { + this.readyState = FakeWebSocket.CLOSED; + this.onclose?.({ code }); + } +} + +function latestWs(): FakeWebSocket { + const ws = FakeWebSocket.instances.at(-1); + if (!ws) throw new Error("no WebSocket was constructed"); + return ws; +} + +describe("streamSessionEventsWs", () => { + beforeEach(() => { + FakeWebSocket.instances = []; + vi.stubGlobal("WebSocket", FakeWebSocket as unknown as typeof WebSocket); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("yields parsed events and marks a clean close on normal (1000) closure", async () => { + const result: WsStreamResult = { sawCleanClose: false }; + const controller = new AbortController(); + const events: string[] = []; + + const iterable = streamSessionEventsWs("conv_a", controller.signal, undefined, result); + // Start consuming before frames arrive; the FakeWebSocket delivers + // synchronously into the generator's queue. + const consumed = (async () => { + for await (const ev of iterable) events.push(ev.type); + })(); + + // Yield a microtask so the generator installs its onmessage/onclose. + await Promise.resolve(); + const ws = latestWs(); + ws.emit({ type: "response.output_text.delta", delta: "hi", item_id: "m1" }); + ws.serverClose(1000); + + await consumed; + expect(events).toContain("text_delta"); + expect(result.sawCleanClose).toBe(true); + }); + + it("marks a drop (not clean) on an abnormal close code", async () => { + const result: WsStreamResult = { sawCleanClose: false }; + const controller = new AbortController(); + + const iterable = streamSessionEventsWs("conv_b", controller.signal, undefined, result); + const consumed = (async () => { + for await (const ev of iterable) void ev; + })(); + + await Promise.resolve(); + // 1001 (going away) is what the server sends on subscriber overflow. + latestWs().serverClose(1001); + + await consumed; + expect(result.sawCleanClose).toBe(false); + }); + + it("closes the socket and ends iteration when the signal aborts", async () => { + const result: WsStreamResult = { sawCleanClose: false }; + const controller = new AbortController(); + + const iterable = streamSessionEventsWs("conv_c", controller.signal, undefined, result); + const consumed = (async () => { + for await (const ev of iterable) void ev; + })(); + + await Promise.resolve(); + const ws = latestWs(); + controller.abort(); + + await consumed; + expect(ws.closeCount).toBeGreaterThanOrEqual(1); + }); + + it("does not open a socket if the signal is already aborted", async () => { + const result: WsStreamResult = { sawCleanClose: false }; + const controller = new AbortController(); + controller.abort(); + + for await (const ev of streamSessionEventsWs("conv_d", controller.signal, undefined, result)) { + void ev; + } + expect(FakeWebSocket.instances).toHaveLength(0); + }); +}); diff --git a/web/src/lib/sessionEventSocket.ts b/web/src/lib/sessionEventSocket.ts new file mode 100644 index 0000000000..683b1cf01a --- /dev/null +++ b/web/src/lib/sessionEventSocket.ts @@ -0,0 +1,173 @@ +// Per-conversation event stream over a WebSocket (`WS /v1/sessions/{id}/stream/ws`). +// +// The functional twin of the SSE stream `openSessionStream` opens, but carried +// over a WebSocket so it rides the browser's separate (effectively unbounded) +// WS connection pool instead of the ~6-per-origin HTTP/1.1 pool. With several +// tabs open the SSE streams alone can exhaust that HTTP pool and stall every +// other request; this transport is the web client's default event stream. +// +// This module owns ONE connection's lifetime only — open, yield parsed events, +// end on close/abort. Reconnect (backoff, presence idle-flip recycle, snapshot +// reconcile) stays in `startStreamPump`, exactly as it does for the SSE path. +// The shape mirrors `parseSseStream(body, result)`: an async iterable of typed +// events plus a `result` out-param whose `sawCleanClose` distinguishes a +// deliberate server close from a transport drop. +// +// Identity rides the transport like the other app WebSockets: the browser +// can't set `X-Forwarded-Email` on a WS handshake, so we rely on the ingress / +// dev proxy to carry the authenticated identity. The server access-checks the +// session id on the connection's user — the id is never trusted for authz. + +import type { StreamEvent } from "@/lib/events"; +import { resolveWebSocketUrl } from "@/lib/host"; +import { parseEvent } from "@/lib/sse"; + +/** + * Out-param filled in as the stream ends, mirroring {@link SseStreamResult}. + * `sawCleanClose` is `true` only when the server closed the socket normally + * (code 1000) — the WS analog of the SSE `[DONE]` sentinel. A drop (abnormal + * close, error, or a going-away code the server sends on subscriber overflow) + * leaves it `false` so the reconnect loop re-subscribes. + */ +export interface WsStreamResult { + sawCleanClose: boolean; +} + +/** WebSocket normal-closure code — the clean-close (`[DONE]`-equivalent) signal. */ +const WS_NORMAL_CLOSURE = 1000; + +/** + * Build the `ws(s)://` URL for a session's event stream, delegating to the + * host seam like the session-updates socket. + * + * @param sessionId - Conversation id to stream. + * @param idle - Connect-time presence idle flag (mirrors the SSE `?idle=`). + * @returns The fully-qualified WebSocket URL. + */ +function buildEventStreamUrl(sessionId: string, idle: boolean): string { + const query = idle ? "?idle=true" : ""; + return resolveWebSocketUrl(`/v1/sessions/${encodeURIComponent(sessionId)}/stream/ws${query}`); +} + +/** + * Open the per-conversation event WebSocket and yield parsed events until the + * socket closes or `signal` aborts. + * + * Consumed by `pumpParsedEvents` in the store exactly like `parseSseStream`'s + * output: the caller reduces the events into blocks and, when this iterable + * ends, reads `result.sawCleanClose` to decide whether to reconnect. + * + * Aborting `signal` (switchTo / unmount / presence idle-flip) closes the + * socket and ends iteration; the pump reads that as `"aborted"`. A network + * drop or a server going-away close ends iteration with + * `sawCleanClose === false`, which the loop treats as reconnectable. + * + * @param sessionId - Conversation id to stream. + * @param signal - Abort signal owned by the caller's connection attempt. + * @param opts - `idle` presence flag, forwarded to the server. + * @param result - Out-param; `sawCleanClose` is set as the stream ends. + * @returns An async iterable of typed {@link StreamEvent}s. + */ +export async function* streamSessionEventsWs( + sessionId: string, + signal: AbortSignal, + opts: { idle?: boolean } | undefined, + result: WsStreamResult, +): AsyncIterable { + if (signal.aborted) return; + + const ws = new WebSocket(buildEventStreamUrl(sessionId, opts?.idle ?? false)); + + // A single-slot handoff between the socket's event callbacks and this + // generator: callbacks push events / a terminal sentinel and wake the + // pending `next()`; the loop below awaits one at a time. + const queue: StreamEvent[] = []; + let ended = false; + let wake: (() => void) | null = null; + + const signalReady = (): void => { + if (wake) { + const w = wake; + wake = null; + w(); + } + }; + const nextReady = (): Promise => + new Promise((resolve) => { + wake = resolve; + }); + + const finish = (): void => { + if (ended) return; + ended = true; + signalReady(); + }; + + ws.onmessage = (event) => { + if (typeof event.data !== "string") return; + let payload: Record; + try { + payload = JSON.parse(event.data) as Record; + } catch { + return; + } + const type = payload["type"]; + if (typeof type !== "string") return; + // The server sends the raw `ServerStreamEvent` payload as one JSON frame + // (no SSE `event:`/`data:` envelope); the discriminant is the `type` + // field, so parse against that — the same `parseEvent` the SSE path uses. + const parsed = parseEvent(type, payload); + if (parsed !== null) { + queue.push(parsed); + signalReady(); + } + }; + ws.onclose = (event) => { + // Normal closure is the server's deliberate end (`[DONE]` analog); any + // other code is a drop the reconnect loop should recover from. + result.sawCleanClose = event.code === WS_NORMAL_CLOSURE; + finish(); + }; + ws.onerror = () => { + // `onerror` precedes `onclose`; let close set the terminal state. Nothing + // to do here beyond ensuring we don't hang if close never fires. + }; + + const onAbort = (): void => { + // Caller tore down (switchTo / idle flip). Close the socket and end + // iteration; `sawCleanClose` stays false so this reads as a drop, but the + // pump maps an aborted attempt to `"aborted"` before that matters. + try { + ws.close(); + } catch { + // Closing an already-closing socket can throw; ignore. + } + finish(); + }; + signal.addEventListener("abort", onAbort); + + try { + while (true) { + while (queue.length > 0) { + yield queue.shift() as StreamEvent; + } + if (ended) return; + // Serial by design: one frame arrives, is drained, then we await the + // next; there is nothing to parallelize. + // eslint-disable-next-line no-await-in-loop + await nextReady(); + } + } finally { + signal.removeEventListener("abort", onAbort); + // Emulate the async-iterator auto-cancel: a consumer that breaks out + // (Stop pressed, session switched) must close the underlying socket + // rather than leave it open until GC. + if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) { + try { + ws.close(); + } catch { + // Ignore close races. + } + } + } +} diff --git a/web/src/store/chatStore.test.ts b/web/src/store/chatStore.test.ts index 1ac50fc3b0..2d1c98936e 100644 --- a/web/src/store/chatStore.test.ts +++ b/web/src/store/chatStore.test.ts @@ -53,6 +53,7 @@ import { isStaleCompletedResponse, reviveStrayCompletedResponse, initChatStore, + pumpParsedEvents, pumpStreamEvents, setPendingInitialPrompt, startStreamPump, @@ -6789,6 +6790,62 @@ describe("chatStore — pumpStreamEvents end reasons", () => { }); }); +describe("chatStore — pumpParsedEvents (WebSocket transport core)", () => { + const setState = useChatStore.setState as unknown as Parameters[3]; + const getState = useChatStore.getState as unknown as Parameters[4]; + const immediate: FrameScheduler = { schedule: (cb) => cb(), cancel: () => {} }; + + /** An async iterable over a fixed list, mimicking the WS event transport. */ + async function* iterableOf(events: StreamEvent[]): AsyncIterable { + for (const ev of events) yield ev; + } + + it("reduces an event iterable into blocks (transport-agnostic)", async () => { + useChatStore.setState({ conversationId: "conv_ws_core", blocks: [] }); + const events: StreamEvent[] = [ + { type: "response_created", response: { id: "r1", status: "in_progress", output: [] } }, + { type: "text_delta", delta: `${"y".repeat(34)} ` }, + ] as unknown as StreamEvent[]; + const reason = await pumpParsedEvents( + "conv_ws_core", + iterableOf(events), + new AbortController(), + setState, + getState, + () => true, // clean close + immediate, + ); + expect(reason).toBe("server_closed"); + expect(useChatStore.getState().blocks.length).toBeGreaterThan(0); + }); + + it("maps a clean close to 'server_closed' and a drop to 'dropped'", async () => { + useChatStore.setState({ conversationId: "conv_ws_close", blocks: [] }); + const clean = await pumpParsedEvents( + "conv_ws_close", + iterableOf([]), + new AbortController(), + setState, + getState, + () => true, + immediate, + ); + expect(clean).toBe("server_closed"); + + useChatStore.setState({ conversationId: "conv_ws_drop", blocks: [] }); + const dropped = await pumpParsedEvents( + "conv_ws_drop", + iterableOf([]), + new AbortController(), + setState, + getState, + () => false, // socket closed without a normal (1000) code + immediate, + ); + expect(dropped).toBe("dropped"); + }); +}); + describe("chatStore — startStreamPump reconnect loop", () => { const setState = useChatStore.setState as unknown as Parameters[2]; const getState = useChatStore.getState as unknown as Parameters[3]; diff --git a/web/src/store/chatStore.ts b/web/src/store/chatStore.ts index 2977174012..2aa153785a 100644 --- a/web/src/store/chatStore.ts +++ b/web/src/store/chatStore.ts @@ -80,6 +80,7 @@ import type { import { createPresenceIdleTracker } from "@/lib/presenceIdle"; import { parseEvent, parseSseStream, type SseStreamResult } from "@/lib/sse"; import { clearSseLog, pushSseEvent } from "@/lib/sseEventLog"; +import { streamSessionEventsWs, type WsStreamResult } from "@/lib/sessionEventSocket"; import { childSessionsQueryKey, type ChildSessionInfo } from "@/hooks/useChildSessions"; import { sessionItemsQueryKey } from "@/hooks/useSessionItems"; import type { Conversation, ConversationsPage } from "@/hooks/useConversations"; @@ -3129,7 +3130,30 @@ if (typeof document !== "undefined") { } /** - * Own the session SSE stream for the lifetime of a bound conversation, + * Whether the event stream should ride a WebSocket instead of the SSE fetch. + * + * WebSocket is the default: it uses the browser's separate, effectively + * unbounded WS connection pool, so many open tabs no longer exhaust the + * ~6-per-origin HTTP/1.1 pool and stall unrelated requests. Set + * `VITE_EVENT_STREAM_TRANSPORT=sse` at build time to force the legacy SSE + * fallback (e.g. a deployment whose ingress can't proxy this WebSocket). + */ +function useEventStreamWebSocket(): boolean { + return import.meta.env.VITE_EVENT_STREAM_TRANSPORT !== "sse"; +} + +/** + * Consecutive WS opens that closed before delivering a single event before we + * give up. The WS handshake collapses a 401/403/404 into an opaque abnormal + * close, so — unlike the SSE preflight, which sees the status directly — the + * only signal of a permanently-broken/forbidden session is "connects, then + * immediately closes empty" repeating. This cap bounds that loop; a transient + * proxy blip recovers well within it. + */ +const MAX_WS_EMPTY_OPENS = 8; + +/** + * Own the session event stream for the lifetime of a bound conversation, * reconnecting transparently across drops. * * One connection at a time: open `/stream`, pump it via @@ -3166,6 +3190,10 @@ export async function startStreamPump( // established stream — failed opens leave it false so a recovered first // connect is still treated as initial, not a reconnect. let hasConnected = false; + // WS-only: consecutive opens that closed before any event (see + // MAX_WS_EMPTY_OPENS). Reset the moment a WS delivers its first event. + let wsEmptyOpens = 0; + const useWebSocket = useEventStreamWebSocket(); // A reconnect loop is inherently sequential — open → pump → reconnect — // so its awaits cannot be parallelized; no-await-in-loop doesn't apply. /* eslint-disable no-await-in-loop */ @@ -3189,6 +3217,76 @@ export async function startStreamPump( presenceAttemptController = attempt; try { const idle = presenceIdle.idleNow(); + + // WebSocket transport: the browser can't preflight a WS handshake for + // a status code the way the SSE fetch does, so open + pump in one step + // and infer a permanently-broken/forbidden session from repeated + // empty opens (MAX_WS_EMPTY_OPENS). Reconnect/backoff/reconcile is + // otherwise identical to the SSE path below. + if (useWebSocket) { + const reconnecting = hasConnected; + const wsResult: WsStreamResult = { sawCleanClose: false }; + let sawEvent = false; + const wsEvents = (async function* () { + for await (const ev of streamSessionEventsWs(id, attempt.signal, { idle }, wsResult)) { + if (!sawEvent) { + sawEvent = true; + // A live connection: mark connected and reset both the + // backoff and the empty-open guard. + hasConnected = true; + failedOpens = 0; + consecutive404s = 0; + wsEmptyOpens = 0; + presenceIdle.noteReported(idle); + if (reconnecting) { + dropEphemeralInFlightBlocks(id, set); + } else { + // Fresh connection (not a reconnect) — clear any stale event + // log from a previous stream bind so the debug panel starts + // clean, same as the SSE path. + clearSseLog(id); + } + } + yield ev; + } + })(); + // Start the pump, then reconcile the snapshot concurrently on a + // reconnect (race-safe via itemId dedup) — same order as the SSE + // path. A snapshot fetch against a session that never reopened is + // harmless (deduped / 404s without effect). + const pumpPromise = pumpParsedEvents( + id, + wsEvents, + controller, + set, + get, + () => wsResult.sawCleanClose, + ); + if (reconnecting) { + await reconcileOnReconnect(id, set, get); + } + let reason = await pumpPromise; + if (reason === "aborted" && !controller.signal.aborted) { + reason = "dropped"; + } + if (reason !== "dropped") break; + if (!sawEvent) { + // Connected (or failed to) and closed without a single event. + // Repeated, this is a permanent failure the WS handshake hid. + wsEmptyOpens += 1; + if (wsEmptyOpens > MAX_WS_EMPTY_OPENS) { + console.warn( + `Session ${id}: event WebSocket closed empty ${wsEmptyOpens}x, giving up`, + ); + finalizeActive(set, "failed", "event stream unavailable", null); + set({ sessionStatus: "failed", status: "idle" }); + break; + } + failedOpens += 1; + } + continue; + } + let streamRes: Response; try { streamRes = await openSessionStream(id, attempt.signal, { idle }); @@ -3674,9 +3772,38 @@ export async function pumpStreamEvents( get: Getter, scheduler: FrameScheduler = createRafScheduler(), ): Promise { - const stream = new BlockStream(); + // SSE wrapper over the transport-agnostic core: parse the byte body into + // typed events, then let the parser's `[DONE]` sentinel decide clean-close + // vs transport-drop. The event WebSocket path (`pumpParsedEvents`) supplies + // its own already-parsed event iterable and clean-close predicate. const sseResult: SseStreamResult = { sawDone: false }; const rawEvents = parseSseStream(body, sseResult); + return pumpParsedEvents(id, rawEvents, controller, set, get, () => sseResult.sawDone, scheduler); +} + +/** + * Transport-agnostic core of the event pump: tap `session.*` side effects, + * reduce blocks through `BlockStream`, and batch-commit to `state.blocks`. + * Shared by {@link pumpStreamEvents} (SSE byte body) and the event WebSocket + * transport, which both produce an `AsyncIterable`. + * + * @param rawEvents - Parsed typed events from either transport. + * @param sawCleanClose - Called once the iterable ends to distinguish a + * deliberate server close (`"server_closed"`, don't reconnect) from a + * transport drop (`"dropped"`, reconnect). For SSE this reflects the + * `[DONE]` sentinel; for WS, a normal (1000) close frame. + * @returns Why the connection ended — see {@link StreamEndReason}. + */ +export async function pumpParsedEvents( + id: string, + rawEvents: AsyncIterable, + controller: AbortController, + set: Setter, + get: Getter, + sawCleanClose: () => boolean, + scheduler: FrameScheduler = createRafScheduler(), +): Promise { + const stream = new BlockStream(); // Tap the raw event stream for `session.*` side effects (sessionStatus, // pending-message promotion, interrupted decoration) before handing it // to the BlockStream reducer. The reducer is intentionally pure @@ -3929,7 +4056,7 @@ export async function pumpStreamEvents( // a deliberate server close (`[DONE]`) or a transport drop without it // (idle proxy disconnect / the Apps ~5-min cap) decides reconnection. flush(); - return sseResult.sawDone ? "server_closed" : "dropped"; + return sawCleanClose() ? "server_closed" : "dropped"; } catch (err) { if (err instanceof Error && err.name === "AbortError") return "aborted"; if (get().conversationId !== id) return "switched"; diff --git a/web/src/test-setup.ts b/web/src/test-setup.ts index c7c9b2cbc0..065417eb1e 100644 --- a/web/src/test-setup.ts +++ b/web/src/test-setup.ts @@ -1,6 +1,15 @@ import "@testing-library/jest-dom/vitest"; import { vi } from "vitest"; +// Default the event-stream transport to SSE in tests. Production defaults to +// the WebSocket (it dodges the ~6-per-origin HTTP/1.1 pool), but the extensive +// `startStreamPump` reconnect/bind suite drives the SSE fetch path by mocking +// `openSessionStream`, and that fallback still ships. The WebSocket transport +// has its own dedicated coverage (`sessionEventSocket.test.ts` + +// `tests/server/routes/test_session_stream_ws.py`). A WS-branch test can +// override this per-case with `vi.stubEnv("VITE_EVENT_STREAM_TRANSPORT", ...)`. +vi.stubEnv("VITE_EVENT_STREAM_TRANSPORT", "sse"); + // The @lobehub icon packages have broken nested-module resolution // under vitest; stub presentational glyphs so component modules that // import them can still load in tests. (The Antigravity glyph additionally From 9d3a66c6a76fb6acc05336fb82257d199b583eff Mon Sep 17 00:00:00 2001 From: Tomu Hirata Date: Wed, 5 Aug 2026 18:14:24 +0900 Subject: [PATCH 2/3] fix(sessions): repair CI for the event WebSocket transport Rebase onto main brought the routes_events explicit-import refactor, so `_iter_session_events` needed adding to that import list and to the sessions package re-exports. CI fixes: - Rename `useEventStreamWebSocket` to `eventStreamUsesWebSocket`: the `use` prefix made oxlint's react-hooks rule treat it as a React Hook called from a non-component function. - Annotate `_iter_session_events` as `AsyncGenerator` so it satisfies `contextlib.aclosing`'s `_SupportsAclose` bound. - Build the e2e SPA with `VITE_EVENT_STREAM_TRANSPORT=sse`. Much of that suite drives the SSE transport directly (intercepting `GET /stream` to inject 404s, monkeypatching `window.fetch` for a controllable ReadableStream), and Playwright's `extra_http_headers` don't apply to a WebSocket handshake, so header-auth'd multi-viewer presence tests can't identify themselves over WS. Add a localStorage override (`omnigent.eventStream.transport`) so a transport can be flipped per-tab without a rebuild, and use it for new coverage of the default transport: an e2e turn over the event WebSocket that also asserts the SSE endpoint is never opened, plus two unit tests for transport selection. Signed-off-by: Tomu Hirata --- omnigent/server/routes/_sessions/helpers.py | 3 +- omnigent/server/routes/sessions/__init__.py | 1 + .../server/routes/sessions/routes_events.py | 1 + .../chat/test_event_stream_websocket.py | 77 +++++++++++++++++++ tests/e2e_ui/conftest.py | 16 +++- web/src/store/chatStore.test.ts | 60 +++++++++++++++ web/src/store/chatStore.ts | 21 ++++- 7 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 tests/e2e_ui/chat/test_event_stream_websocket.py diff --git a/omnigent/server/routes/_sessions/helpers.py b/omnigent/server/routes/_sessions/helpers.py index 6e6c4b977e..da83d37f8a 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/omnigent/server/routes/_sessions/helpers.py @@ -17,6 +17,7 @@ import weakref from collections import deque from collections.abc import ( + AsyncGenerator, AsyncIterator, Awaitable, Callable, @@ -7200,7 +7201,7 @@ async def _iter_session_events( viewer_user_id: str | None = None, viewer_idle: bool = False, presence_root_id: str | None = None, -) -> AsyncIterator[tuple[str, dict[str, Any]]]: +) -> AsyncGenerator[tuple[str, dict[str, Any]], None]: """ Yield validated ``(event_type, payload)`` pairs from the live stream. diff --git a/omnigent/server/routes/sessions/__init__.py b/omnigent/server/routes/sessions/__init__.py index 6ada3cfd8e..be1b48cc94 100644 --- a/omnigent/server/routes/sessions/__init__.py +++ b/omnigent/server/routes/sessions/__init__.py @@ -400,6 +400,7 @@ _invalidate_runner_backed_snapshot_state as _invalidate_runner_backed_snapshot_state, _is_codex_native_subagent as _is_codex_native_subagent, _is_kiro_native_session as _is_kiro_native_session, + _iter_session_events as _iter_session_events, _last_task_error_from_labels as _last_task_error_from_labels, _latest_assistant_text_from_store as _latest_assistant_text_from_store, _latest_message_preview as _latest_message_preview, diff --git a/omnigent/server/routes/sessions/routes_events.py b/omnigent/server/routes/sessions/routes_events.py index e896a1c42d..3747d19065 100644 --- a/omnigent/server/routes/sessions/routes_events.py +++ b/omnigent/server/routes/sessions/routes_events.py @@ -133,6 +133,7 @@ _get_runner_client_for_resource_access, _handle_external_session_todos, _is_codex_native_subagent, + _iter_session_events, _launch_runner_on_host, _persist_external_assistant_message, _persist_external_codex_approval_mode_change, diff --git a/tests/e2e_ui/chat/test_event_stream_websocket.py b/tests/e2e_ui/chat/test_event_stream_websocket.py new file mode 100644 index 0000000000..5dc1481874 --- /dev/null +++ b/tests/e2e_ui/chat/test_event_stream_websocket.py @@ -0,0 +1,77 @@ +"""E2E: the event stream works end-to-end over its WebSocket transport. + +Production defaults the session event stream to +``WS /v1/sessions/{id}/stream/ws`` so many open tabs stop exhausting the +browser's ~6-per-origin HTTP/1.1 connection pool (held-open SSE GETs used +to fill every slot and stall unrelated navigation). The rest of this suite +builds the SPA with ``VITE_EVENT_STREAM_TRANSPORT=sse`` because much of it +drives the SSE transport directly, so these tests opt back in to the +WebSocket via the ``omnigent.eventStream.transport`` localStorage override +and prove the default transport renders a real turn. + +A failure here means one of: + +- The WS route regressed (auth/access gate, snapshot-on-connect, or the + live tail in ``routes_events.stream_session_ws``). +- The client transport regressed (``web/src/lib/sessionEventSocket.ts``) + or ``startStreamPump`` stopped selecting it. +- The shared pump core (``pumpParsedEvents``) stopped reducing events into + blocks for a non-SSE source. +""" + +from __future__ import annotations + +from playwright.sync_api import Page, expect + +_COMPOSER = "Ask the agent anything…" +_ASSISTANT = '[data-testid="message-bubble"][data-role="assistant"]' +# Flip the transport before any app code runs, so the very first stream +# bind uses the WebSocket (an init script runs pre-navigation). +_FORCE_WS = "window.localStorage.setItem('omnigent.eventStream.transport', 'ws')" + + +def test_turn_streams_over_event_websocket( + page: Page, + seeded_session: tuple[str, str], +) -> None: + """A full turn renders with the event stream carried over a WebSocket. + + Also asserts the SSE endpoint is never opened, so a silent fallback to + the old transport can't make this test pass green. + + :param page: Playwright page. + :param seeded_session: ``(base_url, session_id)`` from the fixture. + """ + base_url, session_id = seeded_session + page.add_init_script(_FORCE_WS) + + ws_urls: list[str] = [] + page.on("websocket", lambda ws: ws_urls.append(ws.url)) + sse_opens: list[str] = [] + page.on( + "request", + lambda r: ( + sse_opens.append(r.url) + if r.url.endswith(f"/v1/sessions/{session_id}/stream") + or f"/v1/sessions/{session_id}/stream?" in r.url + else None + ), + ) + + page.goto(f"{base_url}/c/{session_id}") + + composer = page.get_by_placeholder(_COMPOSER) + expect(composer).to_be_visible() + composer.fill("Say hello") + page.get_by_role("button", name="Send").click() + + # The reply can only arrive through the live event stream — the pump has + # no other source for a streamed turn — so a rendered assistant bubble + # proves the WS transport carried it. + expect(page.locator(_ASSISTANT).first).to_be_visible(timeout=60_000) + expect(page.locator(_ASSISTANT).first).not_to_have_text("") + + assert any(f"/v1/sessions/{session_id}/stream/ws" in url for url in ws_urls), ( + f"event WebSocket was never opened (sockets seen: {ws_urls})" + ) + assert not sse_opens, f"SSE stream was opened despite the WS transport: {sse_opens}" diff --git a/tests/e2e_ui/conftest.py b/tests/e2e_ui/conftest.py index 3f5d3b0960..1dbb271474 100644 --- a/tests/e2e_ui/conftest.py +++ b/tests/e2e_ui/conftest.py @@ -757,7 +757,21 @@ def built_spa(request: pytest.FixtureRequest) -> None: # COREPACK_ENABLE_DOWNLOAD_PROMPT=0 keeps a corepack `pnpm` shim # from blocking on its download confirmation under captured # pytest output, which reads as a hung test run. - env = {**os.environ, "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0"} + # Build the SPA against the SSE event-stream transport. Production + # defaults to the event WebSocket (it dodges the ~6-per-origin + # HTTP/1.1 connection cap), but much of this suite drives the SSE + # transport directly — intercepting `GET /stream` to inject 404s or + # monkeypatching `window.fetch` to hand back a controllable + # ReadableStream — and Playwright's `extra_http_headers` do not apply + # to a WebSocket handshake, so header-auth'd multi-viewer tests + # cannot identify themselves over WS. SSE keeps that coverage + # meaningful; the WebSocket transport has its own e2e coverage that + # opts back in via this env var. + env = { + **os.environ, + "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0", + "VITE_EVENT_STREAM_TRANSPORT": "sse", + } subprocess.run( ["pnpm", "install", "--frozen-lockfile", "--filter", "web"], cwd=_REPO_ROOT, diff --git a/web/src/store/chatStore.test.ts b/web/src/store/chatStore.test.ts index 2d1c98936e..220d7f244d 100644 --- a/web/src/store/chatStore.test.ts +++ b/web/src/store/chatStore.test.ts @@ -6790,6 +6790,66 @@ describe("chatStore — pumpStreamEvents end reasons", () => { }); }); +describe("chatStore — event-stream transport selection", () => { + // The suite's SPA/test env pins SSE (see test-setup.ts), so these cases + // assert the localStorage override that lets one tab (or a WS-targeted + // e2e test) pick the other transport without a rebuild. + const TRANSPORT_KEY = "omnigent.eventStream.transport"; + + afterEach(() => { + window.localStorage.removeItem(TRANSPORT_KEY); + }); + + it("uses the SSE fetch when the override selects sse", async () => { + window.localStorage.setItem(TRANSPORT_KEY, "sse"); + seedSession("conv_transport_sse"); + await useChatStore.getState().switchTo("conv_transport_sse"); + // The SSE stream open is an HTTP GET through the mocked fetch; a WS + // transport would never touch it. + expect( + fetchMock.mock.calls.some(([input]) => + String(input).includes("/v1/sessions/conv_transport_sse/stream"), + ), + ).toBe(true); + }); + + it("opens a WebSocket instead of the SSE fetch when the override selects ws", async () => { + window.localStorage.setItem(TRANSPORT_KEY, "ws"); + const sockets: string[] = []; + class RecordingWebSocket { + static readonly OPEN = 1; + static readonly CONNECTING = 0; + readyState = 0; + onopen: (() => void) | null = null; + onmessage: ((e: MessageEvent) => void) | null = null; + onerror: (() => void) | null = null; + onclose: ((e: { code: number }) => void) | null = null; + constructor(url: string) { + sockets.push(url); + } + close(): void { + this.onclose?.({ code: 1000 }); + } + } + vi.stubGlobal("WebSocket", RecordingWebSocket as unknown as typeof WebSocket); + try { + seedSession("conv_transport_ws"); + await useChatStore.getState().switchTo("conv_transport_ws"); + expect(sockets.some((u) => u.includes("/v1/sessions/conv_transport_ws/stream/ws"))).toBe( + true, + ); + // And the constrained HTTP pool is left alone — that's the whole point. + expect( + fetchMock.mock.calls.some(([input]) => + String(input).endsWith("/v1/sessions/conv_transport_ws/stream"), + ), + ).toBe(false); + } finally { + vi.unstubAllGlobals(); + } + }); +}); + describe("chatStore — pumpParsedEvents (WebSocket transport core)", () => { const setState = useChatStore.setState as unknown as Parameters[3]; const getState = useChatStore.getState as unknown as Parameters[4]; diff --git a/web/src/store/chatStore.ts b/web/src/store/chatStore.ts index 2aa153785a..235dddbd83 100644 --- a/web/src/store/chatStore.ts +++ b/web/src/store/chatStore.ts @@ -3129,6 +3129,9 @@ if (typeof document !== "undefined") { ); } +/** localStorage key that overrides the built-in event-stream transport. */ +const EVENT_STREAM_TRANSPORT_KEY = "omnigent.eventStream.transport"; + /** * Whether the event stream should ride a WebSocket instead of the SSE fetch. * @@ -3137,8 +3140,22 @@ if (typeof document !== "undefined") { * ~6-per-origin HTTP/1.1 pool and stall unrelated requests. Set * `VITE_EVENT_STREAM_TRANSPORT=sse` at build time to force the legacy SSE * fallback (e.g. a deployment whose ingress can't proxy this WebSocket). + * + * A `omnigent.eventStream.transport` localStorage value (`"ws"` / `"sse"`) + * overrides the build-time default for the current browser, so a transport + * can be flipped for one tab without a rebuild — used to diagnose a + * suspected transport issue in a deployed build, and by the e2e suite to + * exercise whichever transport a given test targets. */ -function useEventStreamWebSocket(): boolean { +function eventStreamUsesWebSocket(): boolean { + try { + const override = window.localStorage.getItem(EVENT_STREAM_TRANSPORT_KEY); + if (override === "ws") return true; + if (override === "sse") return false; + } catch { + // Storage can throw in private-mode / sandboxed frames; fall through + // to the build-time default. + } return import.meta.env.VITE_EVENT_STREAM_TRANSPORT !== "sse"; } @@ -3193,7 +3210,7 @@ export async function startStreamPump( // WS-only: consecutive opens that closed before any event (see // MAX_WS_EMPTY_OPENS). Reset the moment a WS delivers its first event. let wsEmptyOpens = 0; - const useWebSocket = useEventStreamWebSocket(); + const useWebSocket = eventStreamUsesWebSocket(); // A reconnect loop is inherently sequential — open → pump → reconnect — // so its awaits cannot be parallelized; no-await-in-loop doesn't apply. /* eslint-disable no-await-in-loop */ From 24a9e9a92e273f1381755c0a9f8f47e94280182f Mon Sep 17 00:00:00 2001 From: Tomu Hirata Date: Wed, 5 Aug 2026 18:38:39 +0900 Subject: [PATCH 3/3] ci(e2e-ui): pin the e2e SPA build to the SSE event-stream transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The conftest's own build sets VITE_EVENT_STREAM_TRANSPORT=sse, but every UI workflow builds the SPA in a dedicated step and then runs pytest with --ui-skip-build, so that fixture never ran and the suite kept exercising the WebSocket default — failing the presence tests, which rely on Playwright's extra_http_headers (not applied to a WS handshake). Set the env var on the build step in each workflow that runs the e2e_ui suite: e2e-ui, flake-stress-ui, ui-snapshot, and ui-snapshot-update. The WebSocket default keeps its own coverage: the localStorage override survives an SSE-pinned build, so test_event_stream_websocket.py still opts back in and asserts the SSE endpoint is never opened. Signed-off-by: Tomu Hirata --- .github/workflows/e2e-ui.yml | 10 ++++++++++ .github/workflows/flake-stress-ui.yml | 4 ++++ .github/workflows/ui-snapshot-update.yml | 4 ++++ .github/workflows/ui-snapshot.yml | 4 ++++ 4 files changed, 22 insertions(+) diff --git a/.github/workflows/e2e-ui.yml b/.github/workflows/e2e-ui.yml index 80ccf7e143..745c82902b 100644 --- a/.github/workflows/e2e-ui.yml +++ b/.github/workflows/e2e-ui.yml @@ -227,6 +227,16 @@ jobs: # so never run it under xdist or alongside the live server. env: NPM_CONFIG_REGISTRY: https://registry.npmjs.org/ + # Pin the event-stream transport to SSE, matching the conftest's own + # build (this job passes --ui-skip-build, so that build never runs). + # Much of this suite drives the SSE transport directly — intercepting + # `GET /stream` to inject 404s, monkeypatching `window.fetch` for a + # controllable ReadableStream — and Playwright's extra_http_headers + # do not apply to a WebSocket handshake, so header-auth'd + # multi-viewer presence tests cannot identify themselves over WS. + # The WebSocket default has its own coverage, which opts back in via + # the omnigent.eventStream.transport localStorage override. + VITE_EVENT_STREAM_TRANSPORT: sse run: | pnpm install --frozen-lockfile --filter web pnpm --filter web run build diff --git a/.github/workflows/flake-stress-ui.yml b/.github/workflows/flake-stress-ui.yml index 90f2bfbc12..85123fc8ce 100644 --- a/.github/workflows/flake-stress-ui.yml +++ b/.github/workflows/flake-stress-ui.yml @@ -271,6 +271,10 @@ jobs: # never run it under xdist or alongside the live server. env: NPM_CONFIG_REGISTRY: https://registry.npmjs.org/ + # Pin the event-stream transport to SSE: this suite drives the SSE + # transport directly and Playwright header auth cannot ride a WS + # handshake. See e2e-ui.yml for the full rationale. + VITE_EVENT_STREAM_TRANSPORT: sse run: | pnpm install --frozen-lockfile --filter web pnpm --filter web run build diff --git a/.github/workflows/ui-snapshot-update.yml b/.github/workflows/ui-snapshot-update.yml index 752f117f8f..cd0864f602 100644 --- a/.github/workflows/ui-snapshot-update.yml +++ b/.github/workflows/ui-snapshot-update.yml @@ -115,6 +115,10 @@ jobs: - name: Build web SPA env: NPM_CONFIG_REGISTRY: https://registry.npmjs.org/ + # Pin the event-stream transport to SSE: this suite drives the SSE + # transport directly and Playwright header auth cannot ride a WS + # handshake. See e2e-ui.yml for the full rationale. + VITE_EVENT_STREAM_TRANSPORT: sse run: | pnpm install --frozen-lockfile --filter web pnpm --filter web run build diff --git a/.github/workflows/ui-snapshot.yml b/.github/workflows/ui-snapshot.yml index 83c9d911fe..22810e84e7 100644 --- a/.github/workflows/ui-snapshot.yml +++ b/.github/workflows/ui-snapshot.yml @@ -169,6 +169,10 @@ jobs: # never run it alongside the live server. env: NPM_CONFIG_REGISTRY: https://registry.npmjs.org/ + # Pin the event-stream transport to SSE: this suite drives the SSE + # transport directly and Playwright header auth cannot ride a WS + # handshake. See e2e-ui.yml for the full rationale. + VITE_EVENT_STREAM_TRANSPORT: sse run: | pnpm install --frozen-lockfile --filter web pnpm --filter web run build