test(v2 managed agents): unit tests for GET /v2/sessions/:id

tests/test_litellm/managed_agents/test_sessions.py

@@ -0,0 +1,323 @@
+"""Unit tests for `GET /v2/sessions/:id` (LIT-2923).
+
+Approach:
+- Build a minimal FastAPI app with only `sessions.router` mounted (Wave 3
+  owns proxy_server registration; the test should not depend on it).
+- Stub `litellm.proxy.proxy_server.prisma_client` with a MagicMock whose
+  `db.litellm_managedagentsession.find_first` returns a configurable row.
+- Override `user_api_key_auth` via `app.dependency_overrides` to simulate
+  different callers (and test the no-auth path with no override).
+
+Contract-critical assertions:
+- Response NEVER contains `sandbox_url` or `sandbox_metadata` — these are
+  internal-only (contract §6.2 / §6.3).
+- A session belonging to a different `created_by` returns 404, not 403,
+  to avoid leaking session existence.
+"""
+
+import os
+import sys
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from litellm.managed_agents.endpoints import sessions as sessions_module
+from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_session_row(
+    *,
+    session_id: str = "ses_abc",
+    agent_id: str = "agt_xyz",
+    created_by: Optional[str] = "user_a",
+    status: str = "ready",
+    sandbox_type: str = "opencode",
+    sandbox_size: str = "small",
+    sandbox_timeout_minutes: int = 60,
+    sandbox_idle_timeout_minutes: int = 10,
+    sandbox_image: Optional[str] = "litellm/opencode:latest",
+    sandbox_url: Optional[str] = "http://127.0.0.1:1234",
+    sandbox_metadata: Optional[Dict[str, Any]] = None,
+    repos: Optional[list] = None,
+    env_vars: Optional[Dict[str, Any]] = None,
+    terminated_at: Optional[datetime] = None,
+) -> MagicMock:
+    """Return a MagicMock row that mirrors a Prisma row shape for the
+    LiteLLM_ManagedAgentSession table."""
+    if sandbox_metadata is None:
+        sandbox_metadata = {"opencode_session_id": "oc_sid_xxx"}
+    if repos is None:
+        repos = []
+    if env_vars is None:
+        env_vars = {}
+
+    now = datetime.now(timezone.utc)
+    row = MagicMock()
+    # Make `model_dump()` return the dict the helper actually consumes.
+    row.model_dump.return_value = {
+        "id": session_id,
+        "agent_id": agent_id,
+        "sandbox_type": sandbox_type,
+        "sandbox_size": sandbox_size,
+        "sandbox_timeout_minutes": sandbox_timeout_minutes,
+        "sandbox_idle_timeout_minutes": sandbox_idle_timeout_minutes,
+        "sandbox_image": sandbox_image,
+        "sandbox_url": sandbox_url,
+        "sandbox_metadata": sandbox_metadata,
+        "status": status,
+        "repos": repos,
+        "env_vars": env_vars,
+        "created_by": created_by,
+        "created_at": now,
+        "updated_at": now,
+        "terminated_at": terminated_at,
+    }
+    return row
+
+
+def _make_prisma(find_first_return: Any = None) -> MagicMock:
+    """Build a MagicMock prisma_client with `db.litellm_managedagentsession.find_first`."""
+    prisma = MagicMock()
+    prisma.db = MagicMock()
+    prisma.db.litellm_managedagentsession = MagicMock()
+    prisma.db.litellm_managedagentsession.find_first = AsyncMock(
+        return_value=find_first_return
+    )
+    return prisma
+
+
+def _build_app(auth: Optional[UserAPIKeyAuth] = None) -> FastAPI:
+    """Construct a minimal app with the sessions router mounted.
+
+    If `auth` is provided, override the auth dependency to return it. If
+    omitted, the real `user_api_key_auth` runs (used to test the 401 path).
+    """
+    app = FastAPI()
+    app.include_router(sessions_module.router)
+    if auth is not None:
+        app.dependency_overrides[user_api_key_auth] = lambda: auth
+    return app
+
+
+def _user_auth(user_id: str = "user_a") -> UserAPIKeyAuth:
+    return UserAPIKeyAuth(api_key="sk-test", user_id=user_id)
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_get_session_returns_row_for_owner(monkeypatch):
+    """200 with full SessionRow shape when the session belongs to the caller.
+
+    Contract-critical: response must NOT include `sandbox_url` or
+    `sandbox_metadata` — these are internal-only.
+    """
+    row = _make_session_row(session_id="ses_abc", created_by="user_a")
+    prisma = _make_prisma(find_first_return=row)
+
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    app = _build_app(auth=_user_auth("user_a"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_abc")
+    assert resp.status_code == 200, resp.text
+    body = resp.json()
+
+    # Public fields present.
+    assert body["id"] == "ses_abc"
+    assert body["agent_id"] == "agt_xyz"
+    assert body["status"] == "ready"
+    assert body["created_by"] == "user_a"
+    assert body["terminated_at"] is None
+
+    # Sandbox sub-object correctly reassembled from flat columns.
+    assert body["sandbox"]["type"] == "opencode"
+    assert body["sandbox"]["size"] == "small"
+    assert body["sandbox"]["timeout_minutes"] == 60
+    assert body["sandbox"]["idle_timeout_minutes"] == 10
+    assert body["sandbox"]["image"] == "litellm/opencode:latest"
+
+    # Internal-only fields MUST be stripped.
+    assert "sandbox_url" not in body, "sandbox_url must NOT be returned"
+    assert "sandbox_metadata" not in body, "sandbox_metadata must NOT be returned"
+    # Belt-and-suspenders: also not nested under `sandbox`.
+    assert "sandbox_url" not in body["sandbox"]
+    assert "sandbox_metadata" not in body["sandbox"]
+
+    # The DB lookup was scoped by created_by.
+    prisma.db.litellm_managedagentsession.find_first.assert_awaited_once_with(
+        where={"id": "ses_abc", "created_by": "user_a"}
+    )
+
+
+def test_get_session_returns_repos_array(monkeypatch):
+    """`repos` from the row is mapped onto the response."""
+    row = _make_session_row(
+        session_id="ses_with_repos",
+        created_by="user_a",
+        repos=[
+            {
+                "url": "https://github.com/org/repo-a",
+                "starting_ref": "main",
+                "checked_out_sha": "abc123",
+            }
+        ],
+    )
+    prisma = _make_prisma(find_first_return=row)
+
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    app = _build_app(auth=_user_auth("user_a"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_with_repos")
+    assert resp.status_code == 200, resp.text
+    body = resp.json()
+    assert len(body["repos"]) == 1
+    assert body["repos"][0]["url"] == "https://github.com/org/repo-a"
+    assert body["repos"][0]["starting_ref"] == "main"
+    assert body["repos"][0]["checked_out_sha"] == "abc123"
+
+
+# ---------------------------------------------------------------------------
+# 404 paths
+# ---------------------------------------------------------------------------
+
+
+def test_get_session_returns_404_when_missing(monkeypatch):
+    """If the session does not exist for ANY caller, return 404."""
+    prisma = _make_prisma(find_first_return=None)
+
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    app = _build_app(auth=_user_auth("user_a"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_does_not_exist")
+    assert resp.status_code == 404
+    assert resp.json()["detail"] == "Session ses_does_not_exist not found"
+
+
+def test_get_session_returns_404_for_other_users_session(monkeypatch):
+    """Session exists but belongs to another user → 404, NOT 403.
+
+    The DB query is scoped by `created_by` — so when caller=user_b queries
+    a session owned by user_a, the row is filtered out at the DB level and
+    `find_first` returns None. The handler returns 404 to avoid leaking
+    existence (contract §5).
+    """
+    # Simulate: row exists for user_a, but find_first(where={created_by=user_b})
+    # returns None (Prisma scoping behavior).
+    prisma = _make_prisma(find_first_return=None)
+
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    # Caller is user_b but the session row is owned by user_a.
+    app = _build_app(auth=_user_auth("user_b"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_owned_by_user_a")
+    assert resp.status_code == 404, "Must 404 (not 403) to avoid leaking existence"
+    assert resp.json()["detail"] == "Session ses_owned_by_user_a not found"
+
+    # Verify the scoping: query passed user_b, not user_a.
+    prisma.db.litellm_managedagentsession.find_first.assert_awaited_once_with(
+        where={"id": "ses_owned_by_user_a", "created_by": "user_b"}
+    )
+
+
+# ---------------------------------------------------------------------------
+# Auth + DB-not-connected
+# ---------------------------------------------------------------------------
+
+
+def test_get_session_no_auth_returns_401():
+    """When `user_api_key_auth` rejects (no/invalid bearer token) → 401.
+
+    We can't run the real `user_api_key_auth` in a unit test (it needs a
+    full proxy DB + master-key bootstrap), so we install a dependency
+    override that simulates the auth-failure path: raise HTTPException(401).
+    The handler should never be reached — FastAPI short-circuits at the
+    auth dep.
+    """
+    from fastapi import HTTPException
+
+    def _reject_auth() -> UserAPIKeyAuth:
+        raise HTTPException(
+            status_code=401, detail="Authentication Error, Invalid proxy server token"
+        )
+
+    app = FastAPI()
+    app.include_router(sessions_module.router)
+    app.dependency_overrides[user_api_key_auth] = _reject_auth
+
+    client = TestClient(app)
+    resp = client.get("/v2/sessions/ses_anything")
+    assert resp.status_code == 401, resp.text
+
+
+def test_get_session_returns_500_when_db_not_connected(monkeypatch):
+    """If `prisma_client` is None, return 500 with the canonical error."""
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", None)
+
+    app = _build_app(auth=_user_auth("user_a"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_anything")
+    assert resp.status_code == 500
+    detail = resp.json()["detail"]
+    assert detail["error"] == CommonProxyErrors.db_not_connected_error.value
+
+
+# ---------------------------------------------------------------------------
+# Sanity: status passthrough
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("status", ["provisioning", "ready", "terminated", "error"])
+def test_get_session_passes_through_all_statuses(monkeypatch, status):
+    """The handler does NOT filter by status — all four enum values pass through.
+
+    (Per contract §6.3 the read endpoint just returns the row; status-based
+    fail-closed behavior lives on the messaging endpoints, not here.)
+    """
+    row = _make_session_row(
+        session_id="ses_status_test", created_by="user_a", status=status
+    )
+    prisma = _make_prisma(find_first_return=row)
+
+    import litellm.proxy.proxy_server as ps
+
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    app = _build_app(auth=_user_auth("user_a"))
+    client = TestClient(app)
+
+    resp = client.get("/v2/sessions/ses_status_test")
+    assert resp.status_code == 200, resp.text
+    assert resp.json()["status"] == status