From c9c541a5632e71ae72fed770f0f08e0edc09d4a3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 7 May 2026 18:04:55 -0700 Subject: [PATCH] fix(managed agents): warm pool integration + session key minting (#27438) Co-authored-by: Claude Opus 4.7 (1M context) --- litellm/proxy/_new_secret_config.yaml | 5 + .../endpoints_sessions.py | 230 ++++++++++++++---- .../fargate/bootstrap.py | 54 ++-- .../harnesses/opencode/Dockerfile | 9 +- .../managed_agents_endpoints/lifecycle.py | 20 +- .../proxy/managed_agents_endpoints/types.py | 2 + litellm/proxy/proxy_server.py | 19 ++ .../test_bootstrap.py | 20 +- .../test_endpoints_sessions.py | 174 ++++++++++++- 9 files changed, 442 insertions(+), 91 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 703fe6adc41..7f5150a8243 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -81,3 +81,8 @@ litellm_settings: general_settings: master_key: sk-1234 # REPLACE in production + managed_agents: + enabled: true + pool_enabled: true + pool_min_warm: 1 + aws_region: us-east-1 diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py index bcd6aa6145e..4d3df30a8f8 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py @@ -8,7 +8,10 @@ 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.management_endpoints.key_management_endpoints import ( + delete_verification_tokens, + generate_key_helper_fn, +) from litellm.proxy.managed_agents_endpoints import config_loader as _config_loader from litellm.proxy.managed_agents_endpoints.endpoints import ( _assert_owner_or_admin, @@ -35,6 +38,11 @@ from litellm.proxy.managed_agents_endpoints.types import ( SessionCreateIn, SessionOut, ) +from litellm.proxy.managed_agents_endpoints.warm_pool import ( + post_claim as _warm_pool_post_claim, + schedule_refill as _warm_pool_schedule_refill, + try_claim as _warm_pool_try_claim, +) from litellm.proxy.utils import jsonify_object @@ -57,6 +65,41 @@ def _resolve_region() -> str: return "us-east-1" +async def _mint_session_key( + user_api_key_dict: UserAPIKeyAuth, + agent_id: str, + agent_model: str, + session_id: str, +) -> tuple[str, str]: + """Mint a session-scoped LiteLLM key for a managed-agent sandbox. + + Returns (plaintext_key, token_hash). The hash is what's stored on the + session row; the plaintext is injected as LITELLM_API_KEY into the + container env and never persisted. + """ + key_data = await generate_key_helper_fn( + request_type="key", + duration=None, + models=[agent_model], + user_id=user_api_key_dict.user_id, + team_id=user_api_key_dict.team_id, + agent_id=agent_id, + key_alias=f"managed-agent-session-{session_id}", + metadata={ + "managed_agent_id": agent_id, + "managed_agent_session_id": session_id, + }, + ) + plaintext = key_data.get("token") + token_hash = key_data.get("token_id") + if not plaintext or not token_hash: + raise RuntimeError( + "generate_key_helper_fn returned no token/token_id for session " + f"{session_id}" + ) + return plaintext, token_hash + + def _region_from_arn(arn: Optional[str]) -> Optional[str]: """Extract region from an AWS ARN. @@ -186,61 +229,124 @@ 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": decrypted_key or "", - "LITELLM_API_BASE": metadata.get("litellm_api_base", "") or "", - "LITELLM_DEFAULT_MODEL": agent.model, - "REPO_URL": template.repo_url, - "BRANCH": agent.branch or template.default_branch, - } - git_token = await decrypt_git_token(prisma_client, template.git_credential_id) - if git_token: - env["GIT_TOKEN"] = git_token - if agent.prompt: - env["AGENT_PROMPT"] = agent.prompt - client = httpx.AsyncClient( timeout=httpx.Timeout(connect=10, read=None, write=None, pool=10) ) task_arn: Optional[str] = None + session_key_hash: Optional[str] = None try: - infra = await asyncio.to_thread( - bootstrap_shared_infra, region, aws_overrides, template.container_port - ) - if not infra.subnet_ids: - raise RuntimeError("bootstrap_shared_infra returned no subnets") - subnet = infra.subnet_ids[0] - security_group = infra.security_group_id - - task_arn = await asyncio.to_thread( - run_task_sync, - region=region, - cluster=cluster, - task_def_arn=template.task_def_arn, - container_name="harness", - subnet=subnet, - security_group=security_group, - env=env, - session_id=session_id, + # Mint a session-scoped LiteLLM key for the sandbox container instead + # of passing the user's own key through. The key is scoped to the + # agent's model only, aliased by session_id, and revoked when the + # session stops. Done inside try so a mint failure cleans up the + # session row via _mark_session_failed. + session_key, session_key_hash = await _mint_session_key( + user_api_key_dict=user_api_key_dict, agent_id=agent_id, + agent_model=agent.model, + session_id=session_id, ) await prisma_client.db.litellm_managedagentsessiontable.update( where={"session_id": session_id}, - data={"task_arn": task_arn}, + data={"virtual_key_hash": session_key_hash}, ) - public_ip = await asyncio.to_thread( - wait_running_get_ip_sync, region, cluster, task_arn, 300 - ) + env: Dict[str, str] = { + "LITELLM_API_KEY": session_key, + "LITELLM_API_BASE": metadata.get("litellm_api_base", "") or "", + "LITELLM_DEFAULT_MODEL": agent.model, + "REPO_URL": template.repo_url, + "BRANCH": agent.branch or template.default_branch, + } + git_token = await decrypt_git_token(prisma_client, template.git_credential_id) + if git_token: + env["GIT_TOKEN"] = git_token + if agent.prompt: + env["AGENT_PROMPT"] = agent.prompt + + # Warm-pool fast path: try to claim a pre-spawned Fargate task whose + # shim is already listening. On success we skip RunTask + ENI wait + + # ECR pull entirely (~30-60s saved). On any failure (no slot, claim + # POST 4xx/5xx/timeout) we fall through to the cold path below. + public_ip: Optional[str] = None + cfg = _config_loader.MANAGED_AGENTS_CONFIG + pool_enabled = bool(cfg and getattr(cfg, "pool_enabled", False)) + pool_min_warm = int(getattr(cfg, "pool_min_warm", 1)) if cfg else 1 + if pool_enabled: + slot = await _warm_pool_try_claim(template.template_id) + if slot is not None: + claimed_arn, claimed_ip, claimed_port, claimed_secret = slot + try: + await _warm_pool_post_claim( + public_ip=claimed_ip, + container_port=claimed_port, + secret=claimed_secret, + env=env, + timeout=120.0, + ) + task_arn = claimed_arn + public_ip = claimed_ip + await prisma_client.db.litellm_managedagentsessiontable.update( + where={"session_id": session_id}, + data={"task_arn": task_arn}, + ) + verbose_proxy_logger.info( + f"managed_agents: warm-pool hit template={template.template_id} " + f"session={session_id} ip={public_ip}" + ) + except Exception as e: + verbose_proxy_logger.warning( + "managed_agents: warm-pool claim failed, falling back " + f"(template={template.template_id} arn={claimed_arn}): {e}" + ) + await asyncio.to_thread( + stop_task_sync, + region, + cluster, + claimed_arn, + "warm_pool claim failed", + ) + public_ip = None + task_arn = None + finally: + _warm_pool_schedule_refill( + template=template, + region=region, + aws_overrides=aws_overrides, + cluster=cluster, + min_warm=pool_min_warm, + ) + + if public_ip is None: + infra = await asyncio.to_thread( + bootstrap_shared_infra, region, aws_overrides, template.container_port + ) + if not infra.subnet_ids: + raise RuntimeError("bootstrap_shared_infra returned no subnets") + subnet = infra.subnet_ids[0] + security_group = infra.security_group_id + + task_arn = await asyncio.to_thread( + run_task_sync, + region=region, + cluster=cluster, + task_def_arn=template.task_def_arn, + container_name="harness", + subnet=subnet, + security_group=security_group, + env=env, + session_id=session_id, + agent_id=agent_id, + ) + await prisma_client.db.litellm_managedagentsessiontable.update( + where={"session_id": session_id}, + data={"task_arn": task_arn}, + ) + + 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 @@ -299,6 +405,12 @@ async def create_session( verbose_proxy_logger.warning( f"managed_agents: stop_task after failure raised: {stop_err}" ) + if session_key_hash: + await _revoke_session_key( + session_id=session_id, + token_hash=session_key_hash, + user_api_key_dict=user_api_key_dict, + ) if isinstance(e, HTTPException): raise raise HTTPException(status_code=500, detail=f"session create failed: {e}") @@ -387,9 +499,39 @@ async def delete_session( session_id=session_id, ) + if row.virtual_key_hash: + await _revoke_session_key( + session_id=session_id, + token_hash=row.virtual_key_hash, + user_api_key_dict=user_api_key_dict, + ) + await prisma_client.db.litellm_managedagentsessiontable.update( where={"session_id": session_id}, data={"status": "dead", "stopped_at": _now_utc()}, ) return {"id": session_id, "status": "dead"} + + +async def _revoke_session_key( + session_id: str, + token_hash: str, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Revoke a session-scoped LiteLLM key. Best-effort: a failure here must + not block session teardown.""" + from litellm.proxy.proxy_server import user_api_key_cache + + try: + await delete_verification_tokens( + tokens=[token_hash], + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_proxy_logger.warning( + "managed_agents: failed to revoke session key for session=%s: %s", + session_id, + e, + ) diff --git a/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py b/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py index a7bddd61c9f..5d40b2521a0 100644 --- a/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py +++ b/litellm/proxy/managed_agents_endpoints/fargate/bootstrap.py @@ -138,6 +138,25 @@ def _sg_has_tcp_ingress(sg: dict, port: int) -> bool: return False +def _ensure_tcp_ingress(ec2, sg_id: str, sg_name: str, port: int) -> None: + try: + ec2.authorize_security_group_ingress( + GroupId=sg_id, + IpPermissions=[ + { + "IpProtocol": "tcp", + "FromPort": port, + "ToPort": port, + "IpRanges": [{"CidrIp": "0.0.0.0/0"}], + } + ], + ) + verbose_proxy_logger.info(f"Added ingress for tcp/{port} to {sg_name}") + except ClientError as e: + if "InvalidPermission.Duplicate" not in str(e): + raise + + def ensure_security_group( region: str, sg_name: str, @@ -146,6 +165,7 @@ def ensure_security_group( container_port: int, ) -> str: ec2 = _ec2(region) + shim_port = container_port + 1 r = ec2.describe_security_groups( Filters=[ {"Name": "vpc-id", "Values": [vpc_id]}, @@ -159,30 +179,10 @@ def ensure_security_group( verbose_proxy_logger.debug( f"Security group {sg_name} already exists in VPC {vpc_id}" ) - # 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 + _ensure_tcp_ingress(ec2, sg_id, sg_name, container_port) + if not _sg_has_tcp_ingress(sg, shim_port): + _ensure_tcp_ingress(ec2, sg_id, sg_name, shim_port) return sg_id verbose_proxy_logger.info(f"Creating security group {sg_name} in VPC {vpc_id}") @@ -200,7 +200,13 @@ def ensure_security_group( "FromPort": container_port, "ToPort": container_port, "IpRanges": [{"CidrIp": "0.0.0.0/0"}], - } + }, + { + "IpProtocol": "tcp", + "FromPort": shim_port, + "ToPort": shim_port, + "IpRanges": [{"CidrIp": "0.0.0.0/0"}], + }, ], ) diff --git a/litellm/proxy/managed_agents_endpoints/harnesses/opencode/Dockerfile b/litellm/proxy/managed_agents_endpoints/harnesses/opencode/Dockerfile index ae8065a140c..ed064150723 100644 --- a/litellm/proxy/managed_agents_endpoints/harnesses/opencode/Dockerfile +++ b/litellm/proxy/managed_agents_endpoints/harnesses/opencode/Dockerfile @@ -16,7 +16,7 @@ RUN find /root/.opencode -maxdepth 2 -type d -name 'cache' -exec rm -rf {} + 2>/ FROM debian:bookworm-slim RUN apt-get update && apt-get install -y --no-install-recommends \ - git ca-certificates bash \ + git ca-certificates bash python3 \ && rm -rf /var/lib/apt/lists/* \ && useradd -m -u 1000 -s /bin/bash sandbox \ && mkdir -p /work \ @@ -27,9 +27,10 @@ ENV PATH="/home/sandbox/.opencode/bin:${PATH}" WORKDIR /work COPY --chown=sandbox:sandbox entrypoint.sh /entrypoint.sh -RUN chmod +x /entrypoint.sh +COPY --chown=sandbox:sandbox shim.py /shim.py +RUN chmod +x /entrypoint.sh /shim.py USER sandbox -EXPOSE 4096 -ENTRYPOINT ["/entrypoint.sh"] +EXPOSE 4096 4097 +ENTRYPOINT ["python3", "/shim.py"] diff --git a/litellm/proxy/managed_agents_endpoints/lifecycle.py b/litellm/proxy/managed_agents_endpoints/lifecycle.py index c6762683376..8bf8108e08e 100644 --- a/litellm/proxy/managed_agents_endpoints/lifecycle.py +++ b/litellm/proxy/managed_agents_endpoints/lifecycle.py @@ -27,6 +27,7 @@ from litellm.proxy.managed_agents_endpoints.fargate.tasks import ( list_tagged_task_arns, stop_task_sync, ) +from litellm.proxy.managed_agents_endpoints.warm_pool import POOL_SLOT_PREFIX ALIVE_STATUSES = ("creating", "ready") DEAD_STATUSES = ("dead", "failed", "stopped") @@ -121,16 +122,27 @@ async def reconcile_orphans( if not managed_tasks: return {"scanned": 0, "orphaned_stopped": 0, "stale_creating_stopped": 0} - session_ids = [t["tags"][TAG_SESSION_ID] for t in managed_tasks] - rows = await prisma_client.db.litellm_managedagentsessiontable.find_many( - where={"session_id": {"in": session_ids}} + # Warm-pool tasks are tagged with a sentinel session_id ('pool-warm-...') and + # have no DB row. They are managed in-process by warm_pool.py, not the DB. + real_tasks = [ + t + for t in managed_tasks + if not t["tags"][TAG_SESSION_ID].startswith(POOL_SLOT_PREFIX) + ] + session_ids = [t["tags"][TAG_SESSION_ID] for t in real_tasks] + rows = ( + await prisma_client.db.litellm_managedagentsessiontable.find_many( + where={"session_id": {"in": session_ids}} + ) + if session_ids + else [] ) by_id = {r.session_id: r for r in rows} orphaned = 0 stale = 0 now_ts = time.time() - for task in managed_tasks: + for task in real_tasks: sid = task["tags"][TAG_SESSION_ID] arn = task["taskArn"] row = by_id.get(sid) diff --git a/litellm/proxy/managed_agents_endpoints/types.py b/litellm/proxy/managed_agents_endpoints/types.py index 2f5bc8f0aa1..ff99e376044 100644 --- a/litellm/proxy/managed_agents_endpoints/types.py +++ b/litellm/proxy/managed_agents_endpoints/types.py @@ -33,6 +33,8 @@ class ManagedAgentsConfig(BaseModel): dockerfiles: Dict[str, DockerfileConfig] = Field(default_factory=dict) aws: AwsOverrides = Field(default_factory=AwsOverrides) reconcile_interval_seconds: int = 60 + pool_enabled: bool = False + pool_min_warm: int = 1 class DockerfileOut(BaseModel): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aa7b5ff7796..33389295e11 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -953,6 +953,25 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 f"managed_agents: started Fargate orphan reconciler (cluster={cluster_name})" ) + if managed_agents_config.pool_enabled: + from litellm.proxy.managed_agents_endpoints.warm_pool import ( + warm_pool_startup, + ) + + asyncio.create_task( + warm_pool_startup( + prisma_client=prisma_client, + region=managed_agents_config.aws_region or "us-east-1", + aws_overrides=managed_agents_config.aws, + cluster=cluster_name, + min_warm=managed_agents_config.pool_min_warm, + ) + ) + verbose_proxy_logger.info( + "managed_agents: warm pool startup kicked " + f"(min_warm={managed_agents_config.pool_min_warm})" + ) + # End of startup event yield 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 f115c9ebe5e..609bb2d7c2b 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_bootstrap.py @@ -173,12 +173,15 @@ def test_discover_vpc_subnet_no_public_subnet_raises(mock_ec2): def test_ensure_security_group_existing_with_port_returns_id(mock_ec2): + # Existing SG must already have BOTH the harness port AND the warm-pool + # shim port (container_port + 1) for the no-op path. mock_ec2.describe_security_groups.return_value = { "SecurityGroups": [ { "GroupId": "sg-existing", "IpPermissions": [ - {"IpProtocol": "tcp", "FromPort": 4096, "ToPort": 4096} + {"IpProtocol": "tcp", "FromPort": 4096, "ToPort": 4096}, + {"IpProtocol": "tcp", "FromPort": 4097, "ToPort": 4097}, ], } ] @@ -197,7 +200,7 @@ def test_ensure_security_group_existing_with_port_returns_id(mock_ec2): 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.""" + ingress rules for both the harness port and the shim port (port + 1).""" mock_ec2.describe_security_groups.return_value = { "SecurityGroups": [ { @@ -214,12 +217,13 @@ def test_ensure_security_group_existing_missing_port_authorizes(mock_ec2): ) 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 + # Two calls: one for the harness port (5000), one for the shim port (5001). + assert mock_ec2.authorize_security_group_ingress.call_count == 2 + authorized_ports = sorted( + call.kwargs["IpPermissions"][0]["FromPort"] + for call in mock_ec2.authorize_security_group_ingress.call_args_list + ) + assert authorized_ports == [5000, 5001] def test_ensure_security_group_existing_missing_port_swallows_duplicate(mock_ec2): 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 a77584905c4..b34d15d12e8 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 @@ -91,6 +91,7 @@ def _make_session(session_id="sess-1", **kw): harness_session_id=None, fargate_cluster="litellm-agents", fargate_task_def_arn="arn:td", + virtual_key_hash=None, failure_reason=None, stopped_at=None, last_seen_at=None, @@ -201,8 +202,8 @@ def test_create_session_happy_path_with_initial_prompt(app_factory, user): new=AsyncMock(return_value={"parts": [{"type": "text", "text": "hi"}]}), ), patch( - "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_value_helper", - return_value="sk-x", + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=AsyncMock(return_value=("sk-session", "hash-session")), ), patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", @@ -237,7 +238,7 @@ def test_create_session_happy_path_with_initial_prompt(app_factory, user): # run_task_sync got the env vars from agent.metadata + template _, run_kwargs = mock_run.call_args env = run_kwargs["env"] - assert env["LITELLM_API_KEY"] == "sk-x" + assert env["LITELLM_API_KEY"] == "sk-session" assert env["LITELLM_API_BASE"] == "http://x" assert env["LITELLM_DEFAULT_MODEL"] == "anthropic/claude-sonnet-4-6" assert env["REPO_URL"] == "https://github.com/x/y" @@ -287,6 +288,10 @@ def test_create_session_happy_path_no_initial_prompt(app_factory, user): "litellm.proxy.managed_agents_endpoints.endpoints_sessions.harness_send_message", new=send_mock, ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=AsyncMock(return_value=("sk-session", "hash-session")), + ), patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", new=AsyncMock(return_value=None), @@ -341,6 +346,14 @@ def test_create_session_marks_failed_on_exception(app_factory, user): "litellm.proxy.managed_agents_endpoints.endpoints_sessions.stop_task_sync", return_value=None, ) as mock_stop, + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=AsyncMock(return_value=("sk-session", "hash-session")), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._revoke_session_key", + new=AsyncMock(return_value=None), + ) as mock_revoke, patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", new=AsyncMock(return_value=None), @@ -357,6 +370,8 @@ def test_create_session_marks_failed_on_exception(app_factory, user): assert resp.status_code == 500 assert "session create failed" in resp.json()["detail"] + # Minted key revoked on failure + mock_revoke.assert_awaited_once() # Session row was marked failed update_calls = prisma.db.litellm_managedagentsessiontable.update.call_args_list @@ -700,8 +715,8 @@ def test_create_session_uses_region_from_template_arn(app_factory, user): new=AsyncMock(return_value="harness-sess-1"), ), patch( - "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_value_helper", - return_value="sk-x", + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=AsyncMock(return_value=("sk-session", "hash-session")), ), patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", @@ -761,8 +776,8 @@ def test_create_session_falls_back_to_config_when_arn_malformed(app_factory, use new=AsyncMock(return_value="harness-sess-1"), ), patch( - "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_value_helper", - return_value="sk-x", + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=AsyncMock(return_value=("sk-session", "hash-session")), ), patch( "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", @@ -812,3 +827,148 @@ def test_delete_session_uses_region_from_task_arn(app_factory, user): assert resp.status_code == 200 stop_mock.assert_awaited_once() assert stop_mock.call_args.kwargs["region"] == "us-west-2" + + +# --------------------------------------------------------------------------- +# session-scoped temp key: mint at create, revoke at delete +# --------------------------------------------------------------------------- + + +def test_create_session_mints_key_and_persists_hash(app_factory, user): + """Mint a session-scoped key, store its hash, and pass plaintext to env.""" + client = app_factory(user) + template = _make_template() + agent = _make_agent(template) + prisma = _make_prisma(agent=agent, session=_make_session()) + infra = SimpleNamespace( + cluster_arn="arn:cluster", + task_exec_role_arn="arn:role", + security_group_id="sg-1", + log_group_name="/ecs/x", + vpc_id="vpc-1", + subnet_ids=["subnet-1"], + ) + mint_mock = AsyncMock(return_value=("sk-fresh-token", "hash-abc")) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.bootstrap_shared_infra", + return_value=infra, + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.run_task_sync", + return_value="arn:task/abc", + ) as mock_run, + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.wait_running_get_ip_sync", + return_value="1.2.3.4", + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.wait_http_ready", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.harness_create_session", + new=AsyncMock(return_value="harness-sess-1"), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._mint_session_key", + new=mint_mock, + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.managed_agents_endpoints.config_loader.MANAGED_AGENTS_CONFIG", + SimpleNamespace( + aws_region="us-west-2", + aws=SimpleNamespace(cluster=None), + ), + ), + ): + resp = client.post("/v1/managed_agents/agents/agt-1/session", json={}) + + assert resp.status_code == 200, resp.text + mint_mock.assert_awaited_once() + # Plaintext key flowed into ECS env, not the user's key + env = mock_run.call_args.kwargs["env"] + assert env["LITELLM_API_KEY"] == "sk-fresh-token" + # Hash persisted on session row + update_calls = prisma.db.litellm_managedagentsessiontable.update.call_args_list + hash_updates = [ + c + for c in update_calls + if c.kwargs.get("data", {}).get("virtual_key_hash") == "hash-abc" + ] + assert hash_updates, f"expected virtual_key_hash update, got {update_calls}" + + +def test_delete_session_revokes_key_when_hash_present(app_factory, user): + client = app_factory(user) + sess = _make_session( + session_id="sess-9", + status="ready", + task_arn="arn:aws:ecs:us-west-2:1:task/litellm-agents/xyz", + virtual_key_hash="hash-to-revoke", + ) + prisma = _make_prisma(session=sess) + revoke_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=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._revoke_session_key", + new=revoke_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 == 200 + revoke_mock.assert_awaited_once() + assert revoke_mock.call_args.kwargs["token_hash"] == "hash-to-revoke" + + +def test_delete_session_skips_revoke_when_no_hash(app_factory, user): + client = app_factory(user) + sess = _make_session( + session_id="sess-9", + status="ready", + task_arn="arn:aws:ecs:us-west-2:1:task/litellm-agents/xyz", + virtual_key_hash=None, + ) + prisma = _make_prisma(session=sess) + revoke_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=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions._revoke_session_key", + new=revoke_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 == 200 + revoke_mock.assert_not_called()