diff --git a/litellm/proxy/managed_agents_endpoints/endpoints.py b/litellm/proxy/managed_agents_endpoints/endpoints.py index 4b8dce2cebb..29c28ce36e6 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, Dict, List import boto3 from fastapi import APIRouter, Depends, HTTPException @@ -63,6 +63,31 @@ def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None: raise HTTPException(status_code=403, detail="admin role required") +def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + + +def _assert_owner_or_admin( + user_api_key_dict: UserAPIKeyAuth, + created_by: object, + resource_kind: str, + resource_id: str, +) -> None: + """Reject callers that neither own the resource nor are PROXY_ADMIN. + + A 404 is returned for non-owners (rather than 403) so that resource IDs + cannot be enumerated by unauthorized callers. + """ + if _is_admin(user_api_key_dict): + return + caller = user_api_key_dict.user_id + if caller is not None and created_by == caller: + return + raise HTTPException( + status_code=404, detail=f"{resource_kind} '{resource_id}' not found" + ) + + @router.get("/dockerfiles", response_model=List[DockerfileOut]) async def list_dockerfiles_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -191,6 +216,16 @@ async def create_sandbox_template( return _template_row_to_out(updated_row) +def _template_visible_to(row, user_api_key_dict: UserAPIKeyAuth) -> bool: + """A private template is only visible to its creator and admins.""" + if row.visibility == "public": + return True + if _is_admin(user_api_key_dict): + return True + caller = user_api_key_dict.user_id + return caller is not None and row.created_by == caller + + @router.get("/sandbox-templates", response_model=List[TemplateOut]) async def list_sandbox_templates( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -200,7 +235,18 @@ async def list_sandbox_templates( if prisma_client is None: raise HTTPException(status_code=500, detail="prisma client not available") - rows = await prisma_client.db.litellm_managedagentsandboxtemplatetable.find_many() + if _is_admin(user_api_key_dict): + where_clause = {} + else: + caller = user_api_key_dict.user_id + visibility_filters: List[Dict[str, Any]] = [{"visibility": "public"}] + if caller is not None: + visibility_filters.append({"created_by": caller}) + where_clause = {"OR": visibility_filters} + + rows = await prisma_client.db.litellm_managedagentsandboxtemplatetable.find_many( + where=where_clause + ) return [_template_row_to_out(row) for row in rows] @@ -217,7 +263,7 @@ async def get_sandbox_template( row = await prisma_client.db.litellm_managedagentsandboxtemplatetable.find_unique( where={"template_id": template_id} ) - if row is None: + if row is None or not _template_visible_to(row, user_api_key_dict): raise HTTPException( status_code=404, detail=f"template '{template_id}' not found" ) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py index 12c0b3f7ab2..393d5801623 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py @@ -5,7 +5,11 @@ from typing import List from fastapi import Depends, HTTPException from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.managed_agents_endpoints.endpoints import router +from litellm.proxy.managed_agents_endpoints.endpoints import ( + _assert_owner_or_admin, + _is_admin, + router, +) from litellm.proxy.managed_agents_endpoints.git_validation import ( decrypt_git_token, validate_repo_branch, @@ -81,8 +85,12 @@ async def list_agents( if prisma_client is None: raise HTTPException(status_code=500, detail="prisma client not available") + where: dict = {} + if not _is_admin(user_api_key_dict) and user_api_key_dict.user_id is not None: + where["created_by"] = user_api_key_dict.user_id + rows = await prisma_client.db.litellm_managedagenttable.find_many( - order={"created_at": "desc"} + where=where, order={"created_at": "desc"} ) return [_agent_row_to_out(row) for row in rows] @@ -103,4 +111,6 @@ async def get_agent( if row is None: raise HTTPException(status_code=404, detail=f"agent '{agent_id}' not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "agent", agent_id) + return _agent_row_to_out(row) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_passthrough.py b/litellm/proxy/managed_agents_endpoints/endpoints_passthrough.py index 2ef85a8a85a..ec04da139f5 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_passthrough.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_passthrough.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime, timezone -from typing import Any, Dict +from typing import Any, Dict, Optional import httpx from fastapi import Depends, HTTPException, Request @@ -8,7 +8,10 @@ from fastapi.responses import StreamingResponse from starlette.background import BackgroundTask from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.managed_agents_endpoints.endpoints import router +from litellm.proxy.managed_agents_endpoints.endpoints import ( + _assert_owner_or_admin, + router, +) from litellm.proxy.managed_agents_endpoints.harness_client import ( expand_message, harness_send_message, @@ -26,6 +29,16 @@ HOP_BY_HOP = { "upgrade", } +# Headers that authenticate the caller to the proxy. Must NOT be forwarded +# to the sandbox container (it runs untrusted user code and could exfiltrate +# the caller's LiteLLM API key or session cookie). +SENSITIVE_HEADERS = { + "authorization", + "cookie", + "x-api-key", + "x-litellm-api-key", +} + async def _touch_session(prisma_client, session_id: str) -> None: try: @@ -37,10 +50,34 @@ async def _touch_session(prisma_client, session_id: str) -> None: pass -def _build_http_client() -> httpx.AsyncClient: - return httpx.AsyncClient( - timeout=httpx.Timeout(connect=10, read=None, write=None, pool=10) - ) +# Module-level shared client. Reused across all passthrough requests so we +# get HTTP connection-pool keepalive (one TCP+TLS handshake per sandbox host +# instead of one per chat message). Lazily initialized inside the running +# event loop. Closed on proxy shutdown via close_passthrough_http_client. +_HTTP_CLIENT: Optional[httpx.AsyncClient] = None +_HTTP_CLIENT_LOCK = asyncio.Lock() + + +async def _get_http_client() -> httpx.AsyncClient: + global _HTTP_CLIENT + if _HTTP_CLIENT is None or _HTTP_CLIENT.is_closed: + async with _HTTP_CLIENT_LOCK: + if _HTTP_CLIENT is None or _HTTP_CLIENT.is_closed: + _HTTP_CLIENT = httpx.AsyncClient( + timeout=httpx.Timeout(connect=10, read=None, write=None, pool=10), + limits=httpx.Limits( + max_connections=200, max_keepalive_connections=50 + ), + ) + return _HTTP_CLIENT + + +async def close_passthrough_http_client() -> None: + """Called from proxy shutdown to drain pooled connections.""" + global _HTTP_CLIENT + if _HTTP_CLIENT is not None and not _HTTP_CLIENT.is_closed: + await _HTTP_CLIENT.aclose() + _HTTP_CLIENT = None @router.post("/sessions/{session_id}/message") @@ -60,6 +97,7 @@ async def session_message( ) if row is None: raise HTTPException(status_code=404, detail="session not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "session", session_id) if ( row.status != "ready" or row.sandbox_url is None @@ -72,22 +110,19 @@ async def session_message( parts = expand_message(body.text, body.parts) - client = _build_http_client() + client = await _get_http_client() try: - try: - result = await harness_send_message( - row.sandbox_url, - row.harness_session_id, - client, - model=row.agent.model, - parts=parts, - ) - except httpx.HTTPStatusError as e: - raise HTTPException(e.response.status_code, e.response.text) - except httpx.HTTPError as e: - raise HTTPException(502, f"upstream error: {e}") - finally: - asyncio.create_task(client.aclose()) + result = await harness_send_message( + row.sandbox_url, + row.harness_session_id, + client, + model=row.agent.model, + parts=parts, + ) + except httpx.HTTPStatusError as e: + raise HTTPException(e.response.status_code, e.response.text) + except httpx.HTTPError as e: + raise HTTPException(502, f"upstream error: {e}") asyncio.create_task(_touch_session(prisma_client, session_id)) return result @@ -108,34 +143,30 @@ async def session_events( ) if row is None: raise HTTPException(status_code=404, detail="session not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "session", session_id) if row.status != "ready" or row.sandbox_url is None: raise HTTPException( status_code=409, detail=f"session not ready (status={row.status})", ) - client = _build_http_client() + client = await _get_http_client() req = client.build_request("GET", f"{row.sandbox_url}/event", timeout=None) try: upstream = await client.send(req, stream=True) except httpx.HTTPError as e: - await client.aclose() raise HTTPException(502, f"upstream error: {e}") resp_headers = { k: v for k, v in upstream.headers.items() if k.lower() not in HOP_BY_HOP } - async def _close() -> None: - await upstream.aclose() - await client.aclose() - return StreamingResponse( upstream.aiter_raw(), status_code=upstream.status_code, headers=resp_headers, media_type=upstream.headers.get("content-type", "text/event-stream"), - background=BackgroundTask(_close), + background=BackgroundTask(upstream.aclose), ) @@ -159,6 +190,7 @@ async def session_raw_proxy( ) if row is None: raise HTTPException(status_code=404, detail="session not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "session", session_id) if row.status != "ready" or row.sandbox_url is None: raise HTTPException( status_code=409, @@ -169,11 +201,13 @@ async def session_raw_proxy( fwd_headers = { k: v for k, v in request.headers.items() - if k.lower() not in HOP_BY_HOP and k.lower() != "host" + if k.lower() not in HOP_BY_HOP + and k.lower() not in SENSITIVE_HEADERS + and k.lower() != "host" } body = await request.body() - client = _build_http_client() + client = await _get_http_client() req = client.build_request( method=request.method, url=target, @@ -184,7 +218,6 @@ async def session_raw_proxy( try: upstream = await client.send(req, stream=True) except httpx.HTTPError as e: - await client.aclose() raise HTTPException(502, f"upstream error: {e}") asyncio.create_task(_touch_session(prisma_client, session_id)) @@ -193,13 +226,9 @@ async def session_raw_proxy( k: v for k, v in upstream.headers.items() if k.lower() not in HOP_BY_HOP } - async def _close() -> None: - await upstream.aclose() - await client.aclose() - return StreamingResponse( upstream.aiter_raw(), status_code=upstream.status_code, headers=resp_headers, - background=BackgroundTask(_close), + background=BackgroundTask(upstream.aclose), ) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py index 580142cbbcb..d4a28f348a0 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py @@ -9,7 +9,11 @@ from fastapi import Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.managed_agents_endpoints import config_loader as _config_loader -from litellm.proxy.managed_agents_endpoints.endpoints import router +from litellm.proxy.managed_agents_endpoints.endpoints import ( + _assert_owner_or_admin, + _is_admin, + router, +) from litellm.proxy.managed_agents_endpoints.fargate.bootstrap import ( bootstrap_shared_infra, ) @@ -210,6 +214,10 @@ async def create_session( public_ip = await asyncio.to_thread( wait_running_get_ip_sync, region, cluster, task_arn, 300 ) + # NOTE (v1): proxy↔sandbox traffic is plain HTTP over the task's public IP. + # Tracked for follow-up: route through PrivateLink/VPC-internal addressing + # or terminate TLS on the harness so prompts/responses and the env-injected + # LITELLM_API_KEY do not transit the public internet in cleartext. sandbox_url = f"http://{public_ip}:{template.container_port}" await wait_http_ready(sandbox_url, client, timeout=600) @@ -284,6 +292,8 @@ async def list_sessions( where: Dict[str, Any] = {} if agent_id is not None: where["agent_id"] = agent_id + if not _is_admin(user_api_key_dict) and user_api_key_dict.user_id is not None: + where["created_by"] = user_api_key_dict.user_id rows = await prisma_client.db.litellm_managedagentsessiontable.find_many( where=where, order={"created_at": "desc"} @@ -306,6 +316,7 @@ async def get_session( ) if row is None: raise HTTPException(status_code=404, detail=f"session '{session_id}' not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "session", session_id) return _session_row_to_out(row) @@ -324,6 +335,7 @@ async def delete_session( ) if row is None: raise HTTPException(status_code=404, detail=f"session '{session_id}' not found") + _assert_owner_or_admin(user_api_key_dict, row.created_by, "session", session_id) region = _resolve_region() aws_overrides = _resolve_aws_overrides() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f04e01f5289..aa7b5ff7796 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -964,9 +964,19 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 except asyncio.CancelledError: pass except Exception as e: - verbose_proxy_logger.error( - f"Error stopping managed_agents reconciler: {e}" - ) + verbose_proxy_logger.error(f"Error stopping managed_agents reconciler: {e}") + + # Shutdown event - drain pooled passthrough HTTP client + try: + from litellm.proxy.managed_agents_endpoints.endpoints_passthrough import ( + close_passthrough_http_client, + ) + + await close_passthrough_http_client() + except Exception as e: + verbose_proxy_logger.error( + f"Error closing managed_agents passthrough http client: {e}" + ) # Shutdown event - close shared aiohttp session if shared_aiohttp_session is not None: diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_agents.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_agents.py index 009ad0c3347..5e65fd209bc 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_agents.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_agents.py @@ -1,8 +1,9 @@ """Tests for managed_agents_endpoints/endpoints_agents.py. -Covers GET /v1/managed_agents/agents (list) and GET /v1/managed_agents/agents/{id}. -The POST /agents create flow is exercised end-to-end via the session tests, so -this file focuses on the read endpoints. +Covers GET /v1/managed_agents/agents (list) and GET /v1/managed_agents/agents/{id}, +including the ownership / authorization gate. The POST /agents create flow is +exercised end-to-end via the session tests, so this file focuses on the read +endpoints. """ from datetime import datetime, timezone @@ -28,6 +29,20 @@ def user(): ) +@pytest.fixture +def other_user(): + return UserAPIKeyAuth( + api_key="sk-other", user_id="u2", user_role=LitellmUserRoles.INTERNAL_USER + ) + + +@pytest.fixture +def admin(): + return UserAPIKeyAuth( + api_key="sk-admin", user_id="a1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + @pytest.fixture def app_factory(): def make(auth_user): @@ -39,7 +54,7 @@ def app_factory(): return make -def _make_agent(agent_id="agt-1", **kw): +def _make_agent(agent_id="agt-1", created_by="u1", **kw): base = dict( agent_id=agent_id, agent_name="a", @@ -50,6 +65,7 @@ def _make_agent(agent_id="agt-1", **kw): branch="main", metadata={}, created_at=datetime(2026, 5, 7, tzinfo=timezone.utc), + created_by=created_by, ) base.update(kw) return SimpleNamespace(**base) @@ -109,6 +125,26 @@ def test_list_agents_500_when_prisma_unavailable(app_factory, user): assert "prisma" in resp.json()["detail"].lower() +def test_list_agents_filters_by_owner_for_non_admin(app_factory, user): + client = app_factory(user) + prisma = _make_prisma(agents=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/agents") + assert resp.status_code == 200 + _, kwargs = prisma.db.litellm_managedagenttable.find_many.call_args + assert kwargs["where"] == {"created_by": "u1"} + + +def test_list_agents_admin_no_owner_filter(app_factory, admin): + client = app_factory(admin) + prisma = _make_prisma(agents=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/agents") + assert resp.status_code == 200 + _, kwargs = prisma.db.litellm_managedagenttable.find_many.call_args + assert kwargs["where"] == {} + + # --------------------------------------------------------------------------- # get_agent # --------------------------------------------------------------------------- @@ -116,7 +152,7 @@ def test_list_agents_500_when_prisma_unavailable(app_factory, user): def test_get_agent_happy(app_factory, user): client = app_factory(user) - prisma = _make_prisma(agent=_make_agent(agent_id="agt-9")) + prisma = _make_prisma(agent=_make_agent(agent_id="agt-9", created_by="u1")) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = client.get("/v1/managed_agents/agents/agt-9") assert resp.status_code == 200 @@ -133,3 +169,19 @@ def test_get_agent_404(app_factory, user): resp = client.get("/v1/managed_agents/agents/missing") assert resp.status_code == 404 assert "missing" in resp.json()["detail"] + + +def test_get_agent_returns_404_for_non_owner(app_factory, other_user): + client = app_factory(other_user) + prisma = _make_prisma(agent=_make_agent(agent_id="agt-9", created_by="u1")) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/agents/agt-9") + assert resp.status_code == 404 + + +def test_get_agent_visible_to_admin(app_factory, admin): + client = app_factory(admin) + prisma = _make_prisma(agent=_make_agent(agent_id="agt-9", created_by="u1")) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/agents/agt-9") + assert resp.status_code == 200 diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_sessions.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_sessions.py index 1afd950c1e2..b276b8a3ece 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_sessions.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_sessions.py @@ -464,8 +464,22 @@ def test_delete_session_no_task_arn_skips_stop(app_factory, user): # --------------------------------------------------------------------------- -def test_list_sessions_returns_all_when_no_filter(app_factory, user): - client = app_factory(user) +@pytest.fixture +def other_user(): + return UserAPIKeyAuth( + api_key="sk-other", user_id="u2", user_role=LitellmUserRoles.INTERNAL_USER + ) + + +@pytest.fixture +def admin(): + return UserAPIKeyAuth( + api_key="sk-admin", user_id="a1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + +def test_list_sessions_returns_all_when_no_filter(app_factory, admin): + client = app_factory(admin) rows = [ _make_session(session_id="s1", agent_id="agt-1", status="ready"), _make_session(session_id="s2", agent_id="agt-2", status="dead"), @@ -476,14 +490,14 @@ def test_list_sessions_returns_all_when_no_filter(app_factory, user): assert resp.status_code == 200 body = resp.json() assert [r["id"] for r in body] == ["s1", "s2"] - # No filter: where is empty dict + # Admin caller: no owner filter _, kwargs = prisma.db.litellm_managedagentsessiontable.find_many.call_args assert kwargs["where"] == {} assert kwargs["order"] == {"created_at": "desc"} -def test_list_sessions_filters_by_agent_id(app_factory, user): - client = app_factory(user) +def test_list_sessions_filters_by_agent_id(app_factory, admin): + client = app_factory(admin) rows = [_make_session(session_id="s1", agent_id="agt-1", status="ready")] prisma = _make_prisma(sessions=rows) with patch("litellm.proxy.proxy_server.prisma_client", prisma): @@ -501,3 +515,75 @@ def test_list_sessions_empty(app_factory, user): resp = client.get("/v1/managed_agents/sessions") assert resp.status_code == 200 assert resp.json() == [] + + +def test_list_sessions_filters_by_owner_for_non_admin(app_factory, user): + client = app_factory(user) + prisma = _make_prisma(sessions=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/sessions") + assert resp.status_code == 200 + _, kwargs = prisma.db.litellm_managedagentsessiontable.find_many.call_args + assert kwargs["where"] == {"created_by": "u1"} + + +# --------------------------------------------------------------------------- +# Ownership / authorization on get and delete +# --------------------------------------------------------------------------- + + +def test_get_session_returns_404_for_non_owner(app_factory, other_user): + client = app_factory(other_user) + sess = _make_session( + session_id="sess-9", + status="ready", + sandbox_url="http://1.2.3.4:4096", + created_by="u1", + ) + prisma = _make_prisma(session=sess) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/sessions/sess-9") + assert resp.status_code == 404 + + +def test_get_session_visible_to_admin(app_factory, admin): + client = app_factory(admin) + sess = _make_session( + session_id="sess-9", + status="ready", + sandbox_url="http://1.2.3.4:4096", + created_by="u1", + ) + prisma = _make_prisma(session=sess) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/sessions/sess-9") + assert resp.status_code == 200 + + +def test_delete_session_returns_404_for_non_owner(app_factory, other_user): + client = app_factory(other_user) + sess = _make_session( + session_id="sess-9", + status="ready", + task_arn="arn:task/9", + created_by="u1", + ) + prisma = _make_prisma(session=sess) + stop_mock = AsyncMock(return_value=None) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.stop_session_task", + new=stop_mock, + ), + patch( + "litellm.proxy.managed_agents_endpoints.config_loader.MANAGED_AGENTS_CONFIG", + SimpleNamespace( + aws_region="us-west-2", + aws=SimpleNamespace(cluster=None), + ), + ), + ): + resp = client.delete("/v1/managed_agents/sessions/sess-9") + assert resp.status_code == 404 + stop_mock.assert_not_called() diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_templates.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_templates.py index e7b9b1a5777..22e383e653b 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_templates.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_endpoints_templates.py @@ -163,6 +163,7 @@ def test_template_create_happy_path(app_factory, admin, fake_prisma): container_port=4096, path="/x", context_dir="/x", + build_platform="linux/amd64", ), ), patch("litellm.proxy.managed_agents_endpoints.endpoints.validate_repo_branch"), @@ -200,6 +201,7 @@ def test_template_create_private_requires_token_400(app_factory, admin, fake_pri container_port=4096, path="/x", context_dir="/x", + build_platform="linux/amd64", ), ), ): @@ -207,6 +209,89 @@ def test_template_create_private_requires_token_400(app_factory, admin, fake_pri assert resp.status_code == 400 +def _template_row( + template_id="t1", + visibility="public", + created_by="u1", +): + return SimpleNamespace( + template_id=template_id, + template_name=None, + dockerfile_id="opencode", + container_port=4096, + repo_url="https://github.com/x/y", + default_branch="main", + visibility=visibility, + image_uri=None, + task_def_arn=None, + build_status="ready", + build_error=None, + created_by=created_by, + ) + + +def test_list_templates_filters_private_for_non_owner(app_factory, user, fake_prisma): + """Non-admin caller passes a where-clause that excludes private templates + they don't own; assert the where-clause is correct.""" + client = app_factory(user) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_many = AsyncMock( + return_value=[] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resp = client.get("/v1/managed_agents/sandbox-templates") + assert resp.status_code == 200 + where = fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_many.call_args.kwargs[ + "where" + ] + assert "OR" in where + assert {"visibility": "public"} in where["OR"] + assert {"created_by": "u1"} in where["OR"] + + +def test_list_templates_admin_sees_all(app_factory, admin, fake_prisma): + client = app_factory(admin) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_many = AsyncMock( + return_value=[] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resp = client.get("/v1/managed_agents/sandbox-templates") + assert resp.status_code == 200 + where = fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_many.call_args.kwargs[ + "where" + ] + assert where == {} + + +def test_get_private_template_returns_404_for_non_owner(app_factory, user, fake_prisma): + client = app_factory(user) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_unique = AsyncMock( + return_value=_template_row(visibility="private", created_by="u-other") + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resp = client.get("/v1/managed_agents/sandbox-templates/t1") + assert resp.status_code == 404 + + +def test_get_private_template_visible_to_owner(app_factory, user, fake_prisma): + client = app_factory(user) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_unique = AsyncMock( + return_value=_template_row(visibility="private", created_by="u1") + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resp = client.get("/v1/managed_agents/sandbox-templates/t1") + assert resp.status_code == 200 + + +def test_get_private_template_visible_to_admin(app_factory, admin, fake_prisma): + client = app_factory(admin) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_unique = AsyncMock( + return_value=_template_row(visibility="private", created_by="u-other") + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resp = client.get("/v1/managed_agents/sandbox-templates/t1") + assert resp.status_code == 200 + + def test_template_delete_with_agents_409(app_factory, admin, fake_prisma): client = app_factory(admin) fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_unique = AsyncMock( diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_lifecycle.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_lifecycle.py index 2eeb5e20c39..d9b6d720334 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_lifecycle.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_lifecycle.py @@ -69,15 +69,19 @@ async def test_orphan_no_db_row_stopped(): prisma, update = _fake_prisma([]) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=tasks, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=tasks, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" ) @@ -97,15 +101,19 @@ async def test_orphan_dead_row_stopped(status): prisma, update = _fake_prisma(rows) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=tasks, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=tasks, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" ) @@ -128,15 +136,19 @@ async def test_creating_young_skipped(): prisma, update = _fake_prisma(rows) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=tasks, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=tasks, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" ) @@ -158,15 +170,19 @@ async def test_creating_stale_stopped_and_marked_failed(): prisma, update = _fake_prisma(rows) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=tasks, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=tasks, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" ) @@ -189,15 +205,19 @@ async def test_ready_row_skipped(): prisma, update = _fake_prisma(rows) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=tasks, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=tasks, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" ) @@ -210,21 +230,29 @@ async def test_ready_row_skipped(): @pytest.mark.asyncio async def test_no_managed_tasks_returns_zero(): untagged = [ - {"taskArn": "arn:aws:ecs:us-west-2:123:task/other", "tags": {}, "lastStatus": "RUNNING"} + { + "taskArn": "arn:aws:ecs:us-west-2:123:task/other", + "tags": {}, + "lastStatus": "RUNNING", + } ] arns = [t["taskArn"] for t in untagged] prisma, _ = _fake_prisma([]) - with patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", - return_value=arns, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", - return_value=untagged, - ), patch( - "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" - ) as stop_mock: + with ( + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.list_tagged_task_arns", + return_value=arns, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.describe_tasks_with_tags", + return_value=untagged, + ), + patch( + "litellm.proxy.managed_agents_endpoints.lifecycle.stop_task_sync" + ) as stop_mock, + ): stats = await reconcile_orphans( prisma_client=prisma, region="us-west-2", cluster="test" )