diff --git a/agent_assembly/client/__init__.py b/agent_assembly/client/__init__.py index 93278fa..cd8ee71 100644 --- a/agent_assembly/client/__init__.py +++ b/agent_assembly/client/__init__.py @@ -1,5 +1,6 @@ """Client module for gateway communication.""" +from agent_assembly.client.dispatch import DispatchToolResult from agent_assembly.client.gateway import GatewayClient -__all__ = ["GatewayClient"] +__all__ = ["DispatchToolResult", "GatewayClient"] diff --git a/agent_assembly/client/dispatch.py b/agent_assembly/client/dispatch.py new file mode 100644 index 0000000..e238797 --- /dev/null +++ b/agent_assembly/client/dispatch.py @@ -0,0 +1,28 @@ +"""Result type for `GatewayClient.dispatch_tool` (AAASM-1920 / Secret Injection).""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class DispatchToolResult: + """ + Outcome of a successful ``dispatch_tool`` call. + + The gateway resolves every ``${NAME}`` placeholder in the args via the + registered ``SecretsStore``, emits a placeholder-form audit entry, and + returns this object back to the SDK caller. + + Attributes: + resolved_args: Post-substitution args. Carries the *resolved* + credential values; do not log this or pass it to the LLM. + names_substituted: The placeholder names that were resolved during + this call. Names only — never the resolved values. Echoes the + audit-log shape so callers can correlate dispatches with audit + entries by ``names_substituted`` set. + """ + + resolved_args: dict[str, Any] + names_substituted: list[str] = field(default_factory=list) diff --git a/agent_assembly/client/gateway.py b/agent_assembly/client/gateway.py index c75f7d9..1c7d656 100644 --- a/agent_assembly/client/gateway.py +++ b/agent_assembly/client/gateway.py @@ -2,8 +2,11 @@ from __future__ import annotations +from typing import Any + import httpx +from agent_assembly.client.dispatch import DispatchToolResult from agent_assembly.exceptions import GatewayError @@ -175,3 +178,41 @@ def report_edge( return response.json() except httpx.HTTPError as e: raise GatewayError(f"Failed to report edge: {e}") from e + + async def dispatch_tool(self, tool_name: str, args: dict[str, Any]) -> DispatchToolResult: + """ + Dispatch a tool with placeholder-form args (AAASM-1920 Secret Injection). + + The gateway resolves every ``${NAME}`` placeholder in ``args`` against + its registered ``SecretsStore``, emits a placeholder-form audit entry, + and returns the resolved args plus the list of substituted names. + + The LLM never observes the resolved credential value: the agent code + holds the placeholder, Assembly resolves it on the gateway side, and + the response is forwarded to the tool sink only. + + Args: + tool_name: Name of the tool to dispatch (e.g. ``"call_database"``). + args: Placeholder-form args. May contain ``${NAME}`` tokens that + will be resolved server-side before the tool is invoked. + + Returns: + ``DispatchToolResult`` with the resolved args + the list of + placeholder names that were substituted. + + Raises: + GatewayError: If the request fails for any reason — including a + 422 Unprocessable Entity when ``args`` references a placeholder + that is not registered in the gateway's ``SecretsStore``. + """ + body = {"tool": tool_name, "args": args} + try: + response = self.client.post("/dispatch_tool", json=body) + response.raise_for_status() + data = response.json() + return DispatchToolResult( + resolved_args=data.get("resolved_args", {}), + names_substituted=list(data.get("names_substituted", [])), + ) + except httpx.HTTPError as e: + raise GatewayError(f"Failed to dispatch tool {tool_name}: {e}") from e diff --git a/test/unit/client/test_dispatch_tool.py b/test/unit/client/test_dispatch_tool.py new file mode 100644 index 0000000..b8c16a3 --- /dev/null +++ b/test/unit/client/test_dispatch_tool.py @@ -0,0 +1,97 @@ +"""Unit tests for GatewayClient.dispatch_tool (AAASM-1920 Secret Injection).""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from agent_assembly.client import DispatchToolResult, GatewayClient +from agent_assembly.exceptions import GatewayError + + +def _make_response(status_code: int = 200, payload: dict | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.json.return_value = payload or {} + if status_code >= 400: + resp.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + f"{status_code}", + request=MagicMock(), + response=resp, + ) + ) + else: + resp.raise_for_status = MagicMock() + return resp + + +def _patched_client(client: GatewayClient, mock_post: MagicMock) -> object: + return patch.object( + type(client), + "client", + new_callable=lambda: property(lambda _self: MagicMock(post=mock_post)), + ) + + +@pytest.mark.asyncio +async def test_dispatch_tool_returns_resolved_result_on_success() -> None: + client = GatewayClient(gateway_url="http://gw.test", agent_id="agent-1", api_key="k") + mock_post = MagicMock( + return_value=_make_response( + 200, + { + "resolved_args": {"connection_string": "real-secret-abc"}, + "names_substituted": ["DB_PASSWORD"], + }, + ) + ) + with _patched_client(client, mock_post): + result = await client.dispatch_tool("call_database", {"connection_string": "${DB_PASSWORD}"}) + + assert isinstance(result, DispatchToolResult) + assert result.resolved_args == {"connection_string": "real-secret-abc"} + assert result.names_substituted == ["DB_PASSWORD"] + + # Body sent over the wire is the placeholder-form — never the resolved + # value. Pins the SDK contract complementary to the audit-shape contract + # on the gateway side. + _, kwargs = mock_post.call_args + sent = kwargs.get("json") or {} + assert sent == {"tool": "call_database", "args": {"connection_string": "${DB_PASSWORD}"}} + + +@pytest.mark.asyncio +async def test_dispatch_tool_raises_gateway_error_on_unknown_placeholder() -> None: + """422 from the gateway (unknown placeholder) maps to GatewayError.""" + client = GatewayClient(gateway_url="http://gw.test", agent_id="agent-1", api_key="k") + mock_post = MagicMock(return_value=_make_response(422)) + with _patched_client(client, mock_post), pytest.raises(GatewayError) as excinfo: + await client.dispatch_tool("call_database", {"x": "${UNKNOWN_SECRET}"}) + + assert "Failed to dispatch tool call_database" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_dispatch_tool_raises_gateway_error_on_network_failure() -> None: + """Underlying httpx.ConnectError surfaces as GatewayError.""" + client = GatewayClient(gateway_url="http://gw.test", agent_id="agent-1", api_key="k") + mock_post = MagicMock(side_effect=httpx.ConnectError("connection refused")) + with _patched_client(client, mock_post), pytest.raises(GatewayError) as excinfo: + await client.dispatch_tool("call_database", {"x": "y"}) + + assert "Failed to dispatch tool call_database" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_dispatch_tool_defaults_empty_result_fields_when_server_omits_them() -> None: + """Defensive: server returning an empty body still yields a well-formed result.""" + client = GatewayClient(gateway_url="http://gw.test", agent_id="agent-1", api_key="k") + mock_post = MagicMock(return_value=_make_response(200, {})) + with _patched_client(client, mock_post): + result = await client.dispatch_tool("noop", {"x": "y"}) + + assert result.resolved_args == {} + assert result.names_substituted == []