From f6b321743f9cb66deb491d0bba711ee7bf0d9c94 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 12:32:14 -0700 Subject: [PATCH 1/4] feat(deploy_agent): deliver the role to the agent or fail the step A role the agent card could not take was filtered out silently and still recorded on DeployedAgent, so an agent believed to be an executor or a person called the env as the default role. The step now puts AgentEnv-Role into the MCP registration headers for the gateway env when the card's mcp-config add accepts headers, negotiates it through agent-config when the card lists it, sets AGENT_ENV_ROLE for the generated MCP CLI, and raises before any request when neither path can carry a requested role. The recorded role is therefore always one the agent received. Co-Authored-By: Claude Fable 5.1 --- .../task_step/task_steps/deploy_agent.py | 86 +++++-- .../task_step/test_deploy_agent_mcp_alias.py | 5 + .../test_deploy_agent_role_delivery.py | 238 ++++++++++++++++++ 3 files changed, 308 insertions(+), 21 deletions(-) create mode 100644 tst/unit/task_step/test_deploy_agent_role_delivery.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 6d635f61..affbfa28 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -10,6 +10,7 @@ NamespaceChangelogEnableResponse, ObjectChangelogApplyResponse, ObjectSnapshotLoadResponse, + card_request_accepts, ) from agent_env.a2a_agent import protocol @@ -24,25 +25,32 @@ ) from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config -from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env +from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, DeployedSandboxEnv, Env +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.providers.sandbox_providers.local_sandbox import host_url_for from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url, 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.entity_refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency +from agent_env.task_step.task_steps.mcp_cli_builder.codegen import ROLE_ENV_VAR from agent_env.attribution import deploy_attribution logger = logging.getLogger(__name__) +def _mcp_add_request(mcp_ext: dict) -> dict: + """The request contract the card declares for mcp-config `add`.""" + request = mcp_ext.get("params", {}).get("methods", {}).get("add", {}).get("request", {}) + return request if isinstance(request, dict) else {} + + def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str | None) -> dict: - """Body for the agent's mcp-config `add`; the card name is relayed verbatim when the agent's card lists `name` as optional.""" + """Body for the agent's mcp-config `add`; the card name is relayed verbatim when the agent's card takes `name`.""" body: dict = {"url": url} if headers: body["headers"] = headers - add_request = mcp_ext.get("params", {}).get("methods", {}).get("add", {}).get("request", {}) - if card_name and "name" in (add_request.get("optional") or []): + if card_name and card_request_accepts(_mcp_add_request(mcp_ext), ("url", "name")): body["name"] = card_name return body @@ -150,7 +158,7 @@ def __init__( self.agent_snapshot_target_context_id = agent_snapshot_target_context_id self.metadata = metadata or {} self.network_policy = NetworkPolicy.from_dict(network_policy).to_dict() if network_policy else None - self.role = role + self.role = role or None self.enable_agent_changelog = enable_agent_changelog self.agent_changelog_object_url = agent_changelog_object_url self.agent_changelog_toolcall_position_exclusive = agent_changelog_toolcall_position_exclusive @@ -300,6 +308,10 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas 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 + if self.role is not None: + # The generated MCP CLI reads its role from this variable; without it, CLI traffic from a + # role-bearing agent reaches the gateway as the CLI's default role. + env_vars.setdefault(ROLE_ENV_VAR, self.role) resolved_sandbox_type = user_overrides.get("agent_sandbox") or self.sandbox_type deploy_kwargs = {"env_vars": env_vars if env_vars else None} @@ -328,6 +340,13 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas card = deployed.agent_card logger.info(f"A2A agent '{a2a_agent_id}' deployed at {a2a_url}") + mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) + role_header = ( + {AGENT_ENV_ROLE_HEADER: self.role} + if self.role is not None and mcp_ext and card_request_accepts(_mcp_add_request(mcp_ext), ("url", "headers")) + else None + ) + role_in_headers = False env_mcp_urls: list[tuple[str, str, dict | None, str | None]] = [] # (env_id, url, headers, card_name) for env_id in self.env_ids: deployed_env = env_deployments[env_id] @@ -339,9 +358,44 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas else: url = host_url_for(_live_mcp_url(env_id), deployed.sandbox_type) card_name = None - env_mcp_urls.append((env_id, url, sandbox_request_headers_for_url(url) or None, card_name)) + headers = sandbox_request_headers_for_url(url) + if role_header and isinstance(deployed_env, DeployedGatewayEnv): + headers = {**headers, **role_header} + role_in_headers = True + env_mcp_urls.append((env_id, url, headers or None, card_name)) + + identity_payload: dict[str, str] = {"name": self.agent_name} + if self.agent_description is not None: + identity_payload["description"] = self.agent_description + if self.role is not None: + identity_payload["role"] = self.role + identity = A2AAgent.negotiate_agent_config(card, identity_payload) + role_in_config = identity is not None and "role" in identity.fields + if self.role is not None: + # Decided before the first request to the agent, so a role that cannot reach the env is neither + # registered nor recorded; execute() closes the deployed agent. + if not (role_in_headers or role_in_config): + raise RuntimeError( + f"Agent '{self.agent_name}' cannot carry role '{self.role}': its card takes neither 'headers' on " + f"{A2AAgent.EXT_MCP_CONFIG} add for a gateway env nor 'role' in {A2AAgent.EXT_AGENT_CONFIG}, " + f"so the role would never reach the env" + ) + if not role_in_headers: + logger.warning( + f"Agent '{self.agent_name}' takes role '{self.role}' through {A2AAgent.EXT_AGENT_CONFIG} only; " + f"its own MCP client must send {AGENT_ENV_ROLE_HEADER}" + ) + elif not role_in_config: + logger.warning( + f"Agent '{self.agent_name}' takes role '{self.role}' as {AGENT_ENV_ROLE_HEADER} on its gateway MCP " + f"registration only; {A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" + ) + else: + logger.info( + f"Role '{self.role}' reaches '{self.agent_name}' as {AGENT_ENV_ROLE_HEADER} on its gateway MCP " + f"registration and through {A2AAgent.EXT_AGENT_CONFIG}" + ) - mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) if mcp_ext and env_mcp_urls: endpoint = a2a_url + mcp_ext["params"]["endpoint"] async with httpx.AsyncClient() as client: @@ -378,20 +432,10 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas await protocol.post_agent_config(endpoint, {"system_prompt": system_prompt}) logger.info(f"Set system_prompt on agent '{self.agent_name}'") - identity_payload: dict[str, str] = {"name": self.agent_name} - if self.agent_description is not None: - identity_payload["description"] = self.agent_description - if self.role is not None: - identity_payload["role"] = self.role - config_ext = A2AAgent.find_extension(card, A2AAgent.EXT_AGENT_CONFIG) - if config_ext is not None: - supported = (config_ext.get("params") or {}).get("methods", {}).get("set", {}).get("request", {}).get("supported", []) - identity_payload = {k: v for k, v in identity_payload.items() if k in supported} - if identity_payload: - endpoint = a2a_url + (config_ext.get("params") or {}).get("endpoint", "/ext/agent-config") - await protocol.post_agent_config(endpoint, identity_payload) - card.update(identity_payload) - logger.info(f"Stamped card identity on '{self.agent_name}': {sorted(identity_payload)}") + if identity is not None: + await protocol.post_agent_config(a2a_url + identity.endpoint, identity.fields) + card.update(identity.fields) + logger.info(f"Stamped card identity on '{self.agent_name}': {sorted(identity.fields)}") default_model = config.get_model_for_role("agent") if default_model and not context.default_agent_model: diff --git a/tst/unit/task_step/test_deploy_agent_mcp_alias.py b/tst/unit/task_step/test_deploy_agent_mcp_alias.py index 51cd4880..89fa22ba 100644 --- a/tst/unit/task_step/test_deploy_agent_mcp_alias.py +++ b/tst/unit/task_step/test_deploy_agent_mcp_alias.py @@ -28,6 +28,11 @@ def test_the_card_name_is_relayed_verbatim_whatever_it_is(): assert _mcp_add_body(_NAMING_CARD_EXT, "https://gw/mcp", None, name) == {"url": "https://gw/mcp", "name": name} +def test_a_hand_written_supported_list_also_takes_the_name(): + ext = {"params": {"methods": {"add": {"request": {"supported": ["url", "headers", "name"]}}}}} + assert _mcp_add_body(ext, "https://gw/mcp", None, "crm") == {"url": "https://gw/mcp", "name": "crm"} + + def test_no_card_name_means_the_agent_mints_its_own(): assert _mcp_add_body(_NAMING_CARD_EXT, "https://gw/mcp", None, None) == {"url": "https://gw/mcp"} assert "name" not in _mcp_add_body({"params": {"methods": {"add": {"request": {"optional": None}}}}}, "https://gw/mcp", None, "env") diff --git a/tst/unit/task_step/test_deploy_agent_role_delivery.py b/tst/unit/task_step/test_deploy_agent_role_delivery.py new file mode 100644 index 00000000..5d0a8dab --- /dev/null +++ b/tst/unit/task_step/test_deploy_agent_role_delivery.py @@ -0,0 +1,238 @@ +"""deploy_agent delivers the step's role to the agent, or fails the step before any request reaches it; it never +records a role the agent did not get.""" + +import json +import logging +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.env.env import DeployedGatewayEnv +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps import deploy_agent +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.mcp_cli_builder.codegen import ROLE_ENV_VAR + +_LOGGER = "agent_env.task_step.task_steps.deploy_agent" +_GATEWAY_MCP_URL = "https://gw.example/mcp" +_LIVE_MCP_URL = "https://ext.example/mcp" + + +def _mcp_config(add_request: dict) -> dict: + return { + "uri": A2AAgent.EXT_MCP_CONFIG, + "params": {"endpoint": "/ext/mcp-config", "methods": {"add": {"method": "POST", "request": add_request}}}, + } + + +def _agent_config(supported: list[str]) -> dict: + return { + "uri": A2AAgent.EXT_AGENT_CONFIG, + "params": {"endpoint": "/ext/agent-config", "methods": {"set": {"method": "POST", "request": {"supported": supported}}}}, + } + + +def _card(*extensions: dict) -> dict: + return {"capabilities": {"extensions": list(extensions)}} + + +HEADERS_MCP = _mcp_config({"required": ["url"], "optional": ["headers", "name"]}) # as the protocol renders it +LEGACY_MCP = _mcp_config({"required": ["url"]}) +ROLE_CONFIG = _agent_config(["name", "description", "role"]) +ROLELESS_CONFIG = _agent_config(["name", "description"]) + + +@pytest.mark.asyncio +async def test_a_card_with_both_paths_gets_the_role_on_both(wiring, caplog): + with caplog.at_level(logging.WARNING, logger=_LOGGER): + context = await wiring.deploy(_card(HEADERS_MCP, ROLE_CONFIG), role="executor") + + assert wiring.bodies("/ext/mcp-config") == [ + {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "executor"}, "name": "crm"} + ] + assert wiring.bodies("/ext/agent-config") == [{"name": "solver", "role": "executor"}] + assert not [m for m in caplog.messages if "role" in m] + assert [a.role for a in context.deployed_agents] == ["executor"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "card", [_card(HEADERS_MCP), _card(HEADERS_MCP, ROLELESS_CONFIG)], ids=["no-agent-config", "roleless-agent-config"] +) +async def test_a_header_only_card_gets_the_role_on_the_gateway_registration_and_warns(wiring, caplog, card): + with caplog.at_level(logging.WARNING, logger=_LOGGER): + context = await wiring.deploy(card, role="executor") + + assert wiring.bodies("/ext/mcp-config") == [ + {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "executor"}, "name": "crm"} + ] + assert all("role" not in body for body in wiring.bodies("/ext/agent-config")) + assert any(f"{A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" in m for m in caplog.messages) + assert [a.role for a in context.deployed_agents] == ["executor"] + + +@pytest.mark.asyncio +async def test_a_config_only_card_gets_the_role_through_agent_config_and_warns(wiring, caplog): + with caplog.at_level(logging.WARNING, logger=_LOGGER): + context = await wiring.deploy(_card(LEGACY_MCP, ROLE_CONFIG), role="executor") + + assert wiring.bodies("/ext/mcp-config") == [{"url": _GATEWAY_MCP_URL}] + assert wiring.bodies("/ext/agent-config") == [{"name": "solver", "role": "executor"}] + assert any(f"through {A2AAgent.EXT_AGENT_CONFIG} only" in m for m in caplog.messages) + assert [a.role for a in context.deployed_agents] == ["executor"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "card", + [_card(LEGACY_MCP, ROLELESS_CONFIG), _card(LEGACY_MCP), {}], + ids=["legacy-mcp-and-roleless-config", "legacy-mcp-alone", "empty-card"], +) +async def test_a_card_with_neither_path_fails_before_any_request(wiring, card): + with pytest.raises(RuntimeError, match="cannot carry role 'grader'"): + await wiring.deploy(card, role="grader") + + assert wiring.requests == [] + assert wiring.agent.close.await_count == 1 + assert wiring.context.deployed_agents == [] + + +@pytest.mark.asyncio +async def test_a_header_capable_card_with_no_gateway_env_needs_agent_config(wiring): + with pytest.raises(RuntimeError, match="cannot carry role 'grader'"): + await wiring.deploy(_card(HEADERS_MCP), role="grader", env_ids=["ext"]) + assert wiring.requests == [] + + +@pytest.mark.asyncio +async def test_the_header_lands_on_the_gateway_entry_only(wiring): + context = await wiring.deploy(_card(HEADERS_MCP), role="alice@example.com", env_ids=["crm", "ext"]) + + assert wiring.bodies("/ext/mcp-config") == [ + {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "alice@example.com"}, "name": "crm"}, + {"url": _LIVE_MCP_URL}, + ] + assert wiring.bodies("/ext/agent-config") == [] + assert [a.role for a in context.deployed_agents] == ["alice@example.com"] + + +@pytest.mark.asyncio +async def test_the_role_header_joins_the_sandbox_headers_without_touching_them(wiring, monkeypatch): + sandbox_headers = {"Modal-Key": "k"} + monkeypatch.setattr(deploy_agent, "sandbox_request_headers_for_url", lambda url: sandbox_headers) + + await wiring.deploy(_card(HEADERS_MCP), role="executor") + + assert wiring.bodies("/ext/mcp-config")[0]["headers"] == {"Modal-Key": "k", AGENT_ENV_ROLE_HEADER: "executor"} + assert sandbox_headers == {"Modal-Key": "k"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "add_request, takes_headers", + [ + ({"required": ["url"]}, False), + ({"required": ["url"], "optional": None}, False), + ({"required": ["url"], "optional": ["headers", "name"]}, True), + ({"supported": ["url", "headers"]}, True), + ], + ids=["legacy", "optional-none", "protocol-rendered", "hand-written-supported"], +) +async def test_the_card_decides_whether_the_registration_carries_the_role(wiring, add_request, takes_headers): + await wiring.deploy(_card(_mcp_config(add_request), ROLE_CONFIG), role="executor") + + body = wiring.bodies("/ext/mcp-config")[0] + assert body.get("headers", {}).get(AGENT_ENV_ROLE_HEADER) == ("executor" if takes_headers else None) + + +@pytest.mark.asyncio +async def test_no_role_leaves_every_request_as_before(wiring, caplog): + with caplog.at_level(logging.WARNING, logger=_LOGGER): + context = await wiring.deploy(_card(HEADERS_MCP, ROLE_CONFIG), role=None) + + assert wiring.bodies("/ext/mcp-config") == [{"url": _GATEWAY_MCP_URL, "name": "crm"}] + assert wiring.bodies("/ext/agent-config") == [{"name": "solver"}] + assert not [m for m in caplog.messages if "role" in m] + assert wiring.env_vars is None + assert [a.role for a in context.deployed_agents] == [None] + + +@pytest.mark.asyncio +async def test_an_empty_role_is_no_role(wiring): + assert DeployAgentTaskStep(id="ag", version=None, role="").role is None + assert DeployAgentTaskStep.from_dict({"id": "ag", "type": "deploy_agent", "version": 1, "role": ""}).role is None + + context = await wiring.deploy(_card(LEGACY_MCP), role="") + + assert wiring.bodies("/ext/mcp-config") == [{"url": _GATEWAY_MCP_URL}] + assert wiring.env_vars is None + assert [a.role for a in context.deployed_agents] == [None] + + +@pytest.mark.asyncio +async def test_the_role_is_also_the_generated_clis_role(wiring): + await wiring.deploy(_card(HEADERS_MCP), role="executor") + assert wiring.env_vars == {ROLE_ENV_VAR: "executor"} + + +@pytest.mark.asyncio +async def test_an_explicit_cli_role_env_var_is_kept(wiring): + await wiring.deploy(_card(HEADERS_MCP), role="executor", env_vars={ROLE_ENV_VAR: "cli"}) + assert wiring.env_vars == {ROLE_ENV_VAR: "cli"} + + +class _Wiring: + """One deploy_agent run: the requests the agent received, the agent double, and the step's context.""" + + def __init__(self, monkeypatch): + self.requests: list[httpx.Request] = [] + self.agent = MagicMock(metadata={}) + self.agent.close = AsyncMock() + self.context = TaskStepContext(deployed_envs=[_gateway_env()]) + real_client = httpx.AsyncClient + monkeypatch.setattr( + httpx, "AsyncClient", lambda **kw: real_client(transport=httpx.MockTransport(self._handle), **kw) + ) + + def _handle(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + return httpx.Response(200, json={"ok": True}) + + async def deploy(self, card: dict, *, role: str | None, env_ids=("crm",), **step_kwargs) -> TaskStepContext: + self.agent.deploy = AsyncMock( + return_value=MagicMock(a2a_url="https://agent.example", agent_card=card, sandbox_type=None) + ) + step = DeployAgentTaskStep( + id="ag", version=None, env_ids=list(env_ids), agent_name="solver", role=role, **step_kwargs + ) + with ( + patch.object(A2AAgent, "get", return_value=self.agent), + patch.object(deploy_agent.Env, "get", return_value=SimpleNamespace(mcp_url=_LIVE_MCP_URL)), + patch("agent_env.task_step.task_steps.deploy_agent.get_config") as get_config, + ): + get_config.return_value.get_model_for_role.return_value = None + return await step.execute(self.context) + + def bodies(self, path: str) -> list[dict]: + return [json.loads(r.content) for r in self.requests if r.url.path == path] + + @property + def env_vars(self) -> dict | None: + return self.agent.deploy.await_args.kwargs["env_vars"] + + +@pytest.fixture +def wiring(monkeypatch) -> _Wiring: + return _Wiring(monkeypatch) + + +def _gateway_env() -> DeployedGatewayEnv: + return DeployedGatewayEnv( + env_id="crm", env_version=1, gateway_url="https://gw.example", mcp_url=_GATEWAY_MCP_URL, db_web_url=None, + sandbox_id="sb-1", environment_card_url="https://gw.example/.well-known/agent-env.json", + environment_card={"name": "crm"}, + ) From 37fc96e4d727c5bee74aecbbc1435b02662418b7 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 13:11:41 -0700 Subject: [PATCH 2/4] fix(deploy_agent): judge the role header on the registration the agent actually receives A card that requires `name` on mcp-config add and lists `headers` as optional was refused the role header, because the capability check asked about `url` and `headers` alone while the registration also carried the gateway card's name. The check now covers the fields the body will contain. Co-Authored-By: Claude Fable 5.1 --- .../task_step/task_steps/deploy_agent.py | 17 +++++++++-------- .../test_deploy_agent_role_delivery.py | 4 +++- 2 files changed, 12 insertions(+), 9 deletions(-) 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 affbfa28..a4f59d55 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -341,11 +341,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas logger.info(f"A2A agent '{a2a_agent_id}' deployed at {a2a_url}") mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) - role_header = ( - {AGENT_ENV_ROLE_HEADER: self.role} - if self.role is not None and mcp_ext and card_request_accepts(_mcp_add_request(mcp_ext), ("url", "headers")) - else None - ) + add_request = _mcp_add_request(mcp_ext) if mcp_ext else {} role_in_headers = False env_mcp_urls: list[tuple[str, str, dict | None, str | None]] = [] # (env_id, url, headers, card_name) for env_id in self.env_ids: @@ -359,9 +355,14 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas url = host_url_for(_live_mcp_url(env_id), deployed.sandbox_type) card_name = None headers = sandbox_request_headers_for_url(url) - if role_header and isinstance(deployed_env, DeployedGatewayEnv): - headers = {**headers, **role_header} - role_in_headers = True + if self.role is not None and mcp_ext and isinstance(deployed_env, DeployedGatewayEnv): + # Judge the registration the agent will actually receive: with the card name relayed, a + # card that requires `name` still takes the role header. + carries_name = bool(card_name) and card_request_accepts(add_request, ("url", "name")) + fields = ("url", "headers", "name") if carries_name else ("url", "headers") + if card_request_accepts(add_request, fields): + headers = {**headers, AGENT_ENV_ROLE_HEADER: self.role} + role_in_headers = True env_mcp_urls.append((env_id, url, headers or None, card_name)) identity_payload: dict[str, str] = {"name": self.agent_name} diff --git a/tst/unit/task_step/test_deploy_agent_role_delivery.py b/tst/unit/task_step/test_deploy_agent_role_delivery.py index 5d0a8dab..2dcc397f 100644 --- a/tst/unit/task_step/test_deploy_agent_role_delivery.py +++ b/tst/unit/task_step/test_deploy_agent_role_delivery.py @@ -139,8 +139,10 @@ async def test_the_role_header_joins_the_sandbox_headers_without_touching_them(w ({"required": ["url"], "optional": None}, False), ({"required": ["url"], "optional": ["headers", "name"]}, True), ({"supported": ["url", "headers"]}, True), + ({"required": ["url", "name"], "optional": ["headers"]}, True), + ({"required": ["url", "name"]}, False), ], - ids=["legacy", "optional-none", "protocol-rendered", "hand-written-supported"], + ids=["legacy", "optional-none", "protocol-rendered", "hand-written-supported", "name-required", "name-required-no-headers"], ) async def test_the_card_decides_whether_the_registration_carries_the_role(wiring, add_request, takes_headers): await wiring.deploy(_card(_mcp_config(add_request), ROLE_CONFIG), role="executor") From a7bb965837e882d4e515fa2f22aa31cc6377d7b5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 08:48:55 -0700 Subject: [PATCH 3/4] deploy_agent: put the role header on every MCP registration, not only gateway envs What an env does with the role is the env's business: a gateway forwards it, a protocol server reads it from the header itself, anything else ignores a header it does not know. Judging deliverability by env type refused a header-capable agent on an env that would have honoured the header. Co-Authored-By: Claude Fable 5.1 --- .../task_step/task_steps/deploy_agent.py | 19 ++++++++++--------- .../test_deploy_agent_role_delivery.py | 17 ++++++++++------- 2 files changed, 20 insertions(+), 16 deletions(-) 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 a4f59d55..b0b3e371 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -25,7 +25,7 @@ ) from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config -from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, DeployedSandboxEnv, Env +from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.providers.sandbox_providers.local_sandbox import host_url_for from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url, sandbox_request_headers_for_url @@ -355,9 +355,10 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas url = host_url_for(_live_mcp_url(env_id), deployed.sandbox_type) card_name = None headers = sandbox_request_headers_for_url(url) - if self.role is not None and mcp_ext and isinstance(deployed_env, DeployedGatewayEnv): - # Judge the registration the agent will actually receive: with the card name relayed, a - # card that requires `name` still takes the role header. + if self.role is not None and mcp_ext: + # Every MCP registration carries the role; the env decides what it means (a gateway forwards it, a + # protocol server reads it, anything else ignores an unknown header). Judge the registration the + # agent will actually receive: with the card name relayed, a card that requires `name` still takes it. carries_name = bool(card_name) and card_request_accepts(add_request, ("url", "name")) fields = ("url", "headers", "name") if carries_name else ("url", "headers") if card_request_accepts(add_request, fields): @@ -378,7 +379,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas if not (role_in_headers or role_in_config): raise RuntimeError( f"Agent '{self.agent_name}' cannot carry role '{self.role}': its card takes neither 'headers' on " - f"{A2AAgent.EXT_MCP_CONFIG} add for a gateway env nor 'role' in {A2AAgent.EXT_AGENT_CONFIG}, " + f"{A2AAgent.EXT_MCP_CONFIG} add nor 'role' in {A2AAgent.EXT_AGENT_CONFIG}, " f"so the role would never reach the env" ) if not role_in_headers: @@ -388,13 +389,13 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas ) elif not role_in_config: logger.warning( - f"Agent '{self.agent_name}' takes role '{self.role}' as {AGENT_ENV_ROLE_HEADER} on its gateway MCP " - f"registration only; {A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" + f"Agent '{self.agent_name}' takes role '{self.role}' as {AGENT_ENV_ROLE_HEADER} on its MCP " + f"registrations only; {A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" ) else: logger.info( - f"Role '{self.role}' reaches '{self.agent_name}' as {AGENT_ENV_ROLE_HEADER} on its gateway MCP " - f"registration and through {A2AAgent.EXT_AGENT_CONFIG}" + f"Role '{self.role}' reaches '{self.agent_name}' as {AGENT_ENV_ROLE_HEADER} on its MCP " + f"registrations and through {A2AAgent.EXT_AGENT_CONFIG}" ) if mcp_ext and env_mcp_urls: diff --git a/tst/unit/task_step/test_deploy_agent_role_delivery.py b/tst/unit/task_step/test_deploy_agent_role_delivery.py index 2dcc397f..5bc9ab1b 100644 --- a/tst/unit/task_step/test_deploy_agent_role_delivery.py +++ b/tst/unit/task_step/test_deploy_agent_role_delivery.py @@ -63,7 +63,7 @@ async def test_a_card_with_both_paths_gets_the_role_on_both(wiring, caplog): @pytest.mark.parametrize( "card", [_card(HEADERS_MCP), _card(HEADERS_MCP, ROLELESS_CONFIG)], ids=["no-agent-config", "roleless-agent-config"] ) -async def test_a_header_only_card_gets_the_role_on_the_gateway_registration_and_warns(wiring, caplog, card): +async def test_a_header_only_card_gets_the_role_on_its_mcp_registration_and_warns(wiring, caplog, card): with caplog.at_level(logging.WARNING, logger=_LOGGER): context = await wiring.deploy(card, role="executor") @@ -102,19 +102,22 @@ async def test_a_card_with_neither_path_fails_before_any_request(wiring, card): @pytest.mark.asyncio -async def test_a_header_capable_card_with_no_gateway_env_needs_agent_config(wiring): - with pytest.raises(RuntimeError, match="cannot carry role 'grader'"): - await wiring.deploy(_card(HEADERS_MCP), role="grader", env_ids=["ext"]) - assert wiring.requests == [] +async def test_an_env_that_is_not_behind_a_gateway_gets_the_header_too(wiring): + """deploy_agent does not judge what the env does with the role: a protocol server reads the header itself.""" + context = await wiring.deploy(_card(HEADERS_MCP), role="grader", env_ids=["ext"]) + + assert wiring.bodies("/ext/mcp-config") == [{"url": _LIVE_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "grader"}}] + assert wiring.bodies("/ext/agent-config") == [] + assert [a.role for a in context.deployed_agents] == ["grader"] @pytest.mark.asyncio -async def test_the_header_lands_on_the_gateway_entry_only(wiring): +async def test_the_header_lands_on_every_mcp_registration(wiring): context = await wiring.deploy(_card(HEADERS_MCP), role="alice@example.com", env_ids=["crm", "ext"]) assert wiring.bodies("/ext/mcp-config") == [ {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "alice@example.com"}, "name": "crm"}, - {"url": _LIVE_MCP_URL}, + {"url": _LIVE_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "alice@example.com"}}, ] assert wiring.bodies("/ext/agent-config") == [] assert [a.role for a in context.deployed_agents] == ["alice@example.com"] From 8d9b4aeb21d0914fea7846c19d8b7877bcc681b2 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 09:15:52 -0700 Subject: [PATCH 4/4] deploy_agent: let the add-body builder carry the role; one decision, one log line _mcp_add_body already decides whether the card's `add` takes the env card name; the role header is one more field it may take, judged on the field set actually sent. execute() builds every registration body first, then asks which paths carry the role and raises if none does. The two warnings for a single path go: the one info line names the paths that carried it. Co-Authored-By: Claude Fable 5.1 --- .../task_step/task_steps/deploy_agent.py | 68 +++++++------------ .../test_deploy_agent_role_delivery.py | 12 ++-- 2 files changed, 28 insertions(+), 52 deletions(-) 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 b0b3e371..9555458a 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -45,13 +45,17 @@ def _mcp_add_request(mcp_ext: dict) -> dict: return request if isinstance(request, dict) else {} -def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str | None) -> dict: - """Body for the agent's mcp-config `add`; the card name is relayed verbatim when the agent's card takes `name`.""" +def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str | None, role: str | None = None) -> dict: + """Body for the agent's mcp-config `add`. The card name and the role header ride along only when the card's `add` + takes them, judged on the exact field set sent; what an env does with the role is the env's business.""" + request = _mcp_add_request(mcp_ext) body: dict = {"url": url} if headers: - body["headers"] = headers - if card_name and card_request_accepts(_mcp_add_request(mcp_ext), ("url", "name")): + body["headers"] = dict(headers) + if card_name and card_request_accepts(request, ("url", "name")): body["name"] = card_name + if role and card_request_accepts(request, {*body, "headers"}): + body.setdefault("headers", {})[AGENT_ENV_ROLE_HEADER] = role return body @@ -309,9 +313,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas if resolved_litellm_base_url: env_vars["LITELLM_BASE_URL"] = resolved_litellm_base_url if self.role is not None: - # The generated MCP CLI reads its role from this variable; without it, CLI traffic from a - # role-bearing agent reaches the gateway as the CLI's default role. - env_vars.setdefault(ROLE_ENV_VAR, self.role) + env_vars.setdefault(ROLE_ENV_VAR, self.role) # the generated MCP CLI sends its role from here resolved_sandbox_type = user_overrides.get("agent_sandbox") or self.sandbox_type deploy_kwargs = {"env_vars": env_vars if env_vars else None} @@ -341,9 +343,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas logger.info(f"A2A agent '{a2a_agent_id}' deployed at {a2a_url}") mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) - add_request = _mcp_add_request(mcp_ext) if mcp_ext else {} - role_in_headers = False - env_mcp_urls: list[tuple[str, str, dict | None, str | None]] = [] # (env_id, url, headers, card_name) + registrations: list[tuple[str, dict]] = [] for env_id in self.env_ids: deployed_env = env_deployments[env_id] if deployed_env is not None: @@ -354,17 +354,9 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas else: url = host_url_for(_live_mcp_url(env_id), deployed.sandbox_type) card_name = None - headers = sandbox_request_headers_for_url(url) - if self.role is not None and mcp_ext: - # Every MCP registration carries the role; the env decides what it means (a gateway forwards it, a - # protocol server reads it, anything else ignores an unknown header). Judge the registration the - # agent will actually receive: with the card name relayed, a card that requires `name` still takes it. - carries_name = bool(card_name) and card_request_accepts(add_request, ("url", "name")) - fields = ("url", "headers", "name") if carries_name else ("url", "headers") - if card_request_accepts(add_request, fields): - headers = {**headers, AGENT_ENV_ROLE_HEADER: self.role} - role_in_headers = True - env_mcp_urls.append((env_id, url, headers or None, card_name)) + if mcp_ext: + body = _mcp_add_body(mcp_ext, url, sandbox_request_headers_for_url(url), card_name, self.role) + registrations.append((env_id, body)) identity_payload: dict[str, str] = {"name": self.agent_name} if self.agent_description is not None: @@ -372,38 +364,26 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas if self.role is not None: identity_payload["role"] = self.role identity = A2AAgent.negotiate_agent_config(card, identity_payload) - role_in_config = identity is not None and "role" in identity.fields if self.role is not None: - # Decided before the first request to the agent, so a role that cannot reach the env is neither - # registered nor recorded; execute() closes the deployed agent. - if not (role_in_headers or role_in_config): + carried_by = [via for via, carried in ( + (f"{AGENT_ENV_ROLE_HEADER} on the MCP registrations", + any(AGENT_ENV_ROLE_HEADER in body.get("headers", {}) for _, body in registrations)), + (A2AAgent.EXT_AGENT_CONFIG, identity is not None and "role" in identity.fields), + ) if carried] + if not carried_by: + # Before the first request to the agent, so a role that cannot reach the env is never registered or recorded. raise RuntimeError( f"Agent '{self.agent_name}' cannot carry role '{self.role}': its card takes neither 'headers' on " f"{A2AAgent.EXT_MCP_CONFIG} add nor 'role' in {A2AAgent.EXT_AGENT_CONFIG}, " f"so the role would never reach the env" ) - if not role_in_headers: - logger.warning( - f"Agent '{self.agent_name}' takes role '{self.role}' through {A2AAgent.EXT_AGENT_CONFIG} only; " - f"its own MCP client must send {AGENT_ENV_ROLE_HEADER}" - ) - elif not role_in_config: - logger.warning( - f"Agent '{self.agent_name}' takes role '{self.role}' as {AGENT_ENV_ROLE_HEADER} on its MCP " - f"registrations only; {A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" - ) - else: - logger.info( - f"Role '{self.role}' reaches '{self.agent_name}' as {AGENT_ENV_ROLE_HEADER} on its MCP " - f"registrations and through {A2AAgent.EXT_AGENT_CONFIG}" - ) + logger.info(f"Role '{self.role}' reaches '{self.agent_name}' through {' and '.join(carried_by)}") - if mcp_ext and env_mcp_urls: + if registrations: endpoint = a2a_url + mcp_ext["params"]["endpoint"] async with httpx.AsyncClient() as client: - for env_id, mcp_url, headers, card_name in env_mcp_urls: - body = _mcp_add_body(mcp_ext, mcp_url, headers, card_name) - logger.info(f"Registering MCP for env {env_id} via {mcp_url}") + for env_id, body in registrations: + logger.info(f"Registering MCP for env {env_id} via {body['url']}") resp = await client.post(endpoint, json=body, timeout=180) if resp.status_code >= 400: logger.error(f"Failed to register MCP {env_id}: {resp.status_code} {resp.text}") diff --git a/tst/unit/task_step/test_deploy_agent_role_delivery.py b/tst/unit/task_step/test_deploy_agent_role_delivery.py index 5bc9ab1b..f560f648 100644 --- a/tst/unit/task_step/test_deploy_agent_role_delivery.py +++ b/tst/unit/task_step/test_deploy_agent_role_delivery.py @@ -63,26 +63,22 @@ async def test_a_card_with_both_paths_gets_the_role_on_both(wiring, caplog): @pytest.mark.parametrize( "card", [_card(HEADERS_MCP), _card(HEADERS_MCP, ROLELESS_CONFIG)], ids=["no-agent-config", "roleless-agent-config"] ) -async def test_a_header_only_card_gets_the_role_on_its_mcp_registration_and_warns(wiring, caplog, card): - with caplog.at_level(logging.WARNING, logger=_LOGGER): - context = await wiring.deploy(card, role="executor") +async def test_a_header_only_card_gets_the_role_on_its_mcp_registration(wiring, card): + context = await wiring.deploy(card, role="executor") assert wiring.bodies("/ext/mcp-config") == [ {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "executor"}, "name": "crm"} ] assert all("role" not in body for body in wiring.bodies("/ext/agent-config")) - assert any(f"{A2AAgent.EXT_AGENT_CONFIG} does not list 'role'" in m for m in caplog.messages) assert [a.role for a in context.deployed_agents] == ["executor"] @pytest.mark.asyncio -async def test_a_config_only_card_gets_the_role_through_agent_config_and_warns(wiring, caplog): - with caplog.at_level(logging.WARNING, logger=_LOGGER): - context = await wiring.deploy(_card(LEGACY_MCP, ROLE_CONFIG), role="executor") +async def test_a_config_only_card_gets_the_role_through_agent_config(wiring): + context = await wiring.deploy(_card(LEGACY_MCP, ROLE_CONFIG), role="executor") assert wiring.bodies("/ext/mcp-config") == [{"url": _GATEWAY_MCP_URL}] assert wiring.bodies("/ext/agent-config") == [{"name": "solver", "role": "executor"}] - assert any(f"through {A2AAgent.EXT_AGENT_CONFIG} only" in m for m in caplog.messages) assert [a.role for a in context.deployed_agents] == ["executor"]