diff --git a/tests/test_litellm/managed_agents/test_events.py b/tests/test_litellm/managed_agents/test_events.py new file mode 100644 index 00000000000..a6103ef0aba --- /dev/null +++ b/tests/test_litellm/managed_agents/test_events.py @@ -0,0 +1,434 @@ +"""Unit tests for `GET /v2/sessions/:id/events` (LIT-2920). + +Covers contract §6.6 + §7: +- SSE stream format (one frame per event, blank line separator). +- Required SSE headers (`Cache-Control: no-cache`, `Connection: keep-alive`, + `X-Accel-Buffering: no`, `Content-Type: text/event-stream`). +- Pre-forward checks happen before streaming begins: + - 404 when the session does not exist for this caller. + - 503 when the session is still provisioning. + - 404 when the session is in a terminal state (terminated/error). +- Mid-stream `SandboxUnreachableError` is converted to an inline `error` + SSE event (status code is already 200 by then). + +The handler depends on: +- `user_api_key_auth` (FastAPI dep) — overridden via `app.dependency_overrides`. +- `litellm.proxy.proxy_server.prisma_client` — patched to a sentinel. +- `litellm.managed_agents.endpoints.events.get_session` — patched per test. +- `litellm.managed_agents.endpoints.events.get_adapter` — patched per test. +""" + +from typing import Any, AsyncIterator, Dict, List, Optional, Tuple +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.managed_agents.adapters.base import SandboxUnreachableError +from litellm.managed_agents.endpoints.events import router as events_router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +CALLER_USER_ID = "user_xyz" + + +def _build_app() -> FastAPI: + """Build a FastAPI app with the events router and auth override.""" + app = FastAPI() + app.include_router(events_router) + + async def _fake_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id=CALLER_USER_ID, api_key="sk-test") + + app.dependency_overrides[user_api_key_auth] = _fake_auth + return app + + +def _ready_session_row( + session_id: str = "ses_test", + sandbox_url: str = "http://127.0.0.1:1234", + oc_sid: str = "oc_sid_abc", +) -> Dict[str, Any]: + return { + "id": session_id, + "agent_id": "agt_abc", + "sandbox_type": "opencode", + "sandbox_url": sandbox_url, + "sandbox_metadata": {"opencode_session_id": oc_sid}, + "status": "ready", + "created_by": CALLER_USER_ID, + } + + +class _FakeAdapter: + """Adapter stub matching the SandboxAdapter Protocol's stream_events. + + Yields the supplied list of (event_type, data) tuples in order. If + `raise_after` is set, raises SandboxUnreachableError after that many + events have been yielded — used to test mid-stream error handling. + """ + + def __init__( + self, + events: List[Tuple[str, Dict[str, Any]]], + raise_after: Optional[int] = None, + ) -> None: + self.events = events + self.raise_after = raise_after + + async def stream_events( + self, + sandbox_url: str, + opencode_session_id: str, + our_session_id: str, + ) -> AsyncIterator[Tuple[str, Dict[str, Any]]]: + for i, evt in enumerate(self.events): + if self.raise_after is not None and i >= self.raise_after: + raise SandboxUnreachableError("upstream is gone") + yield evt + if self.raise_after is not None and self.raise_after >= len(self.events): + raise SandboxUnreachableError("upstream is gone") + + +# --------------------------------------------------------------------------- +# Happy path — SSE format + headers +# --------------------------------------------------------------------------- + + +def test_stream_events_happy_path_format_and_headers() -> None: + app = _build_app() + client = TestClient(app) + + fake_adapter = _FakeAdapter( + events=[ + ("connected", {"session_id": "ses_test"}), + ( + "message.started", + {"message_id": "msg_a", "role": "assistant"}, + ), + ( + "message.completed", + { + "message_id": "msg_a", + "content": "hi", + "completed_at": "2026-05-07T15:04:05.123Z", + }, + ), + ] + ) + + async def _fake_get_session(prisma_client, *, session_id, created_by): + assert session_id == "ses_test" + assert created_by == CALLER_USER_ID + return _ready_session_row() + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch( + "litellm.managed_agents.endpoints.events.get_adapter", + return_value=fake_adapter, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + with client.stream("GET", "/v2/sessions/ses_test/events") as r: + assert r.status_code == 200 + # Headers per contract §6.6. + assert r.headers["content-type"].startswith("text/event-stream") + assert r.headers["cache-control"] == "no-cache" + assert r.headers["connection"] == "keep-alive" + assert r.headers["x-accel-buffering"] == "no" + body = b"".join(r.iter_bytes()).decode() + + # First event must be `connected` per contract. + assert body.startswith('event: connected\ndata: {"session_id": "ses_test"}\n\n') + # Each event in order, separated by blank line. + assert ( + "event: message.started\n" + 'data: {"message_id": "msg_a", "role": "assistant"}\n\n' + ) in body + assert "event: message.completed\n" in body + assert body.endswith("\n\n") + + +# --------------------------------------------------------------------------- +# 404 — session missing (status code returned BEFORE streaming starts) +# --------------------------------------------------------------------------- + + +def test_stream_events_404_when_session_missing() -> None: + app = _build_app() + client = TestClient(app) + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return None + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + r = client.get("/v2/sessions/ses_missing/events") + + assert r.status_code == 404 + assert "ses_missing" in r.json()["detail"] + assert r.headers["content-type"].startswith("application/json") + + +# --------------------------------------------------------------------------- +# 503 — session is still provisioning (status code BEFORE streaming starts) +# --------------------------------------------------------------------------- + + +def test_stream_events_503_when_session_provisioning() -> None: + app = _build_app() + client = TestClient(app) + + provisioning_row = _ready_session_row() + provisioning_row["status"] = "provisioning" + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return provisioning_row + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + r = client.get("/v2/sessions/ses_test/events") + + assert r.status_code == 503 + assert r.json()["detail"] == "Session not ready" + # Retry-After hint per contract §7 failure modes table. + assert r.headers.get("retry-after") == "5" + + +# --------------------------------------------------------------------------- +# 404 — terminated session is treated as not-found (per contract §7) +# --------------------------------------------------------------------------- + + +def test_stream_events_404_when_session_terminated() -> None: + app = _build_app() + client = TestClient(app) + + terminated_row = _ready_session_row() + terminated_row["status"] = "terminated" + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return terminated_row + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + r = client.get("/v2/sessions/ses_test/events") + + assert r.status_code == 404 + + +# --------------------------------------------------------------------------- +# Mid-stream SandboxUnreachableError → inline `error` event (status is 200) +# --------------------------------------------------------------------------- + + +def test_stream_events_unreachable_after_stream_started_emits_error_event() -> None: + app = _build_app() + client = TestClient(app) + + fake_adapter = _FakeAdapter( + events=[ + ("connected", {"session_id": "ses_test"}), + ("message.started", {"message_id": "msg_a", "role": "assistant"}), + ], + raise_after=2, # raise immediately after both events stream. + ) + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return _ready_session_row() + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch( + "litellm.managed_agents.endpoints.events.get_adapter", + return_value=fake_adapter, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + with client.stream("GET", "/v2/sessions/ses_test/events") as r: + # Status was already 200 by the time the adapter raised. + assert r.status_code == 200 + body = b"".join(r.iter_bytes()).decode() + + assert "event: connected\n" in body + assert "event: message.started\n" in body + # Error becomes an inline SSE event with the canonical message. + assert 'event: error\ndata: {"error": "Sandbox unreachable"}\n\n' in body + + +# --------------------------------------------------------------------------- +# Pre-stream SandboxUnreachableError → still 200 with error event +# (the adapter raises on the very first iteration, before any events stream). +# --------------------------------------------------------------------------- + + +def test_stream_events_unreachable_on_first_iteration_emits_error_only() -> None: + app = _build_app() + client = TestClient(app) + + # raise_after=0 means: raise before yielding anything. + fake_adapter = _FakeAdapter(events=[], raise_after=0) + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return _ready_session_row() + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch( + "litellm.managed_agents.endpoints.events.get_adapter", + return_value=fake_adapter, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + with client.stream("GET", "/v2/sessions/ses_test/events") as r: + # StreamingResponse always returns 200 once it's been constructed — + # the streaming generator hasn't run yet at that point. + assert r.status_code == 200 + body = b"".join(r.iter_bytes()).decode() + + assert body == 'event: error\ndata: {"error": "Sandbox unreachable"}\n\n' + + +# --------------------------------------------------------------------------- +# 500 — corrupt session row (missing sandbox state) +# --------------------------------------------------------------------------- + + +def test_stream_events_500_when_session_row_missing_sandbox_state() -> None: + app = _build_app() + client = TestClient(app) + + bad_row = _ready_session_row() + bad_row["sandbox_url"] = None # corrupt: ready session must have a url. + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return bad_row + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + r = client.get("/v2/sessions/ses_test/events") + + assert r.status_code == 500 + + +# --------------------------------------------------------------------------- +# JSON-string sandbox_metadata is parsed defensively +# --------------------------------------------------------------------------- + + +def test_stream_events_handles_json_string_sandbox_metadata() -> None: + """Some Prisma clients return JSON columns as strings — handle that.""" + app = _build_app() + client = TestClient(app) + + row = _ready_session_row() + row["sandbox_metadata"] = '{"opencode_session_id": "oc_sid_abc"}' + + fake_adapter = _FakeAdapter(events=[("connected", {"session_id": "ses_test"})]) + + async def _fake_get_session(prisma_client, *, session_id, created_by): + return row + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch( + "litellm.managed_agents.endpoints.events.get_adapter", + return_value=fake_adapter, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + with client.stream("GET", "/v2/sessions/ses_test/events") as r: + assert r.status_code == 200 + body = b"".join(r.iter_bytes()).decode() + + assert "event: connected\n" in body + + +# --------------------------------------------------------------------------- +# 500 — prisma client not initialized +# --------------------------------------------------------------------------- + + +def test_stream_events_500_when_prisma_client_not_initialized() -> None: + app = _build_app() + client = TestClient(app) + + with patch("litellm.proxy.proxy_server.prisma_client", None): + r = client.get("/v2/sessions/ses_test/events") + + assert r.status_code == 500 + + +# --------------------------------------------------------------------------- +# Caller scoping is enforced — _load_ready_session passes user_id through. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_stream_events_scopes_by_caller_user_id() -> None: + """The handler must pass `user_api_key_dict.user_id` as `created_by`. + + Without scoping, callers could read other users' sessions. + """ + app = _build_app() + client = TestClient(app) + + captured: Dict[str, Any] = {} + + async def _fake_get_session(prisma_client, *, session_id, created_by): + captured["session_id"] = session_id + captured["created_by"] = created_by + return None # 404 — but we only care that scoping was used. + + with ( + patch( + "litellm.managed_agents.endpoints.events.get_session", + side_effect=_fake_get_session, + ), + patch("litellm.proxy.proxy_server.prisma_client", object()), + ): + r = client.get("/v2/sessions/ses_test/events") + + assert r.status_code == 404 + assert captured["session_id"] == "ses_test" + assert captured["created_by"] == CALLER_USER_ID