fix: fix trace in webserver (#59)

Author: innsd

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,94 @@ 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:
+                e = HTTPException(status_code=404, detail="Session not found")
+                telemetry.trace_agent_server_finish(
+                    path="/run_sse", func_result="", exception=e
+                )
+                raise e
+
+            # 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="/run_sse", 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",
+            )
+
+        # Move the custom /run_sse route to the beginning of the routes list for priority matching (without deleting the ADK default route)
+        routes = self.app.router.routes
+        for i, r in enumerate(routes):
+            if (
+                    getattr(r, "path", None) == "/run_sse"
+                    and "POST" in getattr(r, "methods", set())
+                    and getattr(r,"endpoint", None) == run_agent_sse
+            ):
+                routes.insert(0, routes.pop(i))
+                break
+
         # Attach ASGI middleware for unified telemetry across all routes
         self.app.add_middleware(AgentkitTelemetryHTTPMiddleware)
 
@@ -163,13 +251,25 @@ async def _invoke_compat(request: Request):
                 for k, v in dict(headers).items()
                 if k.lower() not in {"authorization", "token"}
             }
+            # trace request attributes on current span
+            telemetry.trace_agent_server(
+                func_name="_invoke_compat",
+                span=span,
+                headers=telemetry_headers,
+                text="",
+            )
+
             user_id = headers.get("user_id") or "agentkit_user"
             session_id = headers.get("session_id") or ""
 
             # Determine app_name from loader
             app_names = self.server.agent_loader.list_agents()
             if not app_names:
-                raise HTTPException(status_code=404, detail="No agents configured")
+                exception = HTTPException(status_code=404, detail="No agents configured")
+                telemetry.trace_agent_server_finish(
+                    path="/invoke", func_result="", exception=exception
+                )
+                raise exception
             app_name = app_names[0]
 
             # Parse payload and convert to ADK Content
@@ -193,13 +293,6 @@ async def _invoke_compat(request: Request):
                         text = ""
             content = types.UserContent(parts=[types.Part(text=text or "")])
 
-            # trace request attributes on current span
-            telemetry.trace_agent_server(
-                func_name="_invoke_compat",
-                span=span,
-                headers=telemetry_headers,
-                text=text or "",
-            )
 
             # Ensure session exists
             session = await self.server.session_service.get_session(
@@ -232,10 +325,10 @@ async def event_generator():
                     # finish span on successful end of stream handled by middleware
                     pass
                 except Exception as e:
-                    yield f'data: {{"error": "{str(e)}"}}\n\n'
                     telemetry.trace_agent_server_finish(
                         path="/invoke", func_result="", exception=e
                     )
+                    yield f'data: {{"error": "{str(e)}"}}\n\n'
 
             return StreamingResponse(
                 event_generator(),

agentkit/apps/agent_server_app/telemetry.py

@@ -105,8 +105,10 @@ def trace_agent_server_finish(
             }
             if exception:
                 self.handle_exception(span, exception)
-                attributes["error_type"] = exception.__class__.__name__
-
+                if getattr(exception, "status_code", None):
+                    attributes["error_type"] = f"{exception.__class__.__name__}_{exception.status_code}"
+                else:
+                    attributes["error_type"] = exception.__class__.__name__
             # only record invoke request latency metrics
             if (
                 hasattr(span, "start_time")

agentkit/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.5.2"
+VERSION = "0.5.3"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agentkit-sdk-python"
-version = "0.5.2"
+version = "0.5.3"
 description = "Python SDK for transforming any AI agent into a production-ready application. Framework-agnostic primitives for runtime, memory, authentication, and tools with volcengine-managed infrastructure."
 readme = "README.md"
 requires-python = ">=3.10"