From 8855bb7da427c8c31f4cc8d11d03a5f15758f05d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 7 May 2026 16:28:58 -0700 Subject: [PATCH] fix(managed agents): address Greptile review (iteration 3) - Encrypt litellm_api_key with encrypt_value_helper before storing it in agent.metadata (key: litellm_api_key_encrypted), and decrypt at session spawn before injecting into the harness env. Plaintext keys no longer sit in the metadata JSONB column. - Close the user_id-None bypass on list_sessions and list_agents: non-admin callers without an associated user_id now get an empty list instead of every row in the table. - ensure_security_group now adds an ingress rule for the requested container_port if the shared SG already exists but lacks one. Without this, multi-port deployments produced unreachable harnesses on every template after the first. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../endpoints_agents.py | 13 +++- .../endpoints_sessions.py | 17 +++++- .../fargate/bootstrap.py | 41 ++++++++++++- .../test_bootstrap.py | 59 +++++++++++++++++- .../test_endpoints_agents.py | 60 +++++++++++++++++++ .../test_endpoints_sessions.py | 23 ++++++- 6 files changed, 205 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py index c74229d492b..4ab46f8d287 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_agents.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_agents.py @@ -5,6 +5,7 @@ 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.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.managed_agents_endpoints.endpoints import ( _assert_owner_or_admin, _is_admin, @@ -55,6 +56,10 @@ async def create_agent( git_token = await decrypt_git_token(prisma_client, template.git_credential_id) await asyncio.to_thread(validate_repo_branch, template.repo_url, branch, git_token) + encrypted_key = ( + encrypt_value_helper(body.litellm_api_key) if body.litellm_api_key else None + ) + create_data = jsonify_object( { "agent_name": body.name, @@ -63,7 +68,7 @@ async def create_agent( "tools": json.dumps(body.tools), "branch": branch, "metadata": { - "litellm_api_key": body.litellm_api_key, + "litellm_api_key_encrypted": encrypted_key, "litellm_api_base": body.litellm_api_base, }, "created_by": user_api_key_dict.user_id, @@ -87,7 +92,11 @@ async def list_agents( 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: + if not _is_admin(user_api_key_dict): + # Non-admin callers see only their own rows. If the API key has no + # user_id, treat it as "no rows" rather than exposing every agent. + if user_api_key_dict.user_id is None: + return [] where["created_by"] = user_api_key_dict.user_id rows = await prisma_client.db.litellm_managedagenttable.find_many( diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py index e65817feff1..1c9d453b2bd 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py @@ -8,6 +8,7 @@ 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.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.managed_agents_endpoints import config_loader as _config_loader from litellm.proxy.managed_agents_endpoints.endpoints import ( _assert_owner_or_admin, @@ -168,8 +169,16 @@ async def create_session( ) session_id = row.session_id + encrypted_key = metadata.get("litellm_api_key_encrypted") + decrypted_key = ( + decrypt_value_helper( + encrypted_key, key="litellm_api_key", return_original_value=True + ) + if encrypted_key + else "" + ) env: Dict[str, str] = { - "LITELLM_API_KEY": metadata.get("litellm_api_key", "") or "", + "LITELLM_API_KEY": decrypted_key or "", "LITELLM_API_BASE": metadata.get("litellm_api_base", "") or "", "LITELLM_DEFAULT_MODEL": agent.model, "REPO_URL": template.repo_url, @@ -293,7 +302,11 @@ 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: + if not _is_admin(user_api_key_dict): + # Non-admin callers see only their own rows. If the API key has no + # user_id, treat it as "no rows" rather than exposing every session. + if user_api_key_dict.user_id is None: + return [] where["created_by"] = user_api_key_dict.user_id rows = await prisma_client.db.litellm_managedagentsessiontable.find_many( diff --git a/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py b/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py index c77f905dd48..a7bddd61c9f 100644 --- a/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py +++ b/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py @@ -125,6 +125,19 @@ def discover_vpc_subnet(region: str) -> Tuple[str, str, str]: return vpc_id, vpc_cidr, subnets[0]["SubnetId"] +def _sg_has_tcp_ingress(sg: dict, port: int) -> bool: + for perm in sg.get("IpPermissions", []) or []: + if perm.get("IpProtocol") != "tcp": + continue + if ( + perm.get("FromPort") is not None + and perm.get("ToPort") is not None + and perm["FromPort"] <= port <= perm["ToPort"] + ): + return True + return False + + def ensure_security_group( region: str, sg_name: str, @@ -141,10 +154,36 @@ def ensure_security_group( ) existing = r.get("SecurityGroups", []) if existing: + sg = existing[0] + sg_id = sg["GroupId"] verbose_proxy_logger.debug( f"Security group {sg_name} already exists in VPC {vpc_id}" ) - return existing[0]["GroupId"] + # The SG is shared across all templates. If a new template uses a + # different container_port, the existing SG won't have an ingress + # rule for it and the harness will be unreachable. Add the missing + # rule rather than silently returning. + if not _sg_has_tcp_ingress(sg, container_port): + try: + ec2.authorize_security_group_ingress( + GroupId=sg_id, + IpPermissions=[ + { + "IpProtocol": "tcp", + "FromPort": container_port, + "ToPort": container_port, + "IpRanges": [{"CidrIp": "0.0.0.0/0"}], + } + ], + ) + verbose_proxy_logger.info( + f"Added ingress for tcp/{container_port} to {sg_name}" + ) + except ClientError as e: + # InvalidPermission.Duplicate => another worker added it concurrently + if "InvalidPermission.Duplicate" not in str(e): + raise + return sg_id verbose_proxy_logger.info(f"Creating security group {sg_name} in VPC {vpc_id}") sg_id = ec2.create_security_group( diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py index 114b8c0aa5d..f115c9ebe5e 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py @@ -172,9 +172,16 @@ def test_discover_vpc_subnet_no_public_subnet_raises(mock_ec2): # --------------------------------------------------------------------------- -def test_ensure_security_group_existing_returns_id(mock_ec2): +def test_ensure_security_group_existing_with_port_returns_id(mock_ec2): mock_ec2.describe_security_groups.return_value = { - "SecurityGroups": [{"GroupId": "sg-existing"}] + "SecurityGroups": [ + { + "GroupId": "sg-existing", + "IpPermissions": [ + {"IpProtocol": "tcp", "FromPort": 4096, "ToPort": 4096} + ], + } + ] } with patch.object(bootstrap, "_ec2", return_value=mock_ec2): sg_id = bootstrap.ensure_security_group( @@ -188,6 +195,54 @@ def test_ensure_security_group_existing_returns_id(mock_ec2): mock_ec2.authorize_security_group_egress.assert_not_called() +def test_ensure_security_group_existing_missing_port_authorizes(mock_ec2): + """Existing SG was created for a different container_port; we must add + an ingress rule for the new port so multi-port deployments work.""" + mock_ec2.describe_security_groups.return_value = { + "SecurityGroups": [ + { + "GroupId": "sg-existing", + "IpPermissions": [ + {"IpProtocol": "tcp", "FromPort": 4096, "ToPort": 4096} + ], + } + ] + } + with patch.object(bootstrap, "_ec2", return_value=mock_ec2): + sg_id = bootstrap.ensure_security_group( + "us-west-2", "litellm-sg", "vpc-1", "10.0.0.0/16", 5000 + ) + + assert sg_id == "sg-existing" + mock_ec2.authorize_security_group_ingress.assert_called_once() + perms = mock_ec2.authorize_security_group_ingress.call_args.kwargs["IpPermissions"][ + 0 + ] + assert perms["FromPort"] == 5000 + assert perms["ToPort"] == 5000 + + +def test_ensure_security_group_existing_missing_port_swallows_duplicate(mock_ec2): + mock_ec2.describe_security_groups.return_value = { + "SecurityGroups": [ + { + "GroupId": "sg-existing", + "IpPermissions": [ + {"IpProtocol": "tcp", "FromPort": 4096, "ToPort": 4096} + ], + } + ] + } + mock_ec2.authorize_security_group_ingress.side_effect = _client_error( + "InvalidPermission.Duplicate" + ) + with patch.object(bootstrap, "_ec2", return_value=mock_ec2): + sg_id = bootstrap.ensure_security_group( + "us-west-2", "litellm-sg", "vpc-1", "10.0.0.0/16", 5000 + ) + assert sg_id == "sg-existing" + + def test_ensure_security_group_missing_creates_ingress_revoke_and_authorize(mock_ec2): mock_ec2.describe_security_groups.return_value = {"SecurityGroups": []} mock_ec2.create_security_group.return_value = {"GroupId": "sg-new"} 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 6360df8aa45..591a82b5e0e 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 @@ -206,6 +206,66 @@ def _make_template(template_id="tmpl-1", visibility="public", created_by="u1"): ) +def test_list_agents_returns_empty_when_user_id_none(app_factory): + """Non-admin caller with no user_id must not see other users' rows.""" + user_no_id = UserAPIKeyAuth( + api_key="sk", user_id=None, user_role=LitellmUserRoles.INTERNAL_USER + ) + client = app_factory(user_no_id) + prisma = _make_prisma(agents=[_make_agent(agent_id="x", created_by="someone-else")]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/agents") + assert resp.status_code == 200 + assert resp.json() == [] + prisma.db.litellm_managedagenttable.find_many.assert_not_called() + + +def test_create_agent_encrypts_litellm_api_key(app_factory, user): + """The plaintext API key must NOT be persisted in metadata; the + encrypted form goes under litellm_api_key_encrypted.""" + client = app_factory(user) + template = _make_template(visibility="public", 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(return_value=_make_agent(agent_id="agt-2")) + p.db.litellm_managedagenttable = agent_t + + body = { + "name": "agt", + "model": "anthropic/claude-sonnet-4-6", + "template_id": "tmpl-1", + "litellm_api_key": "sk-secret", + } + + with ( + patch("litellm.proxy.proxy_server.prisma_client", p), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_agents.encrypt_value_helper", + return_value="ENCRYPTED", + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_agents.decrypt_git_token", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_agents.validate_repo_branch" + ), + ): + resp = client.post("/v1/managed_agents/agents", json=body) + + import json as _json + + assert resp.status_code == 200, resp.text + raw_meta = agent_t.create.call_args.kwargs["data"]["metadata"] + metadata = _json.loads(raw_meta) if isinstance(raw_meta, str) else raw_meta + assert metadata["litellm_api_key_encrypted"] == "ENCRYPTED" + assert "litellm_api_key" not in metadata + assert "sk-secret" not in str(metadata) + + 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") 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 30df21b1a5d..a61673a4864 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 @@ -71,7 +71,10 @@ def _make_agent(template, created_by="u1"): tools=[], template_id=template.template_id, branch="main", - metadata={"litellm_api_key": "sk-x", "litellm_api_base": "http://x"}, + metadata={ + "litellm_api_key_encrypted": "encrypted-sk-x", + "litellm_api_base": "http://x", + }, created_by=created_by, ) a.template = template @@ -197,6 +200,10 @@ def test_create_session_happy_path_with_initial_prompt(app_factory, user): "litellm.proxy.managed_agents_endpoints.endpoints_sessions.harness_send_message", new=AsyncMock(return_value={"parts": [{"type": "text", "text": "hi"}]}), ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_value_helper", + return_value="sk-x", + ), patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", new=AsyncMock(return_value=None), @@ -541,6 +548,20 @@ def test_list_sessions_filters_by_owner_for_non_admin(app_factory, user): assert kwargs["where"] == {"created_by": "u1"} +def test_list_sessions_returns_empty_when_user_id_none(app_factory): + """Non-admin caller with no user_id must not see other users' rows.""" + user_no_id = UserAPIKeyAuth( + api_key="sk", user_id=None, user_role=LitellmUserRoles.INTERNAL_USER + ) + client = app_factory(user_no_id) + prisma = _make_prisma(sessions=[_make_session(session_id="s1")]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/managed_agents/sessions") + assert resp.status_code == 200 + assert resp.json() == [] + prisma.db.litellm_managedagentsessiontable.find_many.assert_not_called() + + # --------------------------------------------------------------------------- # Ownership / authorization on get and delete # ---------------------------------------------------------------------------