Feat/agent migration bridge (#136)

Author: wujiaming-ai

agentkit/apps/agent_server_app/agent_server_app.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import json
 import logging
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from typing import Any
 
@@ -62,6 +64,35 @@
 logger = logging.getLogger(__name__)
 
 
+async def _call_lifecycle_handler(handler: Any) -> None:
+    result = handler()
+    if inspect.isawaitable(result):
+        await result
+
+
+@asynccontextmanager
+async def _run_a2a_app_lifespan(a2a_app: Any) -> AsyncIterator[None]:
+    router = getattr(a2a_app, "router", None)
+    if router is None:
+        raise RuntimeError("A2A server app has no router; cannot initialize lifecycle.")
+
+    lifespan_context = getattr(router, "lifespan_context", None)
+    if lifespan_context is not None:
+        async with lifespan_context(a2a_app):
+            yield
+        return
+
+    startup_handlers = tuple(getattr(router, "on_startup", ()) or ())
+    shutdown_handlers = tuple(getattr(router, "on_shutdown", ()) or ())
+    for handler in startup_handlers:
+        await _call_lifecycle_handler(handler)
+    try:
+        yield
+    finally:
+        for handler in shutdown_handlers:
+            await _call_lifecycle_handler(handler)
+
+
 class AgentKitAgentLoader(BaseAgentLoader):
     def __init__(self, agent_or_app: BaseAgent | App) -> None:
         super().__init__()
@@ -157,11 +188,10 @@ def __init__(
         async def lifespan(app: FastAPI):
             # trigger A2A server app startup
             logger.info(
-                "Triggering A2A server app startup within API server..."
+                "Triggering A2A server app lifespan within API server..."
             )
-            for handler in _a2a_server_app.router.on_startup:
-                await handler()
-            yield
+            async with _run_a2a_app_lifespan(_a2a_server_app):
+                yield
 
         resolved_allow_origins = resolve_agentkit_allow_origins(
             allow_origins=allow_origins,
@@ -280,78 +310,6 @@ async def event_generator():
                 routes.insert(0, routes.pop(i))
                 break
 
-        @self.app.post("/run_sse")
-        async def run_agent_sse(req: RunAgentRequest) -> StreamingResponse:
-            print("my run sse !!!")
-            # SSE endpoint
-            session = await self.server.session_service.get_session(
-                app_name=req.app_name,
-                user_id=req.user_id,
-                session_id=req.session_id,
-            )
-            if not session:
-                raise HTTPException(status_code=404, detail="Session not found")
-
-            # Convert the events to properly formatted SSE
-            async def event_generator():
-                try:
-                    stream_mode = (
-                        StreamingMode.SSE
-                        if req.streaming
-                        else StreamingMode.NONE
-                    )
-                    runner = await self.server.get_runner_async(req.app_name)
-                    async with Aclosing(
-                        runner.run_async(
-                            user_id=req.user_id,
-                            session_id=req.session_id,
-                            new_message=req.new_message,
-                            state_delta=req.state_delta,
-                            run_config=RunConfig(streaming_mode=stream_mode),
-                            invocation_id=req.invocation_id,
-                        )
-                    ) as agen:
-                        async for event in agen:
-                            # ADK Web renders artifacts from `actions.artifactDelta`
-                            # during part processing *and* during action processing
-                            # 1) the original event with `artifactDelta` cleared (content)
-                            # 2) a content-less "action-only" event carrying `artifactDelta`
-                            events_to_stream = [event]
-                            if (
-                                event.actions.artifact_delta
-                                and event.content
-                                and event.content.parts
-                            ):
-                                content_event = event.model_copy(deep=True)
-                                content_event.actions.artifact_delta = {}
-                                artifact_event = event.model_copy(deep=True)
-                                artifact_event.content = None
-                                events_to_stream = [
-                                    content_event,
-                                    artifact_event,
-                                ]
-
-                            for event_to_stream in events_to_stream:
-                                sse_event = event_to_stream.model_dump_json(
-                                    exclude_none=True,
-                                    by_alias=True,
-                                )
-                                logger.debug(
-                                    "Generated event in agent run streaming: %s",
-                                    sse_event,
-                                )
-                                yield f"data: {sse_event}\n\n"
-                except Exception as e:
-                    logger.exception("Error in event_generator: %s", e)
-                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
-
-                # Returns a streaming response with the proper media type for SSE
-
-            return StreamingResponse(
-                event_generator(),
-                media_type="text/event-stream",
-            )
-
         # Attach ASGI middleware for unified telemetry across all routes
         self.app.add_middleware(AgentkitTelemetryHTTPMiddleware)
 

agentkit/frameworks/__init__.py

@@ -11,14 +11,20 @@
 
 __all__ = [
     "FrameworkBridgeError",
+    "BedrockAgentCoreAgentkitBridge",
     "LangChainAgentkitBridge",
     "LangGraphAgentkitBridge",
+    "StrandsAgentkitBridge",
     "UnsupportedFrameworkAgentError",
     "load_entry_object",
 ]
 
 
 def __getattr__(name: str) -> Any:
+    if name == "BedrockAgentCoreAgentkitBridge":
+        from agentkit.frameworks.agentcore import BedrockAgentCoreAgentkitBridge
+
+        return BedrockAgentCoreAgentkitBridge
     if name == "LangChainAgentkitBridge":
         from agentkit.frameworks.langchain import LangChainAgentkitBridge
 
@@ -27,6 +33,10 @@ def __getattr__(name: str) -> Any:
         from agentkit.frameworks.langgraph import LangGraphAgentkitBridge
 
         return LangGraphAgentkitBridge
+    if name == "StrandsAgentkitBridge":
+        from agentkit.frameworks.strands import StrandsAgentkitBridge
+
+        return StrandsAgentkitBridge
     if name == "load_entry_object":
         from agentkit.frameworks.migration import load_entry_object