From e4432be886c44729c8cd8cf1d9072c91a77121f5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 7 May 2026 10:06:56 -0700 Subject: [PATCH] test(v2 managed agents): unit tests for messages endpoints --- .../managed_agents/test_messages.py | 558 ++++++++++++++++++ 1 file changed, 558 insertions(+) create mode 100644 tests/test_litellm/managed_agents/test_messages.py diff --git a/tests/test_litellm/managed_agents/test_messages.py b/tests/test_litellm/managed_agents/test_messages.py new file mode 100644 index 00000000000..1211ed62862 --- /dev/null +++ b/tests/test_litellm/managed_agents/test_messages.py @@ -0,0 +1,558 @@ +"""Unit tests for the messaging endpoints (LIT-2920). + +Covers ``POST /v2/sessions/:id/messages`` and ``GET /v2/sessions/:id/messages``. + +Auth pattern: FastAPI ``app.dependency_overrides`` for ``user_api_key_auth``. +Adapter pattern: monkeypatch ``get_adapter`` so the handler resolves to a +mock without touching the network. +DB pattern: stub ``prisma_client`` with in-memory tables (mirrors the +Wave-2b agent test scaffolding). + +The router is mounted onto the proxy ``app`` inside a fixture so this test +file does not require Wave 3's wiring in ``proxy_server.py`` to be present. +""" + +import os +import sys +import types +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +import litellm.proxy.proxy_server as ps +from litellm.managed_agents.adapters.base import ( + SandboxBadGatewayError, + SandboxUnreachableError, +) +from litellm.managed_agents.endpoints.messages import router as messages_router +from litellm.managed_agents.types import MessageRow +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app + +sys.path.insert(0, os.path.abspath("../../../")) + + +# --------------------------------------------------------------------------- +# Shared fixtures and helpers +# --------------------------------------------------------------------------- + +USER_ID = "user_xyz" +SESSION_ID = "ses_test" +AGENT_ID = "agt_test" +SANDBOX_URL = "http://127.0.0.1:1234" +OC_SID = "oc_sid_xxx" +DEFAULT_AGENT_MODEL = "anthropic/claude-opus-4" + + +def _ensure_router_mounted() -> None: + """Mount the messages router onto ``app`` exactly once. + + Wave 3 owns the real registration in ``proxy_server.py``. For tests we + mount it here so we can hit the endpoints via ``TestClient``. + """ + paths = {getattr(r, "path", None) for r in app.router.routes} + if "/v2/sessions/{session_id}/messages" not in paths: + app.include_router(messages_router) + + +def _make_session_row( + *, + session_id: str = SESSION_ID, + agent_id: str = AGENT_ID, + status: str = "ready", + sandbox_url: Optional[str] = SANDBOX_URL, + opencode_session_id: Optional[str] = OC_SID, + sandbox_type: str = "opencode", + created_by: str = USER_ID, +) -> Dict[str, Any]: + metadata: Dict[str, Any] = {} + if opencode_session_id is not None: + metadata["opencode_session_id"] = opencode_session_id + return { + "id": session_id, + "agent_id": agent_id, + "sandbox_type": sandbox_type, + "sandbox_size": "small", + "sandbox_timeout_minutes": 60, + "sandbox_idle_timeout_minutes": 10, + "sandbox_image": None, + "sandbox_url": sandbox_url, + "sandbox_metadata": metadata, + "status": status, + "repos": [], + "env_vars": {}, + "created_by": created_by, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "terminated_at": None, + } + + +def _make_agent_row( + *, + agent_id: str = AGENT_ID, + model: Optional[str] = DEFAULT_AGENT_MODEL, + created_by: str = USER_ID, +) -> Dict[str, Any]: + config: Dict[str, Any] = { + "system_prompt": "You are a code reviewer.", + "tools": ["read"], + "litellm_api_key": "sk-test", + "litellm_base_url": "http://localhost:4000", + } + if model is not None: + config["model"] = model + return { + "id": agent_id, + "name": "code-reviewer", + "config": config, + "created_by": created_by, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + + +class _FakeSessionTable: + def __init__(self, rows: Optional[List[Dict[str, Any]]] = None) -> None: + self.rows: List[Dict[str, Any]] = list(rows or []) + self.find_first = AsyncMock(side_effect=self._find_first) + + async def _find_first( + self, *, where: Dict[str, Any] + ) -> Optional[types.SimpleNamespace]: + for row in self.rows: + if all(row.get(k) == v for k, v in where.items()): + return types.SimpleNamespace(model_dump=lambda r=row: dict(r)) + return None + + +class _FakeAgentTable: + def __init__(self, rows: Optional[List[Dict[str, Any]]] = None) -> None: + self.rows: List[Dict[str, Any]] = list(rows or []) + self.find_first = AsyncMock(side_effect=self._find_first) + + async def _find_first( + self, *, where: Dict[str, Any] + ) -> Optional[types.SimpleNamespace]: + for row in self.rows: + if all(row.get(k) == v for k, v in where.items()): + return types.SimpleNamespace(model_dump=lambda r=row: dict(r)) + return None + + +def _build_fake_prisma( + sessions: Optional[List[Dict[str, Any]]] = None, + agents: Optional[List[Dict[str, Any]]] = None, +) -> MagicMock: + fake_db = types.SimpleNamespace( + litellm_managedagentsession=_FakeSessionTable(sessions), + litellm_managedagent=_FakeAgentTable(agents), + ) + fake_prisma = MagicMock() + fake_prisma.db = fake_db + return fake_prisma + + +def _override_auth(user_id: str = USER_ID) -> UserAPIKeyAuth: + user_key = UserAPIKeyAuth(api_key="sk-test", user_id=user_id) + app.dependency_overrides[user_api_key_auth] = lambda: user_key + return user_key + + +@pytest.fixture(autouse=True) +def _mount_router_and_clear_overrides(): + _ensure_router_mounted() + yield + app.dependency_overrides.clear() + + +@pytest.fixture +def client() -> TestClient: + return TestClient(app) + + +@pytest.fixture +def mock_adapter(monkeypatch): + """Monkeypatch ``get_adapter`` to return a configurable AsyncMock. + + The same adapter instance is returned by ``get_adapter`` for every + call inside the handler — both ``send_message`` and ``list_messages`` + are exposed as AsyncMocks so each test can configure side effects / + return values independently. + """ + adapter = MagicMock() + adapter.send_message = AsyncMock(return_value=None) + adapter.list_messages = AsyncMock(return_value=[]) + monkeypatch.setattr( + "litellm.managed_agents.endpoints.messages.get_adapter", + lambda sandbox_type: adapter, + ) + return adapter + + +# --------------------------------------------------------------------------- +# POST /v2/sessions/:id/messages +# --------------------------------------------------------------------------- + + +class TestSendMessage: + def test_ready_session_returns_202_with_user_message_row( + self, monkeypatch, client, mock_adapter + ): + """Happy path: ready session → 202, adapter.send_message called with + the right args, and the response is a synthesized user MessageRow. + """ + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row()], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello", "model": "anthropic/claude-opus-4"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 202 + body = resp.json() + assert body["session_id"] == SESSION_ID + assert body["role"] == "user" + assert body["content"] == "hello" + assert body["status"] == "in_progress" + assert body["model"] == "anthropic/claude-opus-4" + assert body["completed_at"] is None + assert body["id"].startswith("msg_") + + mock_adapter.send_message.assert_awaited_once_with( + SANDBOX_URL, + OC_SID, + "hello", + "anthropic/claude-opus-4", + ) + + def test_provisioning_session_returns_503(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row(status="provisioning")], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 503 + assert resp.json()["detail"] == "Session not ready" + # contract §7: provisioning carries Retry-After: 5 + assert resp.headers.get("retry-after") == "5" + mock_adapter.send_message.assert_not_awaited() + + def test_terminated_session_returns_404(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row(status="terminated")], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 404 + assert SESSION_ID in resp.json()["detail"] + mock_adapter.send_message.assert_not_awaited() + + def test_missing_session_returns_404(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma(sessions=[]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 404 + mock_adapter.send_message.assert_not_awaited() + + def test_sandbox_unreachable_returns_504(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row()], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + mock_adapter.send_message.side_effect = SandboxUnreachableError("boom") + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 504 + assert resp.json()["detail"] == {"error": "Sandbox unreachable"} + + def test_sandbox_bad_gateway_returns_502(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row()], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + mock_adapter.send_message.side_effect = SandboxBadGatewayError("malformed") + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 502 + assert resp.json()["detail"] == {"error": "Bad gateway"} + + def test_missing_content_returns_422(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row()], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 422 + mock_adapter.send_message.assert_not_awaited() + + def test_optional_model_falls_back_to_agent_config( + self, monkeypatch, client, mock_adapter + ): + """Omitting `model` in the request → handler resolves it from + agent.config["model"] before calling the adapter. + """ + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row(model="anthropic/claude-sonnet-4")], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 202 + assert resp.json()["model"] == "anthropic/claude-sonnet-4" + mock_adapter.send_message.assert_awaited_once_with( + SANDBOX_URL, + OC_SID, + "hello", + "anthropic/claude-sonnet-4", + ) + + def test_missing_auth_returns_401(self, monkeypatch, client, mock_adapter): + """Without an ``Authorization`` header and a master key set, real auth + rejects. The proxy's ``user_api_key_auth`` is permissive when + ``master_key is None`` — set a master key + drop the override so the + real chain runs. + """ + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row()], + agents=[_make_agent_row()], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + monkeypatch.setattr(ps, "master_key", "sk-test-master") + # Do NOT override auth — exercise the real dependency. + app.dependency_overrides.pop(user_api_key_auth, None) + + resp = client.post( + f"/v2/sessions/{SESSION_ID}/messages", + json={"content": "hello"}, + ) + + # The real auth chain rejects unauthenticated calls. Proxy maps the + # raised exception to ProxyException → 401 in most versions, but some + # 4xx is fine — what matters is that it's an auth failure, NOT a + # successful 2xx, and the adapter was never invoked. + assert resp.status_code in (400, 401, 403), resp.text + mock_adapter.send_message.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# GET /v2/sessions/:id/messages +# --------------------------------------------------------------------------- + + +def _make_message_row( + *, + msg_id: str, + role: str, + content: str, + status: str = "completed", +) -> MessageRow: + return MessageRow( + id=msg_id, + session_id=SESSION_ID, + role=role, # type: ignore[arg-type] + content=content, + model=DEFAULT_AGENT_MODEL, + status=status, # type: ignore[arg-type] + created_at=datetime.now(timezone.utc), + completed_at=datetime.now(timezone.utc) if status == "completed" else None, + ) + + +class TestListMessages: + def test_returns_messages_wrapped_in_message_list_shape( + self, monkeypatch, client, mock_adapter + ): + fake_prisma = _build_fake_prisma(sessions=[_make_session_row()]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + msgs = [ + _make_message_row(msg_id="msg_1", role="user", content="hi"), + _make_message_row(msg_id="msg_2", role="assistant", content="hello!"), + ] + mock_adapter.list_messages.return_value = msgs + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 200 + body = resp.json() + assert "data" in body + assert len(body["data"]) == 2 + assert body["data"][0]["id"] == "msg_1" + assert body["data"][1]["id"] == "msg_2" + assert body["next_cursor"] is None + assert body["has_more"] is False + + mock_adapter.list_messages.assert_awaited_once_with( + SANDBOX_URL, + OC_SID, + SESSION_ID, + 50, + ) + + def test_role_filter_narrows_results(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma(sessions=[_make_session_row()]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + mock_adapter.list_messages.return_value = [ + _make_message_row(msg_id="msg_u1", role="user", content="hi"), + _make_message_row(msg_id="msg_a1", role="assistant", content="hello!"), + _make_message_row(msg_id="msg_u2", role="user", content="more"), + ] + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages?role=user", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 200 + body = resp.json() + assert len(body["data"]) == 2 + assert {m["id"] for m in body["data"]} == {"msg_u1", "msg_u2"} + assert all(m["role"] == "user" for m in body["data"]) + + def test_empty_session_returns_empty_data_array( + self, monkeypatch, client, mock_adapter + ): + fake_prisma = _build_fake_prisma(sessions=[_make_session_row()]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + mock_adapter.list_messages.return_value = [] + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["data"] == [] + assert body["next_cursor"] is None + assert body["has_more"] is False + + def test_provisioning_session_returns_503(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma( + sessions=[_make_session_row(status="provisioning")], + ) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 503 + assert resp.headers.get("retry-after") == "5" + mock_adapter.list_messages.assert_not_awaited() + + def test_missing_session_returns_404(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma(sessions=[]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 404 + mock_adapter.list_messages.assert_not_awaited() + + def test_sandbox_unreachable_returns_504(self, monkeypatch, client, mock_adapter): + fake_prisma = _build_fake_prisma(sessions=[_make_session_row()]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + _override_auth() + + mock_adapter.list_messages.side_effect = SandboxUnreachableError("boom") + + resp = client.get( + f"/v2/sessions/{SESSION_ID}/messages", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert resp.status_code == 504 + assert resp.json()["detail"] == {"error": "Sandbox unreachable"} + + def test_missing_auth_returns_401(self, monkeypatch, client, mock_adapter): + """Without an ``Authorization`` header and a master key set, real auth + rejects. See note on the corresponding ``TestSendMessage`` test. + """ + fake_prisma = _build_fake_prisma(sessions=[_make_session_row()]) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + monkeypatch.setattr(ps, "master_key", "sk-test-master") + # Do NOT override auth — exercise the real dependency. + app.dependency_overrides.pop(user_api_key_auth, None) + + resp = client.get(f"/v2/sessions/{SESSION_ID}/messages") + + assert resp.status_code in (400, 401, 403), resp.text + mock_adapter.list_messages.assert_not_awaited()