From 39bc727552fd188b39eac67e420cf0183c589144 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 22:45:36 -0700 Subject: [PATCH 1/8] feat(config): send [model.headers] on every call to the model endpoint, agents' included Some endpoints need a header on every request: Hugging Face's router bills an organization only when each request carries X-HF-Bill-To. [model.params.extra_headers] reached only agent-env's own calls, so every agent had to invent its own variable for it. [model.headers] declares them once, with env:/secret: references resolved per call like api_key. They go only to the configured endpoint, since a header may carry a credential: - agent-env's own calls get them as extra_headers ([model.params.extra_headers] wins per header); - deployed agents (the judge agent included), the unit-test verifier and installed agents that declare the litellm_extra_headers param get them as LITELLM_EXTRA_HEADERS, a JSON object, set next to LITELLM_BASE_URL; the verifier also gets ANTHROPIC_CUSTOM_HEADERS; - an agent or call pointed at another URL gets none, and an explicit LITELLM_EXTRA_HEADERS wins. Core only writes LITELLM_EXTRA_HEADERS, never reads it. The protocol's reference model sends it. Co-Authored-By: Claude Opus 5.5 --- .agentenv/config.example.toml | 7 ++ AGENTS.md | 5 ++ packages/agentenv-protocol/README.md | 5 +- .../examples/reference_model.py | 3 +- .../tests/test_a2a_agent_examples.py | 15 +++- src/agent_env/a2a_agent/a2a_agent.py | 5 ++ src/agent_env/config/model.py | 33 +++++-- src/agent_env/config/runtime.py | 3 + .../task_step/task_steps/install_agent.py | 2 + .../run_container_unit_tests_verifier.py | 11 +++ tst/unit/config/test_model.py | 70 +++++++++++++++ tst/unit/config/test_unread_keys.py | 2 +- .../task_step/local_container_urls_test.py | 88 ++++++++++++++++++- 13 files changed, 234 insertions(+), 15 deletions(-) diff --git a/.agentenv/config.example.toml b/.agentenv/config.example.toml index 30903567..6e2f182d 100644 --- a/.agentenv/config.example.toml +++ b/.agentenv/config.example.toml @@ -21,6 +21,13 @@ default_a2a_agent_id = "a2a-default" # what a deploy_agent step that names no # base_url = "http://localhost:4000" # a model proxy on this machine; local containers reach it as host.docker.internal # # on Linux, bind it to the Docker bridge address and use that: http://172.17.0.1:4000 # api_key = "env:MY_PROXY_KEY" +# +# Hugging Face Inference Providers instead: a fine-grained token with only "Make calls to Inference +# Providers", since this key reaches every agent. The headers go with every call to this endpoint. +# base_url = "https://router.huggingface.co/v1" +# api_key = "env:HF_INFERENCE_TOKEN" +# [model.headers] +# "X-HF-Bill-To" = "my-org" # bill an organization instead of the token's account [runner] impl = "agent_env.runner.local_runner:LocalRunner" diff --git a/AGENTS.md b/AGENTS.md index e8bda252..d2c63bbb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,11 @@ or one of them only where it is used. this machine that local containers call, a model proxy say, has to listen there: bind it to the bridge address and use that address in its URL. Binding `0.0.0.0` also works, but exposes the service to the network. +- `[model.headers]` go with every call to the configured endpoint and nowhere else: as + `extra_headers` on agent-env's own calls, and as `LITELLM_EXTRA_HEADERS` (a JSON object) to every + agent and unit-test verifier pointed at that endpoint, and to an installed agent that declares the + `litellm_extra_headers` param. The verifier also gets them as `ANTHROPIC_CUSTOM_HEADERS`. Core only + writes that variable, never reads it. - A host-mode `install_agent` (no `container_name`) on a `local` sandbox runs the agent's `install_commands_host` on this machine, with your privileges. Every command a local sandbox runs carries `AGENTENV_SANDBOX=`, and teardown stops the processes that carry it, then removes diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 8a4efa7f..03cb6da6 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -679,8 +679,9 @@ multi-operation custom extensions, and advanced ASGI-lifespan plus common SDK-operation override hooks. The agent examples make real OpenAI-compatible model calls. Set -`LITELLM_API_KEY` and, when needed, `LITELLM_BASE_URL`; their typed agent config -selects the model and system prompt for each deployment. Tests inject a fake +`LITELLM_API_KEY` and, when needed, `LITELLM_BASE_URL`; they also send the headers in +`LITELLM_EXTRA_HEADERS`, a JSON object AgentEnv sets when the endpoint needs any. Their typed +agent config selects the model and system prompt for each deployment. Tests inject a fake model client, so the example suite remains offline and deterministic. ### Object transfer diff --git a/packages/agentenv-protocol/examples/reference_model.py b/packages/agentenv-protocol/examples/reference_model.py index 7216d23a..be54dc97 100644 --- a/packages/agentenv-protocol/examples/reference_model.py +++ b/packages/agentenv-protocol/examples/reference_model.py @@ -48,7 +48,8 @@ def _request(self, *, api_key: str | None = None) -> tuple[str, dict[str, str]]: if base_url.endswith("/v1") else f"{base_url}/v1/chat/completions" ) - return endpoint, {"Authorization": f"Bearer {api_key}"} + headers = json.loads(os.getenv("LITELLM_EXTRA_HEADERS") or "{}") + return endpoint, {**headers, "Authorization": f"Bearer {api_key}"} async def complete( self, diff --git a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py index 800bfe89..bd4aa353 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py @@ -64,9 +64,20 @@ def _load_example(name: str) -> ModuleType: @pytest.mark.asyncio -async def test_reference_model_forwards_generic_metadata( +@pytest.mark.parametrize("endpoint_headers, sent_headers", [ + (None, {"Authorization": "Bearer request-scoped-key"}), + ('{"X-Bill-To": "acme", "Authorization": "Bearer theirs"}', + {"X-Bill-To": "acme", "Authorization": "Bearer request-scoped-key"}), +], ids=["no-endpoint-headers", "endpoint-headers"]) +async def test_reference_model_forwards_generic_metadata_and_the_endpoints_headers( monkeypatch: pytest.MonkeyPatch, + endpoint_headers: str | None, + sent_headers: dict[str, str], ) -> None: + if endpoint_headers is None: + monkeypatch.delenv("LITELLM_EXTRA_HEADERS", raising=False) + else: + monkeypatch.setenv("LITELLM_EXTRA_HEADERS", endpoint_headers) captured: dict[str, Any] = {} class Response: def raise_for_status(self) -> None: @@ -107,7 +118,7 @@ async def post(self, url: str, **kwargs: Any) -> Response: "client": {"timeout": 120}, "request": { "url": "https://models.example/v1/chat/completions", - "headers": {"Authorization": "Bearer request-scoped-key"}, + "headers": sent_headers, "json": { "model": "test-model", "messages": [{"role": "user", "content": "hello"}], diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 27527d7f..71bd8efe 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -13,6 +13,7 @@ from agentenv_protocol.a2a_agent import STANDARD_EXTENSIONS from agent_env.config import get_config +from agent_env.config.model import MODEL_HEADERS_ENV from agent_env.entity_refs import EntityRef from agent_env.providers.sandbox_providers.sandbox import ( SANDBOX_LABEL, @@ -436,6 +437,10 @@ async def deploy( resolved_env["LITELLM_API_KEY"] = config.get_litellm_api_key() if "LITELLM_BASE_URL" not in resolved_env and "LITELLM_BASE_URL" not in self.default_env_vars: resolved_env["LITELLM_BASE_URL"] = config.get_litellm_base_url() + if MODEL_HEADERS_ENV not in resolved_env and MODEL_HEADERS_ENV not in self.default_env_vars: + headers = config.get_model_headers(resolved_env.get("LITELLM_BASE_URL")) + if headers: + resolved_env[MODEL_HEADERS_ENV] = json.dumps(headers) if "LITELLM_BASE_URL" in resolved_env: resolved_env["LITELLM_BASE_URL"] = (provider or sandbox).url_from_sandbox(resolved_env["LITELLM_BASE_URL"]) diff --git a/src/agent_env/config/model.py b/src/agent_env/config/model.py index e83dd619..eac55bb5 100644 --- a/src/agent_env/config/model.py +++ b/src/agent_env/config/model.py @@ -1,6 +1,6 @@ -"""Resolution of the ``[model]`` config.toml block: endpoint, default/per-role models, and -provider params. A ``ModelConfig`` value object parsed from the ``[model]`` table; a -``secret_resolver`` is injected per call so secret references resolve lazily.""" +"""Resolution of the ``[model]`` config.toml block: endpoint, default/per-role models, +provider params, and the endpoint's headers. A ``ModelConfig`` value object parsed from the +``[model]`` table; a ``secret_resolver`` is injected per call so secret references resolve lazily.""" from __future__ import annotations @@ -14,6 +14,8 @@ _ENV_BASE_URL = "LITELLM_BASE_URL" _ENV_API_KEY = "LITELLM_API_KEY" +# Agents get ``[model.headers]`` here as a JSON object; unlike its two siblings, agent-env never reads it. +MODEL_HEADERS_ENV = "LITELLM_EXTRA_HEADERS" class ModelParam(StrEnum): @@ -57,13 +59,14 @@ class ModelConfig: default: Optional[str] = None roles: dict = field(default_factory=dict) params: dict = field(default_factory=dict) + headers: dict = field(default_factory=dict) # values may hold unresolved env:/secret: references @classmethod def from_section(cls, section: dict, *, secret_resolver: SecretResolver) -> ModelConfig: """Build from the raw ``[model]`` section: reject unknown keys and reserved ``[model.params]`` keys, then resolve ``env:``/``secret:`` refs — except - ``api_key``, which keeps its raw reference and resolves lazily per call - (an eager secret fetch would make every ``[model]`` read need credentials).""" + ``api_key`` and ``[model.headers]``, which keep their raw references and resolve + lazily per call (an eager secret fetch would make every ``[model]`` read need credentials).""" allowed = {f.name for f in fields(cls)} unknown = set(section) - allowed if unknown: @@ -74,8 +77,11 @@ def from_section(cls, section: dict, *, secret_resolver: SecretResolver) -> Mode f"[model.params] may not set reserved keys {sorted(reserved)} " "(agent-env sets these on the model call)" ) + headers = section.get("headers", {}) + if not isinstance(headers, dict) or not all(isinstance(v, str) for v in headers.values()): + raise ConfigError("[model.headers] maps header names to strings (an env:/secret: reference is fine)") cfg = interpolate( - {k: v for k, v in section.items() if k != "api_key"}, + {k: v for k, v in section.items() if k not in ("api_key", "headers")}, secret_resolver=secret_resolver, ) return cls( @@ -84,6 +90,7 @@ def from_section(cls, section: dict, *, secret_resolver: SecretResolver) -> Mode default=cfg.get("default"), roles=cfg.get("roles", {}), params=cfg.get("params", {}), + headers=dict(headers), ) def model_for_role(self, role: str) -> Optional[str]: @@ -106,6 +113,14 @@ def _resolved_api_key(self, secret_resolver: SecretResolver) -> Optional[str]: return None return interpolate(self.api_key, secret_resolver=secret_resolver) + def headers_for(self, base_url: Optional[str], *, secret_resolver: SecretResolver) -> dict[str, str]: + """``[model.headers]``, resolved, for a call or an agent pointed at ``base_url``: none unless + that is the configured endpoint, since a header may carry a credential.""" + configured_base = os.getenv(_ENV_BASE_URL) or self.base_url + if not self.headers or not base_url or base_url != configured_base: + return {} + return interpolate(dict(self.headers), secret_resolver=secret_resolver) + def resolve_call( self, model: str, @@ -151,4 +166,8 @@ def resolve_call( # URL would receive the configured proxy's credential. if key is None and (api_base is None or api_base == configured_base): key = self._resolved_api_key(secret_resolver) - return ModelCallConfig(api_key=key, api_base=api_base, params=dict(self.params)) + params = dict(self.params) + headers = self.headers_for(api_base, secret_resolver=secret_resolver) + if headers: + params["extra_headers"] = {**headers, **params.get("extra_headers", {})} + return ModelCallConfig(api_key=key, api_base=api_base, params=params) diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 0431042e..98958916 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -407,6 +407,9 @@ def get_model_params(self, overrides: dict | None = None) -> dict: params.update(config_loader.interpolate(overrides, secret_resolver=self._resolve_secret)) return params + def get_model_headers(self, base_url: str | None) -> dict[str, str]: + return self._model_config().headers_for(base_url, secret_resolver=self._resolve_secret) + def resolve_model_call( self, model: str, diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index b38157f4..46974d96 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -96,6 +96,8 @@ def _resolve_container() -> str: # Config-sourced (from agent-env config) "litellm_api_key": lambda: litellm_api_key_override or config.get_litellm_api_key(), "litellm_base_url": lambda: host_url_for(config.get_litellm_base_url(), sandbox_type), + # A JSON object, {} when [model.headers] is unset. + "litellm_extra_headers": lambda: json.dumps(config.get_model_headers(config.get_litellm_base_url())), } diff --git a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py index 654e8e83..5b64dd63 100644 --- a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py @@ -33,6 +33,7 @@ import uuid from typing import Any, ClassVar, Optional +from agent_env.config.model import MODEL_HEADERS_ENV from agent_env.providers.sandbox_providers.local_sandbox import host_url_for from agent_env.providers.sandbox_providers.sandbox_provider import ( all_sandbox_container_env, @@ -65,6 +66,7 @@ _RESERVED_ENV_KEYS = frozenset({ "LITELLM_API_KEY", "ANTHROPIC_API_KEY", "LITELLM_BASE_URL", "ANTHROPIC_BASE_URL", + MODEL_HEADERS_ENV, "ANTHROPIC_CUSTOM_HEADERS", "PATH", "HOME", }) _MODEL_URL_KEYS = frozenset({"LITELLM_BASE_URL", "ANTHROPIC_BASE_URL"}) @@ -255,6 +257,15 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: command, extra_env = self._resolve_command(context) merged_env.update(self.env_vars) merged_env.update(extra_env) + try: + headers = config.get_model_headers(merged_env.get("LITELLM_BASE_URL")) + if headers: + merged_env.setdefault(MODEL_HEADERS_ENV, json.dumps(headers)) + headers = config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")) + if headers: + merged_env.setdefault("ANTHROPIC_CUSTOM_HEADERS", "\n".join(f"{k}: {v}" for k, v in headers.items())) + except (KeyError, ConfigError): + logger.warning("[model.headers] unavailable; the verifier's model calls won't carry them") for key in _MODEL_URL_KEYS & merged_env.keys(): merged_env[key] = host_url_for(merged_env[key], sandbox.type) env_flags = " ".join(f"-e {k}={shlex.quote(v)}" for k, v in merged_env.items()) diff --git a/tst/unit/config/test_model.py b/tst/unit/config/test_model.py index 38f41494..a1e6c4b8 100644 --- a/tst/unit/config/test_model.py +++ b/tst/unit/config/test_model.py @@ -57,6 +57,22 @@ def test_from_section_rejects_unknown_key(): ModelConfig.from_section({"base_ur": "typo"}, secret_resolver=lambda ref: None) +def test_from_section_keeps_header_refs_raw(): + calls = [] + cfg = ModelConfig.from_section( + {"base_url": "https://gw/v1", "headers": {"X-Bill-To": "secret:org", "X-Plain": "v"}}, + secret_resolver=lambda ref: calls.append(ref) or "resolved", + ) + assert cfg.headers == {"X-Bill-To": "secret:org", "X-Plain": "v"} + assert calls == [] + + +@pytest.mark.parametrize("headers", [{"X-Retries": 3}, ["X-Bill-To"]]) +def test_from_section_rejects_headers_that_are_not_strings(headers): + with pytest.raises(ConfigError, match=r"\[model.headers\]"): + ModelConfig.from_section({"headers": headers}, secret_resolver=lambda ref: None) + + # --- model_for_role ----------------------------------------------------------- def test_model_for_role_hit_and_fallback(): @@ -168,3 +184,57 @@ def test_client_kwargs_omits_none_api_base_and_splats_params(): cc = ModelCallConfig(api_key="k", api_base=None, params={"aws_region_name": "us-west-2"}) assert cc.client_kwargs() == {"api_key": "k", "aws_region_name": "us-west-2"} assert ModelCallConfig("k", "https://b", {}).client_kwargs() == {"api_key": "k", "api_base": "https://b"} + + +# --- headers ------------------------------------------------------------------ + +_WITH_HEADERS = ModelConfig(base_url="https://gw/v1", headers={"X-Bill-To": "secret:org", "X-Plain": "v"}) + + +def test_headers_for_the_configured_endpoint_resolve_their_refs(): + assert _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({"org": "acme"})) == { + "X-Bill-To": "acme", "X-Plain": "v"} + + +def test_headers_follow_the_endpoint_the_env_var_configures(monkeypatch): + monkeypatch.setenv("LITELLM_BASE_URL", "https://env-gw/v1") + assert _WITH_HEADERS.headers_for("https://env-gw/v1", secret_resolver=_secrets({"org": "acme"}))["X-Bill-To"] == "acme" + assert _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({})) == {} + + +@pytest.mark.parametrize("base_url", ["https://elsewhere.example/v1", None]) +def test_headers_go_nowhere_but_the_configured_endpoint_and_resolve_nothing(base_url): + assert _WITH_HEADERS.headers_for(base_url, secret_resolver=_secrets({})) == {} + + +def test_headers_without_a_configured_endpoint_go_nowhere(): + assert ModelConfig(headers={"X-Plain": "v"}).headers_for("https://gw/v1", secret_resolver=_secrets({})) == {} + + +def test_an_unresolvable_header_ref_fails_loud(): + with pytest.raises(ConfigError, match="Unresolved secret:org"): + _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({})) + + +def test_resolve_call_sends_the_headers_to_the_configured_endpoint(): + call = _WITH_HEADERS.resolve_call("m", secret_resolver=_secrets({"org": "acme"})) + assert call.params == {"extra_headers": {"X-Bill-To": "acme", "X-Plain": "v"}} + + +def test_resolve_call_sends_no_headers_to_a_foreign_override(): + call = _WITH_HEADERS.resolve_call("m", base_override="https://elsewhere.example", secret_resolver=_secrets({})) + assert call.params == {} + + +def test_resolve_call_sends_no_headers_to_a_model_override(): + call = _WITH_HEADERS.resolve_call( + "m", model_overrides={"m": {"secret_key": "k", "base_url": "https://gw/v1"}}, secret_resolver=_secrets({"k": "sk"})) + assert call.params == {} + + +def test_params_extra_headers_win_per_header(): + cfg = ModelConfig(base_url="https://gw/v1", headers={"X-Bill-To": "acme", "X-Plain": "v"}, + params={"temperature": 0, "extra_headers": {"X-Bill-To": "other"}}) + call = cfg.resolve_call("m", secret_resolver=_secrets({})) + assert call.params == {"temperature": 0, "extra_headers": {"X-Bill-To": "other", "X-Plain": "v"}} + assert cfg.params == {"temperature": 0, "extra_headers": {"X-Bill-To": "other"}} diff --git a/tst/unit/config/test_unread_keys.py b/tst/unit/config/test_unread_keys.py index c49e7304..72b83bc6 100644 --- a/tst/unit/config/test_unread_keys.py +++ b/tst/unit/config/test_unread_keys.py @@ -212,7 +212,7 @@ def test_a_config_key_in_a_free_form_table_shadows_nothing(tmp_path, monkeypatch def test_a_key_beside_config_in_a_table_that_takes_no_config_is_named(tmp_path, monkeypatch): - takes = "[model] takes 'base_url', 'api_key', 'default', 'roles' and 'params'." + takes = "[model] takes 'base_url', 'api_key', 'default', 'roles', 'params' and 'headers'." assert _warnings(tmp_path, monkeypatch, '[model]\nmode = "x"\nconfig = { mode = "y" }\n') == [ f"[model] has a key named 'mode', which agent-env does not read: {takes}", f"[model] has a key named 'config', which agent-env does not read: {takes}", diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index c16d8dd5..b939649f 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -1,11 +1,13 @@ """A URL on this machine that agent-env hands into a container on it reaches the host: an agent's, a unit-test verifier's and an installed agent's model endpoint, a peer's and a trigger executor's A2A URL, and a CLI's gateway -URL. Containers elsewhere get them as they are.""" +URL. Containers elsewhere get them as they are. The model endpoint's ``[model.headers]`` go with it, and only with +it.""" from __future__ import annotations import asyncio import json +import shlex import types import httpx @@ -15,6 +17,7 @@ from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.artifact.artifacts.cli import CliArtifact from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.config.model import ModelConfig from agent_env.env.env import DeployedGatewayEnv from agent_env.providers.sandbox_providers import sandbox_provider from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider @@ -29,6 +32,8 @@ MODEL = "http://localhost:4000/v1" MODEL_IN_CONTAINER = "http://host.docker.internal:4000/v1" +HEADERS = {"X-Bill-To": "acme"} +ELSEWHERE = "http://127.0.0.1:5001" @pytest.fixture @@ -37,6 +42,11 @@ def model_endpoint(monkeypatch): monkeypatch.setenv("LITELLM_API_KEY", "k") +@pytest.fixture +def model_headers(monkeypatch, local_stores, model_endpoint): + monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers=HEADERS)) + + class _Started(Exception): """Raised once a fake has seen the container's env, which is all a test reads.""" @@ -71,6 +81,46 @@ async def test_an_agent_gets_the_model_endpoint_its_container_reaches( await _agent().deploy(sandbox_type="any", env_vars={"LITELLM_BASE_URL": MODEL} if passed else None) assert provider.env["LITELLM_BASE_URL"] == expected + assert "LITELLM_EXTRA_HEADERS" not in provider.env + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_cls", [_LocalProvider, _RemoteProvider]) +@pytest.mark.parametrize("passed", [False, True], ids=["configured", "passed"]) +async def test_an_agent_gets_the_model_endpoints_headers(monkeypatch, model_headers, provider_cls, passed): + provider = provider_cls() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) + + with pytest.raises(_Started): + await _agent().deploy(sandbox_type="any", env_vars={"LITELLM_BASE_URL": MODEL} if passed else None) + + assert json.loads(provider.env["LITELLM_EXTRA_HEADERS"]) == HEADERS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deploy_env, agent_fields", [ + ({"LITELLM_BASE_URL": ELSEWHERE}, {}), + (None, {"default_env_vars": {"LITELLM_BASE_URL": ELSEWHERE}}), +], ids=["passed-elsewhere", "its-own-endpoint"]) +async def test_an_agent_pointed_elsewhere_gets_no_headers(monkeypatch, model_headers, deploy_env, agent_fields): + provider = _RemoteProvider() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) + + with pytest.raises(_Started): + await _agent(**agent_fields).deploy(sandbox_type="any", env_vars=deploy_env) + + assert "LITELLM_EXTRA_HEADERS" not in provider.env + + +@pytest.mark.asyncio +async def test_headers_an_agent_is_given_win(monkeypatch, model_headers): + provider = _RemoteProvider() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) + + with pytest.raises(_Started): + await _agent().deploy(sandbox_type="any", env_vars={"LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}'}) + + assert provider.env["LITELLM_EXTRA_HEADERS"] == '{"X-Bill-To": "mine"}' @pytest.mark.asyncio @@ -122,7 +172,7 @@ async def exec_with_output(self, *args) -> tuple[int, str, str]: return 0, "ok", "" -def _run_verifier(monkeypatch, sandbox_type: str, **step_fields) -> str: +def _run_verifier(monkeypatch, sandbox_type: str, seed: dict | None = None, **step_fields) -> str: """The verifier's test command, as it ran on a sandbox of ``sandbox_type``.""" sandbox = _VerifierSandbox(sandbox_type) provider = types.SimpleNamespace(get_sandbox=lambda sandbox_id: asyncio.sleep(0, result=sandbox)) @@ -132,6 +182,8 @@ def _run_verifier(monkeypatch, sandbox_type: str, **step_fields) -> str: id=artifact_id, version=1, object_url=object_url))) context = TaskStepContext() context.metadata["deployed_docker_containers"] = [{"container_name": "c", "sandbox_name": "h"}] + if seed: + context.metadata["seed"] = seed context.deployed_sandboxes.append( DeployedSandbox(sandbox_name="h", sandbox_id="sb-1", sandbox_mode="vm", sandbox_type=sandbox_type)) @@ -157,6 +209,26 @@ def test_a_model_endpoint_the_verifier_sets_itself_is_reached_through_the_host(m assert "-e LITELLM_BASE_URL=http://host.docker.internal:5001 " in run +def test_a_unit_test_verifier_gets_the_model_endpoints_headers(monkeypatch, model_headers): + run = _run_verifier(monkeypatch, "local") + + assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(HEADERS))} " in run + assert f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote('X-Bill-To: acme')} " in run + + +def test_a_unit_test_verifier_pointed_elsewhere_gets_no_headers(monkeypatch, model_headers): + run = _run_verifier(monkeypatch, "local", env_vars={"LITELLM_BASE_URL": ELSEWHERE, "ANTHROPIC_BASE_URL": ELSEWHERE}) + + assert "LITELLM_EXTRA_HEADERS" not in run and "ANTHROPIC_CUSTOM_HEADERS" not in run + + +def test_a_seed_column_cant_replace_the_verifiers_headers(monkeypatch, model_headers): + run = _run_verifier(monkeypatch, "local", seed={"litellm_extra_headers": "{}", "anthropic_custom_headers": "x"}) + + assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(HEADERS))} " in run + assert f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote('X-Bill-To: acme')} " in run + + @pytest.mark.parametrize("sandbox_type, expected", [ ("local", MODEL_IN_CONTAINER), ("modal_vm", MODEL), (None, MODEL), ], ids=["local-container", "remote-container", "host-install"]) @@ -169,6 +241,18 @@ def test_an_installed_agent_gets_the_model_endpoint_it_reaches(sandbox_type, exp assert resolvers["litellm_base_url"]() == expected +@pytest.mark.parametrize("configured, expected", [(True, HEADERS), (False, {})], ids=["headers", "none"]) +def test_an_installed_agent_gets_the_model_endpoints_headers(monkeypatch, local_stores, model_endpoint, configured, expected): + if configured: + monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers=HEADERS)) + resolvers = _build_param_resolvers( + container_name="c", work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, agent_name="solver", + config=local_stores, sandbox_type="local", + ) + + assert json.loads(resolvers["litellm_extra_headers"]()) == expected + + _PEERS_CARD = {"capabilities": {"extensions": [ {"uri": A2AAgent.EXT_PEER_AGENTS, "params": {"endpoint": "/ext/peer-agents"}}]}} _HUMAN = "http://localhost:18000/api/v1/a2a/human/instance/i-1" From 5aa461ae94b0028a77e553447b2888953de91ed5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 22:56:51 -0700 Subject: [PATCH 2/8] fix(config): scope model headers by the agent's effective URL, case-insensitively; keep secrets out of install logs - A2AAgent.deploy reads the URL an agent will use from its own default_env_vars too, so an agent declaring the configured endpoint itself still gets [model.headers]. - [model.params.extra_headers] overrides a [model.headers] entry whatever the case of its name. - install_agent logs each command's template, not the substituted command: a resolved {litellm_api_key} or {litellm_extra_headers} reached the log before. - The protocol's reference model sends LITELLM_EXTRA_HEADERS only to the endpoint AgentEnv configured, and never lets one replace its own Authorization. Co-Authored-By: Claude Opus 5.5 --- .../examples/reference_model.py | 7 +- .../tests/test_a2a_agent_examples.py | 11 ++-- src/agent_env/a2a_agent/a2a_agent.py | 3 +- src/agent_env/config/model.py | 4 +- .../task_step/task_steps/install_agent.py | 3 +- tst/unit/config/test_model.py | 8 +-- tst/unit/task_step/install_agent_test.py | 65 ++++++++++++++++++- .../task_step/local_container_urls_test.py | 11 ++++ 8 files changed, 98 insertions(+), 14 deletions(-) diff --git a/packages/agentenv-protocol/examples/reference_model.py b/packages/agentenv-protocol/examples/reference_model.py index be54dc97..e308edf5 100644 --- a/packages/agentenv-protocol/examples/reference_model.py +++ b/packages/agentenv-protocol/examples/reference_model.py @@ -42,14 +42,17 @@ def _request(self, *, api_key: str | None = None) -> tuple[str, dict[str, str]]: raise RuntimeError( "LITELLM_BASE_URL must be set before running the agent" ) + # The endpoint's headers are for the endpoint AgentEnv configured, never one this client was pointed at. + headers = (json.loads(os.getenv("LITELLM_EXTRA_HEADERS") or "{}") + if base_url == os.getenv("LITELLM_BASE_URL") else {}) base_url = base_url.rstrip("/") endpoint = ( f"{base_url}/chat/completions" if base_url.endswith("/v1") else f"{base_url}/v1/chat/completions" ) - headers = json.loads(os.getenv("LITELLM_EXTRA_HEADERS") or "{}") - return endpoint, {**headers, "Authorization": f"Bearer {api_key}"} + own = {k: v for k, v in headers.items() if k.lower() != "authorization"} + return endpoint, {**own, "Authorization": f"Bearer {api_key}"} async def complete( self, diff --git a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py index bd4aa353..cf3fe114 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py @@ -64,16 +64,19 @@ def _load_example(name: str) -> ModuleType: @pytest.mark.asyncio -@pytest.mark.parametrize("endpoint_headers, sent_headers", [ - (None, {"Authorization": "Bearer request-scoped-key"}), - ('{"X-Bill-To": "acme", "Authorization": "Bearer theirs"}', +@pytest.mark.parametrize("endpoint_headers, configured_url, sent_headers", [ + (None, "https://models.example/v1", {"Authorization": "Bearer request-scoped-key"}), + ('{"X-Bill-To": "acme", "authorization": "Bearer theirs"}', "https://models.example/v1", {"X-Bill-To": "acme", "Authorization": "Bearer request-scoped-key"}), -], ids=["no-endpoint-headers", "endpoint-headers"]) + ('{"X-Bill-To": "acme"}', "https://configured.example/v1", {"Authorization": "Bearer request-scoped-key"}), +], ids=["no-endpoint-headers", "endpoint-headers", "another-endpoint"]) async def test_reference_model_forwards_generic_metadata_and_the_endpoints_headers( monkeypatch: pytest.MonkeyPatch, endpoint_headers: str | None, + configured_url: str, sent_headers: dict[str, str], ) -> None: + monkeypatch.setenv("LITELLM_BASE_URL", configured_url) if endpoint_headers is None: monkeypatch.delenv("LITELLM_EXTRA_HEADERS", raising=False) else: diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 71bd8efe..ca561544 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -438,7 +438,8 @@ async def deploy( if "LITELLM_BASE_URL" not in resolved_env and "LITELLM_BASE_URL" not in self.default_env_vars: resolved_env["LITELLM_BASE_URL"] = config.get_litellm_base_url() if MODEL_HEADERS_ENV not in resolved_env and MODEL_HEADERS_ENV not in self.default_env_vars: - headers = config.get_model_headers(resolved_env.get("LITELLM_BASE_URL")) + headers = config.get_model_headers( + resolved_env.get("LITELLM_BASE_URL", self.default_env_vars.get("LITELLM_BASE_URL"))) if headers: resolved_env[MODEL_HEADERS_ENV] = json.dumps(headers) if "LITELLM_BASE_URL" in resolved_env: diff --git a/src/agent_env/config/model.py b/src/agent_env/config/model.py index eac55bb5..65dbf8db 100644 --- a/src/agent_env/config/model.py +++ b/src/agent_env/config/model.py @@ -169,5 +169,7 @@ def resolve_call( params = dict(self.params) headers = self.headers_for(api_base, secret_resolver=secret_resolver) if headers: - params["extra_headers"] = {**headers, **params.get("extra_headers", {})} + explicit = params.get("extra_headers", {}) + overridden = {name.lower() for name in explicit} + params["extra_headers"] = {**{k: v for k, v in headers.items() if k.lower() not in overridden}, **explicit} return ModelCallConfig(api_key=key, api_base=api_base, params=params) diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 46974d96..ca7c1eb9 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -292,7 +292,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: f"{commands_key}[{i-1}] references undeclared placeholder {{{e.args[0]}}}; " f"add it to the extension's required_params" ) - logger.info(f" [{i}/{len(install_commands)}] {formatted.splitlines()[0][:140]}") + # The template, not the command: resolved params carry the model key and headers. + logger.info(f" [{i}/{len(install_commands)}] {cmd.splitlines()[0][:140]}") await sandbox.exec_script(formatted) if not host_mode: # copied into the container; a host install may run from it await sandbox.exec_script(f"rm -rf {shlex.quote(work_dir)}") diff --git a/tst/unit/config/test_model.py b/tst/unit/config/test_model.py index a1e6c4b8..40a044b7 100644 --- a/tst/unit/config/test_model.py +++ b/tst/unit/config/test_model.py @@ -232,9 +232,9 @@ def test_resolve_call_sends_no_headers_to_a_model_override(): assert call.params == {} -def test_params_extra_headers_win_per_header(): +def test_params_extra_headers_win_per_header_whatever_its_case(): cfg = ModelConfig(base_url="https://gw/v1", headers={"X-Bill-To": "acme", "X-Plain": "v"}, - params={"temperature": 0, "extra_headers": {"X-Bill-To": "other"}}) + params={"temperature": 0, "extra_headers": {"x-bill-to": "other"}}) call = cfg.resolve_call("m", secret_resolver=_secrets({})) - assert call.params == {"temperature": 0, "extra_headers": {"X-Bill-To": "other", "X-Plain": "v"}} - assert cfg.params == {"temperature": 0, "extra_headers": {"X-Bill-To": "other"}} + assert call.params == {"temperature": 0, "extra_headers": {"X-Plain": "v", "x-bill-to": "other"}} + assert cfg.params == {"temperature": 0, "extra_headers": {"x-bill-to": "other"}} diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index 55d88249..c5440738 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -1,9 +1,17 @@ -"""Unit tests for InstallAgentTaskStep's workspace_dir plumbing (no I/O).""" +"""Unit tests for InstallAgentTaskStep's workspace_dir plumbing, and what its install logs show (no I/O).""" from __future__ import annotations +import asyncio +import logging +from types import SimpleNamespace + import pytest +from agent_env.a2a_agent import A2AAgent +from agent_env.config.model import ModelConfig +from agent_env.providers.sandbox_providers import sandbox_provider +from agent_env.task_step.context import DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.install_agent import ( InstallAgentTaskStep, _build_param_resolvers, @@ -111,3 +119,58 @@ def test_a2a_port_roundtrip(): assert InstallAgentTaskStep.from_dict(step.to_dict()).a2a_port == 8010 step2 = InstallAgentTaskStep(id="s", version=None, sandbox_name="vm", a2a_agent_id="agent") assert InstallAgentTaskStep.from_dict(step2.to_dict()).a2a_port is None + + +class _InstallSandbox: + def __init__(self): + self.scripts: list[str] = [] + + def scoped_name(self, name): + return name + + def host_port(self, port): + return port + + async def load_docker_images(self, images): + return None + + async def load_object_file(self, url, destination): + return None + + async def exec_script(self, script): + self.scripts.append(script) + return "" + + +def _resolved(value): + return lambda *args, **kwargs: asyncio.sleep(0, result=value) + + +@pytest.mark.asyncio +async def test_install_logs_show_the_command_template_not_the_model_key_or_headers(monkeypatch, caplog, local_stores): + monkeypatch.setenv("LITELLM_BASE_URL", "http://localhost:4000/v1") + monkeypatch.setenv("LITELLM_API_KEY", "sk-never-logged") + monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers={"X-Bill-To": "org-never-logged"})) + sandbox = _InstallSandbox() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", + lambda name: SimpleNamespace(get_sandbox=_resolved(sandbox))) + image = SimpleNamespace(image_name="img", build_context_object_url="s3://bucket/ctx.tar.gz") + monkeypatch.setattr(A2AAgent, "get", staticmethod(lambda *args: SimpleNamespace( + id="agent-x", version=1, metadata={}, docker_image_artifact=image))) + step = InstallAgentTaskStep(id="i", version=None, sandbox_name="vm", a2a_agent_id="agent-x", agent_name="solver") + monkeypatch.setattr(step, "_extract_install_extension", _resolved({"params": { + "install_commands_host": ["start --key {litellm_api_key} --headers {litellm_extra_headers}"], + "required_params_host": ["litellm_api_key", "litellm_extra_headers"], "a2a_port": 8000}})) + monkeypatch.setattr(step, "_wait_for_agent_card", _resolved({})) + monkeypatch.setattr(step, "_register_instance", lambda *args: SimpleNamespace(instance_id="solver-1")) + context = TaskStepContext(deployed_sandboxes=[DeployedSandbox( + sandbox_name="vm", sandbox_id="sb-1", sandbox_mode="vm", sandbox_type="local", + tunnel_urls={"8000": "http://127.0.0.1:8000"})]) + + with caplog.at_level(logging.INFO, logger="agent_env.task_step.task_steps.install_agent"): + await step.execute(context) + + [command] = [s for s in sandbox.scripts if s.startswith("start ")] + assert "sk-never-logged" in command and "org-never-logged" in command + assert "start --key {litellm_api_key} --headers {litellm_extra_headers}" in caplog.text + assert "never-logged" not in caplog.text diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index b939649f..d9a05727 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -112,6 +112,17 @@ async def test_an_agent_pointed_elsewhere_gets_no_headers(monkeypatch, model_hea assert "LITELLM_EXTRA_HEADERS" not in provider.env +@pytest.mark.asyncio +async def test_an_agent_whose_own_endpoint_is_the_configured_one_gets_its_headers(monkeypatch, model_headers): + provider = _RemoteProvider() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) + + with pytest.raises(_Started): + await _agent(default_env_vars={"LITELLM_BASE_URL": MODEL}).deploy(sandbox_type="any") + + assert json.loads(provider.env["LITELLM_EXTRA_HEADERS"]) == HEADERS + + @pytest.mark.asyncio async def test_headers_an_agent_is_given_win(monkeypatch, model_headers): provider = _RemoteProvider() From c615e4fe8c6570aa9add1526fd8384b27291e335 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 08:19:21 -0700 Subject: [PATCH 3/8] fix(config): refuse headers HTTP can't carry, mask them in config show, and say when they're held back - A [model.headers] name must be an HTTP token and a value visible ASCII, spaces and tabs: checked at load for a literal and when a reference resolves, so a non-ASCII value fails with a config error instead of inside an agent's client, and a CR or LF can't split one header into two. - config show and explain mask every literal under a `headers` key ([model.headers] and [model.params.extra_headers]): a name like X-Api-Key doesn't say it carries a key. References still show. - Headers held back from a URL that isn't the configured endpoint are logged; a trailing slash on either URL no longer counts as a different endpoint, for the headers and for the configured key alike. - The unit-test verifier fails when configured headers don't resolve, rather than running without them. Co-Authored-By: Claude Opus 5.5 --- .../examples/reference_model.py | 3 +- .../tests/test_a2a_agent_examples.py | 3 +- src/agent_env/config/describe.py | 3 +- src/agent_env/config/model.py | 34 ++++++++++++++++--- .../run_container_unit_tests_verifier.py | 17 +++++----- tst/unit/config/test_describe.py | 22 ++++++++++++ tst/unit/config/test_model.py | 33 ++++++++++++++++-- .../task_step/local_container_urls_test.py | 8 +++++ 8 files changed, 105 insertions(+), 18 deletions(-) diff --git a/packages/agentenv-protocol/examples/reference_model.py b/packages/agentenv-protocol/examples/reference_model.py index e308edf5..aadffd89 100644 --- a/packages/agentenv-protocol/examples/reference_model.py +++ b/packages/agentenv-protocol/examples/reference_model.py @@ -43,8 +43,9 @@ def _request(self, *, api_key: str | None = None) -> tuple[str, dict[str, str]]: "LITELLM_BASE_URL must be set before running the agent" ) # The endpoint's headers are for the endpoint AgentEnv configured, never one this client was pointed at. + configured = os.getenv("LITELLM_BASE_URL") or "" headers = (json.loads(os.getenv("LITELLM_EXTRA_HEADERS") or "{}") - if base_url == os.getenv("LITELLM_BASE_URL") else {}) + if base_url.rstrip("/") == configured.rstrip("/") else {}) base_url = base_url.rstrip("/") endpoint = ( f"{base_url}/chat/completions" diff --git a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py index cf3fe114..2de7408e 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py @@ -68,8 +68,9 @@ def _load_example(name: str) -> ModuleType: (None, "https://models.example/v1", {"Authorization": "Bearer request-scoped-key"}), ('{"X-Bill-To": "acme", "authorization": "Bearer theirs"}', "https://models.example/v1", {"X-Bill-To": "acme", "Authorization": "Bearer request-scoped-key"}), + ('{"X-Bill-To": "acme"}', "https://models.example/v1/", {"X-Bill-To": "acme", "Authorization": "Bearer request-scoped-key"}), ('{"X-Bill-To": "acme"}', "https://configured.example/v1", {"Authorization": "Bearer request-scoped-key"}), -], ids=["no-endpoint-headers", "endpoint-headers", "another-endpoint"]) +], ids=["no-endpoint-headers", "endpoint-headers", "trailing-slash", "another-endpoint"]) async def test_reference_model_forwards_generic_metadata_and_the_endpoints_headers( monkeypatch: pytest.MonkeyPatch, endpoint_headers: str | None, diff --git a/src/agent_env/config/describe.py b/src/agent_env/config/describe.py index d88ba0ec..d0b49a97 100644 --- a/src/agent_env/config/describe.py +++ b/src/agent_env/config/describe.py @@ -59,8 +59,9 @@ # debugging round-trip, under-masking prints a credential into a terminal and a paste. On # that trade `auth` is in even though it also catches `author` — a store `config` table is # where `authorization = "Bearer ..."` lives, and a masked author is the cheaper mistake. +# `headers` masks a whole table of HTTP headers: a name like `X-Api-Key` doesn't say it holds a key. _SECRET_KEY_PARTS = ("password", "passwd", "passphrase", "secret", "token", "credential", - "api_key", "private_key", "auth", "bearer") + "api_key", "private_key", "auth", "bearer", "headers") _URI_USERINFO = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?P[^/@\s]+)@") diff --git a/src/agent_env/config/model.py b/src/agent_env/config/model.py index 65dbf8db..7a297409 100644 --- a/src/agent_env/config/model.py +++ b/src/agent_env/config/model.py @@ -4,18 +4,26 @@ from __future__ import annotations +import logging import os +import re from dataclasses import dataclass, field, fields from enum import StrEnum from typing import Optional +from urllib.parse import urlparse from agent_env.config.errors import ConfigError -from agent_env.config.loader import SecretResolver, interpolate +from agent_env.config.loader import ENV_REF_PREFIX, SECRET_REF_PREFIX, SecretResolver, interpolate + +logger = logging.getLogger(__name__) _ENV_BASE_URL = "LITELLM_BASE_URL" _ENV_API_KEY = "LITELLM_API_KEY" # Agents get ``[model.headers]`` here as a JSON object; unlike its two siblings, agent-env never reads it. MODEL_HEADERS_ENV = "LITELLM_EXTRA_HEADERS" +# RFC 9110: a header name is a token; a value is visible ASCII, spaces and tabs (no CR or LF to split it on). +_HEADER_NAME = re.compile(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+\Z") +_HEADER_VALUE = re.compile(r"[\t\x20-\x7e]*\Z") class ModelParam(StrEnum): @@ -80,6 +88,8 @@ def from_section(cls, section: dict, *, secret_resolver: SecretResolver) -> Mode headers = section.get("headers", {}) if not isinstance(headers, dict) or not all(isinstance(v, str) for v in headers.values()): raise ConfigError("[model.headers] maps header names to strings (an env:/secret: reference is fine)") + for name, value in headers.items(): + _check_header(name, None if value.startswith((ENV_REF_PREFIX, SECRET_REF_PREFIX)) else value) cfg = interpolate( {k: v for k, v in section.items() if k not in ("api_key", "headers")}, secret_resolver=secret_resolver, @@ -116,10 +126,17 @@ def _resolved_api_key(self, secret_resolver: SecretResolver) -> Optional[str]: def headers_for(self, base_url: Optional[str], *, secret_resolver: SecretResolver) -> dict[str, str]: """``[model.headers]``, resolved, for a call or an agent pointed at ``base_url``: none unless that is the configured endpoint, since a header may carry a credential.""" + if not self.headers or not base_url: + return {} configured_base = os.getenv(_ENV_BASE_URL) or self.base_url - if not self.headers or not base_url or base_url != configured_base: + if not configured_base or base_url.rstrip("/") != configured_base.rstrip("/"): + logger.info("Not sending [model.headers] to %s, which is not the configured model endpoint", + urlparse(base_url).hostname or base_url) return {} - return interpolate(dict(self.headers), secret_resolver=secret_resolver) + headers = interpolate(dict(self.headers), secret_resolver=secret_resolver) + for name, value in headers.items(): + _check_header(name, value) + return headers def resolve_call( self, @@ -164,7 +181,7 @@ def resolve_call( # The configured key is scoped to the configured endpoint (and native routing): # a base_override pointing elsewhere must not inherit it, or a caller-supplied # URL would receive the configured proxy's credential. - if key is None and (api_base is None or api_base == configured_base): + if key is None and (api_base is None or api_base.rstrip("/") == (configured_base or "").rstrip("/")): key = self._resolved_api_key(secret_resolver) params = dict(self.params) headers = self.headers_for(api_base, secret_resolver=secret_resolver) @@ -173,3 +190,12 @@ def resolve_call( overridden = {name.lower() for name in explicit} params["extra_headers"] = {**{k: v for k, v in headers.items() if k.lower() not in overridden}, **explicit} return ModelCallConfig(api_key=key, api_base=api_base, params=params) + + +def _check_header(name: str, value: Optional[str]) -> None: + """Refuse a ``[model.headers]`` entry HTTP can't carry (``value`` None checks the name alone): a non-ASCII + value would fail inside an agent's client, and a CR or LF would split one header into two.""" + if not _HEADER_NAME.match(name): + raise ConfigError(f"[model.headers] {name!r} is not an HTTP header name") + if value is not None and not _HEADER_VALUE.match(value): + raise ConfigError(f"[model.headers] {name!r} has a value HTTP can't carry: only visible ASCII, spaces and tabs") diff --git a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py index 5b64dd63..2b91ca32 100644 --- a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py @@ -257,15 +257,14 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: command, extra_env = self._resolve_command(context) merged_env.update(self.env_vars) merged_env.update(extra_env) - try: - headers = config.get_model_headers(merged_env.get("LITELLM_BASE_URL")) - if headers: - merged_env.setdefault(MODEL_HEADERS_ENV, json.dumps(headers)) - headers = config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")) - if headers: - merged_env.setdefault("ANTHROPIC_CUSTOM_HEADERS", "\n".join(f"{k}: {v}" for k, v in headers.items())) - except (KeyError, ConfigError): - logger.warning("[model.headers] unavailable; the verifier's model calls won't carry them") + # Unlike a missing key or URL, configured headers that don't resolve fail the step: a verifier that ran + # without them would grade against an endpoint that answers differently, or bills someone else. + headers = config.get_model_headers(merged_env.get("LITELLM_BASE_URL")) + if headers: + merged_env.setdefault(MODEL_HEADERS_ENV, json.dumps(headers)) + headers = config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")) + if headers: + merged_env.setdefault("ANTHROPIC_CUSTOM_HEADERS", "\n".join(f"{k}: {v}" for k, v in headers.items())) for key in _MODEL_URL_KEYS & merged_env.keys(): merged_env[key] = host_url_for(merged_env[key], sandbox.type) env_flags = " ".join(f"-e {k}={shlex.quote(v)}" for k, v in merged_env.items()) diff --git a/tst/unit/config/test_describe.py b/tst/unit/config/test_describe.py index 939a2fe0..f71b118a 100644 --- a/tst/unit/config/test_describe.py +++ b/tst/unit/config/test_describe.py @@ -613,6 +613,28 @@ def test_a_literal_authorization_header_is_masked(tmp_path, monkeypatch): assert resolved["database"] == "keep" +def test_model_header_values_are_masked_whatever_the_header_is_called(tmp_path, monkeypatch): + """A header's name doesn't say it carries a credential (`X-Api-Key`), so every literal header value is + masked, in [model.headers] and [model.params.extra_headers] alike; a reference still names its secret.""" + path = _write_config(tmp_path, """ +[model] +base_url = "https://example.invalid" + +[model.headers] +"X-Api-Key" = "literal-header-key" +"X-Bill-To" = "secret:bill_to" + +[model.params.extra_headers] +"X-Portkey-Virtual-Key" = "literal-virtual-key" +""") + monkeypatch.setenv("AGENT_ENV_CONFIG", str(path)) + + rendered = render(describe_config()) + + assert "literal-header-key" not in rendered and "literal-virtual-key" not in rendered + assert "secret:bill_to" in rendered + + def test_a_section_name_used_as_a_key_inside_another_table_is_warned_about(tmp_path, monkeypatch): """TOML binds a bare key to the table above it, so `envs = "..."` typed at the foot of a file becomes `[task_steps].envs` — a well-formed document that configures nothing. No diff --git a/tst/unit/config/test_model.py b/tst/unit/config/test_model.py index 40a044b7..c7145333 100644 --- a/tst/unit/config/test_model.py +++ b/tst/unit/config/test_model.py @@ -73,6 +73,16 @@ def test_from_section_rejects_headers_that_are_not_strings(headers): ModelConfig.from_section({"headers": headers}, secret_resolver=lambda ref: None) +@pytest.mark.parametrize("headers, problem", [ + ({"X Bill To": "acme"}, "not an HTTP header name"), + ({"X-Bill-To": "café"}, "value HTTP can't carry"), + ({"X-Bill-To": "acme\r\nX-Injected: 1"}, "value HTTP can't carry"), +]) +def test_from_section_rejects_a_header_http_cant_carry(headers, problem): + with pytest.raises(ConfigError, match=problem): + ModelConfig.from_section({"headers": headers}, secret_resolver=lambda ref: None) + + # --- model_for_role ----------------------------------------------------------- def test_model_for_role_hit_and_fallback(): @@ -156,9 +166,10 @@ def test_base_override_to_a_foreign_url_does_not_inherit_the_configured_key(): assert call.api_base == "https://elsewhere.example" -def test_base_override_equal_to_the_configured_base_keeps_the_key(): +@pytest.mark.parametrize("override", ["https://proxy.example", "https://proxy.example/"]) +def test_base_override_equal_to_the_configured_base_keeps_the_key(override): call = ModelConfig(base_url="https://proxy.example", api_key="secret:k").resolve_call( - "m", base_override="https://proxy.example", secret_resolver=_secrets({"k": "sk-proxy"}) + "m", base_override=override, secret_resolver=_secrets({"k": "sk-proxy"}) ) assert call.api_key == "sk-proxy" @@ -196,6 +207,24 @@ def test_headers_for_the_configured_endpoint_resolve_their_refs(): "X-Bill-To": "acme", "X-Plain": "v"} +def test_headers_ignore_a_trailing_slash_on_either_url(): + assert _WITH_HEADERS.headers_for("https://gw/v1/", secret_resolver=_secrets({"org": "acme"}))["X-Bill-To"] == "acme" + assert ModelConfig(base_url="https://gw/v1/", headers={"X-Plain": "v"}).headers_for( + "https://gw/v1", secret_resolver=_secrets({})) == {"X-Plain": "v"} + + +def test_headers_held_back_from_another_endpoint_are_logged(caplog): + with caplog.at_level("INFO", logger="agent_env.config.model"): + _WITH_HEADERS.headers_for("https://elsewhere.example/v1", secret_resolver=_secrets({})) + + assert "Not sending [model.headers] to elsewhere.example" in caplog.text + + +def test_a_header_ref_that_resolves_to_a_value_http_cant_carry_fails_loud(): + with pytest.raises(ConfigError, match="value HTTP can't carry"): + _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({"org": "acme\nX-Injected: 1"})) + + def test_headers_follow_the_endpoint_the_env_var_configures(monkeypatch): monkeypatch.setenv("LITELLM_BASE_URL", "https://env-gw/v1") assert _WITH_HEADERS.headers_for("https://env-gw/v1", secret_resolver=_secrets({"org": "acme"}))["X-Bill-To"] == "acme" diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index d9a05727..d1c83aa8 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -17,6 +17,7 @@ from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.artifact.artifacts.cli import CliArtifact from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.config.errors import ConfigError from agent_env.config.model import ModelConfig from agent_env.env.env import DeployedGatewayEnv from agent_env.providers.sandbox_providers import sandbox_provider @@ -233,6 +234,13 @@ def test_a_unit_test_verifier_pointed_elsewhere_gets_no_headers(monkeypatch, mod assert "LITELLM_EXTRA_HEADERS" not in run and "ANTHROPIC_CUSTOM_HEADERS" not in run +def test_a_unit_test_verifier_whose_headers_dont_resolve_fails(monkeypatch, local_stores, model_endpoint): + monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers={"X-Bill-To": "secret:fd_no_such_secret"})) + + with pytest.raises(ConfigError, match="fd_no_such_secret"): + _run_verifier(monkeypatch, "local") + + def test_a_seed_column_cant_replace_the_verifiers_headers(monkeypatch, model_headers): run = _run_verifier(monkeypatch, "local", seed={"litellm_extra_headers": "{}", "anthropic_custom_headers": "x"}) From 656273c2842a01523ada941109004325b023cdc6 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 08:33:55 -0700 Subject: [PATCH 4/8] feat(sail_vm): inject [model.headers] at egress like the model key, so their values never enter the Sailbox With inject_model_key on, an agent's Sailbox already never holds its model key: Sail adds it to requests to the model endpoint as they leave. [model.headers] arrived as plain LITELLM_EXTRA_HEADERS, so a credential header undid that guarantee. - Each header value is a Sail secret of its own, AGENTENV_LITELLM_H_, kept like the key's; the saved policy sets each header from its secret next to the key's two. - Every command and file agent-env sends carries LITELLM_EXTRA_HEADERS (and ANTHROPIC_CUSTOM_HEADERS) with the names kept and sail-injected-header as each value, replaced as a whole so a short value can't rewrite other text; the policy overwrites the placeholders. - Reconnect recognises a policy with headers, and recovers their values for scrubbing from this process or from the configured [model.headers]; anyone else's policy is still unknown. - A Sailbox refuses a command or file carrying header values Sail doesn't inject for it, as it does a key. - A header named authorization or x-api-key, which the key is set as, is refused. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 3 +- .../sandbox_providers/sail_vm/README.md | 25 +-- .../sandbox_providers/sail_vm/model_key.py | 150 ++++++++++++++--- .../sandbox_providers/sail_vm/provider.py | 12 +- .../sandbox_providers/sail_vm/sandbox.py | 11 +- .../sail_vm_model_key_test.py | 152 ++++++++++++++++++ 6 files changed, 309 insertions(+), 44 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index d2c63bbb..cd828c96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -73,7 +73,8 @@ or one of them only where it is used. `extra_headers` on agent-env's own calls, and as `LITELLM_EXTRA_HEADERS` (a JSON object) to every agent and unit-test verifier pointed at that endpoint, and to an installed agent that declares the `litellm_extra_headers` param. The verifier also gets them as `ANTHROPIC_CUSTOM_HEADERS`. Core only - writes that variable, never reads it. + writes that variable, never reads it. On `sail_vm` with `inject_model_key`, Sail sets them at egress + like the key, and the Sailbox holds only placeholders. - A host-mode `install_agent` (no `container_name`) on a `local` sandbox runs the agent's `install_commands_host` on this machine, with your privileges. Every command a local sandbox runs carries `AGENTENV_SANDBOX=`, and teardown stops the processes that carry it, then removes diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/README.md b/src/agent_env/providers/sandbox_providers/sail_vm/README.md index fc007ce9..6bedd6ae 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/README.md +++ b/src/agent_env/providers/sandbox_providers/sail_vm/README.md @@ -189,16 +189,20 @@ asyncio.run(main()) - **Reconnect:** `get_sandbox` restores the ports and the applied egress policy. A policy it can't represent leaves `network_policy` unknown, and image loading then fails closed. -## The agent's model key +## The agent's model key and headers -With `inject_model_key` on, an agent's `LITELLM_API_KEY` never enters its Sailbox: +With `inject_model_key` on, an agent's `LITELLM_API_KEY` never enters its Sailbox, and neither do the +values of the endpoint's `[model.headers]`: 1. The key is stored as a Sail secret named `AGENTENV_LITELLM_`, one per distinct key. + Each header value is a secret of its own, `AGENTENV_LITELLM_H_`. 2. The Sailbox is created with a saved egress policy. On requests to `LITELLM_BASE_URL`'s host, the policy - sets the `authorization: Bearer …` and `x-api-key` headers from that secret. Sail adds the key as each - request leaves the box. + sets the `authorization: Bearer …` and `x-api-key` headers from the key's secret, and each header from + its own. Sail adds them as each request leaves the box. 3. Inside the box, every command and file agent-env sends carries the placeholder - `sail-injected-model-key` in place of the key. + `sail-injected-model-key` in place of the key. `LITELLM_EXTRA_HEADERS` (and the unit-test verifier's + `ANTHROPIC_CUSTOM_HEADERS`) keep the header names with `sail-injected-header` as every value, which the + policy overwrites. 4. A `docker` shim on the box gives every container the VM's CA bundle (`SSL_CERT_FILE`, `REQUESTS_CA_BUNDLE`, `NODE_EXTRA_CA_CERTS`, `CURL_CA_BUNDLE`). Sail terminates TLS for the model host with its own CA, so containers need it to trust those requests. @@ -207,12 +211,13 @@ Requirements and behaviour: - **Endpoint:** the model endpoint must be HTTPS and reachable from Sail. - **Key length:** keys shorter than 16 characters are refused, because scrubbing replaces the key - wherever it appears. -- **Other keys are refused:** a Sailbox refuses any command or file that would carry a model key Sail - doesn't inject for it. One example is an agent deployed into a sandbox it didn't create. + wherever it appears. Headers are replaced as a whole, so their values have no minimum. +- **Headers the key sets:** a `[model.headers]` entry named `authorization` or `x-api-key` is refused. +- **Other keys are refused:** a Sailbox refuses any command or file that would carry a model key, or + header values, Sail doesn't inject for it. One example is an agent deployed into a sandbox it didn't create. - **Turning it off:** set `inject_model_key = false` to pass keys in as other providers do. -- **Cleanup:** terminate deletes the Sailbox's policy. The secret is kept, so a concurrent launch with the - same key never loses it. To remove secrets for retired keys, sweep `AGENTENV_LITELLM_*` with +- **Cleanup:** terminate deletes the Sailbox's policy. The secrets are kept, so a concurrent launch with the + same key or headers never loses them. To remove secrets for retired keys, sweep `AGENTENV_LITELLM_*` with `sail secret list`. ## Attribution and cost diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py index e86cc282..f449ea4a 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py @@ -1,7 +1,8 @@ -"""Keeping the model key out of a Sailbox: Sail adds it to the agent's model requests as they leave. +"""Keeping the model key, and the endpoint's ``[model.headers]``, out of a Sailbox: Sail adds them to the +agent's model requests as they leave. The key is stored as a Sail secret named from its SHA-256, so one key is one secret and different keys never -share one. The secret is never deleted, so no launch can remove one another, concurrent or in another +share one; each header value is a secret of its own, named the same way from its name and value. The secret is never deleted, so no launch can remove one another, concurrent or in another process, still needs; a Sailbox's saved policy, which only names it, is deleted with the Sailbox. A saved egress policy on the Sailbox sets it as the auth headers of requests to the model endpoint's host. Inside the Sailbox the key is replaced by a placeholder, and containers start trusting the CA Sail terminates those requests' TLS with. @@ -18,11 +19,18 @@ from typing import Any, Mapping from urllib.parse import urlparse +from agent_env.config.model import MODEL_HEADERS_ENV + SECRET_PREFIX = "AGENTENV_LITELLM_" +HEADER_SECRET_PREFIX = f"{SECRET_PREFIX}H_" POLICY_PREFIX = "agentenv-" PLACEHOLDER = "sail-injected-model-key" +HEADER_PLACEHOLDER = "sail-injected-header" KEY_ENV = "LITELLM_API_KEY" BASE_URL_ENV = "LITELLM_BASE_URL" +HEADERS_ENV = MODEL_HEADERS_ENV +#: The headers the key is set as; a [model.headers] entry may not set them too. +KEY_HEADERS = ("authorization", "x-api-key") #: The Sailbox's CA bundle, which includes the CA Sail intercepts TLS with, and where containers see it. VM_CA_BUNDLE = "/etc/ssl/certs/ca-certificates.crt" @@ -53,28 +61,66 @@ # By secret name: the key, so a handle reconnected here can still scrub its agent's key, and the launches # and Sailboxes in this process that need it, so the key is forgotten once none does. Process memory only. _injected_keys: dict[str, str] = {} +_injected_headers: dict[frozenset[str], dict[str, str]] = {} _holders: dict[str, set[str]] = {} +# The env vars agent-env hands a model endpoint's headers in, and how each writes them. +_HEADER_NAMES = (HEADERS_ENV, "ANTHROPIC_CUSTOM_HEADERS") def secret_name(key: str) -> str: return SECRET_PREFIX + hashlib.sha256(key.encode()).hexdigest()[:32].upper() +def header_secret_name(name: str, value: str) -> str: + return HEADER_SECRET_PREFIX + hashlib.sha256(f"{name.lower()}\0{value}".encode()).hexdigest()[:32].upper() + + +def header_lines(headers: Mapping[str, str]) -> str: + """Headers as ``ANTHROPIC_CUSTOM_HEADERS`` takes them: one ``Name: Value`` per line.""" + return "\n".join(f"{name}: {value}" for name, value in headers.items()) + + +def _is_header_placeholder(name: str, value: str) -> bool: + """Whether ``value``, as the ``name`` env var holds it, carries no header value but the placeholder.""" + if name == HEADERS_ENV: + try: + headers = json.loads(value) + except ValueError: + return False + return isinstance(headers, dict) and all(v == HEADER_PLACEHOLDER for v in headers.values()) + return all(line.partition(": ")[2] == HEADER_PLACEHOLDER for line in value.splitlines() if line) + + def carries_model_key(text: str, *, shell: bool) -> bool: """Whether ``text`` assigns a model-key env var a value other than the placeholder. In a ``shell`` script each word is decoded as the shell would (quotes and escapes, in the name too), and a value that is wholly one unquoted or double-quoted variable reference isn't one; in other text, every value is literal.""" if not shell: - return any(value and value != PLACEHOLDER for value in _LITERAL_ASSIGNMENT.findall(text)) + return (any(value and value != PLACEHOLDER for value in _LITERAL_ASSIGNMENT.findall(text)) + or _carries_headers_literally(text)) for word in _SHELL_WORD.findall(text): try: decoded = "".join(shlex.split(word)) except ValueError: return True name, assigned, value = decoded.partition("=") - if not assigned or name not in _KEY_NAMES or not value or value == PLACEHOLDER: + if not assigned or not value or _REFERENCE.match(_written_value(word)): continue - if not _REFERENCE.match(_written_value(word)): + if name in _KEY_NAMES and value != PLACEHOLDER: + return True + if name in _HEADER_NAMES and not _is_header_placeholder(name, value): + return True + return False + + +def _carries_headers_literally(text: str) -> bool: + """Whether text that isn't a shell script (a file, one argv word) assigns a header env var a value: the rest + of the text after ``NAME=``, which for one argv word is all of it.""" + for name in _HEADER_NAMES: + _, assigned, value = text.partition(f"{name}=") + if name == HEADERS_ENV: + value = value.split("\n", 1)[0].strip().strip("'\"") + if assigned and value and not _REFERENCE.match(value) and not _is_header_placeholder(name, value): return True return False @@ -117,12 +163,17 @@ def docker_shim() -> str: @dataclass class ModelKeyInjection: - """The model endpoint host, the Sail secret holding the key, and the saved policy applying it. - ``key`` is None on a reconnected Sailbox unless the configured key is the one injected.""" + """The model endpoint host, the Sail secrets holding the key and the endpoint's headers, and the saved + policy applying them. ``key`` is None on a reconnected Sailbox unless the configured key is the one + injected, and ``header_values`` unless the configured headers are.""" host: str secret: str key: str | None = field(default=None, repr=False) + #: Each header's lowercased name -> the Sail secret holding its value. + headers: dict[str, str] = field(default_factory=dict) + #: The headers as agent-env hands them to workloads, names as configured: what scrubbing replaces. + header_values: dict[str, str] | None = field(default=None, repr=False) policy_id: str | None = None #: Saved policies a swap may or may not have applied (Sail couldn't say), deleted once it can. unsettled_policy_ids: list[str] = field(default_factory=list) @@ -148,7 +199,15 @@ def for_env(cls, env: Mapping[str, str]) -> ModelKeyInjection | None: f"Sail model-key injection needs a {KEY_ENV} of at least {MIN_KEY_LENGTH} characters, to scrub it " "safely; set inject_model_key = false in [sandbox.providers.sail_vm.config] to pass it in" ) - return cls(host=parsed.hostname, secret=secret_name(key), key=key) + values = json.loads(env.get(HEADERS_ENV) or "{}") + if clash := sorted(name for name in values if name.lower() in KEY_HEADERS): + raise ValueError( + f"[model.headers] sets {clash}, which Sail sets from the model key; drop them, or set " + "inject_model_key = false in [sandbox.providers.sail_vm.config] to pass headers and key in" + ) + return cls(host=parsed.hostname, secret=secret_name(key), key=key, + headers={name.lower(): header_secret_name(name, value) for name, value in values.items()}, + header_values=values or None) @classmethod def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjection | None: @@ -156,35 +215,55 @@ def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjectio rules = document.get("rules") if isinstance(document, dict) else None if not isinstance(rules, dict) or len(rules) != 1: return None - (host,) = rules - names = set(_SECRET_REF.findall(json.dumps(rules))) - if len(names) != 1: + (host, entries), = rules.items() + try: + (entry,) = entries + refs = dict(entry["request"]["set"]["headers"]) + except (TypeError, ValueError, KeyError, AttributeError): + return None + refs.pop("authorization", None) # "Bearer ${secrets.…}", the key again: rules() below checks it + names = {header: _SECRET_REF.fullmatch(ref) if isinstance(ref, str) else None for header, ref in refs.items()} + if "x-api-key" not in names or not all(names.values()): return None - (name,) = names - injection = cls(host=host, secret=name, policy_id=policy_id) - return injection if name.startswith(SECRET_PREFIX) and rules == injection.rules() else None + secret = names.pop("x-api-key").group(1) + headers = {header: match.group(1) for header, match in names.items()} + injection = cls(host=host, secret=secret, headers=headers, policy_id=policy_id) + own = secret.startswith(SECRET_PREFIX) and all(s.startswith(HEADER_SECRET_PREFIX) for s in headers.values()) + return injection if own and rules == injection.rules() else None def rules(self) -> dict[str, Any]: - """Sail egress rules setting the key as both auth headers model endpoints read.""" + """Sail egress rules setting the key as both auth headers model endpoints read, and each header.""" ref = f"${{secrets.{self.secret}}}" - return {self.host: [{"request": {"set": {"headers": {"authorization": f"Bearer {ref}", "x-api-key": ref}}}}]} + headers = {"authorization": f"Bearer {ref}", "x-api-key": ref, + **{name: f"${{secrets.{secret}}}" for name, secret in self.headers.items()}} + return {self.host: [{"request": {"set": {"headers": headers}}}]} + + def header_secrets(self) -> dict[str, str]: + """Each header's Sail secret -> its value, for the launch to store.""" + return {header_secret_name(name, value): value for name, value in (self.header_values or {}).items()} def matches(self, other: ModelKeyInjection | None) -> bool: - return other is not None and (other.host, other.secret) == (self.host, self.secret) + return other is not None and (other.host, other.secret, other.headers) == (self.host, self.secret, self.headers) + + def _headers_id(self) -> frozenset[str]: + return frozenset(self.headers.values()) def hold(self, holder: str) -> None: - """Record ``holder`` (a launch or a Sailbox id) as needing this injection's secret.""" + """Record ``holder`` (a launch or a Sailbox id) as needing this injection's secrets.""" if self.key is not None: _injected_keys[self.secret] = self.key + if self.header_values: + _injected_headers[self._headers_id()] = self.header_values _holders.setdefault(self.secret, set()).add(holder) def release(self, holder: str) -> None: - """Drop ``holder``, forgetting the key once no launch or Sailbox in this process needs it.""" + """Drop ``holder``, forgetting the key and headers once no launch or Sailbox in this process needs them.""" holders = _holders.get(self.secret, set()) holders.discard(holder) if not holders: _holders.pop(self.secret, None) _injected_keys.pop(self.secret, None) + _injected_headers.pop(self._headers_id(), None) def recover_key(self, candidate: str | None) -> None: """Recover the key this injection's secret was named from: one this process injected, else @@ -193,16 +272,35 @@ def recover_key(self, candidate: str | None) -> None: if self.key is None and candidate and secret_name(candidate) == self.secret: self.key = candidate - def _encodings(self) -> list[str]: - """The key as agent-env writes it into commands: shell-quoted, inside single quotes, and raw.""" - return [shlex.quote(self.key), self.key.replace("'", "'\\''"), self.key] + def recover_headers(self, candidate: Mapping[str, str] | None) -> None: + """Recover the headers this injection's secrets were named from: ones this process injected, else + ``candidate`` (the configured headers) when they are those.""" + if not self.headers: + return + self.header_values = _injected_headers.get(self._headers_id()) + if self.header_values is None and candidate and { + name.lower(): header_secret_name(name, value) for name, value in candidate.items() + } == self.headers: + self.header_values = dict(candidate) + + def _encodings(self) -> list[tuple[str, str]]: + """Each value agent-env writes into commands, and the placeholder it becomes, in every encoding: + shell-quoted, inside single quotes, and raw. The headers go whole, as each env var holds them, so a + short value can't rewrite unrelated text.""" + pairs = [(self.key, PLACEHOLDER)] if self.key else [] + if self.header_values: + placeholders = dict.fromkeys(self.header_values, HEADER_PLACEHOLDER) + pairs += [(json.dumps(self.header_values), json.dumps(placeholders)), + (header_lines(self.header_values), header_lines(placeholders))] + return [(encode(value), encode(placeholder)) for value, placeholder in pairs + for encode in (shlex.quote, lambda text: text.replace("'", "'\\''"), lambda text: text)] def scrub(self, text: str) -> str: - for encoded in self._encodings() if self.key else (): - text = text.replace(encoded, PLACEHOLDER) + for encoded, placeholder in self._encodings(): + text = text.replace(encoded, placeholder) return text def scrub_bytes(self, data: bytes) -> bytes: - for encoded in self._encodings() if self.key else (): - data = data.replace(encoded.encode(), PLACEHOLDER.encode()) + for encoded, placeholder in self._encodings(): + data = data.replace(encoded.encode(), placeholder.encode()) return data diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py index 27116548..76170d72 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py @@ -289,6 +289,8 @@ async def release() -> None: if injection is not None: injection.hold(launch) await sdk.Secret.set.aio(injection.secret, injection.key) + for secret, value in injection.header_secrets().items(): + await sdk.Secret.set.aio(secret, value) egress = await create_saved_policy(sdk, policy_name(injection), egress_document(effective_policy, injection)) injection.policy_id = egress.id creating = True @@ -374,7 +376,8 @@ async def create_sandbox( network_policy: NetworkPolicy | None = None, ) -> SailVmSandbox: """A bare VM the caller loads and starts ``image_name`` in, as on the other VM providers. A model key in - ``env`` is injected by Sail rather than passed in (unless ``inject_model_key`` is off).""" + ``env``, and the endpoint's headers, are injected by Sail rather than passed in (unless + ``inject_model_key`` is off).""" del image_name injection = ModelKeyInjection.for_env(env) if self._inject_model_key else None return await self._create( @@ -406,10 +409,15 @@ async def get_sandbox(self, sandbox_id: str) -> SailVmSandbox: applied = getattr(raw, "egress_policy", None) injection = ModelKeyInjection.from_document(getattr(applied, "document", None), getattr(applied, "policy_id", None)) if injection is not None: + config = get_config() try: - injection.recover_key(get_config().get_litellm_api_key()) + injection.recover_key(config.get_litellm_api_key()) except Exception: # noqa: BLE001 - no configured key to recover; the handle just can't scrub it pass + try: + injection.recover_headers(config.get_model_headers(config.get_litellm_base_url())) + except Exception: # noqa: BLE001 - no configured headers to recover; the handle just can't scrub them + pass sandbox = SailVmSandbox( raw, sdk=sdk, tunnel_urls=tunnel_urls, network_policy=None, injection=injection, refuse_model_keys=self._inject_model_key, diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/sandbox.py b/src/agent_env/providers/sandbox_providers/sail_vm/sandbox.py index 8644a777..8bde5e0a 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/sandbox.py @@ -226,8 +226,9 @@ async def wait(self) -> int: class SailVmSandbox(VmSandbox): """A Sailbox from the Docker-capable devbox image; commands run as root. With a model-key ``injection``, - the key's value never enters the Sailbox: every command and file is scrubbed of it. With - ``refuse_model_keys``, one that still carries a model key (one Sail doesn't inject here) is refused.""" + the key's value and the endpoint's header values never enter the Sailbox: every command and file is + scrubbed of them. With ``refuse_model_keys``, one that still carries a model key or header values (ones + Sail doesn't inject here) is refused.""" type = "sail_vm" _DOCKER_PROBE_TIMEOUT = 10 @@ -364,9 +365,9 @@ async def _write_bytes_to_vm_path(self, data: bytes, vm_path: str) -> None: def _refusal(self) -> ModelKeyRefusedError: return ModelKeyRefusedError( - f"Refusing to send a model key into Sailbox {self.sandbox_id}: Sail injects model keys only for an agent " - "deployed on its own Sail sandbox, so this one would land on the Sailbox's disk. Deploy the agent on its " - "own sandbox, or set inject_model_key = false in [sandbox.providers.sail_vm.config] to pass keys in." + f"Refusing to send a model key or [model.headers] into Sailbox {self.sandbox_id}: Sail injects them only " + "for an agent deployed on its own Sail sandbox, so these would land on the Sailbox's disk. Deploy the agent " + "on its own sandbox, or set inject_model_key = false in [sandbox.providers.sail_vm.config] to pass them in." ) async def apply_network_policy(self, policy: NetworkPolicy) -> None: diff --git a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index db731bae..7608f9ee 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py @@ -1,6 +1,7 @@ """Model-key injection on Sailboxes: the key stays at Sail, the Sailbox only ever holds a placeholder.""" import asyncio +import json import logging import shlex import subprocess @@ -13,10 +14,13 @@ from agent_env.providers.sandbox_providers.sail_vm import model_key from agent_env.providers.sandbox_providers.sail_vm.model_key import ( DOCKER_SHIM_PATH, + HEADER_PLACEHOLDER, PLACEHOLDER, ModelKeyInjection, carries_model_key, docker_shim, + header_lines, + header_secret_name, secret_name, ) from agent_env.providers.sandbox_providers.sail_vm.provider import SailVmSandboxProvider @@ -91,6 +95,7 @@ def fresh_key_state(monkeypatch): monkeypatch.setattr(_sdk, "_installed_key", None) monkeypatch.setattr(_sdk, "_apps", {}) monkeypatch.setattr(model_key, "_injected_keys", {}) + monkeypatch.setattr(model_key, "_injected_headers", {}) monkeypatch.setattr(model_key, "_holders", {}) @@ -887,3 +892,150 @@ async def test_an_unconfirmed_lost_create_keeps_its_policy_rather_than_fail_dele saved[0].delete.aio.assert_not_awaited() assert "stays with the unconfirmed Sailbox" in caplog.text assert secret_name(_KEY) not in model_key._injected_keys + + +# --- [model.headers]: each value is a Sail secret the policy sets, and the Sailbox holds placeholders ---------- + +_HEADERS = {"X-HF-Bill-To": "acme-org", "X-Gateway-Key": "it's-a-gateway-secret"} +_ENV_H = {**_ENV, "LITELLM_EXTRA_HEADERS": json.dumps(_HEADERS)} +_PLACEHELD = dict.fromkeys(_HEADERS, HEADER_PLACEHOLDER) + + +def test_header_secrets_are_named_from_the_header_and_its_value(): + assert header_secret_name("X-A", "v") == header_secret_name("x-a", "v") + assert len({header_secret_name("X-A", "v"), header_secret_name("X-A", "w"), header_secret_name("X-B", "v")}) == 3 + assert header_secret_name("X-A", "v").startswith("AGENTENV_LITELLM_H_") + + +def test_the_policy_sets_each_header_from_its_secret_and_never_holds_a_value(): + injection = ModelKeyInjection.for_env(_ENV_H) + (rule,) = injection.rules()["llm.example.com"] + headers = rule["request"]["set"]["headers"] + + assert set(headers) == {"authorization", "x-api-key", "x-hf-bill-to", "x-gateway-key"} + assert headers["x-hf-bill-to"] == f"${{secrets.{header_secret_name('X-HF-Bill-To', 'acme-org')}}}" + assert not any(value in str(injection.rules()) or value in repr(injection) for value in _HEADERS.values()) + + +@pytest.mark.parametrize("name", ["Authorization", "X-Api-Key"]) +def test_a_header_the_key_is_set_as_is_refused(name): + with pytest.raises(ValueError, match="Sail sets from the model key"): + ModelKeyInjection.for_env({**_ENV, "LITELLM_EXTRA_HEADERS": json.dumps({name: "Bearer other"})}) + + +def test_an_injection_with_headers_is_recovered_from_its_saved_policy_only(): + injection = ModelKeyInjection.for_env(_ENV_H) + recovered = ModelKeyInjection.from_document({"rules": injection.rules()}, "ep_1") + assert injection.matches(recovered) and recovered.header_values is None + assert not injection.matches(ModelKeyInjection.for_env(_ENV)) + (rule,) = injection.rules()["llm.example.com"] + widened = {**rule["request"]["set"]["headers"], "x-other": "${secrets.GITHUB_TOKEN}"} + assert ModelKeyInjection.from_document({"rules": {"llm.example.com": [{"request": {"set": {"headers": widened}}}]}}, + "ep_2") is None + + +@pytest.mark.asyncio +async def test_an_agent_sailbox_gets_its_headers_through_secrets_and_its_saved_policy(): + sdk, saved, _ = _fake_sdk(_sailbox()) + + await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox(image_name="agent:1", port=8000, env=_ENV_H) + + stored = {call.args[0]: call.args[1] for call in sdk.Secret.set.aio.await_args_list} + assert stored == {secret_name(_KEY): _KEY, **{header_secret_name(n, v): v for n, v in _HEADERS.items()}} + (policy,) = saved + assert policy.document == {"rules": ModelKeyInjection.for_env(_ENV_H).rules()} + + +@pytest.mark.asyncio +async def test_a_failed_create_forgets_the_headers_too(): + sdk, _, _ = _fake_sdk(_sailbox()) + sdk.Sailbox.create.aio = AsyncMock(side_effect=RuntimeError("no capacity")) + + with pytest.raises(RuntimeError, match="no capacity"): + await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox(image_name="agent:1", port=8000, env=_ENV_H) + + assert model_key._injected_headers == {} + + +@pytest.mark.asyncio +async def test_commands_and_files_carry_header_placeholders_never_the_values(): + sailbox = _sailbox() + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox(image_name="agent:1", port=8000, env=_ENV_H) + written = json.dumps(_HEADERS) + + await sandbox.exec_script(f"docker run -d -e LITELLM_EXTRA_HEADERS='{written.replace(chr(39), chr(39) + chr(92) + chr(39) + chr(39))}' agent:1") + await sandbox.exec_with_output("sudo", "bash", "-c", f"docker exec -e LITELLM_EXTRA_HEADERS={shlex.quote(written)} " + f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_HEADERS))} c pytest") + await sandbox.write_host_file(f"LITELLM_EXTRA_HEADERS={written}\n".encode(), "/opt/agent/.env") + + sent = [arg for call in sailbox.exec.aio.await_args_list for arg in call.args[0]] + sent.append(sailbox.fs.write.aio.await_args.args[1].decode()) + assert not any(value in text for value in _HEADERS.values() for text in sent) + (run,) = [arg for arg in sent if "docker run" in arg] + assert json.dumps(_PLACEHELD) in run + (exec_,) = [arg for arg in sent if "docker exec" in arg] + assert shlex.quote(header_lines(_PLACEHELD)) in exec_ + + +@pytest.mark.asyncio +async def test_reconnect_in_another_process_scrubs_the_configured_headers(monkeypatch): + injection = ModelKeyInjection.for_env(_ENV_H) + sailbox = _sailbox() + sailbox.egress_policy = SimpleNamespace(policy_id="ep_9", document={"rules": injection.rules()}) + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + monkeypatch.setattr("agent_env.providers.sandbox_providers.sail_vm.provider.get_config", lambda: MagicMock( + get_litellm_api_key=lambda: _KEY, get_litellm_base_url=lambda: "https://llm.example.com/v1", + get_model_headers=lambda base_url: dict(_HEADERS))) + + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).get_sandbox("sb_1") + await sandbox.exec_with_output("docker", "exec", "-e", f"LITELLM_EXTRA_HEADERS={json.dumps(_HEADERS)}", "c", "pytest") + + assert sailbox.exec.aio.await_args.args[0][3] == f"LITELLM_EXTRA_HEADERS={json.dumps(_PLACEHELD)}" + + +@pytest.mark.asyncio +async def test_reconnect_with_other_configured_headers_refuses_rather_than_send_them(monkeypatch): + injection = ModelKeyInjection.for_env(_ENV_H) + sailbox = _sailbox() + sailbox.egress_policy = SimpleNamespace(policy_id="ep_9", document={"rules": injection.rules()}) + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + other = {"X-HF-Bill-To": "another-org", "X-Gateway-Key": "another-secret"} + monkeypatch.setattr("agent_env.providers.sandbox_providers.sail_vm.provider.get_config", lambda: MagicMock( + get_litellm_api_key=lambda: _KEY, get_litellm_base_url=lambda: "https://llm.example.com/v1", + get_model_headers=lambda base_url: other)) + + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).get_sandbox("sb_1") + + assert sandbox._injection.header_values is None + with pytest.raises(ModelKeyRefusedError): + await sandbox.exec_with_output("docker", "exec", "-e", f"LITELLM_EXTRA_HEADERS={json.dumps(other)}", "c", "pytest") + + +@pytest.mark.asyncio +async def test_a_sailbox_without_injection_refuses_header_values_but_takes_placeholders(): + sailbox = _sailbox(ports=()) + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_vm(exposed_ports=[]) + + with pytest.raises(ModelKeyRefusedError, match="model.headers"): + await sandbox.exec_with_output("docker", "exec", "-e", f"LITELLM_EXTRA_HEADERS={json.dumps(_HEADERS)}", "c", "true") + await sandbox.exec_with_output("docker", "exec", "-e", f"LITELLM_EXTRA_HEADERS={json.dumps(_PLACEHELD)}", "c", "true") + + +@pytest.mark.parametrize(("text", "shell", "carries"), [ + (f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_HEADERS))}", True, True), + (f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_PLACEHELD))}", True, False), + ('-e LITELLM_EXTRA_HEADERS="$LITELLM_EXTRA_HEADERS"', True, False), + (f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_HEADERS))}", True, True), + (f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_PLACEHELD))}", True, False), + (f"LITELLM_EXTRA_HEADERS={json.dumps(_HEADERS)}", False, True), + (f"LITELLM_EXTRA_HEADERS={json.dumps(_PLACEHELD)}\nOTHER=1", False, False), + ("LITELLM_EXTRA_HEADERS={}", False, False), +]) +def test_a_header_assignment_is_detected_as_the_shell_reads_it(text, shell, carries): + assert carries_model_key(text, shell=shell) is carries From f96322364fee7991484cbfa0d7f2db121c9b05c3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 08:47:28 -0700 Subject: [PATCH 5/8] fix(sail_vm): scope injected headers to the endpoint's path, scrub them as written, count their holders apart - Sail's rule for the headers matches the model endpoint's path prefix (/v1/ for https://host/v1), so another service on the same host never gets them; the key still goes with every request to the host, in a second rule, as before. Reconnect recognises both shapes. - LITELLM_EXTRA_HEADERS a deploy wrote in its own spelling (compact JSON, say) is scrubbed as written too, instead of slipping past the scrub and getting the deploy refused; a value that isn't a JSON object of strings is refused with a message saying so. - The headers a process remembers for scrubbing are released by their own holders, so terminating a Sailbox no longer drops them for another with a different key and the same headers. - The unit-test verifier resolves configured headers only for a variable the step doesn't set itself. Co-Authored-By: Claude Opus 5.5 --- .../sandbox_providers/sail_vm/README.md | 5 +- .../sandbox_providers/sail_vm/model_key.py | 71 ++++++++++++++----- .../run_container_unit_tests_verifier.py | 12 ++-- .../sail_vm_model_key_test.py | 62 ++++++++++++++-- .../task_step/local_container_urls_test.py | 9 +++ 5 files changed, 126 insertions(+), 33 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/README.md b/src/agent_env/providers/sandbox_providers/sail_vm/README.md index 6bedd6ae..e885af5f 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/README.md +++ b/src/agent_env/providers/sandbox_providers/sail_vm/README.md @@ -197,8 +197,9 @@ values of the endpoint's `[model.headers]`: 1. The key is stored as a Sail secret named `AGENTENV_LITELLM_`, one per distinct key. Each header value is a secret of its own, `AGENTENV_LITELLM_H_`. 2. The Sailbox is created with a saved egress policy. On requests to `LITELLM_BASE_URL`'s host, the policy - sets the `authorization: Bearer …` and `x-api-key` headers from the key's secret, and each header from - its own. Sail adds them as each request leaves the box. + sets the `authorization: Bearer …` and `x-api-key` headers from the key's secret. On those under its path + (`/v1/` for `https://host/v1`), it also sets each header from its own, so another service on the same host + never gets them. Sail adds them as each request leaves the box. 3. Inside the box, every command and file agent-env sends carries the placeholder `sail-injected-model-key` in place of the key. `LITELLM_EXTRA_HEADERS` (and the unit-test verifier's `ANTHROPIC_CUSTOM_HEADERS`) keep the header names with `sail-injected-header` as every value, which the diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py index f449ea4a..bc52c2c0 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py @@ -61,8 +61,9 @@ # By secret name: the key, so a handle reconnected here can still scrub its agent's key, and the launches # and Sailboxes in this process that need it, so the key is forgotten once none does. Process memory only. _injected_keys: dict[str, str] = {} -_injected_headers: dict[frozenset[str], dict[str, str]] = {} +_injected_headers: dict[frozenset[str], tuple[dict[str, str], str | None]] = {} _holders: dict[str, set[str]] = {} +_header_holders: dict[frozenset[str], set[str]] = {} # The env vars agent-env hands a model endpoint's headers in, and how each writes them. _HEADER_NAMES = (HEADERS_ENV, "ANTHROPIC_CUSTOM_HEADERS") @@ -174,6 +175,11 @@ class ModelKeyInjection: headers: dict[str, str] = field(default_factory=dict) #: The headers as agent-env hands them to workloads, names as configured: what scrubbing replaces. header_values: dict[str, str] | None = field(default=None, repr=False) + #: ``LITELLM_EXTRA_HEADERS`` exactly as the deploy wrote it, when that isn't the JSON agent-env writes. + header_text: str | None = field(default=None, repr=False) + #: The model endpoint's path as a prefix (``/v1/``) the headers are scoped to, as they are in agent-env; + #: None when it has no path. The key goes with every request to the host either way. + path_prefix: str | None = None policy_id: str | None = None #: Saved policies a swap may or may not have applied (Sail couldn't say), deleted once it can. unsettled_policy_ids: list[str] = field(default_factory=list) @@ -199,15 +205,23 @@ def for_env(cls, env: Mapping[str, str]) -> ModelKeyInjection | None: f"Sail model-key injection needs a {KEY_ENV} of at least {MIN_KEY_LENGTH} characters, to scrub it " "safely; set inject_model_key = false in [sandbox.providers.sail_vm.config] to pass it in" ) - values = json.loads(env.get(HEADERS_ENV) or "{}") + text = env.get(HEADERS_ENV) or "" + try: + values = json.loads(text or "{}") + except ValueError: + values = None + if not isinstance(values, dict) or not all(isinstance(value, str) for value in values.values()): + raise ValueError(f"{HEADERS_ENV} must be a JSON object of header names and values, not {text[:80]!r}") if clash := sorted(name for name in values if name.lower() in KEY_HEADERS): raise ValueError( f"[model.headers] sets {clash}, which Sail sets from the model key; drop them, or set " "inject_model_key = false in [sandbox.providers.sail_vm.config] to pass headers and key in" ) + path = parsed.path.rstrip("/") return cls(host=parsed.hostname, secret=secret_name(key), key=key, headers={name.lower(): header_secret_name(name, value) for name, value in values.items()}, - header_values=values or None) + header_values=values or None, header_text=text if values and text != json.dumps(values) else None, + path_prefix=f"{path}/" if values and path else None) @classmethod def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjection | None: @@ -216,9 +230,11 @@ def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjectio if not isinstance(rules, dict) or len(rules) != 1: return None (host, entries), = rules.items() + if not isinstance(entries, list) or len(entries) not in (1, 2): + return None try: - (entry,) = entries - refs = dict(entry["request"]["set"]["headers"]) + refs = dict(entries[0]["request"]["set"]["headers"]) + prefix = entries[0]["match"]["path"]["prefix"] if len(entries) == 2 else None except (TypeError, ValueError, KeyError, AttributeError): return None refs.pop("authorization", None) # "Bearer ${secrets.…}", the key again: rules() below checks it @@ -227,43 +243,59 @@ def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjectio return None secret = names.pop("x-api-key").group(1) headers = {header: match.group(1) for header, match in names.items()} - injection = cls(host=host, secret=secret, headers=headers, policy_id=policy_id) + injection = cls(host=host, secret=secret, headers=headers, path_prefix=prefix, policy_id=policy_id) own = secret.startswith(SECRET_PREFIX) and all(s.startswith(HEADER_SECRET_PREFIX) for s in headers.values()) return injection if own and rules == injection.rules() else None def rules(self) -> dict[str, Any]: - """Sail egress rules setting the key as both auth headers model endpoints read, and each header.""" + """Sail egress rules setting the key as both auth headers model endpoints read on every request to the + host, and each header on those under the endpoint's path.""" ref = f"${{secrets.{self.secret}}}" - headers = {"authorization": f"Bearer {ref}", "x-api-key": ref, - **{name: f"${{secrets.{secret}}}" for name, secret in self.headers.items()}} - return {self.host: [{"request": {"set": {"headers": headers}}}]} + key = {"request": {"set": {"headers": {"authorization": f"Bearer {ref}", "x-api-key": ref}}}} + if not self.headers: + return {self.host: [key]} + both = {"request": {"set": {"headers": { + **key["request"]["set"]["headers"], **{name: f"${{secrets.{secret}}}" for name, secret in self.headers.items()}, + }}}} + if self.path_prefix is None: + return {self.host: [both]} + return {self.host: [{"match": {"path": {"prefix": self.path_prefix}}, **both}, key]} def header_secrets(self) -> dict[str, str]: """Each header's Sail secret -> its value, for the launch to store.""" return {header_secret_name(name, value): value for name, value in (self.header_values or {}).items()} def matches(self, other: ModelKeyInjection | None) -> bool: - return other is not None and (other.host, other.secret, other.headers) == (self.host, self.secret, self.headers) + return other is not None and (other.host, other.secret, other.headers, other.path_prefix) == ( + self.host, self.secret, self.headers, self.path_prefix) def _headers_id(self) -> frozenset[str]: return frozenset(self.headers.values()) def hold(self, holder: str) -> None: - """Record ``holder`` (a launch or a Sailbox id) as needing this injection's secrets.""" + """Record ``holder`` (a launch or a Sailbox id) as needing this injection's key, and its headers.""" if self.key is not None: _injected_keys[self.secret] = self.key - if self.header_values: - _injected_headers[self._headers_id()] = self.header_values _holders.setdefault(self.secret, set()).add(holder) + if self.headers: + if self.header_values: + _injected_headers[self._headers_id()] = (self.header_values, self.header_text) + _header_holders.setdefault(self._headers_id(), set()).add(holder) def release(self, holder: str) -> None: - """Drop ``holder``, forgetting the key and headers once no launch or Sailbox in this process needs them.""" + """Drop ``holder``, forgetting the key, and the headers, once no launch or Sailbox in this process needs + them; the two are counted apart, since Sailboxes with different keys can share headers.""" holders = _holders.get(self.secret, set()) holders.discard(holder) if not holders: _holders.pop(self.secret, None) _injected_keys.pop(self.secret, None) - _injected_headers.pop(self._headers_id(), None) + if self.headers: + held = _header_holders.get(self._headers_id(), set()) + held.discard(holder) + if not held: + _header_holders.pop(self._headers_id(), None) + _injected_headers.pop(self._headers_id(), None) def recover_key(self, candidate: str | None) -> None: """Recover the key this injection's secret was named from: one this process injected, else @@ -277,8 +309,9 @@ def recover_headers(self, candidate: Mapping[str, str] | None) -> None: ``candidate`` (the configured headers) when they are those.""" if not self.headers: return - self.header_values = _injected_headers.get(self._headers_id()) - if self.header_values is None and candidate and { + if remembered := _injected_headers.get(self._headers_id()): + self.header_values, self.header_text = remembered + elif candidate and { name.lower(): header_secret_name(name, value) for name, value in candidate.items() } == self.headers: self.header_values = dict(candidate) @@ -292,6 +325,8 @@ def _encodings(self) -> list[tuple[str, str]]: placeholders = dict.fromkeys(self.header_values, HEADER_PLACEHOLDER) pairs += [(json.dumps(self.header_values), json.dumps(placeholders)), (header_lines(self.header_values), header_lines(placeholders))] + if self.header_text: + pairs.append((self.header_text, json.dumps(placeholders))) return [(encode(value), encode(placeholder)) for value, placeholder in pairs for encode in (shlex.quote, lambda text: text.replace("'", "'\\''"), lambda text: text)] diff --git a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py index 2b91ca32..83cd6775 100644 --- a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py @@ -259,12 +259,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: merged_env.update(extra_env) # Unlike a missing key or URL, configured headers that don't resolve fail the step: a verifier that ran # without them would grade against an endpoint that answers differently, or bills someone else. - headers = config.get_model_headers(merged_env.get("LITELLM_BASE_URL")) - if headers: - merged_env.setdefault(MODEL_HEADERS_ENV, json.dumps(headers)) - headers = config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")) - if headers: - merged_env.setdefault("ANTHROPIC_CUSTOM_HEADERS", "\n".join(f"{k}: {v}" for k, v in headers.items())) + if MODEL_HEADERS_ENV not in merged_env: + if headers := config.get_model_headers(merged_env.get("LITELLM_BASE_URL")): + merged_env[MODEL_HEADERS_ENV] = json.dumps(headers) + if "ANTHROPIC_CUSTOM_HEADERS" not in merged_env: + if headers := config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")): + merged_env["ANTHROPIC_CUSTOM_HEADERS"] = "\n".join(f"{k}: {v}" for k, v in headers.items()) for key in _MODEL_URL_KEYS & merged_env.keys(): merged_env[key] = host_url_for(merged_env[key], sandbox.type) env_flags = " ".join(f"-e {k}={shlex.quote(v)}" for k, v in merged_env.items()) diff --git a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index 7608f9ee..55bdf2fd 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py @@ -96,6 +96,7 @@ def fresh_key_state(monkeypatch): monkeypatch.setattr(_sdk, "_apps", {}) monkeypatch.setattr(model_key, "_injected_keys", {}) monkeypatch.setattr(model_key, "_injected_headers", {}) + monkeypatch.setattr(model_key, "_header_holders", {}) monkeypatch.setattr(model_key, "_holders", {}) @@ -907,16 +908,31 @@ def test_header_secrets_are_named_from_the_header_and_its_value(): assert header_secret_name("X-A", "v").startswith("AGENTENV_LITELLM_H_") -def test_the_policy_sets_each_header_from_its_secret_and_never_holds_a_value(): +def test_the_policy_sets_each_header_from_its_secret_under_the_endpoint_path_and_the_key_everywhere(): injection = ModelKeyInjection.for_env(_ENV_H) - (rule,) = injection.rules()["llm.example.com"] - headers = rule["request"]["set"]["headers"] + scoped, everywhere = injection.rules()["llm.example.com"] + assert scoped["match"] == {"path": {"prefix": "/v1/"}} + headers = scoped["request"]["set"]["headers"] assert set(headers) == {"authorization", "x-api-key", "x-hf-bill-to", "x-gateway-key"} assert headers["x-hf-bill-to"] == f"${{secrets.{header_secret_name('X-HF-Bill-To', 'acme-org')}}}" + assert everywhere == ModelKeyInjection.for_env(_ENV).rules()["llm.example.com"][0] assert not any(value in str(injection.rules()) or value in repr(injection) for value in _HEADERS.values()) +def test_an_endpoint_without_a_path_gets_its_headers_with_every_request_to_the_host(): + injection = ModelKeyInjection.for_env({**_ENV_H, "LITELLM_BASE_URL": "https://llm.example.com"}) + + (rule,) = injection.rules()["llm.example.com"] + assert "match" not in rule and "x-hf-bill-to" in rule["request"]["set"]["headers"] + + +@pytest.mark.parametrize("text", ["not json", '["X-A"]', '{"X-A": 1}']) +def test_headers_that_arent_a_json_object_of_strings_are_refused(text): + with pytest.raises(ValueError, match="must be a JSON object of header names and values"): + ModelKeyInjection.for_env({**_ENV, "LITELLM_EXTRA_HEADERS": text}) + + @pytest.mark.parametrize("name", ["Authorization", "X-Api-Key"]) def test_a_header_the_key_is_set_as_is_refused(name): with pytest.raises(ValueError, match="Sail sets from the model key"): @@ -928,10 +944,12 @@ def test_an_injection_with_headers_is_recovered_from_its_saved_policy_only(): recovered = ModelKeyInjection.from_document({"rules": injection.rules()}, "ep_1") assert injection.matches(recovered) and recovered.header_values is None assert not injection.matches(ModelKeyInjection.for_env(_ENV)) - (rule,) = injection.rules()["llm.example.com"] - widened = {**rule["request"]["set"]["headers"], "x-other": "${secrets.GITHUB_TOKEN}"} - assert ModelKeyInjection.from_document({"rules": {"llm.example.com": [{"request": {"set": {"headers": widened}}}]}}, - "ep_2") is None + scoped, everywhere = injection.rules()["llm.example.com"] + widened = {**scoped, "request": {"set": {"headers": {**scoped["request"]["set"]["headers"], + "x-other": "${secrets.GITHUB_TOKEN}"}}}} + assert ModelKeyInjection.from_document({"rules": {"llm.example.com": [widened, everywhere]}}, "ep_2") is None + moved = {**scoped, "match": {"path": {"prefix": "/"}}} + assert not injection.matches(ModelKeyInjection.from_document({"rules": {"llm.example.com": [moved, everywhere]}}, "ep_3")) @pytest.mark.asyncio @@ -1039,3 +1057,33 @@ async def test_a_sailbox_without_injection_refuses_header_values_but_takes_place ]) def test_a_header_assignment_is_detected_as_the_shell_reads_it(text, shell, carries): assert carries_model_key(text, shell=shell) is carries + + +@pytest.mark.asyncio +async def test_headers_a_deploy_writes_in_its_own_spelling_are_scrubbed_too(): + compact = '{"X-HF-Bill-To":"acme-org","X-Gateway-Key":"it\'s-a-gateway-secret"}' + sailbox = _sailbox() + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox( + image_name="agent:1", port=8000, env={**_ENV, "LITELLM_EXTRA_HEADERS": compact}) + + await sandbox.exec_script(f"docker run -d -e LITELLM_EXTRA_HEADERS={shlex.quote(compact)} agent:1") + + (run,) = [arg for call in sailbox.exec.aio.await_args_list for arg in call.args[0] if "docker run" in arg] + assert json.dumps(_PLACEHELD) in run and not any(value in run for value in _HEADERS.values()) + + +@pytest.mark.asyncio +async def test_sailboxes_sharing_headers_but_not_a_key_keep_each_others_headers_to_scrub(): + sdk, _, _ = _fake_sdk(_sailbox()) + provider = SailVmSandboxProvider(api_key="sail-key", sdk=sdk) + first = await provider.create_sandbox(image_name="agent:1", port=8000, env=_ENV_H) + second = ModelKeyInjection.for_env({**_ENV_H, "LITELLM_API_KEY": "sk-another-model-key-9876543210"}) + second.hold("sb_2") + + await first.terminate() + + reconnected = ModelKeyInjection.from_document({"rules": second.rules()}, "ep_9") + reconnected.recover_headers(None) + assert reconnected.header_values == _HEADERS diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index d1c83aa8..97fe45c7 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -241,6 +241,15 @@ def test_a_unit_test_verifier_whose_headers_dont_resolve_fails(monkeypatch, loca _run_verifier(monkeypatch, "local") +def test_a_unit_test_verifier_given_its_own_headers_never_resolves_the_configured_ones(monkeypatch, local_stores, model_endpoint): + monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers={"X-Bill-To": "secret:fd_no_such_secret"})) + + run = _run_verifier(monkeypatch, "local", env_vars={ + "LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}', "ANTHROPIC_CUSTOM_HEADERS": "X-Bill-To: mine"}) + + assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps({'X-Bill-To': 'mine'}))} " in run + + def test_a_seed_column_cant_replace_the_verifiers_headers(monkeypatch, model_headers): run = _run_verifier(monkeypatch, "local", seed={"litellm_extra_headers": "{}", "anthropic_custom_headers": "x"}) From 1a5b44baedbc1f49d71fa2fd291180f5ca0df36e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 08:56:53 -0700 Subject: [PATCH 6/8] fix(sail_vm): never quote headers in an error, and remember every spelling a Sailbox's headers were written in - A LITELLM_EXTRA_HEADERS that isn't a JSON object of strings is refused without quoting it: it may hold a credential. - The headers a process remembers for scrubbing keep every spelling the deploys sharing them wrote, so a Sailbox reconnected after another one held the same headers in another spelling still scrubs its own. Co-Authored-By: Claude Opus 5.5 --- .../sandbox_providers/sail_vm/model_key.py | 19 ++++++++-------- .../sail_vm_model_key_test.py | 22 ++++++++++++++++--- 2 files changed, 29 insertions(+), 12 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py index bc52c2c0..d9c54fee 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py @@ -61,7 +61,7 @@ # By secret name: the key, so a handle reconnected here can still scrub its agent's key, and the launches # and Sailboxes in this process that need it, so the key is forgotten once none does. Process memory only. _injected_keys: dict[str, str] = {} -_injected_headers: dict[frozenset[str], tuple[dict[str, str], str | None]] = {} +_injected_headers: dict[frozenset[str], tuple[dict[str, str], frozenset[str]]] = {} _holders: dict[str, set[str]] = {} _header_holders: dict[frozenset[str], set[str]] = {} # The env vars agent-env hands a model endpoint's headers in, and how each writes them. @@ -175,8 +175,8 @@ class ModelKeyInjection: headers: dict[str, str] = field(default_factory=dict) #: The headers as agent-env hands them to workloads, names as configured: what scrubbing replaces. header_values: dict[str, str] | None = field(default=None, repr=False) - #: ``LITELLM_EXTRA_HEADERS`` exactly as the deploy wrote it, when that isn't the JSON agent-env writes. - header_text: str | None = field(default=None, repr=False) + #: ``LITELLM_EXTRA_HEADERS`` exactly as each deploy wrote it, where that isn't the JSON agent-env writes. + header_texts: frozenset[str] = field(default=frozenset(), repr=False) #: The model endpoint's path as a prefix (``/v1/``) the headers are scoped to, as they are in agent-env; #: None when it has no path. The key goes with every request to the host either way. path_prefix: str | None = None @@ -211,7 +211,7 @@ def for_env(cls, env: Mapping[str, str]) -> ModelKeyInjection | None: except ValueError: values = None if not isinstance(values, dict) or not all(isinstance(value, str) for value in values.values()): - raise ValueError(f"{HEADERS_ENV} must be a JSON object of header names and values, not {text[:80]!r}") + raise ValueError(f"{HEADERS_ENV} must be a JSON object of header names and values") if clash := sorted(name for name in values if name.lower() in KEY_HEADERS): raise ValueError( f"[model.headers] sets {clash}, which Sail sets from the model key; drop them, or set " @@ -220,7 +220,8 @@ def for_env(cls, env: Mapping[str, str]) -> ModelKeyInjection | None: path = parsed.path.rstrip("/") return cls(host=parsed.hostname, secret=secret_name(key), key=key, headers={name.lower(): header_secret_name(name, value) for name, value in values.items()}, - header_values=values or None, header_text=text if values and text != json.dumps(values) else None, + header_values=values or None, + header_texts=frozenset({text}) if values and text != json.dumps(values) else frozenset(), path_prefix=f"{path}/" if values and path else None) @classmethod @@ -279,7 +280,8 @@ def hold(self, holder: str) -> None: _holders.setdefault(self.secret, set()).add(holder) if self.headers: if self.header_values: - _injected_headers[self._headers_id()] = (self.header_values, self.header_text) + _, texts = _injected_headers.get(self._headers_id(), (None, frozenset())) + _injected_headers[self._headers_id()] = (self.header_values, texts | self.header_texts) _header_holders.setdefault(self._headers_id(), set()).add(holder) def release(self, holder: str) -> None: @@ -310,7 +312,7 @@ def recover_headers(self, candidate: Mapping[str, str] | None) -> None: if not self.headers: return if remembered := _injected_headers.get(self._headers_id()): - self.header_values, self.header_text = remembered + self.header_values, self.header_texts = remembered elif candidate and { name.lower(): header_secret_name(name, value) for name, value in candidate.items() } == self.headers: @@ -325,8 +327,7 @@ def _encodings(self) -> list[tuple[str, str]]: placeholders = dict.fromkeys(self.header_values, HEADER_PLACEHOLDER) pairs += [(json.dumps(self.header_values), json.dumps(placeholders)), (header_lines(self.header_values), header_lines(placeholders))] - if self.header_text: - pairs.append((self.header_text, json.dumps(placeholders))) + pairs += [(text, json.dumps(placeholders)) for text in sorted(self.header_texts)] return [(encode(value), encode(placeholder)) for value, placeholder in pairs for encode in (shlex.quote, lambda text: text.replace("'", "'\\''"), lambda text: text)] diff --git a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index 55bdf2fd..885f97aa 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py @@ -927,11 +927,13 @@ def test_an_endpoint_without_a_path_gets_its_headers_with_every_request_to_the_h assert "match" not in rule and "x-hf-bill-to" in rule["request"]["set"]["headers"] -@pytest.mark.parametrize("text", ["not json", '["X-A"]', '{"X-A": 1}']) -def test_headers_that_arent_a_json_object_of_strings_are_refused(text): - with pytest.raises(ValueError, match="must be a JSON object of header names and values"): +@pytest.mark.parametrize("text", ["not json", '["X-A"]', '{"X-Gateway-Key": "private-token", "X-Retries": 3}']) +def test_headers_that_arent_a_json_object_of_strings_are_refused_without_quoting_them(text): + with pytest.raises(ValueError, match="must be a JSON object of header names and values") as refused: ModelKeyInjection.for_env({**_ENV, "LITELLM_EXTRA_HEADERS": text}) + assert "private-token" not in str(refused.value) + @pytest.mark.parametrize("name", ["Authorization", "X-Api-Key"]) def test_a_header_the_key_is_set_as_is_refused(name): @@ -1087,3 +1089,17 @@ async def test_sailboxes_sharing_headers_but_not_a_key_keep_each_others_headers_ reconnected = ModelKeyInjection.from_document({"rules": second.rules()}, "ep_9") reconnected.recover_headers(None) assert reconnected.header_values == _HEADERS + + +def test_sailboxes_writing_the_same_headers_differently_each_keep_their_spelling_to_scrub(): + spellings = ['{"X-HF-Bill-To":"acme-org","X-Gateway-Key":"it\'s-a-gateway-secret"}', + '{ "X-HF-Bill-To" : "acme-org" , "X-Gateway-Key" : "it\'s-a-gateway-secret" }'] + for holder, text in zip(("sb_1", "sb_2"), spellings): + ModelKeyInjection.for_env({**_ENV, "LITELLM_EXTRA_HEADERS": text}).hold(holder) + + reconnected = ModelKeyInjection.from_document({"rules": ModelKeyInjection.for_env(_ENV_H).rules()}, "ep_9") + reconnected.recover_headers(None) + + for text in spellings: + assert reconnected.scrub(f"LITELLM_EXTRA_HEADERS={shlex.quote(text)}") == ( + f"LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_PLACEHELD))}") From 1bfee23a8af6befb56924f532a4358c371392e84 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 09:36:43 -0700 Subject: [PATCH 7/8] refactor(config): one header variable, per-holder header memory on Sail, table-driven scoping tests - The unit-test verifier no longer sets ANTHROPIC_CUSTOM_HEADERS: nothing in this repository reads it, and its "Name: Value" lines were a second encoding Sail had to scrub and detect. Agents and the verifier get [model.headers] as LITELLM_EXTRA_HEADERS alone. - Sail remembers each holder's headers in its own entry, dropped with the holder, instead of reference- counting header sets: Sailboxes sharing headers can't drop each other's, whatever spelling each wrote. - Header assignments in a file or argv word are found wherever they are, not only the first. - The scoping tests for headers_for, resolve_call, deploys and the verifier are tables. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 9 +- .../sandbox_providers/sail_vm/README.md | 5 +- .../sandbox_providers/sail_vm/model_key.py | 82 +++++++---------- .../run_container_unit_tests_verifier.py | 5 +- tst/unit/config/test_model.py | 90 +++++++++---------- .../sail_vm_model_key_test.py | 20 ++--- .../task_step/local_container_urls_test.py | 86 ++++++------------ 7 files changed, 113 insertions(+), 184 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index cd828c96..96f27c3f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,12 +69,9 @@ or one of them only where it is used. this machine that local containers call, a model proxy say, has to listen there: bind it to the bridge address and use that address in its URL. Binding `0.0.0.0` also works, but exposes the service to the network. -- `[model.headers]` go with every call to the configured endpoint and nowhere else: as - `extra_headers` on agent-env's own calls, and as `LITELLM_EXTRA_HEADERS` (a JSON object) to every - agent and unit-test verifier pointed at that endpoint, and to an installed agent that declares the - `litellm_extra_headers` param. The verifier also gets them as `ANTHROPIC_CUSTOM_HEADERS`. Core only - writes that variable, never reads it. On `sail_vm` with `inject_model_key`, Sail sets them at egress - like the key, and the Sailbox holds only placeholders. +- `[model.headers]` go with every call to the configured endpoint and nowhere else: `extra_headers` on + agent-env's own calls, `LITELLM_EXTRA_HEADERS` (JSON, written by core, never read) for agents and the + unit-test verifier. On `sail_vm` with `inject_model_key`, Sail sets them at egress like the key. - A host-mode `install_agent` (no `container_name`) on a `local` sandbox runs the agent's `install_commands_host` on this machine, with your privileges. Every command a local sandbox runs carries `AGENTENV_SANDBOX=`, and teardown stops the processes that carry it, then removes diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/README.md b/src/agent_env/providers/sandbox_providers/sail_vm/README.md index e885af5f..a0c2159b 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/README.md +++ b/src/agent_env/providers/sandbox_providers/sail_vm/README.md @@ -201,9 +201,8 @@ values of the endpoint's `[model.headers]`: (`/v1/` for `https://host/v1`), it also sets each header from its own, so another service on the same host never gets them. Sail adds them as each request leaves the box. 3. Inside the box, every command and file agent-env sends carries the placeholder - `sail-injected-model-key` in place of the key. `LITELLM_EXTRA_HEADERS` (and the unit-test verifier's - `ANTHROPIC_CUSTOM_HEADERS`) keep the header names with `sail-injected-header` as every value, which the - policy overwrites. + `sail-injected-model-key` in place of the key. `LITELLM_EXTRA_HEADERS` keeps the header names with + `sail-injected-header` as every value, which the policy overwrites. 4. A `docker` shim on the box gives every container the VM's CA bundle (`SSL_CERT_FILE`, `REQUESTS_CA_BUNDLE`, `NODE_EXTRA_CA_CERTS`, `CURL_CA_BUNDLE`). Sail terminates TLS for the model host with its own CA, so containers need it to trust those requests. diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py index d9c54fee..02acc610 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py @@ -28,7 +28,6 @@ HEADER_PLACEHOLDER = "sail-injected-header" KEY_ENV = "LITELLM_API_KEY" BASE_URL_ENV = "LITELLM_BASE_URL" -HEADERS_ENV = MODEL_HEADERS_ENV #: The headers the key is set as; a [model.headers] entry may not set them too. KEY_HEADERS = ("authorization", "x-api-key") @@ -61,11 +60,12 @@ # By secret name: the key, so a handle reconnected here can still scrub its agent's key, and the launches # and Sailboxes in this process that need it, so the key is forgotten once none does. Process memory only. _injected_keys: dict[str, str] = {} -_injected_headers: dict[frozenset[str], tuple[dict[str, str], frozenset[str]]] = {} _holders: dict[str, set[str]] = {} -_header_holders: dict[frozenset[str], set[str]] = {} -# The env vars agent-env hands a model endpoint's headers in, and how each writes them. -_HEADER_NAMES = (HEADERS_ENV, "ANTHROPIC_CUSTOM_HEADERS") +# By holder: the headers it was launched or reconnected with, under their secrets, so a handle reconnected here +# can scrub them; each holder drops its own, so Sailboxes sharing headers never drop one another's. +_held_headers: dict[str, tuple[frozenset[str], dict[str, str], frozenset[str]]] = {} +# In text that isn't a shell script: the headers' JSON, to the end of its line. +_LITERAL_HEADERS = re.compile(rf"\b{MODEL_HEADERS_ENV}['\"]?(?:=|\s*:\s*)(.*)") def secret_name(key: str) -> str: @@ -76,20 +76,13 @@ def header_secret_name(name: str, value: str) -> str: return HEADER_SECRET_PREFIX + hashlib.sha256(f"{name.lower()}\0{value}".encode()).hexdigest()[:32].upper() -def header_lines(headers: Mapping[str, str]) -> str: - """Headers as ``ANTHROPIC_CUSTOM_HEADERS`` takes them: one ``Name: Value`` per line.""" - return "\n".join(f"{name}: {value}" for name, value in headers.items()) - - -def _is_header_placeholder(name: str, value: str) -> bool: - """Whether ``value``, as the ``name`` env var holds it, carries no header value but the placeholder.""" - if name == HEADERS_ENV: - try: - headers = json.loads(value) - except ValueError: - return False - return isinstance(headers, dict) and all(v == HEADER_PLACEHOLDER for v in headers.values()) - return all(line.partition(": ")[2] == HEADER_PLACEHOLDER for line in value.splitlines() if line) +def _is_header_placeholder(value: str) -> bool: + """Whether the headers' JSON ``value`` carries no header value but the placeholder.""" + try: + headers = json.loads(value) + except ValueError: + return False + return isinstance(headers, dict) and all(v == HEADER_PLACEHOLDER for v in headers.values()) def carries_model_key(text: str, *, shell: bool) -> bool: @@ -109,19 +102,16 @@ def carries_model_key(text: str, *, shell: bool) -> bool: continue if name in _KEY_NAMES and value != PLACEHOLDER: return True - if name in _HEADER_NAMES and not _is_header_placeholder(name, value): + if name == MODEL_HEADERS_ENV and not _is_header_placeholder(value): return True return False def _carries_headers_literally(text: str) -> bool: - """Whether text that isn't a shell script (a file, one argv word) assigns a header env var a value: the rest - of the text after ``NAME=``, which for one argv word is all of it.""" - for name in _HEADER_NAMES: - _, assigned, value = text.partition(f"{name}=") - if name == HEADERS_ENV: - value = value.split("\n", 1)[0].strip().strip("'\"") - if assigned and value and not _REFERENCE.match(value) and not _is_header_placeholder(name, value): + """Whether text that isn't a shell script (a file, one argv word) assigns the headers' env var a value.""" + for value in _LITERAL_HEADERS.findall(text): + value = value.strip().strip("'\"") + if value and not _REFERENCE.match(value) and not _is_header_placeholder(value): return True return False @@ -205,13 +195,13 @@ def for_env(cls, env: Mapping[str, str]) -> ModelKeyInjection | None: f"Sail model-key injection needs a {KEY_ENV} of at least {MIN_KEY_LENGTH} characters, to scrub it " "safely; set inject_model_key = false in [sandbox.providers.sail_vm.config] to pass it in" ) - text = env.get(HEADERS_ENV) or "" + text = env.get(MODEL_HEADERS_ENV) or "" try: values = json.loads(text or "{}") except ValueError: values = None if not isinstance(values, dict) or not all(isinstance(value, str) for value in values.values()): - raise ValueError(f"{HEADERS_ENV} must be a JSON object of header names and values") + raise ValueError(f"{MODEL_HEADERS_ENV} must be a JSON object of header names and values") if clash := sorted(name for name in values if name.lower() in KEY_HEADERS): raise ValueError( f"[model.headers] sets {clash}, which Sail sets from the model key; drop them, or set " @@ -278,26 +268,18 @@ def hold(self, holder: str) -> None: if self.key is not None: _injected_keys[self.secret] = self.key _holders.setdefault(self.secret, set()).add(holder) - if self.headers: - if self.header_values: - _, texts = _injected_headers.get(self._headers_id(), (None, frozenset())) - _injected_headers[self._headers_id()] = (self.header_values, texts | self.header_texts) - _header_holders.setdefault(self._headers_id(), set()).add(holder) + if self.header_values: + _held_headers[holder] = (self._headers_id(), self.header_values, self.header_texts) def release(self, holder: str) -> None: - """Drop ``holder``, forgetting the key, and the headers, once no launch or Sailbox in this process needs - them; the two are counted apart, since Sailboxes with different keys can share headers.""" + """Drop ``holder``, forgetting the key once no launch or Sailbox in this process needs it, and the + headers ``holder`` held.""" holders = _holders.get(self.secret, set()) holders.discard(holder) if not holders: _holders.pop(self.secret, None) _injected_keys.pop(self.secret, None) - if self.headers: - held = _header_holders.get(self._headers_id(), set()) - held.discard(holder) - if not held: - _header_holders.pop(self._headers_id(), None) - _injected_headers.pop(self._headers_id(), None) + _held_headers.pop(holder, None) def recover_key(self, candidate: str | None) -> None: """Recover the key this injection's secret was named from: one this process injected, else @@ -311,8 +293,10 @@ def recover_headers(self, candidate: Mapping[str, str] | None) -> None: ``candidate`` (the configured headers) when they are those.""" if not self.headers: return - if remembered := _injected_headers.get(self._headers_id()): - self.header_values, self.header_texts = remembered + held = [(values, texts) for headers, values, texts in _held_headers.values() if headers == self._headers_id()] + if held: + self.header_values = held[0][0] + self.header_texts = frozenset().union(*(texts for _, texts in held)) elif candidate and { name.lower(): header_secret_name(name, value) for name, value in candidate.items() } == self.headers: @@ -320,14 +304,12 @@ def recover_headers(self, candidate: Mapping[str, str] | None) -> None: def _encodings(self) -> list[tuple[str, str]]: """Each value agent-env writes into commands, and the placeholder it becomes, in every encoding: - shell-quoted, inside single quotes, and raw. The headers go whole, as each env var holds them, so a - short value can't rewrite unrelated text.""" + shell-quoted, inside single quotes, and raw. The headers go whole, as their JSON in each spelling, so + a short value can't rewrite unrelated text.""" pairs = [(self.key, PLACEHOLDER)] if self.key else [] if self.header_values: - placeholders = dict.fromkeys(self.header_values, HEADER_PLACEHOLDER) - pairs += [(json.dumps(self.header_values), json.dumps(placeholders)), - (header_lines(self.header_values), header_lines(placeholders))] - pairs += [(text, json.dumps(placeholders)) for text in sorted(self.header_texts)] + placeholders = json.dumps(dict.fromkeys(self.header_values, HEADER_PLACEHOLDER)) + pairs += [(text, placeholders) for text in [json.dumps(self.header_values), *sorted(self.header_texts)]] return [(encode(value), encode(placeholder)) for value, placeholder in pairs for encode in (shlex.quote, lambda text: text.replace("'", "'\\''"), lambda text: text)] diff --git a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py index 83cd6775..e7c34531 100644 --- a/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py @@ -66,7 +66,7 @@ _RESERVED_ENV_KEYS = frozenset({ "LITELLM_API_KEY", "ANTHROPIC_API_KEY", "LITELLM_BASE_URL", "ANTHROPIC_BASE_URL", - MODEL_HEADERS_ENV, "ANTHROPIC_CUSTOM_HEADERS", + MODEL_HEADERS_ENV, "PATH", "HOME", }) _MODEL_URL_KEYS = frozenset({"LITELLM_BASE_URL", "ANTHROPIC_BASE_URL"}) @@ -262,9 +262,6 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if MODEL_HEADERS_ENV not in merged_env: if headers := config.get_model_headers(merged_env.get("LITELLM_BASE_URL")): merged_env[MODEL_HEADERS_ENV] = json.dumps(headers) - if "ANTHROPIC_CUSTOM_HEADERS" not in merged_env: - if headers := config.get_model_headers(merged_env.get("ANTHROPIC_BASE_URL")): - merged_env["ANTHROPIC_CUSTOM_HEADERS"] = "\n".join(f"{k}: {v}" for k, v in headers.items()) for key in _MODEL_URL_KEYS & merged_env.keys(): merged_env[key] = host_url_for(merged_env[key], sandbox.type) env_flags = " ".join(f"-e {k}={shlex.quote(v)}" for k, v in merged_env.items()) diff --git a/tst/unit/config/test_model.py b/tst/unit/config/test_model.py index c7145333..77ba2750 100644 --- a/tst/unit/config/test_model.py +++ b/tst/unit/config/test_model.py @@ -199,18 +199,33 @@ def test_client_kwargs_omits_none_api_base_and_splats_params(): # --- headers ------------------------------------------------------------------ -_WITH_HEADERS = ModelConfig(base_url="https://gw/v1", headers={"X-Bill-To": "secret:org", "X-Plain": "v"}) - - -def test_headers_for_the_configured_endpoint_resolve_their_refs(): - assert _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({"org": "acme"})) == { - "X-Bill-To": "acme", "X-Plain": "v"} - - -def test_headers_ignore_a_trailing_slash_on_either_url(): - assert _WITH_HEADERS.headers_for("https://gw/v1/", secret_resolver=_secrets({"org": "acme"}))["X-Bill-To"] == "acme" - assert ModelConfig(base_url="https://gw/v1/", headers={"X-Plain": "v"}).headers_for( - "https://gw/v1", secret_resolver=_secrets({})) == {"X-Plain": "v"} +_HEADERS = {"X-Bill-To": "secret:org", "X-Plain": "v"} +_RESOLVED = {"X-Bill-To": "acme", "X-Plain": "v"} +_WITH_HEADERS = ModelConfig(base_url="https://gw/v1", headers=_HEADERS) + + +@pytest.mark.parametrize("configured, env_base, asked, expected", [ + ("https://gw/v1", None, "https://gw/v1", _RESOLVED), + ("https://gw/v1", None, "https://gw/v1/", _RESOLVED), + ("https://gw/v1/", None, "https://gw/v1", _RESOLVED), + ("https://gw/v1", "https://env-gw/v1", "https://env-gw/v1", _RESOLVED), + ("https://gw/v1", "https://env-gw/v1", "https://gw/v1", {}), + ("https://gw/v1", None, "https://elsewhere.example/v1", {}), + ("https://gw/v1", None, None, {}), + (None, None, "https://gw/v1", {}), +], ids=["configured", "trailing-slash-asked", "trailing-slash-configured", "env-configured", "file-shadowed-by-env", + "elsewhere", "no-url", "no-endpoint"]) +def test_headers_go_only_to_the_configured_endpoint_and_resolve_only_for_it(monkeypatch, configured, env_base, asked, + expected): + if env_base: + monkeypatch.setenv("LITELLM_BASE_URL", env_base) + resolved = [] + + headers = ModelConfig(base_url=configured, headers=_HEADERS).headers_for( + asked, secret_resolver=lambda key: resolved.append(key) or {"org": "acme"}.get(key)) + + assert headers == expected + assert bool(resolved) is bool(expected) def test_headers_held_back_from_another_endpoint_are_logged(caplog): @@ -220,45 +235,22 @@ def test_headers_held_back_from_another_endpoint_are_logged(caplog): assert "Not sending [model.headers] to elsewhere.example" in caplog.text -def test_a_header_ref_that_resolves_to_a_value_http_cant_carry_fails_loud(): - with pytest.raises(ConfigError, match="value HTTP can't carry"): - _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({"org": "acme\nX-Injected: 1"})) - - -def test_headers_follow_the_endpoint_the_env_var_configures(monkeypatch): - monkeypatch.setenv("LITELLM_BASE_URL", "https://env-gw/v1") - assert _WITH_HEADERS.headers_for("https://env-gw/v1", secret_resolver=_secrets({"org": "acme"}))["X-Bill-To"] == "acme" - assert _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({})) == {} - - -@pytest.mark.parametrize("base_url", ["https://elsewhere.example/v1", None]) -def test_headers_go_nowhere_but_the_configured_endpoint_and_resolve_nothing(base_url): - assert _WITH_HEADERS.headers_for(base_url, secret_resolver=_secrets({})) == {} - - -def test_headers_without_a_configured_endpoint_go_nowhere(): - assert ModelConfig(headers={"X-Plain": "v"}).headers_for("https://gw/v1", secret_resolver=_secrets({})) == {} - - -def test_an_unresolvable_header_ref_fails_loud(): - with pytest.raises(ConfigError, match="Unresolved secret:org"): - _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets({})) - - -def test_resolve_call_sends_the_headers_to_the_configured_endpoint(): - call = _WITH_HEADERS.resolve_call("m", secret_resolver=_secrets({"org": "acme"})) - assert call.params == {"extra_headers": {"X-Bill-To": "acme", "X-Plain": "v"}} - - -def test_resolve_call_sends_no_headers_to_a_foreign_override(): - call = _WITH_HEADERS.resolve_call("m", base_override="https://elsewhere.example", secret_resolver=_secrets({})) - assert call.params == {} +@pytest.mark.parametrize("secrets, problem", [ + ({}, "Unresolved secret:org"), + ({"org": "acme\nX-Injected: 1"}, "value HTTP can't carry"), +], ids=["unresolvable", "resolves-to-a-value-http-cant-carry"]) +def test_a_header_ref_that_cant_be_sent_fails_loud(secrets, problem): + with pytest.raises(ConfigError, match=problem): + _WITH_HEADERS.headers_for("https://gw/v1", secret_resolver=_secrets(secrets)) -def test_resolve_call_sends_no_headers_to_a_model_override(): - call = _WITH_HEADERS.resolve_call( - "m", model_overrides={"m": {"secret_key": "k", "base_url": "https://gw/v1"}}, secret_resolver=_secrets({"k": "sk"})) - assert call.params == {} +@pytest.mark.parametrize("call, expected", [ + ({}, {"extra_headers": _RESOLVED}), + ({"base_override": "https://elsewhere.example"}, {}), + ({"model_overrides": {"m": {"secret_key": "k", "base_url": "https://gw/v1"}}}, {}), +], ids=["configured", "foreign-override", "model-override"]) +def test_resolve_call_sends_the_headers_only_to_the_configured_endpoint(call, expected): + assert _WITH_HEADERS.resolve_call("m", secret_resolver=_secrets({"org": "acme", "k": "sk"}), **call).params == expected def test_params_extra_headers_win_per_header_whatever_its_case(): diff --git a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index 885f97aa..a9735b7d 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py @@ -19,7 +19,6 @@ ModelKeyInjection, carries_model_key, docker_shim, - header_lines, header_secret_name, secret_name, ) @@ -95,8 +94,7 @@ def fresh_key_state(monkeypatch): monkeypatch.setattr(_sdk, "_installed_key", None) monkeypatch.setattr(_sdk, "_apps", {}) monkeypatch.setattr(model_key, "_injected_keys", {}) - monkeypatch.setattr(model_key, "_injected_headers", {}) - monkeypatch.setattr(model_key, "_header_holders", {}) + monkeypatch.setattr(model_key, "_held_headers", {}) monkeypatch.setattr(model_key, "_holders", {}) @@ -974,7 +972,7 @@ async def test_a_failed_create_forgets_the_headers_too(): with pytest.raises(RuntimeError, match="no capacity"): await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox(image_name="agent:1", port=8000, env=_ENV_H) - assert model_key._injected_headers == {} + assert model_key._held_headers == {} @pytest.mark.asyncio @@ -986,17 +984,15 @@ async def test_commands_and_files_carry_header_placeholders_never_the_values(): written = json.dumps(_HEADERS) await sandbox.exec_script(f"docker run -d -e LITELLM_EXTRA_HEADERS='{written.replace(chr(39), chr(39) + chr(92) + chr(39) + chr(39))}' agent:1") - await sandbox.exec_with_output("sudo", "bash", "-c", f"docker exec -e LITELLM_EXTRA_HEADERS={shlex.quote(written)} " - f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_HEADERS))} c pytest") + await sandbox.exec_with_output("sudo", "bash", "-c", f"docker exec -e LITELLM_EXTRA_HEADERS={shlex.quote(written)} c pytest") await sandbox.write_host_file(f"LITELLM_EXTRA_HEADERS={written}\n".encode(), "/opt/agent/.env") sent = [arg for call in sailbox.exec.aio.await_args_list for arg in call.args[0]] sent.append(sailbox.fs.write.aio.await_args.args[1].decode()) assert not any(value in text for value in _HEADERS.values() for text in sent) - (run,) = [arg for arg in sent if "docker run" in arg] - assert json.dumps(_PLACEHELD) in run - (exec_,) = [arg for arg in sent if "docker exec" in arg] - assert shlex.quote(header_lines(_PLACEHELD)) in exec_ + for command in ("docker run", "docker exec"): + (script,) = [arg for arg in sent if command in arg] + assert json.dumps(_PLACEHELD) in script @pytest.mark.asyncio @@ -1051,10 +1047,10 @@ async def test_a_sailbox_without_injection_refuses_header_values_but_takes_place (f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_HEADERS))}", True, True), (f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_PLACEHELD))}", True, False), ('-e LITELLM_EXTRA_HEADERS="$LITELLM_EXTRA_HEADERS"', True, False), - (f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_HEADERS))}", True, True), - (f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote(header_lines(_PLACEHELD))}", True, False), (f"LITELLM_EXTRA_HEADERS={json.dumps(_HEADERS)}", False, True), (f"LITELLM_EXTRA_HEADERS={json.dumps(_PLACEHELD)}\nOTHER=1", False, False), + (f"LITELLM_EXTRA_HEADERS={json.dumps(_PLACEHELD)}\nLITELLM_EXTRA_HEADERS={json.dumps(_HEADERS)}", False, True), + (f"LITELLM_EXTRA_HEADERS: '{json.dumps(_HEADERS)}'", False, True), ("LITELLM_EXTRA_HEADERS={}", False, False), ]) def test_a_header_assignment_is_detected_as_the_shell_reads_it(text, shell, carries): diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index 97fe45c7..27363434 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -87,52 +87,25 @@ async def test_an_agent_gets_the_model_endpoint_its_container_reaches( @pytest.mark.asyncio @pytest.mark.parametrize("provider_cls", [_LocalProvider, _RemoteProvider]) -@pytest.mark.parametrize("passed", [False, True], ids=["configured", "passed"]) -async def test_an_agent_gets_the_model_endpoints_headers(monkeypatch, model_headers, provider_cls, passed): +@pytest.mark.parametrize("passed, own, expected", [ + (None, None, HEADERS), + ({"LITELLM_BASE_URL": MODEL}, None, HEADERS), + (None, {"LITELLM_BASE_URL": MODEL}, HEADERS), + ({"LITELLM_BASE_URL": ELSEWHERE}, None, None), + (None, {"LITELLM_BASE_URL": ELSEWHERE}, None), + ({"LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}'}, None, {"X-Bill-To": "mine"}), +], ids=["configured", "passed-configured", "own-configured", "passed-elsewhere", "own-elsewhere", "given-its-own"]) +async def test_an_agent_gets_the_headers_only_for_the_configured_endpoint( + monkeypatch, model_headers, provider_cls, passed, own, expected, +): provider = provider_cls() monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) with pytest.raises(_Started): - await _agent().deploy(sandbox_type="any", env_vars={"LITELLM_BASE_URL": MODEL} if passed else None) - - assert json.loads(provider.env["LITELLM_EXTRA_HEADERS"]) == HEADERS - - -@pytest.mark.asyncio -@pytest.mark.parametrize("deploy_env, agent_fields", [ - ({"LITELLM_BASE_URL": ELSEWHERE}, {}), - (None, {"default_env_vars": {"LITELLM_BASE_URL": ELSEWHERE}}), -], ids=["passed-elsewhere", "its-own-endpoint"]) -async def test_an_agent_pointed_elsewhere_gets_no_headers(monkeypatch, model_headers, deploy_env, agent_fields): - provider = _RemoteProvider() - monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) - - with pytest.raises(_Started): - await _agent(**agent_fields).deploy(sandbox_type="any", env_vars=deploy_env) - - assert "LITELLM_EXTRA_HEADERS" not in provider.env - - -@pytest.mark.asyncio -async def test_an_agent_whose_own_endpoint_is_the_configured_one_gets_its_headers(monkeypatch, model_headers): - provider = _RemoteProvider() - monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) - - with pytest.raises(_Started): - await _agent(default_env_vars={"LITELLM_BASE_URL": MODEL}).deploy(sandbox_type="any") - - assert json.loads(provider.env["LITELLM_EXTRA_HEADERS"]) == HEADERS - - -@pytest.mark.asyncio -async def test_headers_an_agent_is_given_win(monkeypatch, model_headers): - provider = _RemoteProvider() - monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: provider) - - with pytest.raises(_Started): - await _agent().deploy(sandbox_type="any", env_vars={"LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}'}) + await _agent(default_env_vars=own).deploy(sandbox_type="any", env_vars=passed) - assert provider.env["LITELLM_EXTRA_HEADERS"] == '{"X-Bill-To": "mine"}' + sent = provider.env.get("LITELLM_EXTRA_HEADERS") + assert (json.loads(sent) if sent else None) == expected @pytest.mark.asyncio @@ -221,17 +194,18 @@ def test_a_model_endpoint_the_verifier_sets_itself_is_reached_through_the_host(m assert "-e LITELLM_BASE_URL=http://host.docker.internal:5001 " in run -def test_a_unit_test_verifier_gets_the_model_endpoints_headers(monkeypatch, model_headers): - run = _run_verifier(monkeypatch, "local") - - assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(HEADERS))} " in run - assert f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote('X-Bill-To: acme')} " in run - - -def test_a_unit_test_verifier_pointed_elsewhere_gets_no_headers(monkeypatch, model_headers): - run = _run_verifier(monkeypatch, "local", env_vars={"LITELLM_BASE_URL": ELSEWHERE, "ANTHROPIC_BASE_URL": ELSEWHERE}) +@pytest.mark.parametrize("step_env, seed, expected", [ + ({}, None, HEADERS), + ({"LITELLM_BASE_URL": ELSEWHERE}, None, None), + ({}, {"litellm_extra_headers": "{}"}, HEADERS), +], ids=["configured", "elsewhere", "a-seed-column-cant-replace-them"]) +def test_a_unit_test_verifier_gets_the_headers_only_for_the_configured_endpoint( + monkeypatch, model_headers, step_env, seed, expected, +): + run = _run_verifier(monkeypatch, "local", seed=seed, env_vars=step_env) - assert "LITELLM_EXTRA_HEADERS" not in run and "ANTHROPIC_CUSTOM_HEADERS" not in run + flag = None if expected is None else f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(expected))} " + assert (flag in run) if flag else ("LITELLM_EXTRA_HEADERS" not in run) def test_a_unit_test_verifier_whose_headers_dont_resolve_fails(monkeypatch, local_stores, model_endpoint): @@ -244,19 +218,11 @@ def test_a_unit_test_verifier_whose_headers_dont_resolve_fails(monkeypatch, loca def test_a_unit_test_verifier_given_its_own_headers_never_resolves_the_configured_ones(monkeypatch, local_stores, model_endpoint): monkeypatch.setattr(local_stores, "_model_cfg", ModelConfig(headers={"X-Bill-To": "secret:fd_no_such_secret"})) - run = _run_verifier(monkeypatch, "local", env_vars={ - "LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}', "ANTHROPIC_CUSTOM_HEADERS": "X-Bill-To: mine"}) + run = _run_verifier(monkeypatch, "local", env_vars={"LITELLM_EXTRA_HEADERS": '{"X-Bill-To": "mine"}'}) assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps({'X-Bill-To': 'mine'}))} " in run -def test_a_seed_column_cant_replace_the_verifiers_headers(monkeypatch, model_headers): - run = _run_verifier(monkeypatch, "local", seed={"litellm_extra_headers": "{}", "anthropic_custom_headers": "x"}) - - assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(HEADERS))} " in run - assert f"-e ANTHROPIC_CUSTOM_HEADERS={shlex.quote('X-Bill-To: acme')} " in run - - @pytest.mark.parametrize("sandbox_type, expected", [ ("local", MODEL_IN_CONTAINER), ("modal_vm", MODEL), (None, MODEL), ], ids=["local-container", "remote-container", "host-install"]) From 08829f6b45dfbffca7b5ff999d6e1b29a183774b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 11:23:04 -0700 Subject: [PATCH 8/8] fix(sail_vm): scrub the headers' JSON before the key, so a header holding the key stays whole A [model.headers] value that is the model key itself (a gateway token, say) was rewritten by the key's scrub first, so the headers' JSON no longer matched and the deploy was refused. The headers go first now. The secret names' hash length is one named constant. Co-Authored-By: Claude Opus 5.5 --- .../sandbox_providers/sail_vm/model_key.py | 13 +++++++++---- .../sandbox_providers/sail_vm_model_key_test.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py index 02acc610..9dada3ed 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/model_key.py @@ -22,6 +22,8 @@ from agent_env.config.model import MODEL_HEADERS_ENV SECRET_PREFIX = "AGENTENV_LITELLM_" +#: Hex characters of a value's SHA-256 a secret name keeps. +SECRET_HASH_CHARS = 32 HEADER_SECRET_PREFIX = f"{SECRET_PREFIX}H_" POLICY_PREFIX = "agentenv-" PLACEHOLDER = "sail-injected-model-key" @@ -69,11 +71,12 @@ def secret_name(key: str) -> str: - return SECRET_PREFIX + hashlib.sha256(key.encode()).hexdigest()[:32].upper() + return SECRET_PREFIX + hashlib.sha256(key.encode()).hexdigest()[:SECRET_HASH_CHARS].upper() def header_secret_name(name: str, value: str) -> str: - return HEADER_SECRET_PREFIX + hashlib.sha256(f"{name.lower()}\0{value}".encode()).hexdigest()[:32].upper() + digest = hashlib.sha256(f"{name.lower()}\0{value}".encode()).hexdigest() + return HEADER_SECRET_PREFIX + digest[:SECRET_HASH_CHARS].upper() def _is_header_placeholder(value: str) -> bool: @@ -305,11 +308,13 @@ def recover_headers(self, candidate: Mapping[str, str] | None) -> None: def _encodings(self) -> list[tuple[str, str]]: """Each value agent-env writes into commands, and the placeholder it becomes, in every encoding: shell-quoted, inside single quotes, and raw. The headers go whole, as their JSON in each spelling, so - a short value can't rewrite unrelated text.""" - pairs = [(self.key, PLACEHOLDER)] if self.key else [] + a short value can't rewrite unrelated text, and first, so a header holding the key is still whole.""" + pairs = [] if self.header_values: placeholders = json.dumps(dict.fromkeys(self.header_values, HEADER_PLACEHOLDER)) pairs += [(text, placeholders) for text in [json.dumps(self.header_values), *sorted(self.header_texts)]] + if self.key: + pairs.append((self.key, PLACEHOLDER)) return [(encode(value), encode(placeholder)) for value, placeholder in pairs for encode in (shlex.quote, lambda text: text.replace("'", "'\\''"), lambda text: text)] diff --git a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index a9735b7d..58cb7b1f 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py @@ -1099,3 +1099,19 @@ def test_sailboxes_writing_the_same_headers_differently_each_keep_their_spelling for text in spellings: assert reconnected.scrub(f"LITELLM_EXTRA_HEADERS={shlex.quote(text)}") == ( f"LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(_PLACEHELD))}") + + +@pytest.mark.asyncio +async def test_a_header_holding_the_model_key_is_scrubbed_whole_and_not_refused(): + headers = {"X-Gateway-Token": _KEY, "X-HF-Bill-To": "acme-org"} + sailbox = _sailbox() + sailbox.exec.aio = AsyncMock(return_value=_process()) + sdk, _, _ = _fake_sdk(sailbox) + sandbox = await SailVmSandboxProvider(api_key="sail-key", sdk=sdk).create_sandbox( + image_name="agent:1", port=8000, env={**_ENV, "LITELLM_EXTRA_HEADERS": json.dumps(headers)}) + + await sandbox.exec_script(f"docker run -d -e LITELLM_API_KEY={_KEY} " + f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps(headers))} agent:1") + + (run,) = [arg for call in sailbox.exec.aio.await_args_list for arg in call.args[0] if "docker run" in arg] + assert _KEY not in run and json.dumps(dict.fromkeys(headers, HEADER_PLACEHOLDER)) in run