diff --git a/litellm/proxy/managed_agents_endpoints/endpoints.py b/litellm/proxy/managed_agents_endpoints/endpoints.py index 1eef857116d..65b4ad3c5df 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints.py @@ -59,6 +59,12 @@ def _resolve_aws_overrides() -> AwsOverrides: return AwsOverrides() +def _deregister_task_def_sync(region: str, task_def_arn: str) -> None: + """Synchronous boto3 wrapper, intended to be called via asyncio.to_thread.""" + ecs = boto3.client("ecs", region_name=region) + ecs.deregister_task_definition(taskDefinition=task_def_arn) + + def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="admin role required") @@ -325,8 +331,7 @@ async def delete_sandbox_template( if row.task_def_arn: try: - ecs = boto3.client("ecs", region_name=region) - ecs.deregister_task_definition(taskDefinition=row.task_def_arn) + await asyncio.to_thread(_deregister_task_def_sync, region, row.task_def_arn) except Exception as e: verbose_proxy_logger.warning( "managed_agents: deregister_task_definition failed for arn=%s: %s", diff --git a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py index 1c9d453b2bd..bcd6aa6145e 100644 --- a/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py +++ b/litellm/proxy/managed_agents_endpoints/endpoints_sessions.py @@ -57,6 +57,20 @@ def _resolve_region() -> str: return "us-east-1" +def _region_from_arn(arn: Optional[str]) -> Optional[str]: + """Extract region from an AWS ARN. + + ARN format: arn::::: + Returns None for malformed input so callers can fall back to global config. + """ + if not arn: + return None + parts = arn.split(":", 5) + if len(parts) < 4 or not parts[3]: + return None + return parts[3] + + def _resolve_aws_overrides() -> AwsOverrides: cfg = _config_loader.MANAGED_AGENTS_CONFIG if cfg is not None: @@ -147,9 +161,12 @@ async def create_session( detail=f"template '{template.template_id}' has no task_def_arn", ) - region = _resolve_region() aws_overrides = _resolve_aws_overrides() cluster = _resolve_cluster(aws_overrides) + # Region must match the template's task-def ARN. Falling back to global + # config produces "Invalid Region in ARN" when a template was built in a + # different region than the proxy's current default. + region = _region_from_arn(template.task_def_arn) or _resolve_region() metadata = _coerce_metadata(getattr(agent, "metadata", None)) @@ -351,9 +368,16 @@ async def delete_session( 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() cluster = row.fargate_cluster or _resolve_cluster(aws_overrides) + # Stop the task in the same region it runs in — the session's task ARN + # encodes that region. The template's task-def ARN is the next-best fallback + # for sessions created before task_arn was persisted. + region = ( + _region_from_arn(row.task_arn) + or _region_from_arn(getattr(row, "fargate_task_def_arn", None)) + or _resolve_region() + ) if row.task_arn: await stop_session_task( diff --git a/litellm/proxy/managed_agents_endpoints/git_validation.py b/litellm/proxy/managed_agents_endpoints/git_validation.py index 2cecc431b33..2697af6dd9a 100644 --- a/litellm/proxy/managed_agents_endpoints/git_validation.py +++ b/litellm/proxy/managed_agents_endpoints/git_validation.py @@ -1,10 +1,10 @@ """Git repo + branch validation utilities for managed-agent template create.""" +import base64 import os import subprocess import uuid from typing import Any, Optional -from urllib.parse import urlparse, urlunparse from fastapi import HTTPException @@ -16,36 +16,43 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( def authed_repo_url(repo_url: str, git_token: Optional[str]) -> str: + """Return the URL unchanged. The git token is now injected via + `http.extraheader` (passed through GIT_CONFIG_* env vars) rather than + embedded in the URL netloc — embedding in argv leaks the token to + `/proc//cmdline` and `ps aux`.""" + return repo_url + + +def _git_auth_env(git_token: Optional[str]) -> dict: + """Build the environment for a `git ls-remote` call. + + When a token is provided we set `http.extraheader = AUTHORIZATION: Basic + ` via GIT_CONFIG_COUNT/KEY_*/VALUE_*, which keeps the credential + out of argv. Falls back to no auth when no token is provided. + """ + env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"} if not git_token: - return repo_url - parsed = urlparse(repo_url) - if parsed.scheme != "https": - return repo_url - if not parsed.hostname: - return repo_url - netloc = f"x-access-token:{git_token}@{parsed.hostname}" - if parsed.port: - netloc += f":{parsed.port}" - return urlunparse( - ( - parsed.scheme, - netloc, - parsed.path, - parsed.params, - parsed.query, - parsed.fragment, - ) + return env + auth_value = base64.b64encode(f"x-access-token:{git_token}".encode("utf-8")).decode( + "ascii" ) + env.update( + { + "GIT_CONFIG_COUNT": "1", + "GIT_CONFIG_KEY_0": "http.extraheader", + "GIT_CONFIG_VALUE_0": f"Authorization: Basic {auth_value}", + } + ) + return env def validate_repo_branch( repo_url: str, branch: str, git_token: Optional[str] = None ) -> None: - env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"} - url = authed_repo_url(repo_url, git_token) + env = _git_auth_env(git_token) try: result = subprocess.run( - ["git", "ls-remote", "--heads", "--tags", url, branch], + ["git", "ls-remote", "--heads", "--tags", repo_url, branch], capture_output=True, text=True, timeout=15, @@ -62,8 +69,6 @@ def validate_repo_branch( if result.returncode != 0: msg_lines = (result.stderr or result.stdout).strip().splitlines() tail = msg_lines[-1] if msg_lines else "unknown error" - # Scrub authed URL so token is not echoed back to clients/logs. - tail = tail.replace(url, repo_url) raise HTTPException( status_code=400, detail=f"git ls-remote failed for {repo_url}: {tail}", 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 a61673a4864..a77584905c4 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 @@ -622,3 +622,193 @@ def test_delete_session_returns_404_for_non_owner(app_factory, other_user): resp = client.delete("/v1/managed_agents/sessions/sess-9") assert resp.status_code == 404 stop_mock.assert_not_called() + + +# --------------------------------------------------------------------------- +# region resolution: ARN-derived per-template instead of global config +# --------------------------------------------------------------------------- + + +def test_region_from_arn_extracts_region_segment(): + from litellm.proxy.managed_agents_endpoints.endpoints_sessions import ( + _region_from_arn, + ) + + assert ( + _region_from_arn("arn:aws:ecs:us-west-2:888602223428:task-definition/litellm:3") + == "us-west-2" + ) + assert ( + _region_from_arn("arn:aws:ecs:eu-central-1:1:task/litellm-agents/abc") + == "eu-central-1" + ) + + +def test_region_from_arn_returns_none_for_malformed(): + from litellm.proxy.managed_agents_endpoints.endpoints_sessions import ( + _region_from_arn, + ) + + assert _region_from_arn(None) is None + assert _region_from_arn("") is None + assert _region_from_arn("arn:td") is None + assert _region_from_arn("arn:aws:ecs::1:task-definition/x") is None + + +def test_create_session_uses_region_from_template_arn(app_factory, user): + """Region must be derived from template.task_def_arn, not global config. + + Reproduces the production bug where a template built in us-west-2 plus a + proxy defaulting to us-east-1 raised "Invalid Region in ARN" on RunTask. + """ + client = app_factory(user) + template = _make_template() + template.task_def_arn = ( + "arn:aws:ecs:us-west-2:888602223428:task-definition/litellm-x:1" + ) + 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"], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.bootstrap_shared_infra", + return_value=infra, + ) as mock_boot, + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.run_task_sync", + return_value="arn:aws:ecs:us-west-2:1:task/litellm-agents/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.decrypt_value_helper", + return_value="sk-x", + ), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.decrypt_git_token", + new=AsyncMock(return_value=None), + ), + # Global config says us-east-1 — should be ignored in favor of ARN. + patch( + "litellm.proxy.managed_agents_endpoints.config_loader.MANAGED_AGENTS_CONFIG", + SimpleNamespace( + aws_region="us-east-1", + aws=SimpleNamespace(cluster=None), + ), + ), + ): + resp = client.post("/v1/managed_agents/agents/agt-1/session", json={}) + + assert resp.status_code == 200, resp.text + assert mock_boot.call_args.args[0] == "us-west-2" + assert mock_run.call_args.kwargs["region"] == "us-west-2" + + +def test_create_session_falls_back_to_config_when_arn_malformed(app_factory, user): + client = app_factory(user) + template = _make_template() + template.task_def_arn = "garbage" + 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"], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.bootstrap_shared_infra", + return_value=infra, + ) as mock_boot, + patch( + "litellm.proxy.managed_agents_endpoints.endpoints_sessions.run_task_sync", + return_value="arn:task/abc", + ), + 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.decrypt_value_helper", + return_value="sk-x", + ), + 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="ap-south-1", + aws=SimpleNamespace(cluster=None), + ), + ), + ): + resp = client.post("/v1/managed_agents/agents/agt-1/session", json={}) + + assert resp.status_code == 200, resp.text + assert mock_boot.call_args.args[0] == "ap-south-1" + + +def test_delete_session_uses_region_from_task_arn(app_factory, user): + """stop_task must run in the region the task was created in.""" + 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", + fargate_cluster="my-cluster", + ) + 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-east-1", + aws=SimpleNamespace(cluster=None), + ), + ), + ): + resp = client.delete("/v1/managed_agents/sessions/sess-9") + + assert resp.status_code == 200 + stop_mock.assert_awaited_once() + assert stop_mock.call_args.kwargs["region"] == "us-west-2" diff --git a/tests/test_litellm/proxy/managed_agents_endpoints/test_git_validation.py b/tests/test_litellm/proxy/managed_agents_endpoints/test_git_validation.py index 00770556b28..d0f22ec7313 100644 --- a/tests/test_litellm/proxy/managed_agents_endpoints/test_git_validation.py +++ b/tests/test_litellm/proxy/managed_agents_endpoints/test_git_validation.py @@ -32,27 +32,16 @@ from litellm.proxy.managed_agents_endpoints.git_validation import ( # --------------------------------------------------------------------------- -def test_authed_repo_url_https_with_token_injects_netloc(): - url = authed_repo_url("https://github.com/org/repo.git", "secret-token") - assert "x-access-token:secret-token@github.com" in url - assert url.endswith("/org/repo.git") - - -def test_authed_repo_url_https_no_token_returns_input_unchanged(): +def test_authed_repo_url_returns_input_unchanged(): + """authed_repo_url no longer rewrites the URL — auth is now passed via + GIT_CONFIG_VALUE_0 env var to keep the token out of argv.""" original = "https://github.com/org/repo.git" + assert authed_repo_url(original, "secret-token") == original assert authed_repo_url(original, None) == original assert authed_repo_url(original, "") == original - - -def test_authed_repo_url_ssh_returns_input_unchanged(): - original = "git@github.com:org/repo.git" - assert authed_repo_url(original, "secret-token") == original - - -def test_authed_repo_url_malformed_returns_input_unchanged(): - # No scheme/host → not https, returned untouched. - original = "not-a-real-url" - assert authed_repo_url(original, "secret-token") == original + assert authed_repo_url("git@github.com:org/repo.git", "secret-token") == ( + "git@github.com:org/repo.git" + ) # --------------------------------------------------------------------------- @@ -88,27 +77,26 @@ def test_validate_repo_branch_empty_stdout_raises_400_branch_not_found(): assert "not found" in exc_info.value.detail -def test_validate_repo_branch_nonzero_exit_raises_400_and_scrubs_token(): +def test_validate_repo_branch_token_not_in_argv(): + """The git token must not appear as a subprocess argument — it should + flow through the GIT_CONFIG_* env vars instead.""" repo_url = "https://github.com/org/repo.git" token = "super-secret-token" - stderr = ( - "remote: Repository not found.\n" - "fatal: repository 'https://x-access-token:super-secret-token@github.com/org/repo.git/' not found" - ) with patch( "litellm.proxy.managed_agents_endpoints.git_validation.subprocess.run", - return_value=_completed(128, stdout="", stderr=stderr), - ): - with pytest.raises(HTTPException) as exc_info: - validate_repo_branch(repo_url, "main", git_token=token) + return_value=_completed(0, stdout="abc\trefs/heads/main\n"), + ) as run_mock: + validate_repo_branch(repo_url, "main", git_token=token) - assert exc_info.value.status_code == 400 - detail = exc_info.value.detail - # Plain repo URL should appear in the error. - assert repo_url in detail - # Authed URL (with token) must NOT leak. - assert token not in detail - assert "x-access-token" not in detail + args, kwargs = run_mock.call_args + cmd = args[0] + assert token not in " ".join(cmd) + env = kwargs["env"] + assert env["GIT_CONFIG_COUNT"] == "1" + assert env["GIT_CONFIG_KEY_0"] == "http.extraheader" + # value is `Authorization: Basic )>` + assert env["GIT_CONFIG_VALUE_0"].startswith("Authorization: Basic ") + assert token not in env["GIT_CONFIG_VALUE_0"] def test_validate_repo_branch_file_not_found_raises_500():