feat(protocol): a canonical attribution-probe/v1 extension

packages/agentenv-protocol/README.md

@@ -474,6 +474,19 @@ tools with (agent-env relays the env card's name: the MultiEnv's declared name,
 `env` + 4 random digits, giving e.g. `mcp__env4821__<tool>`); when absent the agent mints
 `mcp_<8 hex>`.
 
+`ATTRIBUTION_PROBE_V1.probe` lets a caller check that an agent forwards attribution
+upstream. Attribution is an open map of string dimensions, like AgentEnv's cost
+attribution: the SDK assigns no keys, and each deployment uses whatever dimensions
+its model gateway understands (the AgentEnv control plane configures `project_id`
+and `task_id` on agents that declare those config fields). An agent records the
+dimensions it last sent with a model request and returns them from a bodyless
+`POST /ext/attribution-probe` as an `AttributionProbeResponse`:
+`{"last_seen_attribution": {<dimension>: <value>, ...}, "last_seen_at_utc": ...}`.
+The SDK does not record attribution itself, because only the agent knows what it sent.
+The answer is agent-wide rather than per caller, and like every extension route the
+probe carries no access check of its own, so expose an agent only to the control plane
+that drives it.
+
 Runnable, self-contained reference agents live in [`examples/`](https://github.com/scaleapi/agentenv-framework/tree/main/packages/agentenv-protocol/examples):
 the normal, streaming, and multimodal `run(request)` paths, single- and
 multi-operation custom extensions, and advanced ASGI-lifespan plus

packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py

@@ -4,6 +4,7 @@
 
 from .extensions import (
     AGENT_CONFIG_V1,
+    ATTRIBUTION_PROBE_V1,
     INSTALL_V1,
     MCP_CONFIG_V1,
     PEER_AGENTS_V1,
@@ -12,6 +13,7 @@
     STANDARD_EXTENSIONS,
     TRAJECTORY_V1,
     TRIGGERS_V1,
+    AttributionProbeResponse,
     BundleSkillRequest,
     ChangelogIncrement,
     ContextObjectTrajectoryRequest,
@@ -108,6 +110,7 @@
 
 __all__ = [
     "AGENT_CONFIG_V1",
+    "ATTRIBUTION_PROBE_V1",
     "INSTALL_V1",
     "MCP_CONFIG_V1",
     "PEER_AGENTS_V1",
@@ -122,6 +125,7 @@
     "AgentEnvAgent",
     "AgentIdentity",
     "AgentRunResult",
+    "AttributionProbeResponse",
     "BundleSkillRequest",
     "ChangelogIncrement",
     "ContextObjectTrajectoryRequest",

packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py

@@ -231,6 +231,11 @@ class PeerAgentsSetRequest(ExtensionRequest):
     peers: list[PeerAgent]
 
 
+class AttributionProbeResponse(ExtensionResponse):
+    last_seen_attribution: dict[str, str]
+    last_seen_at_utc: str | None = None
+
+
 class TriggerRegisterRequest(ExtensionRequest):
     triggers: list[dict[str, Any]]
 
@@ -1037,6 +1042,21 @@ def _configure_install(**declaration: Any) -> ExtensionConfiguration:
     },
 )
 
+ATTRIBUTION_PROBE_V1 = ExtensionDefinition(
+    uri="urn:agentenv:attribution-probe/v1",
+    description="Report the attribution dimensions last sent with a model request.",
+    endpoint="/ext/attribution-probe",
+    core_operations={
+        "probe": _op(
+            "probe",
+            "POST",
+            "/ext/attribution-probe",
+            RUNTIME,
+            response_model=AttributionProbeResponse,
+        ),
+    },
+)
+
 TRIGGERS_V1 = ExtensionDefinition(
     uri="urn:agentenv:triggers/v1",
     description="Register triggers and decide deterministic responses.",
@@ -1084,6 +1104,7 @@ def _configure_install(**declaration: Any) -> ExtensionConfiguration:
             TRAJECTORY_V1,
             SNAPSHOT_V1,
             PEER_AGENTS_V1,
+            ATTRIBUTION_PROBE_V1,
             TRIGGERS_V1,
             INSTALL_V1,
         )

packages/agentenv-protocol/tests/test_a2a_agent.py

@@ -30,6 +30,7 @@
 )
 from agentenv_protocol.a2a_agent import (
     AGENT_CONFIG_V1,
+    ATTRIBUTION_PROBE_V1,
     MCP_CONFIG_V1,
     PEER_AGENTS_V1,
     SKILL_CONFIG_V1,
@@ -40,6 +41,7 @@
     AgentConfig,
     AgentEnvAgent,
     AgentIdentity,
+    AttributionProbeResponse,
     BundleSkillRequest,
     ContextObjectTrajectoryRequest,
     ContextTrajectoryRequest,
@@ -1348,6 +1350,86 @@ async def list_peers(self):
         }
 
 
+def test_attribution_probe_reports_what_the_agent_last_sent() -> None:
+    @a2a_agent(
+        identity=AgentIdentity(name="attributing-agent", description="test", version="1")
+    )
+    class Agent(AgentEnvAgent):
+        def __init__(self) -> None:
+            self.sent: dict[str, str] = {}
+
+        async def run(self, request: TaskRequest) -> TaskResult:
+            self.sent = {"project_id": "project-1", "task_id": "task-1", "team": "evals"}
+            return TaskResult.text("ok")
+
+        @extension(ATTRIBUTION_PROBE_V1.probe)
+        async def probe(self):
+            return AttributionProbeResponse(
+                last_seen_attribution=self.sent, last_seen_at_utc="2026-01-01T00:00:00Z"
+            )
+
+    agent = Agent()
+    with TestClient(agent.create_app()) as client:
+        card = client.get("/.well-known/agent.json").json()
+        assert _extension(card, ATTRIBUTION_PROBE_V1.uri)["params"]["endpoint"] == (
+            "/ext/attribution-probe"
+        )
+        assert _operation(client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {}).json() == {
+            "last_seen_attribution": {},
+            "last_seen_at_utc": "2026-01-01T00:00:00Z",
+        }
+        response = client.post("/a2a", json=_message_request())
+        assert response.json()["result"]["status"]["state"] == "completed"
+        assert _operation(client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {}).json() == {
+            "last_seen_attribution": {
+                "project_id": "project-1",
+                "task_id": "task-1",
+                "team": "evals",
+            },
+            "last_seen_at_utc": "2026-01-01T00:00:00Z",
+        }
+        assert (
+            _operation(
+                client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {"project_id": "x"}
+            ).status_code
+            == 400
+        )
+
+
+def test_attribution_probe_rejects_an_undeclared_response_field() -> None:
+    @a2a_agent(
+        identity=AgentIdentity(name="attributing-agent", description="test", version="1")
+    )
+    class Agent(AgentEnvAgent):
+        async def run(self, request: TaskRequest) -> TaskResult:
+            return TaskResult.text("ok")
+
+        @extension(ATTRIBUTION_PROBE_V1.probe)
+        async def probe(self):
+            return {"last_seen_attribution": {}, "headers": {"x-secret": "value"}}
+
+    with TestClient(Agent().create_app(), raise_server_exceptions=False) as client:
+        card = client.get("/.well-known/agent.json").json()
+        response = _operation(client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {})
+        assert response.status_code == 500
+        assert "x-secret" not in response.text
+
+
+def test_attribution_probe_is_only_declared_by_a_handler() -> None:
+    @a2a_agent(
+        identity=AgentIdentity(name="plain-agent", description="test", version="1")
+    )
+    class Agent(AgentEnvAgent):
+        async def run(self, request: TaskRequest) -> TaskResult:
+            return TaskResult.text("ok")
+
+    with TestClient(Agent().create_app()) as client:
+        card = client.get("/.well-known/agent.json").json()
+        uris = {item["uri"] for item in card["capabilities"]["extensions"]}
+        assert ATTRIBUTION_PROBE_V1.uri not in uris
+        assert client.post("/ext/attribution-probe", json={}).status_code in (404, 405)
+
+
 def test_handler_advertises_the_canonical_request_model() -> None:
     class Agent:
         @extension(SNAPSHOT_V1.save)