mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(managed agents): address Greptile P1/P2 security review
- Strip Authorization/cookie/api-key headers from raw passthrough so the caller's LiteLLM key never reaches the untrusted sandbox container. - Enforce ownership on get_session, delete_session, get_agent, and the three passthrough endpoints (message, events, raw); non-owners get 404 to avoid ID enumeration. PROXY_ADMIN bypasses the check. - Filter private templates from list/get unless caller is admin or creator. - Pool a module-level httpx client for the passthrough endpoints instead of opening a new one per request; close on proxy shutdown. - Document the v1 plain-HTTP-over-public-IP sandbox transport as a known follow-up. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia
parent
9b7f4ac59a
commit
57b4a9ca20
9 changed files with 468 additions and 110 deletions
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue