diff --git a/src/agent_env/task_step/__init__.py b/src/agent_env/task_step/__init__.py index 626cea566..dd249d93b 100644 --- a/src/agent_env/task_step/__init__.py +++ b/src/agent_env/task_step/__init__.py @@ -1,7 +1,7 @@ """TaskStep module for AgentEnv.""" from agent_env.env.env import DeployedEnv -from .context import DeployedAgent, PromptResponse, TaskStepContext +from .context import DeployedAgent, PromptResponse, StepContextError, TaskStepContext from .task_step import TaskStep from .task_steps import AddSkillsTaskStep, AgentPromptResponseVerifierTaskStep, BuildMcpCliTaskStep, CombineUniverseVerdictsStep, DeployAgentTaskStep, DeployEnvTaskStep, EnvOutcomeVerifierTaskStep, LoadArtifactTaskStep, ModifyEnvToolAccessStep, PromptAgentTaskStep, RubricsVerifierTaskStep, ScoreAggregator, ValidationGateAggregatorStep, VerifyCoreA2AProtocolStep, VerifyCoreEnvironmentProtocolStep, VerifyEnvironmentCardStep, VerifyMCPEnvAssessmentStep, VerifyMCPToolSchemaTaskStep, VerifySpecConformanceTaskStep, VerifyUniverseLoadExportRoundtripStep, VerifyA2AAgentCardStep, VerifyA2AAgentMCPStep, VerifyA2ASkillConfigStep, VerifyA2ATrajectoryStep from .task_steps.verifiers.judge_utils.judge_output_format import JudgeOutputFormat, get_judge_output_format_spec @@ -14,6 +14,7 @@ "get_judge_output_format_spec", "PromptResponse", "ScoreAggregator", + "StepContextError", "TaskStepContext", "TaskStep", "AddSkillsTaskStep", diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index e6d541bae..082a3d736 100644 --- a/src/agent_env/task_step/context.py +++ b/src/agent_env/task_step/context.py @@ -14,6 +14,12 @@ HEARTBEAT_KEY = "_heartbeat_fn" + +class StepContextError(RuntimeError): + """The run's context can't satisfy a step: something the step needs from an earlier step is missing, or + something it would create is already there. Running the step again on the same context fails the same + way, so a runner should not retry it.""" + _REDACTED_KEYS = { "litellm_api_key", "judge_litellm_api_key", "usersim_api_key", "remote_tokens", "cf_access_client_secret", HEARTBEAT_KEY, diff --git a/src/agent_env/task_step/task_steps/add_skills.py b/src/agent_env/task_step/task_steps/add_skills.py index f8fccc3d5..019b02a27 100644 --- a/src/agent_env/task_step/task_steps/add_skills.py +++ b/src/agent_env/task_step/task_steps/add_skills.py @@ -15,7 +15,7 @@ from agentenv_protocol.client import children, find_extension_method from agent_env.config import get_config -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -261,7 +261,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise ValueError("; ".join(problems)) agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents") if agent.instance_id is None: raise RuntimeError(f"Agent '{self.agent_name}' is missing instance_id; cannot resolve deployment") diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index c55923a6a..b440a5ce0 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -80,7 +80,7 @@ from agent_env.config import get_config from agent_env.env.gateway.constants import EXT_STEP_URI from agent_env.store.ids import derive_path_id, is_local_id, validate_local_id -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.sandbox_utils.sandbox_utils import find_agent_container @@ -650,7 +650,7 @@ def _resolve_cua_env(self, context): DeployedEnv metadata, so its presence authoritatively marks a CUA env.""" deployed_env = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed_env is None: - raise RuntimeError( + raise StepContextError( f"Env '{self.env_id}' not found in context.deployed_envs — cannot collect artifacts via controller" ) if _CUA_ENV_METADATA_MARKER not in (deployed_env.metadata or {}): @@ -758,7 +758,7 @@ async def _collect_via_agent_container(self, context, store, artifact_id, versio agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents") if not agent.sandbox_id: raise RuntimeError(f"Agent '{self.agent_name}' has no sandbox_id — cannot collect artifacts") @@ -804,7 +804,7 @@ async def _collect_via_vm_host(self, context, store, artifact_id, version): (sb for sb in context.deployed_sandboxes if sb.sandbox_name == self.sandbox_name), None ) if ds is None: - raise RuntimeError( + raise StepContextError( f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes" ) @@ -862,7 +862,7 @@ async def _collect_via_sandbox_container(self, context, store, artifact_id, vers (s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None ) if deployed is None: - raise RuntimeError( + raise StepContextError( f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes" ) diff --git a/src/agent_env/task_step/task_steps/configure_scheduled_checks.py b/src/agent_env/task_step/task_steps/configure_scheduled_checks.py index aa46f4f5c..e46bdc979 100644 --- a/src/agent_env/task_step/task_steps/configure_scheduled_checks.py +++ b/src/agent_env/task_step/task_steps/configure_scheduled_checks.py @@ -18,7 +18,7 @@ from agent_env.env.gateway.constants import EXT_NOTIFICATIONS_URI from agent_env.env.gateway.notifications import DEFAULT_MAX_ACTIVE, DEFAULT_MIN_INTERVAL from agent_env.entity.refs import EntityRef -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency logger = logging.getLogger(__name__) @@ -77,7 +77,7 @@ def from_dict(cls, data: dict) -> ConfigureScheduledChecksTaskStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") if not deployed.supports(EXT_NOTIFICATIONS_URI, "config"): raise RuntimeError(f"Env '{self.env_id}' runs a gateway without {EXT_NOTIFICATIONS_URI}; redeploy it on a " "gateway image that has it.") 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 cbacdc804..c95b198ff 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -40,7 +40,7 @@ sandbox_request_headers_for_url, ) from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -356,13 +356,13 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas existing = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if existing is not None: - raise RuntimeError(f"Agent with name '{self.agent_name}' is already deployed") + raise StepContextError(f"Agent with name '{self.agent_name}' is already deployed") linked_sandbox = None if self.sandbox_name: ds = next((s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None) if ds is None: - raise RuntimeError( + raise StepContextError( f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes; " f"deploy it with DeploySandboxTaskStep before this step" ) 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 1ff00eaa3..96a669f65 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -11,7 +11,7 @@ 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 +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -155,7 +155,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: existing = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if existing is not None: - raise RuntimeError(f"Env '{self.env_id}' is already deployed") + raise StepContextError(f"Env '{self.env_id}' is already deployed") env = Env.get(self.env_id, self.env_version) user_overrides = context.metadata.get("user_overrides", {}) diff --git a/src/agent_env/task_step/task_steps/deploy_human_agent.py b/src/agent_env/task_step/task_steps/deploy_human_agent.py index a262ebc79..b7719d1c9 100644 --- a/src/agent_env/task_step/task_steps/deploy_human_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_human_agent.py @@ -10,7 +10,7 @@ import httpx -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -60,7 +60,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: "only one human agent is allowed per task" ) if any(a.agent_name == self.agent_name for a in context.deployed_agents): - raise RuntimeError(f"Agent with name '{self.agent_name}' is already deployed") + raise StepContextError(f"Agent with name '{self.agent_name}' is already deployed") if self.a2a_url: url = self.a2a_url diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index 6e8bce2f4..bb833ab40 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -8,7 +8,7 @@ from agent_env.artifact.artifacts.docker_image import pull_problem from agent_env.providers.sandbox_providers.chained_sandbox_provider import skip_refusing from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy -from agent_env.task_step.context import DeployedSandbox, TaskStepContext +from agent_env.task_step.context import DeployedSandbox, StepContextError, TaskStepContext from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution @@ -113,7 +113,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_sandbox_provider if any(s.sandbox_name == self.sandbox_name for s in context.deployed_sandboxes): - raise RuntimeError(f"Sandbox '{self.sandbox_name}' is already deployed") + raise StepContextError(f"Sandbox '{self.sandbox_name}' is already deployed") user_overrides = context.metadata.get("user_overrides") or {} # Moves only deploy_sandbox steps: deploy_env and deploy_agent read env_sandbox and agent_sandbox. diff --git a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_card.py b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_card.py index 890048177..63bd1e7c2 100644 --- a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_card.py +++ b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_card.py @@ -5,7 +5,7 @@ import logging from typing import ClassVar, Optional -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -48,7 +48,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") card = None error = None diff --git a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py index 20d25ea43..20a92ff0b 100644 --- a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py +++ b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py @@ -9,7 +9,7 @@ from agentenv_protocol import METHOD_ADD, METHOD_GET, METHOD_RESET, RPC_PATH from agentenv_protocol import client as protocol_v1 -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -54,7 +54,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") card = None try: 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 6fdd517a9..edebd8b2b 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -33,7 +33,7 @@ trust_env, ) from agent_env.store.base import NotFoundError -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -226,7 +226,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: existing = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if existing is not None: - raise RuntimeError(f"Agent with name '{self.agent_name}' is already deployed") + raise StepContextError(f"Agent with name '{self.agent_name}' is already deployed") if self.container_name is not None: containers = context.metadata.get("deployed_docker_containers", []) @@ -246,7 +246,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: ds = next((s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None) if ds is None: - raise RuntimeError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") + raise StepContextError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") if ds.sandbox_mode != SANDBOX_MODE_VM: raise RuntimeError( f"InstallAgent requires a VM-mode sandbox; '{self.sandbox_name}' is mode={ds.sandbox_mode!r}" diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index e2f4a8c55..1762baaf0 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -18,7 +18,7 @@ from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandbox from agent_env.providers.sandbox_providers.sandbox import into_container from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.task_step.preflight import stored_env from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope @@ -322,7 +322,7 @@ def _resolve_step_artifact(context: TaskStepContext, step_id: str) -> Optional[d continue ref = extract(entry) return {"id": ref["id"], "version": ref.get("version")} if ref else None - raise RuntimeError( + raise StepContextError( f"No artifact recorded by step '{step_id}'. Ensure a step that produces " f"one (e.g. `collect_artifacts`, `snapshot_env`) ran with that id before " f"this load step." @@ -447,7 +447,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if self.env_id is not None: deployed = self._deployed_env(context) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") inputs = self._resolve_inputs(context) resolved_artifacts = inputs.artifacts @@ -474,7 +474,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: ) ds = next((s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None) if ds is None: - raise RuntimeError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") + raise StepContextError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") provider = ( build_sandbox_provider(ds.sandbox_type) if ds.sandbox_type else get_sandbox_provider() ) @@ -545,7 +545,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None or not agent.instance_id: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" files = await A2AAgent.load_file_artifact_universe( @@ -582,7 +582,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None or not agent.sandbox_id: - raise RuntimeError( + raise StepContextError( f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing sandbox_id)" ) provider = ( @@ -641,7 +641,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError("env_id is required when loading a CliArtifact (for gateway_url)") agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None or not agent.instance_id: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) gateway_url = require_gateway_url(deployed, "Installing a CliArtifact") # A local sandbox's host is this machine, which reaches the gateway as it was issued. @@ -664,7 +664,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None or not agent.sandbox_id: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing sandbox_id)") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing sandbox_id)") provider = ( build_sandbox_provider(agent.sandbox_type) if agent.sandbox_type diff --git a/src/agent_env/task_step/task_steps/mcp_cli_builder/build_mcp_cli.py b/src/agent_env/task_step/task_steps/mcp_cli_builder/build_mcp_cli.py index 23c0f22d3..afab39ad8 100644 --- a/src/agent_env/task_step/task_steps/mcp_cli_builder/build_mcp_cli.py +++ b/src/agent_env/task_step/task_steps/mcp_cli_builder/build_mcp_cli.py @@ -19,7 +19,7 @@ from agent_env.env import Env from agent_env.env.env import require_gateway_url, website_child from agent_env.store.ids import derive_id -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef, RefRole from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.thread_work import finish_on_thread @@ -160,7 +160,7 @@ def from_dict(cls, data: dict) -> BuildMcpCliTaskStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") gateway_url = require_gateway_url(deployed, "build_mcp_cli") # the CLI it generates talks to the gateway env = Env.get(self.env_id, deployed.env_version) diff --git a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_env_assessment.py b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_env_assessment.py index 6030ceaef..72ba54990 100644 --- a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_env_assessment.py +++ b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_env_assessment.py @@ -6,7 +6,7 @@ import logging from typing import Any, ClassVar, Optional -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -45,7 +45,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: prompt_response = next((pr for pr in context.prompt_responses if pr.prompt_id == self.prompt_id), None) if prompt_response is None: - raise RuntimeError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") + raise StepContextError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") if prompt_response.error_type: logger.warning(f"Skipping assessment: prompt had error_type={prompt_response.error_type}") @@ -55,7 +55,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") env = Env.get(self.env_id, deployed.env_version) env.update_metadata({**env.metadata, "mcp_tool_correctness_validation": result}) diff --git a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_tool_schema.py b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_tool_schema.py index c5e5c42e6..1a7df4f06 100644 --- a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_tool_schema.py +++ b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_tool_schema.py @@ -9,7 +9,7 @@ from mcp import ClientSession from mcp.client.streamable_http import streamable_http_client -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -122,7 +122,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") mcp_url = deployed.mcp_url logger.info(f"Connecting to MCP at {mcp_url} to verify tool schemas") diff --git a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_spec_conformance.py b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_spec_conformance.py index 8a61e84dc..9787c3ae5 100644 --- a/src/agent_env/task_step/task_steps/mcp_env_validator/verify_spec_conformance.py +++ b/src/agent_env/task_step/task_steps/mcp_env_validator/verify_spec_conformance.py @@ -27,7 +27,7 @@ from mcp.client.streamable_http import streamable_http_client from agent_env.env import legacy_protocol -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.mcp_env_validator.verify_mcp_tool_schema import ( @@ -405,7 +405,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") env = Env.get(self.env_id, deployed.env_version) environment_name = env.environment_name diff --git a/src/agent_env/task_step/task_steps/modify_env_tool_access.py b/src/agent_env/task_step/task_steps/modify_env_tool_access.py index 589da671a..4e71d2f21 100644 --- a/src/agent_env/task_step/task_steps/modify_env_tool_access.py +++ b/src/agent_env/task_step/task_steps/modify_env_tool_access.py @@ -7,7 +7,7 @@ from agent_env.env.gateway import TOOL_DISABLE_ACTION, TOOL_ENABLE_ACTION from agent_env.env.gateway.constants import EXT_DISABLE_TOOL_URI, EXT_ENABLE_TOOL_URI -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -62,7 +62,7 @@ def from_dict(cls, data: dict) -> ModifyEnvToolAccessStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") body = await deployed.invoke(_TOOL_ACCESS_URIS[self.action], self.action, {"role": self.role, "tools": self.tools}) logger.info(f"{self.action} role={self.role!r} tools={self.tools} on env={self.env_id}: {body}") context.metadata.setdefault("tool_access_changes", []).append({ diff --git a/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py b/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py index 9e8274beb..585c2a945 100644 --- a/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py +++ b/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py @@ -12,7 +12,7 @@ from agent_env.env.env_artifact_store import EnvArtifactType, get_env_artifact_store from agent_env.artifact.store import artifact_write_lock from agent_env.store.ids import derive_id, fs_safe, is_local_id -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -101,7 +101,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # Phase 1: Resolve deployed env and universe artifact deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") env = Env.get(deployed.env_id, deployed.env_version) env = await type(env).from_deployed_env(deployed) universe = EnvironmentUniverseArtifact.get(self.universe_artifact_id, self.universe_artifact_version) 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 c11fb716b..c0b986b00 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -35,7 +35,7 @@ ) from agent_env.store import DuplicateKeyError, get_config from agent_env.store.object_store import ObjectStore -from agent_env.task_step.context import DeployedAgent, PromptResponse, TaskStepContext +from agent_env.task_step.context import DeployedAgent, PromptResponse, StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import RetryConfig, TaskStep, TaskStepDependency @@ -452,7 +452,7 @@ async def _converse( ) -> TaskStepContext: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None: - raise RuntimeError(f"Agent with name '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent with name '{self.agent_name}' not found in context.deployed_agents") target_url = agent.a2a_url or agent.api_url user_sim = next((a for a in context.deployed_agents if a.agent_name == self.user_agent_name), None) user_url = self.user_a2a_url or ( diff --git a/src/agent_env/task_step/task_steps/register_agent_triggers.py b/src/agent_env/task_step/task_steps/register_agent_triggers.py index 64074b983..05171c26f 100644 --- a/src/agent_env/task_step/task_steps/register_agent_triggers.py +++ b/src/agent_env/task_step/task_steps/register_agent_triggers.py @@ -9,7 +9,7 @@ import httpx from agent_env.a2a_agent.a2a_agent import A2AAgent -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -76,7 +76,7 @@ def from_dict(cls, data: dict) -> RegisterAgentTriggersStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents") card = agent.a2a_card or {} ext = A2AAgent.find_extension(card, A2AAgent.EXT_TRIGGERS) if ext is None: diff --git a/src/agent_env/task_step/task_steps/register_env_triggers.py b/src/agent_env/task_step/task_steps/register_env_triggers.py index 237f77259..6229a59e4 100644 --- a/src/agent_env/task_step/task_steps/register_env_triggers.py +++ b/src/agent_env/task_step/task_steps/register_env_triggers.py @@ -23,7 +23,7 @@ TRIGGER_WEBSITE, TRIGGER_WEBSITE_OP, ) -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -131,7 +131,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError("; ".join(problems)) deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") body: dict = {"triggers": self.triggers} if self.watch_roles is not None: body["watch_roles"] = self.watch_roles @@ -227,7 +227,7 @@ def _await_failure(self, seen: dict[str, Optional[str]], state: dict, stuck: lis def _executor_config(self, context: TaskStepContext, deployed: DeployedEnv, agent_name: str) -> dict: agent = next((a for a in context.deployed_agents if a.agent_name == agent_name), None) if agent is None: - raise RuntimeError(f"Executor agent '{agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Executor agent '{agent_name}' not found in context.deployed_agents") if not agent.role: raise RuntimeError( f"Executor agent '{agent_name}' was deployed without an explicit role; " diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 6695a0e18..2652f4f01 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -42,7 +42,7 @@ from agent_env.env.env import DeployedSandboxEnv, require_sandbox from agent_env.env.envs.computer.utils import inside from agent_env.providers.sandbox_providers.sandbox import into_container -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -406,7 +406,7 @@ def _deployed_env(self, context: TaskStepContext) -> DeployedSandboxEnv: if deployed is not None: deployed = require_sandbox(deployed, "run_code on the env's host") if deployed is None or not deployed.sandbox_id: - raise RuntimeError( + raise StepContextError( f"Env '{self.env_id}' not found in context.deployed_envs " "(or missing sandbox_id)" ) @@ -457,7 +457,7 @@ async def _sandbox_target(self, context: TaskStepContext) -> _ExecTarget: (s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None ) if deployed is None: - raise RuntimeError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") + raise StepContextError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") provider = ( build_sandbox_provider(deployed.sandbox_type) if deployed.sandbox_type @@ -487,7 +487,7 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: (a for a in context.deployed_agents if a.agent_name == self.agent_name), None ) if agent is None or not agent.sandbox_id: - raise RuntimeError( + raise StepContextError( f"Agent '{self.agent_name}' not found in context.deployed_agents " "(or missing sandbox_id)" ) diff --git a/src/agent_env/task_step/task_steps/run_docker_container.py b/src/agent_env/task_step/task_steps/run_docker_container.py index 15280d4db..6d8269170 100644 --- a/src/agent_env/task_step/task_steps/run_docker_container.py +++ b/src/agent_env/task_step/task_steps/run_docker_container.py @@ -11,7 +11,7 @@ from urllib.parse import urlparse from agent_env.artifact.artifacts.docker_image import DockerImageArtifact -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, port_bindings @@ -177,7 +177,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: None, ) if ds is None: - raise RuntimeError( + raise StepContextError( f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes; " f"deploy it with DeploySandboxTaskStep before this step" ) diff --git a/src/agent_env/task_step/task_steps/sync_env_clock.py b/src/agent_env/task_step/task_steps/sync_env_clock.py index 040c958c2..b4c92793e 100644 --- a/src/agent_env/task_step/task_steps/sync_env_clock.py +++ b/src/agent_env/task_step/task_steps/sync_env_clock.py @@ -19,7 +19,7 @@ from agentenv_protocol import client as protocol_v1 from agent_env.env.gateway.constants import EXT_CLOCK_URI -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -124,7 +124,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") try: await deployed.invoke(EXT_CLOCK_URI, "set_time", {"virtual_time": self.armed_virtual_time, "virtual_seconds_per_real_second": self.virtual_seconds_per_real_second}, diff --git a/src/agent_env/task_step/task_steps/verifiers/agent_prompt_response_verifier.py b/src/agent_env/task_step/task_steps/verifiers/agent_prompt_response_verifier.py index a4fb779f2..7d868f330 100644 --- a/src/agent_env/task_step/task_steps/verifiers/agent_prompt_response_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/agent_prompt_response_verifier.py @@ -21,7 +21,7 @@ import uuid from typing import ClassVar, Optional -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.verifiers.scoring import ScoreAggregator, aggregate_score @@ -81,7 +81,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: None, ) if prompt_response is None: - raise RuntimeError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") + raise StepContextError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") if prompt_response.error_type: logger.warning( diff --git a/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py b/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py index 51f0c7213..5e45774af 100644 --- a/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py @@ -13,7 +13,7 @@ from typing import ClassVar, Optional from agent_env.config import get_config -from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.context import PromptResponse, StepContextError, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.verifiers.run_context import load_run_context @@ -108,7 +108,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") artifact = FileArtifact.get(self.file_artifact_id, self.file_artifact_version) verifier_bytes = await asyncio.to_thread(artifact.load) @@ -151,7 +151,7 @@ def _graded_prompt(self, context: TaskStepContext) -> Optional[PromptResponse]: return context.prompt_responses[-1] if context.prompt_responses else None prompt_response = next((pr for pr in context.prompt_responses if pr.prompt_id == self.prompt_id), None) if prompt_response is None: - raise RuntimeError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") + raise StepContextError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") return prompt_response 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 443e86103..915fca57e 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 @@ -20,7 +20,7 @@ from agent_env.config import get_config from agent_env.config.model import ModelParam from agent_env.providers.sandbox_providers.sandbox import into_container -from agent_env.task_step.context import PromptResponse, TaskStepContext, dual_keyed +from agent_env.task_step.context import PromptResponse, StepContextError, TaskStepContext, dual_keyed from agent_env.task_step.task_steps.verifiers.judge_utils.judge_output_format import ( SCREENSHOT_GROUNDING, JudgeOutputFormat, @@ -636,7 +636,7 @@ def _resolved_criteria(self, context: TaskStepContext) -> list[dict]: async def execute(self, context: TaskStepContext) -> TaskStepContext: prompt_response = next((pr for pr in context.prompt_responses if pr.prompt_id == self.prompt_id), None) if prompt_response is None: - raise RuntimeError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") + raise StepContextError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context") overrides_for_step = self.step_param_overrides(context) grade_up_to_turn = ( _validated_turn(overrides_for_step["grade_up_to_turn"]) @@ -744,7 +744,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if self.agent_name is not None: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents") if not agent.sandbox_id: raise RuntimeError(f"Agent '{self.agent_name}' has no sandbox_id") judge_a2a_url = agent.a2a_url or agent.api_url 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 61c3c4838..ab8cb2e96 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 @@ -50,7 +50,7 @@ registered_sandbox_provider_classes, ) from agent_env.store.ids import derive_id, is_local_id, key_segment, validate_local_id -from agent_env.task_step.context import TaskStepContext, dual_keyed +from agent_env.task_step.context import StepContextError, TaskStepContext, dual_keyed from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.run_code import _DriverTarget @@ -395,7 +395,7 @@ async def _container_sandbox(self, context: TaskStepContext): ) ds = next((s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None) if ds is None: - raise RuntimeError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") + raise StepContextError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") provider = ( build_sandbox_provider(ds.sandbox_type) if ds.sandbox_type else get_sandbox_provider() @@ -439,7 +439,7 @@ def _computer(self, context: TaskStepContext) -> _DriverTarget: timeout.""" deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") deployed = require_sandbox(deployed, "Running a verifier in an env") if missing := inside.missing(deployed): raise RuntimeError(f"env '{self.env_id}' doesn't serve {', '.join(missing)}, which running a verifier in " diff --git a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py index 9144fbc58..5d5776b34 100644 --- a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py +++ b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py @@ -35,7 +35,7 @@ get_agent_sandbox_provider, get_sandbox_provider, ) -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.task_steps.sandbox_utils.sandbox_utils import find_agent_container from agent_env.task_step.task_steps.verifiers.scoring import ScoreAggregator, aggregate_score @@ -194,7 +194,7 @@ async def _agent_sandbox(self, context: TaskStepContext): None, ) if agent is None: - raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents") + raise StepContextError(f"Agent '{self.agent_name}' not found in context.deployed_agents") if not agent.sandbox_id: raise RuntimeError(f"Agent '{self.agent_name}' has no sandbox_id") provider = ( @@ -210,7 +210,7 @@ async def _deployed_sandbox(self, context: TaskStepContext): None, ) if deployed is None: - raise RuntimeError( + raise StepContextError( f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes " f"(deployed: {[s.sandbox_name for s in context.deployed_sandboxes]})" ) diff --git a/tst/unit/task_step/deploy_human_agent_test.py b/tst/unit/task_step/deploy_human_agent_test.py index 8d604af2b..98ec3c308 100644 --- a/tst/unit/task_step/deploy_human_agent_test.py +++ b/tst/unit/task_step/deploy_human_agent_test.py @@ -7,7 +7,7 @@ import pytest from agent_env.config.errors import ConfigError -from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.context import DeployedAgent, StepContextError, TaskStepContext from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_steps.deploy_human_agent import DeployHumanAgentTaskStep @@ -94,7 +94,7 @@ async def test_execute_explicit_a2a_url_is_verbatim(): async def test_execute_rejects_duplicate_agent_name(): context = TaskStepContext(deployed_agents=[DeployedAgent(agent_name="human", api_url="http://x")]) step = DeployHumanAgentTaskStep(id="t.dh", version=None, agent_name="human") - with pytest.raises(RuntimeError, match="already deployed"): + with pytest.raises(StepContextError, match="already deployed"): await step.execute(context) diff --git a/tst/unit/task_step/test_load_artifact_overrides.py b/tst/unit/task_step/test_load_artifact_overrides.py index bc8bd98c8..9d0ad2f0d 100644 --- a/tst/unit/task_step/test_load_artifact_overrides.py +++ b/tst/unit/task_step/test_load_artifact_overrides.py @@ -10,7 +10,7 @@ import pytest -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep @@ -119,7 +119,7 @@ def test_collected_missing_entry_still_raises(self): step = LoadArtifactTaskStep( id="load", version=None, agent_name="agent", collected_artifacts_step_id="collect", ) - with pytest.raises(RuntimeError, match="No artifact recorded by step 'collect'"): + with pytest.raises(StepContextError, match="No artifact recorded by step 'collect'"): step._resolve_inputs(_ctx()) diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index a972dd17a..b21f10b7f 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -25,7 +25,7 @@ from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, DeployedSandboxEnv, EnvNeedsSandbox from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandbox from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox -from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext +from agent_env.task_step.context import DeployedAgent, DeployedSandbox, StepContextError, TaskStepContext from agent_env.task_step.task_steps import run_code as _run_code_module from tst.unit.computer_home import CARD, TOKEN, Computer from tst.unit.event_loop_probe import on_event_loop @@ -492,7 +492,7 @@ async def test_sandbox_target_honors_the_sandboxs_type(): @pytest.mark.asyncio async def test_sandbox_target_names_a_sandbox_the_run_never_deployed(): - with pytest.raises(RuntimeError, match="^Sandbox 'ghost' not found in context.deployed_sandboxes$"): + with pytest.raises(StepContextError, match="^Sandbox 'ghost' not found in context.deployed_sandboxes$"): await _step(sandbox_name="ghost")._sandbox_target(_sandbox_context()) diff --git a/tst/unit/task_step/test_step_context_error.py b/tst/unit/task_step/test_step_context_error.py new file mode 100644 index 000000000..0767cc44e --- /dev/null +++ b/tst/unit/task_step/test_step_context_error.py @@ -0,0 +1,62 @@ +"""A step that finds the run's context can't satisfy it raises StepContextError, which a runner does not +retry: running the step again on the same context fails the same way.""" + +import ast +import re +from pathlib import Path + +import pytest + +import agent_env.task_step.task_steps as task_steps + +_CONTEXT_FAILURE = re.compile( + r"not found in context\.deployed_(envs|agents|sandboxes)|' is already deployed" + r"|^PromptResponse with prompt_id=.* not found in context|^No artifact recorded by step" +) + + +def _message_text(node: ast.expr) -> str: + """The literal text of a message, with each formatted value as ``{}``.""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + if isinstance(node, ast.JoinedStr): + return "".join(_message_text(part) if isinstance(part, ast.Constant) else "{}" for part in node.values) + return "" + + +def _plain_context_raises(source: str) -> list[int]: + """The lines of ``source`` raising a context failure as a plain RuntimeError.""" + return [ + node.lineno + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.Raise) + and isinstance(node.exc, ast.Call) + and isinstance(node.exc.func, ast.Name) + and node.exc.func.id == "RuntimeError" + and node.exc.args + and _CONTEXT_FAILURE.search(_message_text(node.exc.args[0])) + ] + + +@pytest.mark.parametrize( + "source", + [ + """raise RuntimeError(f"Env '{env_id}' not found in context.deployed_envs")""", + """raise RuntimeError(f'Agent {name!r} not found in context.deployed_agents')""", + """raise RuntimeError(\n f"Sandbox '{name}' not found "\n "in context.deployed_sandboxes"\n)""", + """raise RuntimeError("No artifact recorded by step 'collect'")""", + ], + ids=["double-quoted", "single-quoted", "split", "plain-string"], +) +def test_the_guard_sees_a_context_failure_however_it_is_written(source): + assert _plain_context_raises(source) == [1] + + +def test_no_step_raises_a_context_failure_as_a_plain_runtime_error(): + root = Path(task_steps.__file__).parent + found = [ + f"{path.relative_to(root)}:{line}" + for path in sorted(root.rglob("*.py")) + for line in _plain_context_raises(path.read_text()) + ] + assert found == [], "raise StepContextError (agent_env.task_step.context) for these instead"