fix: override run_sse endpoint

(cherry picked from commit 71cb03ab3cad01bef50986fef7ddf162f8f3d2ef)

Author: feng-95

agentkit/apps/agent_server_app/agent_server_app.py

@@ -31,7 +31,7 @@
 from google.adk.auth.credential_service.in_memory_credential_service import (
     InMemoryCredentialService,
 )
-from google.adk.cli.adk_web_server import AdkWebServer
+from google.adk.cli.adk_web_server import AdkWebServer, RunAgentRequest
 from google.adk.cli.utils.base_agent_loader import BaseAgentLoader
 from google.adk.evaluation.local_eval_set_results_manager import (
     LocalEvalSetResultsManager,
@@ -149,6 +149,79 @@ async def lifespan(app: FastAPI):
 
         self.app = self.server.get_fast_api_app(lifespan=lifespan)
 
+        @self.app.post("/run_sse")
+        async def run_agent_sse(req: RunAgentRequest) -> StreamingResponse:
+            logger.info("Overriding run_agent_sse endpoint...")
+            # 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)
+                    telemetry.trace_agent_server_finish(
+                        path="/invoke", func_result="", exception=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)