adk_server/rest/
mod.rs

1pub mod controllers;
2mod routes;
3
4pub use controllers::{
5    A2aController, AppsController, ArtifactsController, DebugController, RuntimeController,
6    SessionController,
7};
8
9use crate::{
10    ServerConfig,
11    auth_bridge::{RequestContext, RequestContextError, RequestContextExtractor},
12    web_ui,
13};
14use axum::{
15    Json, Router,
16    body::Body,
17    extract::{DefaultBodyLimit, State},
18    http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode, header},
19    middleware::{self, Next},
20    response::{IntoResponse, Response},
21    routing::{get, post},
22};
23use serde::Serialize;
24use std::sync::Arc;
25use tokio_util::sync::CancellationToken;
26use tower::ServiceBuilder;
27use tower_http::{
28    cors::{AllowOrigin, CorsLayer},
29    set_header::SetResponseHeaderLayer,
30    timeout::TimeoutLayer,
31    trace::TraceLayer,
32};
33
34const REQUEST_ID_HEADER: &str = "x-request-id";
35
36#[derive(Clone)]
37struct HealthController {
38    session_service: Arc<dyn adk_session::SessionService>,
39    artifact_service: Option<Arc<dyn adk_artifact::ArtifactService>>,
40    memory_service: Option<Arc<dyn adk_core::Memory>>,
41}
42
43impl HealthController {
44    fn new(config: &ServerConfig) -> Self {
45        Self {
46            session_service: config.session_service.clone(),
47            artifact_service: config.artifact_service.clone(),
48            memory_service: config.memory_service.clone(),
49        }
50    }
51}
52
53#[derive(Clone, Debug)]
54struct RequestId(String);
55
56impl RequestId {
57    fn as_str(&self) -> &str {
58        &self.0
59    }
60}
61
62#[derive(Serialize)]
63#[serde(rename_all = "camelCase")]
64struct HealthResponse {
65    status: &'static str,
66    components: HealthComponents,
67}
68
69#[derive(Serialize)]
70#[serde(rename_all = "camelCase")]
71struct HealthComponents {
72    session: ComponentHealth,
73    memory: ComponentHealth,
74    artifact: ComponentHealth,
75}
76
77#[derive(Serialize)]
78#[serde(rename_all = "camelCase")]
79struct ComponentHealth {
80    status: &'static str,
81    #[serde(skip_serializing_if = "Option::is_none")]
82    error: Option<String>,
83}
84
85impl ComponentHealth {
86    fn healthy() -> Self {
87        Self { status: "healthy", error: None }
88    }
89
90    fn unhealthy(error: impl Into<String>) -> Self {
91        Self { status: "unhealthy", error: Some(error.into()) }
92    }
93
94    fn not_configured() -> Self {
95        Self { status: "not_configured", error: None }
96    }
97}
98
99/// Build CORS layer based on security configuration
100fn build_cors_layer(config: &ServerConfig) -> CorsLayer {
101    let cors = CorsLayer::new()
102        .allow_methods([Method::GET, Method::POST, Method::PUT, Method::DELETE, Method::OPTIONS])
103        .allow_headers([
104            header::CONTENT_TYPE,
105            header::AUTHORIZATION,
106            HeaderName::from_static(REQUEST_ID_HEADER),
107            HeaderName::from_static("x-adk-ui-protocol"),
108            HeaderName::from_static("x-adk-ui-transport"),
109        ]);
110
111    if config.security.allowed_origins.is_empty() {
112        cors.allow_origin(AllowOrigin::any())
113    } else {
114        let origins: Vec<HeaderValue> = config
115            .security
116            .allowed_origins
117            .iter()
118            .filter_map(|origin| origin.parse().ok())
119            .collect();
120        cors.allow_origin(origins)
121    }
122}
123
124fn validate_request_id(headers: &HeaderMap) -> Option<String> {
125    let value = headers.get(REQUEST_ID_HEADER)?;
126    let raw = value.to_str().ok()?;
127    if raw.len() > 128 {
128        return None;
129    }
130    uuid::Uuid::parse_str(raw).ok()?;
131    Some(raw.to_string())
132}
133
134async fn request_id_middleware(mut request: Request<Body>, next: Next) -> Response {
135    let request_id =
136        validate_request_id(request.headers()).unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
137
138    request.extensions_mut().insert(RequestId(request_id.clone()));
139
140    let mut response = next.run(request).await;
141    if let Ok(value) = HeaderValue::from_str(&request_id) {
142        response.headers_mut().insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
143    }
144    response
145}
146
147async fn auth_middleware(
148    request: Request<Body>,
149    next: Next,
150    extractor: Option<Arc<dyn RequestContextExtractor>>,
151) -> Response {
152    let (mut parts, body) = request.into_parts();
153
154    let request_context = match extractor {
155        Some(extractor) => match extractor.extract(&parts).await {
156            Ok(context) => Some(context),
157            Err(RequestContextError::MissingAuth) => {
158                return (
159                    StatusCode::UNAUTHORIZED,
160                    Json(serde_json::json!({ "error": "missing authorization" })),
161                )
162                    .into_response();
163            }
164            Err(RequestContextError::InvalidToken(message)) => {
165                return (
166                    StatusCode::UNAUTHORIZED,
167                    Json(serde_json::json!({ "error": format!("invalid token: {message}") })),
168                )
169                    .into_response();
170            }
171            Err(RequestContextError::ExtractionFailed(message)) => {
172                return (
173                    StatusCode::INTERNAL_SERVER_ERROR,
174                    Json(serde_json::json!({
175                        "error": format!("auth extraction failed: {message}")
176                    })),
177                )
178                    .into_response();
179            }
180        },
181        None => None,
182    };
183
184    parts.extensions.insert::<Option<RequestContext>>(request_context);
185    next.run(Request::from_parts(parts, body)).await
186}
187
188async fn health_check(State(controller): State<HealthController>) -> impl IntoResponse {
189    let session = match controller.session_service.health_check().await {
190        Ok(()) => ComponentHealth::healthy(),
191        Err(error) => ComponentHealth::unhealthy(error.to_string()),
192    };
193
194    let memory = match controller.memory_service.as_ref() {
195        Some(service) => match service.health_check().await {
196            Ok(()) => ComponentHealth::healthy(),
197            Err(error) => ComponentHealth::unhealthy(error.to_string()),
198        },
199        None => ComponentHealth::not_configured(),
200    };
201
202    let artifact = match controller.artifact_service.as_ref() {
203        Some(service) => match service.health_check().await {
204            Ok(()) => ComponentHealth::healthy(),
205            Err(error) => ComponentHealth::unhealthy(error.to_string()),
206        },
207        None => ComponentHealth::not_configured(),
208    };
209
210    let healthy = session.status == "healthy"
211        && memory.status != "unhealthy"
212        && artifact.status != "unhealthy";
213
214    (
215        if healthy { StatusCode::OK } else { StatusCode::SERVICE_UNAVAILABLE },
216        Json(HealthResponse {
217            status: if healthy { "healthy" } else { "unhealthy" },
218            components: HealthComponents { session, memory, artifact },
219        }),
220    )
221}
222
223/// Create the server application with optional A2A support
224pub fn create_app(config: ServerConfig) -> Router {
225    create_app_with_a2a(config, None)
226}
227
228/// Start hot reload watchers for configured YAML agent directories.
229///
230/// For each directory in `config.yaml_agent_dirs`, creates an
231/// [`AgentConfigLoader`](crate::yaml_agent::AgentConfigLoader) and
232/// [`HotReloadWatcher`](crate::yaml_agent::HotReloadWatcher), performs
233/// the initial load, and spawns a background watcher task.
234///
235/// Returns the list of active watchers so route handlers can look up
236/// YAML-defined agents.
237#[cfg(feature = "yaml-agent")]
238async fn start_yaml_agent_watchers(
239    dirs: &[std::path::PathBuf],
240) -> Vec<Arc<crate::yaml_agent::HotReloadWatcher>> {
241    use crate::yaml_agent::{AgentConfigLoader, HotReloadWatcher};
242
243    let mut watchers = Vec::new();
244
245    for dir in dirs {
246        // Create a minimal tool registry (no pre-registered tools) and a
247        // placeholder model factory. Real deployments should configure these
248        // via ServerConfig extensions; for now we use empty defaults so the
249        // watcher can start and load YAML definitions.
250        let registry: Arc<dyn adk_core::ToolRegistry> = Arc::new(EmptyToolRegistry);
251        let factory: Arc<dyn crate::yaml_agent::ModelFactory> = Arc::new(NoOpModelFactory);
252        let loader = Arc::new(AgentConfigLoader::new(registry, factory));
253        let watcher = Arc::new(HotReloadWatcher::new(loader));
254
255        match watcher.watch(dir).await {
256            Ok(handle) => {
257                tracing::info!("started YAML agent hot reload watcher for {}", dir.display());
258                // Detach the watcher task — it runs until the server shuts down.
259                drop(handle);
260                watchers.push(watcher);
261            }
262            Err(e) => {
263                tracing::warn!("failed to start YAML agent watcher for {}: {e}", dir.display());
264            }
265        }
266    }
267
268    watchers
269}
270
271/// Empty tool registry used as default when no tools are pre-registered.
272#[cfg(feature = "yaml-agent")]
273struct EmptyToolRegistry;
274
275#[cfg(feature = "yaml-agent")]
276impl adk_core::ToolRegistry for EmptyToolRegistry {
277    fn resolve(&self, _tool_name: &str) -> Option<Arc<dyn adk_core::Tool>> {
278        None
279    }
280
281    fn available_tools(&self) -> Vec<String> {
282        vec![]
283    }
284}
285
286/// Placeholder model factory that returns an error for any provider.
287///
288/// Real deployments should provide a proper `ModelFactory` via
289/// `ServerConfig` extensions. This exists so the watcher can start
290/// even when no model factory is explicitly configured.
291#[cfg(feature = "yaml-agent")]
292struct NoOpModelFactory;
293
294#[cfg(feature = "yaml-agent")]
295#[async_trait::async_trait]
296impl crate::yaml_agent::ModelFactory for NoOpModelFactory {
297    async fn create_model(
298        &self,
299        provider: &str,
300        model_id: &str,
301    ) -> adk_core::Result<Arc<dyn adk_core::Llm>> {
302        Err(adk_core::AdkError::config(format!(
303            "no model factory configured for YAML agent loading \
304             (requested provider='{provider}', model_id='{model_id}'). \
305             Configure a ModelFactory on the server to enable YAML agent model creation."
306        )))
307    }
308}
309
310/// Create the server application with A2A support at the specified base URL
311pub fn create_app_with_a2a(config: ServerConfig, a2a_base_url: Option<&str>) -> Router {
312    let session_controller = SessionController::new(config.session_service.clone());
313    let runtime_controller = RuntimeController::new(config.clone());
314    let apps_controller = AppsController::new(config.clone());
315    let artifacts_controller = ArtifactsController::new(config.clone());
316    let debug_controller = DebugController::new(config.clone());
317    let health_controller = HealthController::new(&config);
318
319    // Start YAML agent hot reload watchers if configured.
320    #[cfg(feature = "yaml-agent")]
321    {
322        let dirs = config.yaml_agent_dirs.clone();
323        if !dirs.is_empty() {
324            tokio::spawn(async move {
325                let _watchers = start_yaml_agent_watchers(&dirs).await;
326                // Watchers are kept alive for the lifetime of this task.
327                // They run background filesystem watch loops internally.
328                // We hold them here so they aren't dropped.
329                std::future::pending::<()>().await;
330            });
331        }
332    }
333
334    let auth_layer = middleware::from_fn({
335        let extractor = config.request_context_extractor.clone();
336        move |request: Request<Body>, next: Next| {
337            let extractor = extractor.clone();
338            async move { auth_middleware(request, next, extractor).await }
339        }
340    });
341
342    let health_router =
343        Router::new().route("/health", get(health_check)).with_state(health_controller);
344
345    let ui_api_router = Router::new()
346        .route("/apps", get(controllers::apps::list_apps))
347        .route("/list-apps", get(controllers::apps::list_apps_compat))
348        .with_state(apps_controller)
349        .route("/ui/capabilities", get(controllers::ui::ui_capabilities))
350        .route("/ui/initialize", post(controllers::ui::ui_initialize))
351        .route("/ui/message", post(controllers::ui::ui_message))
352        .route("/ui/update-model-context", post(controllers::ui::ui_update_model_context))
353        .route("/ui/notifications/poll", post(controllers::ui::ui_poll_notifications))
354        .route(
355            "/ui/notifications/resources-list-changed",
356            post(controllers::ui::ui_notify_resources_list_changed),
357        )
358        .route(
359            "/ui/notifications/tools-list-changed",
360            post(controllers::ui::ui_notify_tools_list_changed),
361        )
362        .route("/ui/resources", get(controllers::ui::list_ui_resources))
363        .route("/ui/resources/read", get(controllers::ui::read_ui_resource))
364        .route("/ui/resources/register", post(controllers::ui::register_ui_resource))
365        // These routes mutate and read shared bridge and resource state, so they
366        // carry the same authentication as the session, artifact, and debug routers.
367        .layer(auth_layer.clone());
368
369    let session_router = Router::new()
370        .route("/sessions", post(controllers::session::create_session))
371        .route(
372            "/sessions/{app_name}/{user_id}/{session_id}",
373            get(controllers::session::get_session).delete(controllers::session::delete_session),
374        )
375        .route(
376            "/apps/{app_name}/users/{user_id}/sessions",
377            get(controllers::session::list_sessions)
378                .post(controllers::session::create_session_from_path),
379        )
380        .route(
381            "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
382            get(controllers::session::get_session_from_path)
383                .post(controllers::session::create_session_from_path)
384                .delete(controllers::session::delete_session_from_path),
385        )
386        .with_state(session_controller)
387        .layer(auth_layer.clone());
388
389    let runtime_router = Router::new()
390        .route("/run/{app_name}/{user_id}/{session_id}", post(controllers::runtime::run_sse))
391        .route("/run_sse", post(controllers::runtime::run_sse_compat))
392        .with_state(runtime_controller);
393
394    let artifacts_router = Router::new()
395        .route(
396            "/sessions/{app_name}/{user_id}/{session_id}/artifacts",
397            get(controllers::artifacts::list_artifacts),
398        )
399        .route(
400            "/sessions/{app_name}/{user_id}/{session_id}/artifacts/{artifact_name}",
401            get(controllers::artifacts::get_artifact),
402        )
403        .with_state(artifacts_controller)
404        .layer(auth_layer.clone());
405
406    let mut debug_router = Router::new()
407        .route("/debug/trace/session/{session_id}", get(controllers::debug::get_session_traces))
408        .route(
409            "/debug/graph/{app_name}/{user_id}/{session_id}/{event_id}",
410            get(controllers::debug::get_graph),
411        )
412        .route(
413            "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
414            get(controllers::debug::get_graph),
415        )
416        .route("/apps/{app_name}/eval_sets", get(controllers::debug::get_eval_sets))
417        .route(
418            "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}",
419            get(controllers::debug::get_event),
420        );
421
422    if config.request_context_extractor.is_none() || config.security.expose_admin_debug {
423        debug_router = debug_router
424            .route("/debug/trace/{event_id}", get(controllers::debug::get_trace_by_event_id));
425    }
426
427    let debug_router = debug_router.with_state(debug_controller.clone()).layer(auth_layer.clone());
428
429    let api_router = Router::new()
430        .merge(health_router)
431        .merge(ui_api_router)
432        .merge(session_router)
433        .merge(runtime_router)
434        .merge(artifacts_router)
435        .merge(debug_router);
436
437    let ui_router = Router::new()
438        .route("/", get(web_ui::root_redirect))
439        .route("/ui/", get(web_ui::serve_ui_index))
440        .route("/ui/assets/config/runtime-config.json", get(web_ui::serve_runtime_config))
441        .with_state(config.clone())
442        .route("/ui/{*path}", get(web_ui::serve_ui_assets));
443
444    let mut app = Router::new().nest("/api", api_router).merge(ui_router);
445
446    if let Some(base_url) = a2a_base_url {
447        let a2a_controller = A2aController::new(config.clone(), base_url);
448        // Discovery stays public — an agent card is meant to be fetched by unknown peers.
449        // The JSON-RPC routes execute agent and tool work, so they carry the same
450        // authentication as every other mutation surface. They were previously merged at the
451        // root, outside the layer applied to `/api`, so anyone who could reach the port could
452        // drive the agent and incur its costs.
453        let a2a_discovery = Router::new()
454            .route("/.well-known/agent.json", get(controllers::a2a::get_agent_card))
455            .with_state(a2a_controller.clone());
456        let a2a_rpc = Router::new()
457            .route("/a2a", post(controllers::a2a::handle_jsonrpc))
458            .route("/a2a/stream", post(controllers::a2a::handle_jsonrpc_stream))
459            .with_state(a2a_controller)
460            .layer(auth_layer.clone());
461        app = app.merge(a2a_discovery).merge(a2a_rpc);
462    }
463
464    let cors_layer = build_cors_layer(&config);
465    let trace_layer = TraceLayer::new_for_http().make_span_with(|request: &Request<Body>| {
466        let request_id =
467            request.extensions().get::<RequestId>().map(RequestId::as_str).unwrap_or("");
468        tracing::info_span!(
469            "http.request",
470            request.id = %request_id,
471            http.method = %request.method(),
472            http.path = %request.uri().path()
473        )
474    });
475
476    app.layer(
477        ServiceBuilder::new()
478            .layer(middleware::from_fn(request_id_middleware))
479            .layer(trace_layer)
480            .layer(TimeoutLayer::with_status_code(
481                StatusCode::REQUEST_TIMEOUT,
482                config.security.request_timeout,
483            ))
484            .layer(DefaultBodyLimit::max(config.security.max_body_size))
485            .layer(cors_layer)
486            .layer(SetResponseHeaderLayer::if_not_present(
487                header::X_CONTENT_TYPE_OPTIONS,
488                HeaderValue::from_static("nosniff"),
489            ))
490            .layer(SetResponseHeaderLayer::if_not_present(
491                header::X_FRAME_OPTIONS,
492                HeaderValue::from_static("DENY"),
493            ))
494            .layer(SetResponseHeaderLayer::if_not_present(
495                header::X_XSS_PROTECTION,
496                HeaderValue::from_static("1; mode=block"),
497            )),
498    )
499}
500
501// ---------------------------------------------------------------------------
502// ServerBuilder — extensible server construction with custom routes
503// ---------------------------------------------------------------------------
504
505/// Builder for constructing an ADK server with custom routes.
506///
507/// `ServerBuilder` allows registering additional Axum routers alongside the
508/// built-in REST, A2A, and UI routes. Custom routes benefit from the same
509/// middleware stack (auth, CORS, tracing, timeout, security headers) as the
510/// built-in routes.
511///
512/// # Example
513///
514/// ```rust,ignore
515/// use adk_server::{ServerBuilder, ServerConfig};
516/// use axum::{Router, routing::get};
517///
518/// let config = ServerConfig::new(agent, session_service);
519///
520/// let app = ServerBuilder::new(config)
521///     .add_api_routes(
522///         Router::new()
523///             .route("/projects", get(list_projects))
524///             .route("/projects/{id}", get(get_project))
525///     )
526///     .add_api_routes(
527///         Router::new()
528///             .route("/automations", get(list_automations))
529///     )
530///     .with_a2a("http://localhost:8080")
531///     .build();
532///
533/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
534/// axum::serve(listener, app).await?;
535/// ```
536pub struct ServerBuilder {
537    config: ServerConfig,
538    a2a_base_url: Option<String>,
539    api_routes: Vec<Router>,
540    root_routes: Vec<Router>,
541    shutdown_endpoint: bool,
542}
543
544impl ServerBuilder {
545    /// Create a new server builder with the given configuration.
546    pub fn new(config: ServerConfig) -> Self {
547        Self {
548            config,
549            a2a_base_url: None,
550            api_routes: Vec::new(),
551            root_routes: Vec::new(),
552            shutdown_endpoint: false,
553        }
554    }
555
556    /// Add custom routes nested under `/api`.
557    ///
558    /// These routes are merged into the API router and benefit from the auth
559    /// middleware layer. Multiple calls accumulate routes.
560    ///
561    /// # Example
562    ///
563    /// ```rust,ignore
564    /// builder.add_api_routes(
565    ///     Router::new()
566    ///         .route("/projects", get(list_projects))
567    ///         .route("/projects/{id}", get(get_project))
568    /// )
569    /// ```
570    pub fn add_api_routes(mut self, routes: Router) -> Self {
571        self.api_routes.push(routes);
572        self
573    }
574
575    /// Add custom routes at the root level (not nested under `/api`).
576    ///
577    /// These routes are merged at the top level of the application, alongside
578    /// the UI and A2A routes. They receive the full middleware stack (CORS,
579    /// tracing, timeout, security headers) but NOT the auth middleware.
580    ///
581    /// Use this for routes that need their own auth handling or public endpoints.
582    pub fn add_root_routes(mut self, routes: Router) -> Self {
583        self.root_routes.push(routes);
584        self
585    }
586
587    /// Enable A2A protocol support at the specified base URL.
588    ///
589    /// The base URL is used to construct the agent card's endpoint URL.
590    pub fn with_a2a(mut self, base_url: impl Into<String>) -> Self {
591        self.a2a_base_url = Some(base_url.into());
592        self
593    }
594
595    /// Enable the `POST /api/shutdown` endpoint for graceful shutdown.
596    ///
597    /// When enabled, the server exposes a shutdown endpoint that triggers
598    /// graceful shutdown: stops accepting new connections, completes in-flight
599    /// requests, and then exits. Use [`build_with_shutdown`](Self::build_with_shutdown)
600    /// to get the [`ShutdownHandle`] for wiring into `axum::serve().with_graceful_shutdown()`.
601    ///
602    /// The endpoint is protected by the auth middleware when a
603    /// `RequestContextExtractor` is configured.
604    pub fn enable_shutdown_endpoint(mut self) -> Self {
605        self.shutdown_endpoint = true;
606        self
607    }
608
609    /// Build the final Axum router with all routes and middleware applied.
610    pub fn build(self) -> Router {
611        self.build_inner().0
612    }
613
614    /// Build the final Axum router and return a [`ShutdownHandle`].
615    ///
616    /// Use this when [`enable_shutdown_endpoint()`](Self::enable_shutdown_endpoint) is set.
617    /// Pass the handle's signal to `axum::serve().with_graceful_shutdown()`.
618    ///
619    /// # Example
620    ///
621    /// ```rust,ignore
622    /// let (app, shutdown_handle) = ServerBuilder::new(config)
623    ///     .enable_shutdown_endpoint()
624    ///     .build_with_shutdown();
625    ///
626    /// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
627    /// axum::serve(listener, app)
628    ///     .with_graceful_shutdown(shutdown_handle.signal())
629    ///     .await?;
630    /// ```
631    pub fn build_with_shutdown(self) -> (Router, ShutdownHandle) {
632        let (router, handle) = self.build_inner();
633        (router, handle.expect("build_with_shutdown requires enable_shutdown_endpoint()"))
634    }
635
636    fn build_inner(self) -> (Router, Option<ShutdownHandle>) {
637        let config = &self.config;
638        let session_controller = SessionController::new(config.session_service.clone());
639        let runtime_controller = RuntimeController::new(config.clone());
640        let apps_controller = AppsController::new(config.clone());
641        let artifacts_controller = ArtifactsController::new(config.clone());
642        let debug_controller = DebugController::new(config.clone());
643        let health_controller = HealthController::new(config);
644
645        // Start YAML agent hot reload watchers if configured.
646        #[cfg(feature = "yaml-agent")]
647        {
648            let dirs = config.yaml_agent_dirs.clone();
649            if !dirs.is_empty() {
650                tokio::spawn(async move {
651                    let _watchers = start_yaml_agent_watchers(&dirs).await;
652                    std::future::pending::<()>().await;
653                });
654            }
655        }
656
657        let auth_layer = middleware::from_fn({
658            let extractor = config.request_context_extractor.clone();
659            move |request: Request<Body>, next: Next| {
660                let extractor = extractor.clone();
661                async move { auth_middleware(request, next, extractor).await }
662            }
663        });
664
665        let health_router =
666            Router::new().route("/health", get(health_check)).with_state(health_controller);
667
668        let ui_api_router = Router::new()
669            .route("/apps", get(controllers::apps::list_apps))
670            .route("/list-apps", get(controllers::apps::list_apps_compat))
671            .with_state(apps_controller)
672            .route("/ui/capabilities", get(controllers::ui::ui_capabilities))
673            .route("/ui/initialize", post(controllers::ui::ui_initialize))
674            .route("/ui/message", post(controllers::ui::ui_message))
675            .route("/ui/update-model-context", post(controllers::ui::ui_update_model_context))
676            .route("/ui/notifications/poll", post(controllers::ui::ui_poll_notifications))
677            .route(
678                "/ui/notifications/resources-list-changed",
679                post(controllers::ui::ui_notify_resources_list_changed),
680            )
681            .route(
682                "/ui/notifications/tools-list-changed",
683                post(controllers::ui::ui_notify_tools_list_changed),
684            )
685            .route("/ui/resources", get(controllers::ui::list_ui_resources))
686            .route("/ui/resources/read", get(controllers::ui::read_ui_resource))
687            .route("/ui/resources/register", post(controllers::ui::register_ui_resource))
688            // These routes mutate and read shared bridge and resource state, so they
689            // carry the same authentication as the session, artifact, and debug routers.
690            .layer(auth_layer.clone());
691
692        let session_router = Router::new()
693            .route("/sessions", post(controllers::session::create_session))
694            .route(
695                "/sessions/{app_name}/{user_id}/{session_id}",
696                get(controllers::session::get_session).delete(controllers::session::delete_session),
697            )
698            .route(
699                "/apps/{app_name}/users/{user_id}/sessions",
700                get(controllers::session::list_sessions)
701                    .post(controllers::session::create_session_from_path),
702            )
703            .route(
704                "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
705                get(controllers::session::get_session_from_path)
706                    .post(controllers::session::create_session_from_path)
707                    .delete(controllers::session::delete_session_from_path),
708            )
709            .with_state(session_controller)
710            .layer(auth_layer.clone());
711
712        let runtime_router = Router::new()
713            .route("/run/{app_name}/{user_id}/{session_id}", post(controllers::runtime::run_sse))
714            .route("/run_sse", post(controllers::runtime::run_sse_compat))
715            .with_state(runtime_controller);
716
717        let artifacts_router = Router::new()
718            .route(
719                "/sessions/{app_name}/{user_id}/{session_id}/artifacts",
720                get(controllers::artifacts::list_artifacts),
721            )
722            .route(
723                "/sessions/{app_name}/{user_id}/{session_id}/artifacts/{artifact_name}",
724                get(controllers::artifacts::get_artifact),
725            )
726            .with_state(artifacts_controller)
727            .layer(auth_layer.clone());
728
729        let mut debug_router = Router::new()
730            .route("/debug/trace/session/{session_id}", get(controllers::debug::get_session_traces))
731            .route(
732                "/debug/graph/{app_name}/{user_id}/{session_id}/{event_id}",
733                get(controllers::debug::get_graph),
734            )
735            .route(
736                "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
737                get(controllers::debug::get_graph),
738            )
739            .route("/apps/{app_name}/eval_sets", get(controllers::debug::get_eval_sets))
740            .route(
741                "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}",
742                get(controllers::debug::get_event),
743            );
744
745        if config.request_context_extractor.is_none() || config.security.expose_admin_debug {
746            debug_router = debug_router
747                .route("/debug/trace/{event_id}", get(controllers::debug::get_trace_by_event_id));
748        }
749
750        let debug_router =
751            debug_router.with_state(debug_controller.clone()).layer(auth_layer.clone());
752
753        // Assemble the API router with built-in + custom routes
754        let mut api_router = Router::new()
755            .merge(health_router)
756            .merge(ui_api_router)
757            .merge(session_router)
758            .merge(runtime_router)
759            .merge(artifacts_router)
760            .merge(debug_router);
761
762        // Merge custom API routes — these get the same /api prefix and auth middleware
763        for custom_routes in self.api_routes {
764            api_router = api_router.merge(custom_routes.layer(auth_layer.clone()));
765        }
766
767        // Add shutdown endpoint if enabled
768        let shutdown_handle = if self.shutdown_endpoint {
769            let handle = ShutdownHandle::new();
770            let shutdown_router = Router::new()
771                .route("/shutdown", post(handle_shutdown))
772                .with_state(handle.token.clone())
773                .layer(auth_layer.clone());
774            api_router = api_router.merge(shutdown_router);
775            Some(handle)
776        } else {
777            None
778        };
779
780        let ui_router = Router::new()
781            .route("/", get(web_ui::root_redirect))
782            .route("/ui/", get(web_ui::serve_ui_index))
783            .route("/ui/assets/config/runtime-config.json", get(web_ui::serve_runtime_config))
784            .with_state(config.clone())
785            .route("/ui/{*path}", get(web_ui::serve_ui_assets));
786
787        let mut app = Router::new().nest("/api", api_router).merge(ui_router);
788
789        // Merge custom root routes
790        for custom_routes in self.root_routes {
791            app = app.merge(custom_routes);
792        }
793
794        if let Some(base_url) = &self.a2a_base_url {
795            let a2a_controller = A2aController::new(config.clone(), base_url);
796            // Same split as `create_app_with_a2a`: discovery is public, RPC is authenticated.
797            let a2a_discovery = Router::new()
798                .route("/.well-known/agent.json", get(controllers::a2a::get_agent_card))
799                .with_state(a2a_controller.clone());
800            let a2a_rpc = Router::new()
801                .route("/a2a", post(controllers::a2a::handle_jsonrpc))
802                .route("/a2a/stream", post(controllers::a2a::handle_jsonrpc_stream))
803                .with_state(a2a_controller)
804                .layer(auth_layer.clone());
805            app = app.merge(a2a_discovery).merge(a2a_rpc);
806        }
807
808        let cors_layer = build_cors_layer(config);
809        let trace_layer = TraceLayer::new_for_http().make_span_with(|request: &Request<Body>| {
810            let request_id =
811                request.extensions().get::<RequestId>().map(RequestId::as_str).unwrap_or("");
812            tracing::info_span!(
813                "http.request",
814                request.id = %request_id,
815                http.method = %request.method(),
816                http.path = %request.uri().path()
817            )
818        });
819
820        (
821            app.layer(
822                ServiceBuilder::new()
823                    .layer(middleware::from_fn(request_id_middleware))
824                    .layer(trace_layer)
825                    .layer(TimeoutLayer::with_status_code(
826                        StatusCode::REQUEST_TIMEOUT,
827                        config.security.request_timeout,
828                    ))
829                    .layer(DefaultBodyLimit::max(config.security.max_body_size))
830                    .layer(cors_layer)
831                    .layer(SetResponseHeaderLayer::if_not_present(
832                        header::X_CONTENT_TYPE_OPTIONS,
833                        HeaderValue::from_static("nosniff"),
834                    ))
835                    .layer(SetResponseHeaderLayer::if_not_present(
836                        header::X_FRAME_OPTIONS,
837                        HeaderValue::from_static("DENY"),
838                    ))
839                    .layer(SetResponseHeaderLayer::if_not_present(
840                        header::X_XSS_PROTECTION,
841                        HeaderValue::from_static("1; mode=block"),
842                    )),
843            ),
844            shutdown_handle,
845        )
846    }
847}
848
849/// Wait for a process shutdown signal.
850pub async fn shutdown_signal() {
851    let ctrl_c = async {
852        let _ = tokio::signal::ctrl_c().await;
853    };
854
855    #[cfg(unix)]
856    let terminate = async {
857        if let Ok(mut signal) =
858            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
859        {
860            let _ = signal.recv().await;
861        }
862    };
863
864    #[cfg(not(unix))]
865    let terminate = std::future::pending::<()>();
866
867    tokio::select! {
868        _ = ctrl_c => {}
869        _ = terminate => {}
870    }
871}
872
873// ---------------------------------------------------------------------------
874// ShutdownHandle — programmatic graceful shutdown trigger
875// ---------------------------------------------------------------------------
876
877/// Handle for triggering graceful server shutdown.
878///
879/// Returned by [`ServerBuilder::build_with_shutdown`]. Pass the future from
880/// [`ShutdownHandle::signal()`] to `axum::serve(...).with_graceful_shutdown()`
881/// to enable both OS signal-based and HTTP endpoint-based shutdown.
882///
883/// # Example
884///
885/// ```rust,ignore
886/// use adk_server::{ServerBuilder, ServerConfig};
887///
888/// let (app, shutdown_handle) = ServerBuilder::new(config)
889///     .enable_shutdown_endpoint()
890///     .build_with_shutdown();
891///
892/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
893/// axum::serve(listener, app)
894///     .with_graceful_shutdown(shutdown_handle.signal())
895///     .await?;
896/// ```
897#[derive(Clone)]
898pub struct ShutdownHandle {
899    token: CancellationToken,
900}
901
902impl ShutdownHandle {
903    /// Create a new shutdown handle.
904    fn new() -> Self {
905        Self { token: CancellationToken::new() }
906    }
907
908    /// Trigger graceful shutdown programmatically.
909    ///
910    /// This has the same effect as calling `POST /api/shutdown` — the server
911    /// stops accepting new connections and completes in-flight requests.
912    pub fn shutdown(&self) {
913        tracing::info!("graceful shutdown triggered programmatically");
914        self.token.cancel();
915    }
916
917    /// Returns a future that resolves when shutdown is triggered.
918    ///
919    /// Combines OS signals (Ctrl+C, SIGTERM) with the programmatic/HTTP trigger.
920    /// Pass this to `axum::serve(...).with_graceful_shutdown()`.
921    pub async fn signal(self) {
922        let token = self.token.clone();
923
924        let ctrl_c = async {
925            let _ = tokio::signal::ctrl_c().await;
926        };
927
928        #[cfg(unix)]
929        let terminate = async {
930            if let Ok(mut signal) =
931                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
932            {
933                let _ = signal.recv().await;
934            }
935        };
936
937        #[cfg(not(unix))]
938        let terminate = std::future::pending::<()>();
939
940        tokio::select! {
941            _ = ctrl_c => {
942                tracing::info!("received Ctrl+C, initiating graceful shutdown");
943            }
944            _ = terminate => {
945                tracing::info!("received SIGTERM, initiating graceful shutdown");
946            }
947            _ = token.cancelled() => {
948                // Shutdown triggered via POST /api/shutdown or programmatic call
949            }
950        }
951    }
952
953    /// Returns whether shutdown has been triggered.
954    pub fn is_shutdown(&self) -> bool {
955        self.token.is_cancelled()
956    }
957}
958
959/// Handler for `POST /api/shutdown`.
960///
961/// Triggers graceful shutdown: the server stops accepting new connections,
962/// completes in-flight requests, and then exits.
963async fn handle_shutdown(State(token): State<CancellationToken>) -> impl IntoResponse {
964    tracing::info!("POST /api/shutdown received, initiating graceful shutdown");
965    token.cancel();
966    (StatusCode::OK, Json(serde_json::json!({ "status": "shutting_down" })))
967}
adk_server/rest/mod.rs

adk_server/rest/
mod.rs
