fix(server): send proxy errors as a2a errors

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

apps/agentstack-server/src/agentstack_server/exceptions.py

@@ -69,6 +69,9 @@ def __init__(
         super().__init__("Insufficient permissions", status_code)
 
 
+class InvalidProviderCallError(PlatformError): ...
+
+
 class InvalidVectorDimensionError(PlatformError): ...
 
 

apps/agentstack-server/src/agentstack_server/service_layer/services/a2a.py

@@ -5,7 +5,7 @@
 import inspect
 import logging
 import uuid
-from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from datetime import timedelta
 from typing import NamedTuple, cast
@@ -24,6 +24,7 @@
     AgentCard,
     DeleteTaskPushNotificationConfigParams,
     GetTaskPushNotificationConfigParams,
+    InternalError,
     InvalidRequestError,
     ListTaskPushNotificationConfigParams,
     Message,
@@ -47,9 +48,10 @@
     NetworkProviderLocation,
     Provider,
     ProviderDeploymentState,
+    UnmanagedState,
 )
 from agentstack_server.domain.models.user import User
-from agentstack_server.exceptions import EntityNotFoundError, ForbiddenUpdateError
+from agentstack_server.exceptions import EntityNotFoundError, ForbiddenUpdateError, InvalidProviderCallError
 from agentstack_server.service_layer.deployment_manager import (
     IProviderDeploymentManager,
 )
@@ -116,6 +118,10 @@ async def _fn(*args, **kwargs):
             raise ServerError(error=InvalidRequestError(message=str(e))) from e
         except A2AClientJSONRPCError as e:
             raise ServerError(error=e.error) from e
+        except InvalidProviderCallError as e:
+            raise ServerError(error=InvalidRequestError(message=f"Invalid request to agent: {e!r}")) from e
+        except Exception as e:
+            raise ServerError(error=InternalError(message=f"Internal error: {e!r}")) from e
 
     @functools.wraps(fn)
     async def _fn_iter(*args, **kwargs):
@@ -131,18 +137,28 @@ async def _fn_iter(*args, **kwargs):
 class ProxyRequestHandler(RequestHandler):
     def __init__(
         self,
-        agent_card: AgentCard,
+        *,
         provider_id: UUID,
         uow: IUnitOfWorkFactory,
         user: User,
+        # Calling the factory have side-effects, such as rotating the agent
+        agent_card_factory: Callable[[], Awaitable[AgentCard]] | None = None,
+        agent_card: AgentCard | None = None,
     ):
+        if agent_card_factory is None and agent_card is None:
+            raise ValueError("One of agent_card_factory or agent_card must be provided")
+        self._agent_card_factory = agent_card_factory
         self._agent_card = agent_card
         self._provider_id = provider_id
         self._user = user
         self._uow = uow
 
     @asynccontextmanager
     async def _client_transport(self) -> AsyncIterator[ClientTransport]:
+        if self._agent_card is None:
+            assert self._agent_card_factory is not None
+            self._agent_card = await self._agent_card_factory()
+
         async with httpx.AsyncClient(follow_redirects=True, timeout=timedelta(hours=1).total_seconds()) as httpx_client:
             client: BaseClient = cast(
                 BaseClient,
@@ -306,10 +322,13 @@ def __init__(
         self._expire_requests_after = timedelta(days=configuration.a2a_proxy.requests_expire_after_days)
 
     async def get_request_handler(self, *, provider: Provider, user: User) -> RequestHandler:
-        url = await self.ensure_agent(provider_id=provider.id)
-        agent_card = create_deployment_agent_card(provider.agent_card, deployment_base=str(url))
+        async def agent_card_factory() -> AgentCard:
+            # Delay ensure_agent to the handler so that errors are wrapped properly
+            url = await self.ensure_agent(provider_id=provider.id)
+            return create_deployment_agent_card(provider.agent_card, deployment_base=str(url))
+
         return ProxyRequestHandler(
-            agent_card=agent_card,
+            agent_card_factory=agent_card_factory,
             provider_id=provider.id,
             uow=self._uow,
             user=user,
@@ -332,6 +351,11 @@ async def ensure_agent(self, *, provider_id: UUID) -> HttpUrl:
                 await uow.commit()
 
             if not provider.managed:
+                if provider.unmanaged_state is UnmanagedState.OFFLINE:
+                    raise InvalidProviderCallError(
+                        f"Cannot send message to provider {provider_id}: provider is offline"
+                    )
+
                 assert isinstance(provider.source, NetworkProviderLocation)
                 return provider.source.a2a_url
 
@@ -340,7 +364,9 @@ async def ensure_agent(self, *, provider_id: UUID) -> HttpUrl:
             should_wait = False
             match state:
                 case ProviderDeploymentState.ERROR:
-                    raise RuntimeError("Provider is in an error state")
+                    raise InvalidProviderCallError(
+                        f"Cannot send message to provider {provider_id}: provider is in an error state"
+                    )
                 case (
                     ProviderDeploymentState.MISSING
                     | ProviderDeploymentState.RUNNING

apps/agentstack-server/template.env

@@ -36,7 +36,7 @@ OCI_REGISTRY__AGENTSTACK-REGISTRY-SVC.DEFAULT:5001__INSECURE=true
 # AUTH__BASIC__ENABLED=true
 # AUTH__BASIC__ADMIN_PASSWORD=test-password
 # AUTH__BASIC__ADMIN_PASSWORD=test-password
-# AUTH__BASIC__JWT_SECRET_KEY=test-key
+# AUTH__JWT_SECRET_KEY=test-key
 
 # AUTH__OIDC__ENABLED=true
 # AUTH__OIDC__ADMIN_EMAILS='["abc@ibm.com"]'

apps/agentstack-server/tests/e2e/routes/test_a2a_proxy.py

@@ -1195,3 +1195,36 @@ async def test_task_and_context_both_specified_single_query(client: Client, hand
         {"context_id": "dual-context-456"},
     )
     assert context_result.fetchone() is not None
+
+
+async def test_invalid_request_raises_a2a_error(client: Client, handler: mock.AsyncMock, db_transaction):
+    """Test that an invalid request to an offline provider returns an A2A error."""
+
+    # set provider as offline
+    provider_id = str(client.base_url).rstrip("/").split("/")[-1]
+    await db_transaction.execute(
+        text("UPDATE providers SET unmanaged_state = 'offline' WHERE id = :provider_id"),
+        {"provider_id": provider_id},
+    )
+    await db_transaction.commit()
+
+    message_data = {
+        "jsonrpc": "2.0",
+        "id": "123",
+        "method": "message/send",
+        "params": {
+            "message": {
+                "role": "agent",
+                "parts": [{"kind": "text", "text": "Hello"}],
+                "messageId": "111",
+                "kind": "message",
+            }
+        },
+    }
+    response = client.post("/", json=message_data)
+    assert response.status_code == 200
+    data = response.json()
+    assert data["id"] == "123"
+    assert "error" in data
+    assert data["error"]["code"] == InvalidRequestError().code
+    assert "provider is offline" in data["error"]["message"]