From 768defe90b1304922b7acb6741687fdcebc02dce Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 7 May 2026 16:18:01 -0700 Subject: [PATCH] fix(managed agents): address Greptile review (iteration 2) - delete_sandbox_template defaulted cluster to "litellm-managed-agents" while every other call site uses "litellm-agents"; this caused stop_sessions_for_template + ECS deregister to target a nonexistent cluster, leaving orphan Fargate tasks on template delete. Use the shared default. - create_sandbox_template called validate_repo_branch synchronously, blocking the asyncio event loop for up to 15s per request. Wrap in asyncio.to_thread to mirror the agent path. - create_agent skipped the visibility check, so any authenticated caller who knew a private template UUID could attach an agent (and spawn Fargate tasks) against it. Apply _template_visible_to after lookup. - create_session now also asserts the caller owns the agent before spawning a Fargate task on it. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../managed_agents_endpoints/endpoints.py | 7 +- .../endpoints_agents.py | 3 +- .../endpoints_sessions.py | 1 + .../test_endpoints_agents.py | 43 +++++++++ .../test_endpoints_sessions.py | 16 +++- .../test_endpoints_templates.py | 91 +++++++++++++++++++ 6 files changed, 157 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints.py b/litellm/proxy/managed_agents_endpoints/endpoints.py index 29c28ce36e6..1eef857116d 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints.py @@ -1,3 +1,4 @@ +import asyncio from typing import Any, Dict, List import boto3 @@ -133,7 +134,9 @@ async def create_sandbox_template( detail="visibility=public must not include git_token", ) - validate_repo_branch(body.repo_url, body.default_branch, body.git_token) + await asyncio.to_thread( + validate_repo_branch, body.repo_url, body.default_branch, body.git_token + ) git_credential_id = None if body.git_token: @@ -304,7 +307,7 @@ async def delete_sandbox_template( region = _resolve_region() aws_overrides = _resolve_aws_overrides() - cluster = aws_overrides.cluster or "litellm-managed-agents" + cluster = aws_overrides.cluster or "litellm-agents" try: await stop_sessions_for_template( diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py index 393d5801623..c74229d492b 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py @@ -8,6 +8,7 @@ from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_au from litellm.proxy.managed_agents_endpoints.endpoints import ( _assert_owner_or_admin, _is_admin, + _template_visible_to, router, ) from litellm.proxy.managed_agents_endpoints.git_validation import ( @@ -44,7 +45,7 @@ async def create_agent( where={"template_id": body.template_id} ) ) - if template is None: + if template is None or not _template_visible_to(template, user_api_key_dict): raise HTTPException( status_code=404, detail=f"template '{body.template_id}' not found" ) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py index d4a28f348a0..e65817feff1 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py @@ -122,6 +122,7 @@ async def create_session( ) if agent is None: raise HTTPException(status_code=404, detail=f"agent '{agent_id}' not found") + _assert_owner_or_admin(user_api_key_dict, agent.created_by, "agent", agent_id) template = getattr(agent, "template", None) if template is 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 5e65fd209bc..6360df8aa45 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 @@ -185,3 +185,46 @@ def test_get_agent_visible_to_admin(app_factory, admin): with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = client.get("/v1/managed_agents/agents/agt-9") assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# create_agent — template visibility enforcement +# --------------------------------------------------------------------------- + + +def _make_template(template_id="tmpl-1", visibility="public", created_by="u1"): + return SimpleNamespace( + template_id=template_id, + template_name="t", + dockerfile_id="opencode", + container_port=4096, + repo_url="https://github.com/x/y", + default_branch="main", + visibility=visibility, + git_credential_id=None, + created_by=created_by, + ) + + +def test_create_agent_rejects_private_template_for_non_owner(app_factory, other_user): + client = app_factory(other_user) + template = _make_template(visibility="private", created_by="u1") + p = MagicMock() + template_t = MagicMock() + template_t.find_unique = AsyncMock(return_value=template) + p.db.litellm_managedagentsandboxtemplatetable = template_t + agent_t = MagicMock() + agent_t.create = AsyncMock() + p.db.litellm_managedagenttable = agent_t + + body = { + "name": "agt", + "model": "anthropic/claude-sonnet-4-6", + "template_id": "tmpl-1", + } + + with patch("litellm.proxy.proxy_server.prisma_client", p): + resp = client.post("/v1/managed_agents/agents", json=body) + + assert resp.status_code == 404 + agent_t.create.assert_not_called() 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 b276b8a3ece..30df21b1a5d 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 @@ -62,7 +62,7 @@ def _make_template(build_status="ready"): ) -def _make_agent(template): +def _make_agent(template, created_by="u1"): a = SimpleNamespace( agent_id="agt-1", agent_name="a", @@ -72,6 +72,7 @@ def _make_agent(template): template_id=template.template_id, branch="main", metadata={"litellm_api_key": "sk-x", "litellm_api_base": "http://x"}, + created_by=created_by, ) a.template = template return a @@ -130,6 +131,19 @@ def test_create_session_404_when_agent_missing(app_factory, user): assert "missing" in resp.json()["detail"] +def test_create_session_404_when_caller_does_not_own_agent(app_factory): + other = UserAPIKeyAuth( + api_key="sk-other", user_id="u2", user_role=LitellmUserRoles.INTERNAL_USER + ) + client = app_factory(other) + template = _make_template() + agent = _make_agent(template, created_by="u1") + prisma = _make_prisma(agent=agent) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.post("/v1/managed_agents/agents/agt-1/session", json={}) + assert resp.status_code == 404 + + def test_create_session_409_when_template_not_ready(app_factory, user): client = app_factory(user) template = _make_template(build_status="pending") 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 22e383e653b..8a031c7db56 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 @@ -314,3 +314,94 @@ def test_template_delete_with_agents_409(app_factory, admin, fake_prisma): with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): resp = client.delete("/v1/managed_agents/sandbox-templates/t1") assert resp.status_code == 409 + + +def test_template_delete_uses_litellm_agents_cluster_default( + app_factory, admin, fake_prisma +): + """Default cluster name on template delete must match _resolve_cluster + elsewhere ('litellm-agents'); a mismatch leaves orphan Fargate tasks.""" + client = app_factory(admin) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.find_unique = AsyncMock( + return_value=_template_row(template_id="t1") + ) + fake_prisma.db.litellm_managedagenttable.count = AsyncMock(return_value=0) + stop_mock = AsyncMock(return_value=0) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints.stop_sessions_for_template", + new=stop_mock, + ), + patch( + "litellm.proxy.managed_agents_endpoints.config_loader.MANAGED_AGENTS_CONFIG", + SimpleNamespace(aws_region="us-west-2", aws=SimpleNamespace(cluster=None)), + ), + patch("boto3.client", return_value=MagicMock()), + ): + resp = client.delete("/v1/managed_agents/sandbox-templates/t1") + assert resp.status_code == 200 + _, kwargs = stop_mock.call_args + assert kwargs["cluster"] == "litellm-agents" + + +def test_template_create_validate_repo_branch_runs_off_thread( + app_factory, admin, fake_prisma +): + """validate_repo_branch must be wrapped in asyncio.to_thread to avoid + blocking the event loop on the 15s subprocess.""" + import asyncio as _asyncio + + client = app_factory(admin) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.create = AsyncMock( + return_value=_template_row(template_id="t1", visibility="public") + ) + fake_prisma.db.litellm_managedagentsandboxtemplatetable.update = AsyncMock( + return_value=_template_row(template_id="t1", visibility="public") + ) + provisioned = SimpleNamespace( + image_uri="img:abc", + task_def_arn="arn:td", + image_hash="abc", + ) + + captured: dict = {"called_in_thread": None} + + def fake_validate(*args, **kwargs): + try: + running_loop = _asyncio.get_running_loop() + except RuntimeError: + running_loop = None + captured["called_in_thread"] = running_loop is None + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints.get_dockerfile", + return_value=SimpleNamespace( + dockerfile_id="opencode", + container_port=4096, + path="/x", + context_dir="/x", + build_platform="linux/amd64", + ), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints.validate_repo_branch", + side_effect=fake_validate, + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints.provision_template", + AsyncMock(return_value=provisioned), + ), + patch( + "litellm.proxy.managed_agents_endpoints.config_loader.MANAGED_AGENTS_CONFIG", + SimpleNamespace(aws_region="us-west-2", aws=SimpleNamespace(cluster=None)), + ), + patch("boto3.client", return_value=MagicMock()), + ): + resp = client.post("/v1/managed_agents/sandbox-templates", json=_public_body()) + assert resp.status_code == 200, resp.text + # validate_repo_branch was invoked from a worker thread (no running loop) + assert captured["called_in_thread"] is True