chore(server): check task and context ownership for a2a proxy requests

Signed-off-by: Radek Ježek <[email protected]>

Author: jezekra1

apps/beeai-cli/src/beeai_cli/api.py

@@ -112,7 +112,7 @@ async def a2a_client(agent_card: AgentCard, use_auth: bool = True) -> AsyncItera
         follow_redirects=True,
         timeout=timedelta(hours=1).total_seconds(),
     ) as httpx_client:
-        yield ClientFactory(ClientConfig(httpx_client=httpx_client)).create(card=agent_card)
+        yield ClientFactory(ClientConfig(httpx_client=httpx_client, use_client_preference=True)).create(card=agent_card)
 
 
 @asynccontextmanager

apps/beeai-cli/uv.lock

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "IBM Corp." }]
 requires-python = ">=3.11,<3.14"
 dependencies = [
-    "a2a-sdk==0.3.7",
+    "a2a-sdk==0.3.9",
     "objprint>=0.3.0",
     "uvicorn>=0.35.0",
     "asyncclick>=8.1.8",

apps/beeai-sdk-py/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "IBM Corp." }]
 requires-python = "==3.12.*"
 dependencies = [
-    "a2a-sdk~=0.3.5",
+    "a2a-sdk~=0.3.9",
     "aiohttp>=3.11.11",
     "anyio>=4.9.0",
     "asgiref>=3.8.1",
@@ -42,6 +42,8 @@ dependencies = [
     "openai>=1.97.0",
     "authlib>=1.6.4",
     "async-lru>=2.0.5",
+    "starlette>=0.48.0",
+    "sse-starlette>=3.0.2",
     "exceptiongroup>=1.3.0",
 ]
 

apps/beeai-sdk-py/uv.lock

@@ -2,14 +2,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from typing import Annotated
-from urllib.parse import urljoin, urlparse
+from urllib.parse import urljoin
 from uuid import UUID
 
 import fastapi
 import fastapi.responses
-from a2a.types import AgentCard, TransportProtocol
+from a2a.server.apps import A2AFastAPIApplication
+from a2a.server.apps.rest.rest_adapter import RESTAdapter
+from a2a.types import AgentCard, AgentInterface, TransportProtocol
 from a2a.utils import AGENT_CARD_WELL_KNOWN_PATH
-from fastapi import Depends, Request
+from fastapi import Depends, HTTPException, Request
 
 from beeai_server.api.dependencies import (
     A2AProxyServiceDependency,
@@ -19,45 +21,21 @@
 from beeai_server.domain.models.permissions import AuthorizedUser
 from beeai_server.service_layer.services.a2a import A2AServerResponse
 
-_SUPPORTED_TRANSPORTS = [TransportProtocol.jsonrpc, TransportProtocol.http_json]
-
-
 router = fastapi.APIRouter()
 
 
-def _create_proxy_url(url: str, *, proxy_base: str) -> str:
-    return urljoin(proxy_base, urlparse(url).path.lstrip("/"))
-
-
 def create_proxy_agent_card(agent_card: AgentCard, *, provider_id: UUID, request: Request) -> AgentCard:
-    proxy_base = str(request.url_for(proxy_request.__name__, provider_id=provider_id, path=""))
-    proxy_interfaces = (
-        [
-            interface.model_copy(update={"url": _create_proxy_url(interface.url, proxy_base=proxy_base)})
-            for interface in agent_card.additional_interfaces
-            if interface.transport in _SUPPORTED_TRANSPORTS
-        ]
-        if agent_card.additional_interfaces is not None
-        else None
+    proxy_base = str(request.url_for(a2a_proxy_jsonrpc_transport.__name__, provider_id=provider_id))
+    return agent_card.model_copy(
+        update={
+            "preferred_transport": TransportProtocol.jsonrpc,
+            "url": proxy_base,
+            "additional_interfaces": [
+                AgentInterface(transport=TransportProtocol.http_json, url=urljoin(proxy_base, "http")),
+                AgentInterface(transport=TransportProtocol.jsonrpc, url=proxy_base),
+            ],
+        }
     )
-    if agent_card.preferred_transport in _SUPPORTED_TRANSPORTS:
-        return agent_card.model_copy(
-            update={
-                "url": _create_proxy_url(agent_card.url, proxy_base=proxy_base),
-                "additional_interfaces": proxy_interfaces,
-            }
-        )
-    elif proxy_interfaces:
-        interface = proxy_interfaces[0]
-        return agent_card.model_copy(
-            update={
-                "url": interface.url,
-                "preferred_transport": interface.transport,
-                "additional_interfaces": proxy_interfaces,
-            }
-        )
-    else:
-        raise RuntimeError("Provider doesn't have any transport supported by the proxy.")
 
 
 def _to_fastapi(response: A2AServerResponse):
@@ -79,15 +57,45 @@ async def get_agent_card(
     return create_proxy_agent_card(provider.agent_card, provider_id=provider.id, request=request)
 
 
-@router.api_route("/{provider_id}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
-@router.api_route("/{provider_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
-async def proxy_request(
+@router.post("/{provider_id}")
+@router.post("/{provider_id}/")
+async def a2a_proxy_jsonrpc_transport(
+    provider_id: UUID,
+    request: fastapi.requests.Request,
+    a2a_proxy: A2AProxyServiceDependency,
+    provider_service: ProviderServiceDependency,
+    user: Annotated[AuthorizedUser, Depends(RequiresPermissions(a2a_proxy={"*"}))],
+):
+    provider = await provider_service.get_provider(provider_id=provider_id)
+    agent_card = create_proxy_agent_card(provider.agent_card, provider_id=provider.id, request=request)
+
+    handler = await a2a_proxy.get_request_handler(provider=provider, user=user.user)
+    app = A2AFastAPIApplication(agent_card=agent_card, http_handler=handler)
+    return await app._handle_requests(request)
+
+
+@router.api_route("/{provider_id}/http", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
+@router.api_route(
+    "/{provider_id}/http/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]
+)
+async def a2a_proxy_http_transport(
     provider_id: UUID,
     request: fastapi.requests.Request,
     a2a_proxy: A2AProxyServiceDependency,
-    _: Annotated[AuthorizedUser, Depends(RequiresPermissions(a2a_proxy={"*"}))],
+    provider_service: ProviderServiceDependency,
+    user: Annotated[AuthorizedUser, Depends(RequiresPermissions(a2a_proxy={"*"}))],
     path: str = "",
 ):
-    client = await a2a_proxy.get_proxy_client(provider_id=provider_id)
-    response = await client.send_request(method=request.method, url=f"/{path}", content=request.stream())
-    return _to_fastapi(response)
+    provider = await provider_service.get_provider(provider_id=provider_id)
+    agent_card = create_proxy_agent_card(provider.agent_card, provider_id=provider.id, request=request)
+
+    handler = await a2a_proxy.get_request_handler(provider=provider, user=user.user)
+    adapter = RESTAdapter(agent_card=agent_card, http_handler=handler)
+
+    if not (handler := adapter.routes().get((f"/{path.rstrip('/')}", request.method), None)):
+        raise HTTPException(status_code=404, detail="Not found")
+
+    return await handler(request)
+
+
+# TODO: extra a2a routes are not supported

apps/beeai-server/pyproject.toml

@@ -230,6 +230,11 @@ class ContextConfiguration(BaseModel):
     resource_expire_after_days: int = 7  # Expires files and vector_stores attached to a context
 
 
+class A2AProxyConfiguration(BaseModel):
+    # Expires a2a_request_tasks and a2a_request_contexts (WARNING: has security implications!)
+    requests_expire_after_days: int = 14
+
+
 class FeatureConfiguration(BaseModel):
     generate_conversation_title: bool = True
     provider_builds: bool = True
@@ -265,6 +270,7 @@ class Configuration(BaseSettings):
     vector_stores: VectorStoresConfiguration = Field(default_factory=VectorStoresConfiguration)
     text_extraction: DoclingExtractionConfiguration = Field(default_factory=DoclingExtractionConfiguration)
     context: ContextConfiguration = Field(default_factory=ContextConfiguration)
+    a2a_proxy: A2AProxyConfiguration = Field(default_factory=A2AProxyConfiguration)
     k8s_namespace: str | None = None
     k8s_kubeconfig: Path | None = None
 

apps/beeai-server/src/beeai_server/api/routes/a2a.py

@@ -0,0 +1,14 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from uuid import UUID
+
+from pydantic import AwareDatetime, BaseModel
+
+
+class A2ARequestTask(BaseModel):
+    task_id: str
+    created_by: UUID
+    provider_id: UUID
+    created_at: AwareDatetime
+    last_accessed_at: AwareDatetime

apps/beeai-server/src/beeai_server/configuration.py

@@ -0,0 +1,24 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+from datetime import timedelta
+from typing import Protocol, runtime_checkable
+from uuid import UUID
+
+from beeai_server.domain.models.a2a_request import A2ARequestTask
+
+
+@runtime_checkable
+class IA2ARequestRepository(Protocol):
+    async def track_request_ids_ownership(
+        self,
+        user_id: UUID,
+        provider_id: UUID,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        allow_task_creation: bool = False,
+    ) -> None: ...
+
+    async def get_task(self, *, task_id: str, user_id: UUID) -> A2ARequestTask: ...
+
+    async def delete_tasks(self, *, older_than: timedelta) -> int: ...
+    async def delete_contexts(self, *, older_than: timedelta) -> int: ...

apps/beeai-server/src/beeai_server/domain/models/a2a_request.py

@@ -55,6 +55,20 @@ def __init__(
         super().__init__(f"{entity} with {attribute} {id} not found", status_code)
 
 
+class ForbiddenUpdateError(PlatformError):
+    entity: str
+    id: UUID | str
+    attribute: str
+
+    def __init__(
+        self, entity: str, id: UUID | str, status_code: int = status.HTTP_404_NOT_FOUND, attribute: str = "id"
+    ):
+        self.entity = entity
+        self.id = id
+        self.attribute = attribute
+        super().__init__("Insufficient permissions", status_code)
+
+
 class InvalidVectorDimensionError(PlatformError): ...