feat(task-steps): per-run step_overrides for the solver agent, model and effort

src/agent_env/task_step/task_steps/deploy_agent.py

@@ -286,6 +286,19 @@ def from_dict(cls, data: dict) -> DeployAgentTaskStep:
             enable_docker=data.get("enable_docker", False),
         )
 
+    def _resolve_agent_ref(self, context: TaskStepContext) -> tuple[str, Optional[int]]:
+        """Agent id and version for this run: this step's ``step_overrides`` win over the run-wide override and the stored values; an overridden id takes the override's version (latest when unset), never the stored agent's pin."""
+        step_overrides = self.step_param_overrides(context)
+        user_overrides = context.metadata.get("user_overrides", {})
+        if "a2a_agent_id" in step_overrides:
+            return step_overrides["a2a_agent_id"], step_overrides.get("a2a_agent_version")
+        agent_id = user_overrides.get("a2a_agent_id") or self.a2a_agent_id or get_config().get_default_a2a_agent_id()
+        if "a2a_agent_version" in step_overrides:
+            return agent_id, step_overrides["a2a_agent_version"]
+        if "a2a_agent_version" in user_overrides:
+            return agent_id, user_overrides["a2a_agent_version"]
+        return agent_id, self.a2a_agent_version
+
     async def execute(self, context: TaskStepContext) -> TaskStepContext:
         # An agent is recorded in the context only once it is configured, so one that fails on the way there
         # (a snapshot or changelog that can't be loaded, say) is closed here, or no teardown would find it.
@@ -335,12 +348,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas
 
         config = get_config()
         user_overrides = context.metadata.get("user_overrides", {})
-        a2a_agent_id = user_overrides.get("a2a_agent_id") or self.a2a_agent_id or config.get_default_a2a_agent_id()
-        a2a_agent_version = (
-            user_overrides["a2a_agent_version"]
-            if "a2a_agent_version" in user_overrides
-            else self.a2a_agent_version
-        )
+        a2a_agent_id, a2a_agent_version = self._resolve_agent_ref(context)
         agent = A2AAgent.get(a2a_agent_id, a2a_agent_version)
         requested_disk = user_overrides.get("agent_disk_size_gb") or self.disk_size_gb
         min_disk = agent.metadata.get("min_disk_size_gb")

src/agent_env/task_step/task_steps/prompt_agent.py

@@ -350,6 +350,14 @@ def _warn_on_snapshot_budget(self) -> None:
                 f"after turn {cfg.max_snapshots} and the rest are recorded as skipped"
             )
 
+    def _resolve_model_and_effort(self, context: TaskStepContext) -> tuple[Optional[str], Optional[str]]:
+        """Model and effort for this run: this step's ``step_overrides`` win over the run-wide ``agent_model`` / ``agent_effort`` and the stored values."""
+        step_overrides = self.step_param_overrides(context)
+        overrides = context.metadata.get("user_overrides", {})
+        model = step_overrides.get("model") or context.agent_model or self.model or context.default_agent_model
+        effort = step_overrides.get("effort") or overrides.get("agent_effort") or self.effort
+        return model, effort
+
     async def execute(self, context: TaskStepContext) -> TaskStepContext:
         # A changelog staged on the agent moves into the store while it works, and the rest when it is done.
         async with draining(staged_changelogs(context.metadata, agent_name=self.agent_name)):
@@ -411,10 +419,9 @@ async def _execute_conversation(
         initial_parts = self._apply_seed(self._effective_parts(), seed)
         prompt_text = initial_parts[0]["text"] if self.prompt is not None else None
 
-        model = context.agent_model or self.model or context.default_agent_model
+        model, effort = self._resolve_model_and_effort(context)
         harness = context.agent_harness or self.harness
         overrides = context.metadata.get("user_overrides", {})
-        effort = overrides.get("agent_effort") or self.effort
         max_thinking_tokens = overrides.get("agent_max_thinking_tokens") or self.max_thinking_tokens
         max_turns = overrides.get("agent_max_turns") or self.max_turns
 

tst/unit/task_step/test_step_overrides_agent_and_model.py

@@ -0,0 +1,88 @@
+"""Unit tests for per-run ``step_overrides`` of the solver on DeployAgentTaskStep and PromptAgentTaskStep.
+
+One task version can then run under a different A2A agent, model or effort per run while the other
+agents in the task (user simulators, judges) keep their stored settings.
+"""
+
+from __future__ import annotations
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep
+from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep
+
+
+def _ctx(step_params: dict | None = None, run_overrides: dict | None = None, **fields) -> TaskStepContext:
+    user_overrides = dict(run_overrides or {})
+    if step_params is not None:
+        user_overrides["step_params"] = step_params
+    metadata = {"user_overrides": user_overrides} if user_overrides else {}
+    return TaskStepContext(metadata=metadata, **fields)
+
+
+def _deploy_step(**kw) -> DeployAgentTaskStep:
+    kw.setdefault("id", "deploy-solver")
+    kw.setdefault("version", None)
+    kw.setdefault("agent_name", "solver")
+    kw.setdefault("a2a_agent_id", "stored-agent")
+    kw.setdefault("a2a_agent_version", 7)
+    return DeployAgentTaskStep(**kw)
+
+
+def _prompt_step(**kw) -> PromptAgentTaskStep:
+    kw.setdefault("id", "prompt")
+    kw.setdefault("version", None)
+    kw.setdefault("agent_name", "solver")
+    kw.setdefault("prompt", "go")
+    kw.setdefault("model", "stored-model")
+    kw.setdefault("effort", "stored-effort")
+    return PromptAgentTaskStep(**kw)
+
+
+class TestDeployAgentRef:
+    def test_stored_values_without_overrides(self):
+        assert _deploy_step()._resolve_agent_ref(_ctx()) == ("stored-agent", 7)
+
+    def test_step_override_swaps_agent_and_drops_stored_pin(self):
+        ctx = _ctx({"deploy-solver": {"a2a_agent_id": "other-agent"}})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("other-agent", None)
+
+    def test_step_override_with_version(self):
+        ctx = _ctx({"deploy-solver": {"a2a_agent_id": "other-agent", "a2a_agent_version": 3}})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("other-agent", 3)
+
+    def test_step_override_version_only_keeps_agent(self):
+        ctx = _ctx({"deploy-solver": {"a2a_agent_version": 9}})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("stored-agent", 9)
+
+    def test_step_override_wins_over_run_wide_override(self):
+        ctx = _ctx({"deploy-solver": {"a2a_agent_id": "other-agent"}}, {"a2a_agent_id": "run-wide", "a2a_agent_version": 2})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("other-agent", None)
+
+    def test_run_wide_override_still_applies(self):
+        ctx = _ctx(run_overrides={"a2a_agent_id": "run-wide", "a2a_agent_version": 2})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("run-wide", 2)
+
+    def test_other_steps_overrides_are_ignored(self):
+        ctx = _ctx({"deploy-judge": {"a2a_agent_id": "other-agent"}})
+        assert _deploy_step()._resolve_agent_ref(ctx) == ("stored-agent", 7)
+
+
+class TestPromptModelAndEffort:
+    def test_stored_values_without_overrides(self):
+        assert _prompt_step()._resolve_model_and_effort(_ctx()) == ("stored-model", "stored-effort")
+
+    def test_step_override_wins_over_run_wide_and_stored(self):
+        ctx = _ctx({"prompt": {"model": "step-model", "effort": "xhigh"}}, {"agent_effort": "low"}, agent_model="run-model")
+        assert _prompt_step()._resolve_model_and_effort(ctx) == ("step-model", "xhigh")
+
+    def test_run_wide_overrides_still_apply(self):
+        ctx = _ctx(run_overrides={"agent_effort": "low"}, agent_model="run-model")
+        assert _prompt_step()._resolve_model_and_effort(ctx) == ("run-model", "low")
+
+    def test_partial_step_override_leaves_the_other_field(self):
+        ctx = _ctx({"prompt": {"model": "step-model"}}, agent_model="run-model")
+        assert _prompt_step()._resolve_model_and_effort(ctx) == ("step-model", "stored-effort")
+
+    def test_default_model_is_the_last_resort(self):
+        ctx = _ctx(default_agent_model="fallback")
+        assert _prompt_step(model=None)._resolve_model_and_effort(ctx) == ("fallback", "stored-effort")