refactor: refine a2a ping route

(cherry picked from commit 386d53feca3bcbc69de1d177d9811354bd47d2c1)

Author: innsd

agentkit/apps/a2a_app/a2a_app.py

@@ -26,10 +26,10 @@
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import AgentCard
-from starlette.requests import Request
 from starlette.applications import Starlette
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
+from starlette.requests import Request
 
 from agentkit.apps.a2a_app.telemetry import telemetry
 from agentkit.apps.base_app import BaseAgentkitApp
@@ -43,9 +43,7 @@ async def wrapper(*args, **kwargs):
         context: RequestContext = args[1]
         event_queue: EventQueue = args[2]
 
-        with telemetry.tracer.start_as_current_span(
-            name="a2a_invocation"
-        ) as span:
+        with telemetry.tracer.start_as_current_span(name="a2a_invocation") as span:
             exception = None
             try:
                 result = await execute_func(
@@ -89,9 +87,7 @@ def wrapper(cls: type) -> type[AgentExecutor]:
                 )
 
             if self._agent_executor:
-                raise RuntimeError(
-                    "An executor is already bound to this app instance."
-                )
+                raise RuntimeError("An executor is already bound to this app instance.")
 
             # Wrap the execute method for intercepting context and event_queue
             cls.execute = _wrap_agent_executor_execute_func(cls.execute)
@@ -122,23 +118,6 @@ def wrapper(cls: type) -> type[TaskStore]:
 
         return wrapper
 
-    def add_env_detect_route(self, app: Starlette):
-        def is_agentkit_runtime() -> bool:
-            if os.getenv("RUNTIME_IAM_ROLE_TRN", ""):
-                return True
-            else:
-                return False
-
-        route = Route(
-            "/env",
-            endpoint=lambda request: JSONResponse(
-                {"env": "agentkit" if is_agentkit_runtime() else "veadk"}
-            ),
-            methods=["GET"],
-            name="env_detect",
-        )
-        app.routes.append(route)
-
     def ping(self, func: Callable) -> Callable:
         """Register a zero-argument health check function and expose it via GET /ping.
 
@@ -163,6 +142,46 @@ def _format_ping_status(self, result: str | dict) -> dict:
         )
         return {"status": "error", "message": "Invalid response type."}
 
+    async def ping_endpoint(self, request: Request) -> Response:
+        if not self._ping_func:
+            logger.error("Ping handler function is not set")
+            return Response(status_code=404)
+
+        try:
+            result = (
+                await self._ping_func()
+                if inspect.iscoroutinefunction(self._ping_func)
+                else self._ping_func()
+            )
+            payload = self._format_ping_status(result)
+            return JSONResponse(content=payload)
+        except Exception as e:
+            logger.exception("Ping handler function failed: %s", e)
+            return JSONResponse(
+                content={"status": "error", "message": str(e)},
+                status_code=500,
+            )
+
+    def add_env_detect_route(self, app: Starlette):
+        def is_agentkit_runtime() -> bool:
+            if os.getenv("RUNTIME_IAM_ROLE_TRN", ""):
+                return True
+            else:
+                return False
+
+        route = Route(
+            "/env",
+            endpoint=lambda request: JSONResponse(
+                {"env": "agentkit" if is_agentkit_runtime() else "veadk"}
+            ),
+            methods=["GET"],
+            name="env_detect",
+        )
+        app.routes.append(route)
+
+    def add_ping_route(self, app: Starlette):
+        app.add_route("/ping", self.ping_endpoint, methods=["GET"])
+
     @override
     def run(self, agent_card: AgentCard, host: str, port: int = 8000):
         if not self._agent_executor:
@@ -182,30 +201,8 @@ def run(self, agent_card: AgentCard, host: str, port: int = 8000):
             ),
         ).build()
 
-        # Register /ping route consistent with SimpleApp behavior
-        async def _ping_handler(request: Request):
-            if not self._ping_func:
-                logger.error("Ping handler function is not set")
-                return Response(status_code=404)
-
-            try:
-                result = (
-                    await self._ping_func()
-                    if inspect.iscoroutinefunction(self._ping_func)
-                    else self._ping_func()
-                )
-                payload = self._format_ping_status(result)
-                return JSONResponse(content=payload, media_type="application/json")
-            except Exception as e:
-                logger.exception("Ping handler function failed: %s", e)
-                return JSONResponse(
-                    content={"status": "error", "message": str(e)},
-                    media_type="application/json",
-                    status_code=500,
-                )
-
-        a2a_app.add_route("/ping", _ping_handler, methods=["GET"])
-
+        # Register routes in the same style
+        self.add_ping_route(a2a_app)
         self.add_env_detect_route(a2a_app)
 
         uvicorn.run(a2a_app, host=host, port=port)

agentkit/apps/mcp_app/mcp_app.py

@@ -40,9 +40,7 @@ def tool(self, func: Callable) -> Callable:
             @wraps(func)
             async def async_wrapper(*args, **kwargs) -> Any:
                 # with tracer.start_as_current_span("tool") as span:
-                with telemetry.tracer.start_as_current_span(
-                    name="tool"
-                ) as span:
+                with telemetry.tracer.start_as_current_span(name="tool") as span:
                     exception = None
                     try:
                         result = await func(*args, **kwargs)
@@ -70,9 +68,7 @@ async def async_wrapper(*args, **kwargs) -> Any:
             @wraps(func)
             def sync_wrapper(*args, **kwargs) -> Any:
                 # with tracer.start_as_current_span("tool") as span:
-                with telemetry.tracer.start_as_current_span(
-                    name="tool"
-                ) as span:
+                with telemetry.tracer.start_as_current_span(name="tool") as span:
                     exception = None
                     try:
                         result = func(*args, **kwargs)
@@ -100,9 +96,7 @@ def agent_as_a_tool(self, func: Callable) -> Callable:
 
             @wraps(func)
             async def async_wrapper(*args, **kwargs) -> Any:
-                with telemetry.tracer.start_as_current_span(
-                    name="tool"
-                ) as span:
+                with telemetry.tracer.start_as_current_span(name="tool") as span:
                     exception = None
                     try:
                         result = await func(*args, **kwargs)
@@ -126,9 +120,7 @@ async def async_wrapper(*args, **kwargs) -> Any:
 
             @wraps(func)
             def sync_wrapper(*args, **kwargs) -> Any:
-                with telemetry.tracer.start_as_current_span(
-                    name="tool"
-                ) as span:
+                with telemetry.tracer.start_as_current_span(name="tool") as span:
                     exception = None
                     try:
                         result = func(*args, **kwargs)

agentkit/toolkit/resources/samples/a2a.py

@@ -74,8 +74,8 @@ def ping() -> str:
         capabilities=AgentCapabilities(streaming=True),
         description=agent.description,
         name=agent.name,
-        defaultInputModes=["text"],
-        defaultOutputModes=["text"],
+        default_input_modes=["text"],
+        default_output_modes=["text"],
         provider=AgentProvider(organization="veadk", url=""),
         skills=[AgentSkill(id="0", name="chat", description="Chat", tags=["chat"])],
         url="http://0.0.0.0:8000",