From ce1b42fa665e0991bc290ab38b856a9ed6b64490 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 17:27:38 -0700 Subject: [PATCH] feat(task-steps): per-run step_overrides for the solver agent, model and effort DeployAgentTaskStep reads a2a_agent_id / a2a_agent_version and PromptAgentTaskStep reads model / effort from the run's step_overrides for their own step id, ahead of the run-wide a2a_agent_id / agent_model / agent_effort overrides and the stored values. An overridden agent id takes the override's version (latest when unset), never the stored agent's pin. One task version can now run under a different solver harness and model per run while the user simulators and judges in the same task keep their stored agents; the run-wide override swaps every deploy_agent step at once. Co-Authored-By: Claude Fable 5.1 --- .../task_step/task_steps/deploy_agent.py | 20 +++-- .../task_step/task_steps/prompt_agent.py | 11 ++- .../test_step_overrides_agent_and_model.py | 88 +++++++++++++++++++ 3 files changed, 111 insertions(+), 8 deletions(-) create mode 100644 tst/unit/task_step/test_step_overrides_agent_and_model.py 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 b4e1c6e8..8e6a6213 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/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") diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index 1b0c3062..eb2ddb37 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/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 diff --git a/tst/unit/task_step/test_step_overrides_agent_and_model.py b/tst/unit/task_step/test_step_overrides_agent_and_model.py new file mode 100644 index 00000000..480db671 --- /dev/null +++ b/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")