adk_server/rest/controllers/
debug.rs

1use crate::ServerConfig;
2use axum::{
3    Json,
4    extract::{Path, State},
5    http::StatusCode,
6};
7use serde::Serialize;
8use std::collections::HashMap;
9
10#[derive(Clone)]
11pub struct DebugController {
12    config: ServerConfig,
13}
14
15impl DebugController {
16    pub fn new(config: ServerConfig) -> Self {
17        Self { config }
18    }
19}
20
21#[derive(Serialize)]
22pub struct GraphResponse {
23    #[serde(rename = "dotSrc")]
24    pub dot_src: String,
25}
26
27// ADK-Go compatible trace response (attributes map)
28pub async fn get_trace_by_event_id(
29    State(controller): State<DebugController>,
30    Path(event_id): Path<String>,
31) -> Result<Json<HashMap<String, String>>, StatusCode> {
32    if let Some(exporter) = &controller.config.span_exporter {
33        // First try direct lookup by event_id
34        if let Some(attributes) = exporter.get_trace_by_event_id(&event_id) {
35            return Ok(Json(attributes));
36        }
37
38        // If not found, search through all spans for matching session event ID
39        let trace_dict = exporter.get_trace_dict();
40        for (_, attributes) in trace_dict.iter() {
41            // Check if any span has this event_id in its attributes
42            if attributes.values().any(|v| v == &event_id) {
43                return Ok(Json(attributes.clone()));
44            }
45        }
46    }
47
48    Err(StatusCode::NOT_FOUND)
49}
50
51// Convert ADK exporter format to UI-compatible SpanData format
52// Field names must match adk-web Trace.ts interface exactly
53fn convert_to_span_data(attributes: &HashMap<String, String>) -> serde_json::Value {
54    let start_time: u64 = attributes.get("start_time").and_then(|s| s.parse().ok()).unwrap_or(0);
55    let end_time: u64 = attributes.get("end_time").and_then(|s| s.parse().ok()).unwrap_or(0);
56
57    // Build JSON object - omit parent_span_id entirely to prevent nesting
58    let mut obj = serde_json::json!({
59        "name": attributes.get("span_name").map_or("unknown", |v| v.as_str()),
60        "span_id": attributes.get("span_id").map_or("", |v| v.as_str()),
61        "trace_id": attributes.get("trace_id").map_or("", |v| v.as_str()),
62        "start_time": start_time,
63        "end_time": end_time,
64        "attributes": attributes,
65        "invoc_id": attributes.get("gcp.vertex.agent.invocation_id").map_or("", |v| v.as_str())
66    });
67
68    // Add LLM request/response if present (for UI display)
69    if let Some(llm_req) = attributes.get("gcp.vertex.agent.llm_request") {
70        obj["gcp.vertex.agent.llm_request"] = serde_json::Value::String(llm_req.clone());
71    }
72    if let Some(llm_resp) = attributes.get("gcp.vertex.agent.llm_response") {
73        obj["gcp.vertex.agent.llm_response"] = serde_json::Value::String(llm_resp.clone());
74    }
75
76    obj
77}
78
79// Get all spans for a session (UI-compatible format)
80pub async fn get_session_traces(
81    State(controller): State<DebugController>,
82    Path(session_id): Path<String>,
83) -> Result<Json<Vec<serde_json::Value>>, StatusCode> {
84    if let Some(exporter) = &controller.config.span_exporter {
85        let traces = exporter.get_session_trace(&session_id);
86        let span_data: Vec<serde_json::Value> = traces.iter().map(convert_to_span_data).collect();
87        return Ok(Json(span_data));
88    }
89
90    Ok(Json(Vec::new()))
91}
92
93pub async fn get_graph(
94    State(_controller): State<DebugController>,
95    Path((_app_name, _user_id, _session_id, _event_id)): Path<(String, String, String, String)>,
96) -> Result<Json<GraphResponse>, StatusCode> {
97    // Stub: Return a simple DOT graph
98    let dot_src = "digraph G { Agent -> User [label=\"response\"]; }".to_string();
99    Ok(Json(GraphResponse { dot_src }))
100}
101
102/// Get evaluation sets for an app (stub - returns empty array)
103pub async fn get_eval_sets(
104    State(_controller): State<DebugController>,
105    Path(_app_name): Path<String>,
106) -> Result<Json<Vec<serde_json::Value>>, StatusCode> {
107    // Stub: Return empty array - eval sets not yet implemented
108    Ok(Json(Vec::new()))
109}
110
111/// Get event data by event_id - returns event with invocationId for trace linking
112pub async fn get_event(
113    State(controller): State<DebugController>,
114    Path((app_name, _user_id, session_id, event_id)): Path<(String, String, String, String)>,
115) -> Result<Json<serde_json::Value>, StatusCode> {
116    // Try to find trace data for this event_id
117    if let Some(exporter) = &controller.config.span_exporter {
118        let traces = exporter.get_session_trace(&session_id);
119
120        // Find a trace with matching event_id
121        for attrs in traces {
122            if let Some(stored_event_id) = attrs.get("gcp.vertex.agent.event_id") {
123                if stored_event_id == &event_id {
124                    // Found matching trace - return event-like structure
125                    let invocation_id =
126                        attrs.get("gcp.vertex.agent.invocation_id").cloned().unwrap_or_default();
127
128                    return Ok(Json(serde_json::json!({
129                        "id": event_id,
130                        "invocationId": invocation_id,
131                        "appName": app_name,
132                        "sessionId": session_id,
133                        "attributes": attrs,
134                        "gcp.vertex.agent.llm_request": attrs.get("gcp.vertex.agent.llm_request"),
135                        "gcp.vertex.agent.llm_response": attrs.get("gcp.vertex.agent.llm_response")
136                    })));
137                }
138            }
139        }
140    }
141
142    // Event not found - return a minimal stub to prevent UI errors
143    Ok(Json(serde_json::json!({
144        "id": event_id,
145        "invocationId": "",
146        "appName": app_name,
147        "sessionId": session_id
148    })))
149}
adk_server/rest/controllers/debug.rs

adk_server/rest/controllers/
debug.rs
