Added tests

Author: vladkol

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -71,6 +71,7 @@
         from sse_starlette.sse import EventSourceResponse
         from starlette.applications import Starlette
         from starlette.authentication import BaseUser
+        from starlette.datastructures import URL
         from starlette.exceptions import HTTPException
         from starlette.requests import Request
         from starlette.responses import JSONResponse, Response
@@ -488,6 +489,12 @@ async def event_generator(
         )
 
     def _modify_rpc_url(self, agent_card: AgentCard, request: Request):
+        """Modifies Agent's RPC URL based on the AgentCard request.
+
+        Args:
+            agent_card (AgentCard): Original AgentCard
+            request (Request): AgentCard request
+        """
         rpc_url = URL(agent_card.url)
         rpc_path = rpc_url.path
         port = None
@@ -499,6 +506,7 @@ def _modify_rpc_url(self, agent_card: AgentCard, request: Request):
 
         if "X-Forwarded-Proto" in request.headers:
             scheme = request.headers["X-Forwarded-Proto"]
+            port = None
         else:
             scheme = request.url.scheme
         if not scheme:
@@ -526,23 +534,14 @@ def _modify_rpc_url(self, agent_card: AgentCard, request: Request):
                     new_path = new_path.rstrip("/")
                 rpc_path = new_path
 
-        if port:
-            agent_card.url = str(
-                rpc_url.replace(
-                    hostname=host,
-                    port=port,
-                    scheme=scheme,
-                    path=rpc_path
-                )
-            )
-        else:
-            agent_card.url = str(
-                rpc_url.replace(
-                    hostname=host,
-                    scheme=scheme,
-                    path=rpc_path
-                )
+        agent_card.url = str(
+            rpc_url.replace(
+                hostname=host,
+                port=port,
+                scheme=scheme,
+                path=rpc_path
             )
+        )
 
     async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
         """Handles GET requests for the agent card endpoint.

src/a2a/server/apps/rest/rest_adapter.py

@@ -7,6 +7,7 @@
 
 if TYPE_CHECKING:
     from sse_starlette.sse import EventSourceResponse
+    from starlette.datastructures import URL
     from starlette.requests import Request
     from starlette.responses import JSONResponse, Response
 
@@ -15,6 +16,7 @@
 else:
     try:
         from sse_starlette.sse import EventSourceResponse
+        from starlette.datastructures import URL
         from starlette.requests import Request
         from starlette.responses import JSONResponse, Response
 
@@ -119,7 +121,8 @@ async def handle_get_agent_card(self, request: Request) -> JSONResponse:
             A JSONResponse containing the agent card data.
         """
         # The public agent card is a direct serialization of the agent_card
-        # provided at initialization.
+        # provided at initialization except for the RPC URL.
+        self._modify_rpc_url(self.agent_card, request)
         return JSONResponse(
             self.agent_card.model_dump(mode='json', exclude_none=True)
         )
@@ -145,9 +148,65 @@ async def handle_authenticated_agent_card(
                     message='Authenticated card not supported'
                 )
             )
+        self._modify_rpc_url(self.agent_card, request)
         return JSONResponse(
             self.agent_card.model_dump(mode='json', exclude_none=True)
         )
+    
+    def _modify_rpc_url(self, agent_card: AgentCard, request: Request):
+        """Modifies Agent's RPC URL based on the AgentCard request.
+
+        Args:
+            agent_card (AgentCard): Original AgentCard
+            request (Request): AgentCard request
+        """
+        rpc_url = URL(agent_card.url)
+        rpc_path = rpc_url.path
+        port = None
+        if "X-Forwarded-Host" in request.headers:
+            host = request.headers["X-Forwarded-Host"]
+        else:
+            host = request.url.hostname
+            port = request.url.port
+            
+        if "X-Forwarded-Proto" in request.headers:
+            scheme = request.headers["X-Forwarded-Proto"]
+            port = None
+        else:
+            scheme = request.url.scheme
+        if not scheme:
+            scheme = "http"
+        if ":" in host: # type: ignore
+            comps = host.rsplit(":", 1) # type: ignore
+            host = comps[0]
+            port = comps[1]
+
+        # Handle URL maps,
+        # e.g. "agents/my-agent/.well-known/agent-card.json"
+        if "X-Forwarded-Path" in request.headers:
+            forwarded_path = request.headers["X-Forwarded-Path"].strip()
+            if (
+                forwarded_path and
+                request.url.path != forwarded_path
+                and forwarded_path.endswith(request.url.path)
+            ):
+                # "agents/my-agent" for "agents/my-agent/.well-known/agent-card.json"
+                extra_path = forwarded_path[:-len(request.url.path)]
+                new_path = extra_path + rpc_path
+                # If original path was just "/",
+                # we remove trailing "/" in the the extended one
+                if len(new_path) > 1 and rpc_path == "/":
+                    new_path = new_path.rstrip("/")
+                rpc_path = new_path
+
+        agent_card.url = str(
+            rpc_url.replace(
+                hostname=host,
+                port=port,
+                scheme=scheme,
+                path=rpc_path
+            )
+        )
 
     def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]:
         """Constructs a dictionary of API routes and their corresponding handlers.

tests/server/apps/jsonrpc/test_jsonrpc_app.py

@@ -26,6 +26,7 @@
     RequestHandler,
 )  # For mock spec
 from a2a.types import (
+    AgentCapabilities,
     AgentCard,
     Message,
     MessageSendParams,
@@ -36,7 +37,7 @@
     SendMessageSuccessResponse,
     TextPart,
 )
-
+from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
 
 # --- StarletteUserProxy Tests ---
 
@@ -356,5 +357,58 @@ def side_effect(request, context: ServerCallContext):
         }
 
 
+class TestAgentCardHandler:
+        @pytest.fixture
+        def agent_card(self):
+            return AgentCard(
+                name='APIKeyAgent',
+                description='An agent that uses API Key auth.',
+                url='http://localhost:8000',
+                version='1.0.0',
+                capabilities=AgentCapabilities(),
+                default_input_modes=['text/plain'],
+                default_output_modes=['text/plain'],
+                skills=[],
+            )
+
+        def test_agent_card_url_rewriting(
+            self, agent_card: AgentCard,
+        ):
+            """
+            Tests that the A2AStarletteApplication endpoint correctly handles Agent URL rewriting.
+            """
+            handler = AsyncMock()
+            app_instance = A2AStarletteApplication(agent_card, handler)
+            client = TestClient(
+                app_instance.build(),
+                base_url="https://my-agents.com:5000"
+            )
+
+            response = client.get(AGENT_CARD_WELL_KNOWN_PATH)
+            response.raise_for_status()
+            assert response.json()["url"] == "https://my-agents.com:5000"
+
+            response = client.get(
+                AGENT_CARD_WELL_KNOWN_PATH,
+                headers={
+                    "X-Forwarded-Host": "my-great-agents.com:5678",
+                    "X-Forwarded-Proto": "http",
+                    "X-Forwarded-Path": 
+                        "/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH
+                }
+            )
+            assert response.json()["url"] == "http://my-great-agents.com:5678/agents/my-agent"
+
+            client = TestClient(
+                app_instance.build(
+                    agent_card_url="/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH
+                ),
+                base_url="https://my-mighty-agents.com"
+            )
+
+            response = client.get("/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH)
+            assert response.json()["url"] == "https://my-mighty-agents.com/agents/my-agent"
+
+
 if __name__ == '__main__':
     pytest.main([__file__])

tests/server/apps/rest/test_rest_fastapi_app.py

@@ -1,7 +1,7 @@
 import logging
 
 from typing import Any
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -15,6 +15,7 @@
 from a2a.server.apps.rest.rest_adapter import RESTAdapter
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.types import (
+    AgentCapabilities,
     AgentCard,
     Message,
     Part,
@@ -24,7 +25,7 @@
     TaskStatus,
     TextPart,
 )
-
+from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
 
 logger = logging.getLogger(__name__)
 
@@ -222,5 +223,62 @@ async def test_send_message_success_task(
     assert expected_response == actual_response
 
 
+class TestAgentCardHandler:
+        @pytest.fixture
+        def agent_card(self):
+            return AgentCard(
+                name='APIKeyAgent',
+                description='An agent that uses API Key auth.',
+                url='http://localhost:8000',
+                version='1.0.0',
+                capabilities=AgentCapabilities(),
+                default_input_modes=['text/plain'],
+                default_output_modes=['text/plain'],
+                skills=[],
+            )
+        
+        @pytest.mark.anyio
+        async def test_agent_card_url_rewriting(
+            self, agent_card: AgentCard,
+        ):
+            """
+            Tests that the REST endpoint correctly handles Agent URL rewriting.
+            """
+            app_instance = A2ARESTFastAPIApplication(agent_card, AsyncMock())
+            app =  app_instance.build(
+                agent_card_url=AGENT_CARD_WELL_KNOWN_PATH
+            )
+            client = AsyncClient(
+                transport=ASGITransport(app=app),
+                base_url="https://my-agents.com:5000"
+            )
+
+            response = await client.get(AGENT_CARD_WELL_KNOWN_PATH)
+            response.raise_for_status()
+            assert response.json()["url"] == "https://my-agents.com:5000"
+
+            response = await client.get(
+                AGENT_CARD_WELL_KNOWN_PATH,
+                headers={
+                    "X-Forwarded-Host": "my-great-agents.com:5678",
+                    "X-Forwarded-Proto": "http",
+                    "X-Forwarded-Path": 
+                        "/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH
+                }
+            )
+            assert response.json()["url"] == "http://my-great-agents.com:5678/agents/my-agent"
+
+            app = app_instance.build(
+                agent_card_url="/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH
+            )
+            client = AsyncClient(
+                transport=ASGITransport(app=app),
+                base_url="https://my-mighty-agents.com"
+            )
+
+            response = await client.get("/agents/my-agent" + AGENT_CARD_WELL_KNOWN_PATH)
+            assert response.json()["url"] == "https://my-mighty-agents.com/agents/my-agent"
+
+
 if __name__ == '__main__':
     pytest.main([__file__])