Merge pull request #18 from volcengine/feat/env-route

feat: support env detect routes for a2a and mcp server

Author: innsd

agentkit/apps/a2a_app/a2a_app.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 from typing import Callable, override
 
 import uvicorn
@@ -24,9 +25,12 @@
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import AgentCard
+from starlette.applications import Starlette
+from starlette.responses import JSONResponse
+from starlette.routing import Route
 
-from agentkit.apps.base_app import BaseAgentkitApp
 from agentkit.apps.a2a_app.telemetry import telemetry
+from agentkit.apps.base_app import BaseAgentkitApp
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +41,9 @@ 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(
@@ -80,7 +86,9 @@ 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)
@@ -111,6 +119,23 @@ 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)
+
     @override
     def run(self, agent_card: AgentCard, host: str, port: int = 8000):
         if not self._agent_executor:
@@ -130,4 +155,6 @@ def run(self, agent_card: AgentCard, host: str, port: int = 8000):
             ),
         ).build()
 
+        self.add_env_detect_route(a2a_app)
+
         uvicorn.run(a2a_app, host=host, port=port)

agentkit/apps/mcp_app/mcp_app.py

@@ -15,6 +15,7 @@
 
 import inspect
 import logging
+import os
 from functools import wraps
 from typing import Any, Callable, override
 
@@ -39,7 +40,9 @@ 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)
@@ -67,7 +70,9 @@ 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)
@@ -95,7 +100,9 @@ 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)
@@ -119,7 +126,9 @@ 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)
@@ -142,8 +151,25 @@ def sync_wrapper(*args, **kwargs) -> Any:
 
         return func
 
+    def add_env_detect_tool(self):
+        def is_agentkit_runtime() -> bool:
+            if os.getenv("RUNTIME_IAM_ROLE_TRN", ""):
+                return True
+            else:
+                return False
+
+        def get_env() -> dict:
+            return {"env": "agentkit" if is_agentkit_runtime() else "veadk"}
+
+        self._mcp_server.tool(get_env)
+
     @override
     def run(
-        self, host: str, port: int = 8000, transport: Transport = "streamable-http"
+        self,
+        host: str,
+        port: int = 8000,
+        transport: Transport = "streamable-http",
     ) -> None:
+        self.add_env_detect_tool()
+
         self._mcp_server.run(host=host, port=port, transport=transport)