diff --git a/AGENTS.md b/AGENTS.md index 5531965d..e11b9828 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -182,6 +182,13 @@ orchestrators call it the same way. What callers rely on: Both write each prompt response through `PromptResponse.to_dict()`, which adds the S3-named key of every renamed trajectory field, since stored documents and raw-doc readers use them; serialize a context with `TaskStepContext.to_dict()`, not `dataclasses.asdict`. +- A per-run secret given as a bare `secret:KEY` or `env:NAME` reference names the secret rather + than holding it, so redaction keeps it, and the step that uses it resolves it with + `Config.resolve_reference()`: a runner whose steps run on another host passes per-run keys this + way. A reference with a `?default` is stripped like a secret. The run index strips the same keys + from a run's overrides. A reference resolves against the environment and `[stores.secret]` of the + host running the step, so a runner that takes runs from callers who may not read those refuses + references it did not create. - A run is resumable: pass a previously persisted `context`, a `start_step` and the existing `instance_id`; completed steps are seeded as SUCCESS. - A runner that must hear from a long step puts a zero-argument callable at diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index f6bc8f28..30262d78 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -416,8 +416,8 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: @click.option("--agent-artifact-id", default=None, type=str, help="Override the agent image artifact for deploy_agent steps") @click.option("--start-step", default=0, type=int, help="0-based step index to start execution from") @click.option("--context-json", default=None, type=click.Path(exists=True, dir_okay=False), help="Path to a JSON file to seed the TaskStepContext") -@click.option("--litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY for this run") -@click.option("--judge-litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY used by rubrics_verifier judge agents (falls back to --litellm-api-key, then to config secret)") +@click.option("--litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY for this run; a secret:KEY or env:NAME reference is resolved where it is used") +@click.option("--judge-litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY used by rubrics_verifier judge agents (falls back to --litellm-api-key, then to config secret); a secret:KEY or env:NAME reference is resolved where it is used") @click.option("--apply-trajectory-filter/--no-trajectory-filter", "apply_trajectory_filter", default=None, help="Force trajectory filtering on/off for rubrics_verifier (overrides the step's trajectory_filter setting)") @click.option("--a2a-agent-id", default=None, type=str, help="Override A2A agent id for deploy_agent steps") @click.option("--agent-sandbox", default=None, @@ -558,8 +558,8 @@ def _seed_universe_id(task, seed: dict) -> str | None: @click.option("--user-model", default=None, type=str, help="Override the model the user simulator runs on") @click.option("--judge-model", default=None, type=str, help="Override the model rubrics_verifier judges grade with") @click.option("--agent-artifact-id", default=None, type=str, help="Override the agent image artifact for deploy_agent steps") -@click.option("--litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY for all runs") -@click.option("--judge-litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY used by rubrics_verifier judge agents") +@click.option("--litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY for all runs; a secret:KEY or env:NAME reference is resolved where it is used") +@click.option("--judge-litellm-api-key", default=None, type=str, help="Override LITELLM_API_KEY used by rubrics_verifier judge agents; a secret:KEY or env:NAME reference is resolved where it is used") @click.option("--apply-trajectory-filter/--no-trajectory-filter", "apply_trajectory_filter", default=None, help="Force trajectory filtering on/off for rubrics_verifier") @click.option("--agent-sandbox", default=None, help="Override agent sandbox backend(s) for deploy_agent steps; comma-separated for fallback chain (e.g. 'modal,local')") diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 6628176a..6c7c5c52 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -315,6 +315,11 @@ def _get_secret(self) -> Mapping[str, Any]: self._secret = bundle return self._secret + def resolve_reference(self, value: Any) -> Any: + """``value`` with an ``env:NAME`` / ``secret:KEY`` reference resolved, the latter through ``[stores.secret]``; + any other value as it is.""" + return config_loader.interpolate(value, secret_resolver=self._resolve_secret) + def get_litellm_api_key(self) -> str: return self._model_config().resolve_api_key(secret_resolver=self._resolve_secret) diff --git a/src/agent_env/runner/store.py b/src/agent_env/runner/store.py index d30b0b6c..5eaf2698 100644 --- a/src/agent_env/runner/store.py +++ b/src/agent_env/runner/store.py @@ -13,6 +13,7 @@ from agent_env.config import get_config from agent_env.runner.runner import RunRecord, RunStatus from agent_env.store import Eq, Filter, In, Sort, UpdateSpec +from agent_env.task_step.context import _strip_redacted_keys RUNS_COLLECTION = "runs" @@ -38,7 +39,9 @@ def ensure_indexes() -> None: def insert_run(record: RunRecord) -> RunRecord: record.created_at_utc = record.created_at_utc or _now() - _docs().insert(RUNS_COLLECTION, record.to_dict()) + doc = record.to_dict() + doc["overrides"] = _strip_redacted_keys(doc["overrides"]) + _docs().insert(RUNS_COLLECTION, doc) return record diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index 936d00dd..48a0629e 100644 --- a/src/agent_env/task_step/context.py +++ b/src/agent_env/task_step/context.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from typing import Any +from agent_env.config.loader import ENV_REF_PREFIX, SECRET_REF_PREFIX from agent_env.env.env import DeployedEnv from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, host_url_for from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url @@ -19,6 +20,14 @@ } +def _is_redacted(key: str, value: Any) -> bool: + """Whether ``key`` holds a secret kept out of storage. A bare ``env:`` / ``secret:`` reference names one rather + than holding it, so it is kept; one with a ``?default`` may carry the secret itself, so it is not.""" + return key in _REDACTED_KEYS and not ( + isinstance(value, str) and value.startswith((ENV_REF_PREFIX, SECRET_REF_PREFIX)) and "?" not in value + ) + + def _holds_redacted_key(value: Any) -> bool: """Whether ``value`` carries a ``_REDACTED_KEYS`` entry at any depth.""" return isinstance(value, dict) and any( @@ -51,7 +60,7 @@ def _strip_redacted_keys(value: Any) -> Any: never reach Mongo regardless of which code path persists them. """ if isinstance(value, dict): - return {k: _strip_redacted_keys(v) for k, v in value.items() if k not in _REDACTED_KEYS} + return {k: _strip_redacted_keys(v) for k, v in value.items() if not _is_redacted(k, v)} if isinstance(value, list): return [_strip_redacted_keys(v) for v in value] return value diff --git a/src/agent_env/task_step/context_ops.py b/src/agent_env/task_step/context_ops.py index e5301847..9649b9ae 100644 --- a/src/agent_env/task_step/context_ops.py +++ b/src/agent_env/task_step/context_ops.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints -from agent_env.task_step.context import PromptResponse, TaskStepContext, _REDACTED_KEYS, _strip_redacted_keys +from agent_env.task_step.context import PromptResponse, TaskStepContext, _is_redacted, _strip_redacted_keys if TYPE_CHECKING: from agent_env.store.document_store import UpdateSpec @@ -191,12 +191,12 @@ def _diff_dict( return for k in pre.keys() - post.keys(): - if k in _REDACTED_KEYS: + if _is_redacted(k, pre[k]): continue ops.unsets.add(f"{path}.{k}") for k, post_v in post.items(): - if k in _REDACTED_KEYS: + if _is_redacted(k, post_v): continue sub_path = f"{path}.{k}" if k not in pre: diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 87696db1..f2b3728b 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -368,7 +368,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas env_vars = dict(self.env_vars) if self.env_vars else {} override_litellm_key = user_overrides.get("litellm_api_key") if override_litellm_key: - env_vars["LITELLM_API_KEY"] = override_litellm_key + env_vars["LITELLM_API_KEY"] = config.resolve_reference(override_litellm_key) resolved_litellm_base_url = user_overrides.get("litellm_base_url") or self.litellm_base_url if resolved_litellm_base_url: env_vars["LITELLM_BASE_URL"] = resolved_litellm_base_url diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index cb2d6229..d31618bc 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -6,6 +6,7 @@ import logging from typing import ClassVar, Optional +from agent_env.config import get_config from agent_env.env.gateway import GatewayMode from agent_env.env.store import update_env_instance_metadata from agent_env.task_step.context import TaskStepContext @@ -170,7 +171,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) model_api_key = user_overrides.get("litellm_api_key") if model_api_key and ("litellm_api_key" in params or accepts_kwargs): - deploy_kwargs["litellm_api_key"] = model_api_key + deploy_kwargs["litellm_api_key"] = get_config().resolve_reference(model_api_key) if self.artifact_id is not None and ("artifact_id" in params or accepts_kwargs): deploy_kwargs["artifact_id"] = self.artifact_id if self.artifact_version is not None and ("artifact_version" in params or accepts_kwargs): 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 5d5d10b5..768d9788 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -96,7 +96,8 @@ def _resolve_container() -> str: "agent_name": lambda: agent_name, "workspace_dir": _resolve_workspace_dir, # Config-sourced (from agent-env config) - "litellm_api_key": lambda: litellm_api_key_override or config.get_litellm_api_key(), + "litellm_api_key": lambda: (config.resolve_reference(litellm_api_key_override) if litellm_api_key_override + else 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/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index 43427dca..f0eda557 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -626,8 +626,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: config = get_config() api_key = ( - overrides.get("judge_litellm_api_key") - or overrides.get("litellm_api_key") + config.resolve_reference(overrides.get("judge_litellm_api_key") or overrides.get("litellm_api_key")) or config.get_litellm_api_key() ) litellm_base_url = ( @@ -993,7 +992,9 @@ async def _prompt_llm_judge( config = get_config() call_config = config.resolve_model_call( model, - default_api_key=overrides.get("judge_litellm_api_key") or overrides.get("litellm_api_key"), + default_api_key=config.resolve_reference( + overrides.get("judge_litellm_api_key") or overrides.get("litellm_api_key") + ), base_override=overrides.get("judge_litellm_base_url") or overrides.get("litellm_base_url") or self.default_model_api_base, diff --git a/tst/unit/config/test_resolve_reference.py b/tst/unit/config/test_resolve_reference.py new file mode 100644 index 00000000..b6e48645 --- /dev/null +++ b/tst/unit/config/test_resolve_reference.py @@ -0,0 +1,11 @@ +"""Config.resolve_reference: the resolver a step applies to a per-run secret given by reference.""" + +from agent_env.config import get_config + + +def test_an_env_reference_resolves_and_anything_else_passes_through(monkeypatch): + monkeypatch.setenv("PER_RUN_KEY", "sk-from-env") + config = get_config() + assert config.resolve_reference("env:PER_RUN_KEY") == "sk-from-env" + assert config.resolve_reference("sk-plain") == "sk-plain" + assert config.resolve_reference(None) is None diff --git a/tst/unit/runner/run_store_test.py b/tst/unit/runner/run_store_test.py index 5c1eff2a..eab918dc 100644 --- a/tst/unit/runner/run_store_test.py +++ b/tst/unit/runner/run_store_test.py @@ -50,6 +50,16 @@ def test_set_running_advances_only_a_queued_run(docs): assert run_store.get_run("c").status == RunStatus.CANCELED +def test_the_index_never_holds_a_per_run_key_but_keeps_a_reference(docs): + record = _record("r", RunStatus.QUEUED) + record.overrides = {"metadata": {"user_overrides": { + "litellm_api_key": "sk-raw", "judge_litellm_api_key": "secret:judge", "agent_sandbox": "local"}}} + run_store.insert_run(record) + assert run_store.get_run("r").overrides["metadata"]["user_overrides"] == { + "judge_litellm_api_key": "secret:judge", "agent_sandbox": "local"} + assert record.overrides["metadata"]["user_overrides"]["litellm_api_key"] == "sk-raw" + + def test_count_runs_counts_one_task_or_every_task(docs): run_store.insert_run(_record("a", RunStatus.QUEUED)) run_store.insert_run(RunRecord(run_id="b", runner="other", task_id="u", task_version=1, instance_id="j")) diff --git a/tst/unit/task_step/context_ops_journal_test.py b/tst/unit/task_step/context_ops_journal_test.py index 3fad38f7..5a8c6677 100644 --- a/tst/unit/task_step/context_ops_journal_test.py +++ b/tst/unit/task_step/context_ops_journal_test.py @@ -85,6 +85,15 @@ def test_the_judge_key_never_reaches_the_store_but_survives_a_rebuild(): assert live.metadata["user_overrides"] == {"judge_litellm_api_key": "sk-judge", "agent_effort": "high"} +def test_a_per_run_key_by_reference_is_persisted_but_not_one_carrying_a_default(): + live = TaskStepContext(metadata={"user_overrides": { + "litellm_api_key": "secret:team_key", "judge_litellm_api_key": "env:JUDGE_KEY?sk-raw-default"}}) + stored = live.to_safe_dict() + assert stored["metadata"]["user_overrides"] == {"litellm_api_key": "secret:team_key"} + journal = repr(build_context_update_ops(None, live).to_journal_dict()) + assert "secret:team_key" in journal and "sk-raw-default" not in journal + + def test_the_runner_heartbeat_is_never_persisted_but_survives_a_rebuild(): def beat(): return None diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index d36e38ac..67231e0b 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -11,6 +11,7 @@ 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.store.secret_store.local_secret_store import LocalSecretStore from agent_env.task_step.context import DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.install_agent import ( InstallAgentTaskStep, @@ -98,6 +99,9 @@ class _FakeConfig: def get_litellm_api_key(self): return "config-key" + def resolve_reference(self, value): + return value + def test_litellm_key_resolver_prefers_per_run_override(): r = _build_param_resolvers( @@ -178,6 +182,52 @@ async def test_install_logs_show_the_command_template_not_the_model_key_or_heade assert "never-logged" not in caplog.text +@pytest.mark.asyncio +async def test_a_per_run_key_given_by_reference_is_resolved_through_the_secret_store(monkeypatch, local_stores): + monkeypatch.setenv("LITELLM_BASE_URL", "http://localhost:4000/v1") + local_stores.set_secret_store(LocalSecretStore(values={"team_key": "sk-from-the-store"}, use_env=False)) + 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}"], + "required_params_host": ["litellm_api_key"], "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"})], + metadata={"user_overrides": {"litellm_api_key": "secret:team_key"}}, + ) + + await step.execute(context) + + assert [s for s in sandbox.scripts if s.startswith("start ")] == ["start --key sk-from-the-store"] + + +def test_a_per_run_reference_is_resolved_only_when_the_install_asks_for_the_key(): + resolved = [] + + class _Config(_FakeConfig): + def resolve_reference(self, value): + resolved.append(value) + return "sk-resolved" + + r = _build_param_resolvers( + container_name="c", work_dir="/tmp/stage", agent_ctx_tar="/tmp/stage/ctx.tar.gz", + a2a_port=8000, config=_Config(), agent_name="default-agent", workspace_dir=None, + litellm_api_key_override="secret:team_key", + ) + assert resolved == [] + assert r["litellm_api_key"]() == "sk-resolved" + assert resolved == ["secret:team_key"] + + @pytest.mark.asyncio async def test_a_step_naming_no_agent_installs_the_configured_default(monkeypatch, local_stores): asked = []