feat(deploy_agent): deliver the role to the agent or fail the step

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
@@ -25,25 +26,36 @@
 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.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_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."""
+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, 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
-    add_request = mcp_ext.get("params", {}).get("methods", {}).get("add", {}).get("request", {})
-    if card_name and "name" in (add_request.get("optional") or []):
+        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
 
 
@@ -150,7 +162,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 +312,8 @@ 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:
+            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}
@@ -328,7 +342,8 @@ 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}")
 
-        env_mcp_urls: list[tuple[str, str, dict | None, str | None]] = []  # (env_id, url, headers, card_name)
+        mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG)
+        registrations: list[tuple[str, dict]] = []
         for env_id in self.env_ids:
             deployed_env = env_deployments[env_id]
             if deployed_env is not None:
@@ -339,15 +354,36 @@ 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))
+            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))
 
-        mcp_ext = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG)
-        if mcp_ext and env_mcp_urls:
+        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)
+        if self.role is not None:
+            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"
+                )
+            logger.info(f"Role '{self.role}' reaches '{self.agent_name}' through {' and '.join(carried_by)}")
+
+        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}")
@@ -378,20 +414,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:

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")

tst/unit/task_step/test_deploy_agent_role_delivery.py

@@ -0,0 +1,239 @@
+"""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_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 [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(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 [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_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_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, "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"]
+
+
+@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),
+        ({"required": ["url", "name"], "optional": ["headers"]}, True),
+        ({"required": ["url", "name"]}, False),
+    ],
+    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")
+
+    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"},
+    )