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..96f27c3f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +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: `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/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..aadffd89 100644 --- a/packages/agentenv-protocol/examples/reference_model.py +++ b/packages/agentenv-protocol/examples/reference_model.py @@ -42,13 +42,18 @@ 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. + configured = os.getenv("LITELLM_BASE_URL") or "" + headers = (json.loads(os.getenv("LITELLM_EXTRA_HEADERS") or "{}") + if base_url.rstrip("/") == configured.rstrip("/") else {}) base_url = base_url.rstrip("/") endpoint = ( f"{base_url}/chat/completions" if base_url.endswith("/v1") else f"{base_url}/v1/chat/completions" ) - return endpoint, {"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 800bfe89..2de7408e 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py @@ -64,9 +64,24 @@ def _load_example(name: str) -> ModuleType: @pytest.mark.asyncio -async def test_reference_model_forwards_generic_metadata( +@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"}), + ('{"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", "trailing-slash", "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: + monkeypatch.setenv("LITELLM_EXTRA_HEADERS", endpoint_headers) captured: dict[str, Any] = {} class Response: def raise_for_status(self) -> None: @@ -107,7 +122,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..ca561544 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,11 @@ 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", 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: resolved_env["LITELLM_BASE_URL"] = (provider or sandbox).url_from_sandbox(resolved_env["LITELLM_BASE_URL"]) 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 e83dd619..7a297409 100644 --- a/src/agent_env/config/model.py +++ b/src/agent_env/config/model.py @@ -1,19 +1,29 @@ -"""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 +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): @@ -57,13 +67,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 +85,13 @@ 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)") + 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 != "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 +100,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 +123,21 @@ 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.""" + if not self.headers or not base_url: + return {} + configured_base = os.getenv(_ENV_BASE_URL) or self.base_url + 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 {} + 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, model: str, @@ -149,6 +181,21 @@ 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) - 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: + 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) + + +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/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/providers/sandbox_providers/sail_vm/README.md b/src/agent_env/providers/sandbox_providers/sail_vm/README.md index fc007ce9..a0c2159b 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. 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. + `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. @@ -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..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 @@ -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,19 @@ from typing import Any, Mapping from urllib.parse import urlparse +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" +HEADER_PLACEHOLDER = "sail-injected-header" KEY_ENV = "LITELLM_API_KEY" BASE_URL_ENV = "LITELLM_BASE_URL" +#: 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" @@ -54,10 +63,29 @@ # and Sailboxes in this process that need it, so the key is forgotten once none does. Process memory only. _injected_keys: dict[str, str] = {} _holders: dict[str, set[str]] = {} +# 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: - 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: + 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: + """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: @@ -65,16 +93,28 @@ def carries_model_key(text: str, *, shell: bool) -> bool: 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 == 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 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 @@ -117,12 +157,22 @@ 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) + #: ``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 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 +198,24 @@ 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) + 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"{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 " + "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_texts=frozenset({text}) if values and text != json.dumps(values) else frozenset(), + path_prefix=f"{path}/" if values and path else None) @classmethod def from_document(cls, document: Any, policy_id: str | None) -> ModelKeyInjection | None: @@ -156,35 +223,66 @@ 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() + if not isinstance(entries, list) or len(entries) not in (1, 2): + return None + try: + 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 + 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, 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.""" + """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}}}" - return {self.host: [{"request": {"set": {"headers": {"authorization": f"Bearer {ref}", "x-api-key": ref}}}}]} + 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) == (self.host, self.secret) + 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 secret.""" + """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 _holders.setdefault(self.secret, 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 once no launch or Sailbox in this process needs it.""" + """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) + _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 @@ -193,16 +291,39 @@ 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 + 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: + 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 their JSON in each spelling, so + 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)] 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/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index b38157f4..ca7c1eb9 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())), } @@ -290,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/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..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 @@ -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, "PATH", "HOME", }) _MODEL_URL_KEYS = frozenset({"LITELLM_BASE_URL", "ANTHROPIC_BASE_URL"}) @@ -255,6 +257,11 @@ 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) + # 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. + 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) 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 38f41494..77ba2750 100644 --- a/tst/unit/config/test_model.py +++ b/tst/unit/config/test_model.py @@ -57,6 +57,32 @@ 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) + + +@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(): @@ -140,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" @@ -168,3 +195,67 @@ 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 ------------------------------------------------------------------ + +_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): + 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 + + +@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)) + + +@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(): + 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-Plain": "v", "x-bill-to": "other"}} + 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/providers/sandbox_providers/sail_vm_model_key_test.py b/tst/unit/providers/sandbox_providers/sail_vm_model_key_test.py index db731bae..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 @@ -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,12 @@ 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_secret_name, secret_name, ) from agent_env.providers.sandbox_providers.sail_vm.provider import SailVmSandboxProvider @@ -91,6 +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, "_held_headers", {}) monkeypatch.setattr(model_key, "_holders", {}) @@ -887,3 +891,227 @@ 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_under_the_endpoint_path_and_the_key_everywhere(): + injection = ModelKeyInjection.for_env(_ENV_H) + 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-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): + 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)) + 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 +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._held_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)} 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) + 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 +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"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): + 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 + + +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))}") + + +@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 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 c16d8dd5..27363434 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,8 @@ 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 from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider @@ -29,6 +33,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 +43,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 +82,30 @@ 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, 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(default_env_vars=own).deploy(sandbox_type="any", env_vars=passed) + + sent = provider.env.get("LITELLM_EXTRA_HEADERS") + assert (json.loads(sent) if sent else None) == expected @pytest.mark.asyncio @@ -122,7 +157,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 +167,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 +194,35 @@ 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 +@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) + + 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): + 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_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"}'}) + + assert f"-e LITELLM_EXTRA_HEADERS={shlex.quote(json.dumps({'X-Bill-To': 'mine'}))} " 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 +235,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"