feat(task_step): pass per-run secrets by reference, and keep raw ones out of the run index

AGENTS.md

@@ -174,6 +174,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.
 - An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task in which

src/agent_env/cli/task/run.py

@@ -414,8 +414,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,
@@ -550,8 +550,8 @@ def _seed_universe_id(task, seed: dict) -> str | None:
 @click.option("--output-dir", default=None, type=click.Path(file_okay=False), help="Directory to write context JSON outputs")
 @click.option("--agent-model", default=None, type=str, help="Override the model used by prompt_agent steps")
 @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')")

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)
 

src/agent_env/runner/store.py

@@ -14,6 +14,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"
 
@@ -39,7 +40,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
 
 

src/agent_env/task_step/context.py

@@ -6,6 +6,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
@@ -15,6 +16,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(
@@ -47,7 +56,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

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:

src/agent_env/task_step/task_steps/deploy_agent.py

@@ -356,7 +356,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

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
@@ -167,7 +168,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):

src/agent_env/task_step/task_steps/install_agent.py

@@ -94,7 +94,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())),

src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py

@@ -619,8 +619,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 = (
@@ -980,7 +979,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,

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

tst/unit/runner/run_store_test.py

@@ -48,3 +48,13 @@ def test_set_running_advances_only_a_queued_run(docs):
     run_store.insert_run(_record("c", RunStatus.CANCELED))
     run_store.set_running("c")
     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"

tst/unit/task_step/context_ops_journal_test.py

@@ -83,3 +83,12 @@ def test_the_judge_key_never_reaches_the_store_but_survives_a_rebuild():
     assert "sk-judge" not in repr(build_context_update_ops(None, live).to_journal_dict())
     _rebuild_context(live, stored)
     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

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,
@@ -96,6 +97,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(
@@ -174,3 +178,49 @@ async def test_install_logs_show_the_command_template_not_the_model_key_or_heade
     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
+
+
+@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"]