Skip to main content

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        .route("/ui/agents/{name}", get(controllers::apps::get_agent_details))
349        .with_state(apps_controller)
350        .route("/ui/capabilities", get(controllers::ui::ui_capabilities))
351        .route("/ui/initialize", post(controllers::ui::ui_initialize))
352        .route("/ui/message", post(controllers::ui::ui_message))
353        .route("/ui/update-model-context", post(controllers::ui::ui_update_model_context))
354        .route("/ui/notifications/poll", post(controllers::ui::ui_poll_notifications))
355        .route(
356            "/ui/notifications/resources-list-changed",
357            post(controllers::ui::ui_notify_resources_list_changed),
358        )
359        .route(
360            "/ui/notifications/tools-list-changed",
361            post(controllers::ui::ui_notify_tools_list_changed),
362        )
363        .route("/ui/resources", get(controllers::ui::list_ui_resources))
364        .route("/ui/resources/read", get(controllers::ui::read_ui_resource))
365        .route("/ui/resources/register", post(controllers::ui::register_ui_resource))
366        // These routes mutate and read shared bridge and resource state, so they
367        // carry the same authentication as the session, artifact, and debug routers.
368        .layer(auth_layer.clone());
369
370    let session_router = Router::new()
371        .route("/sessions", post(controllers::session::create_session))
372        .route(
373            "/sessions/{app_name}/{user_id}/{session_id}",
374            get(controllers::session::get_session).delete(controllers::session::delete_session),
375        )
376        .route(
377            "/apps/{app_name}/users/{user_id}/sessions",
378            get(controllers::session::list_sessions)
379                .post(controllers::session::create_session_from_path),
380        )
381        .route(
382            "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
383            get(controllers::session::get_session_from_path)
384                .post(controllers::session::create_session_from_path)
385                .delete(controllers::session::delete_session_from_path),
386        )
387        .with_state(session_controller)
388        .layer(auth_layer.clone());
389
390    let runtime_router = Router::new()
391        .route("/run", post(controllers::runtime::run_collect))
392        .route("/run/{app_name}/{user_id}/{session_id}", post(controllers::runtime::run_sse))
393        .route("/run_sse", post(controllers::runtime::run_sse_compat))
394        .with_state(runtime_controller);
395
396    let artifacts_router = Router::new()
397        .route(
398            "/sessions/{app_name}/{user_id}/{session_id}/artifacts",
399            get(controllers::artifacts::list_artifacts),
400        )
401        .route(
402            "/sessions/{app_name}/{user_id}/{session_id}/artifacts/{artifact_name}",
403            get(controllers::artifacts::get_artifact),
404        )
405        .with_state(artifacts_controller)
406        .layer(auth_layer.clone());
407
408    let mut debug_router = Router::new()
409        .route("/debug/trace/session/{session_id}", get(controllers::debug::get_session_traces))
410        .route(
411            "/debug/graph/{app_name}/{user_id}/{session_id}/{event_id}",
412            get(controllers::debug::get_graph),
413        )
414        .route(
415            "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
416            get(controllers::debug::get_graph),
417        )
418        .route("/apps/{app_name}/eval_sets", get(controllers::debug::get_eval_sets))
419        .route(
420            "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}",
421            get(controllers::debug::get_event),
422        );
423
424    if config.request_context_extractor.is_none() || config.security.expose_admin_debug {
425        debug_router = debug_router
426            .route("/debug/trace/{event_id}", get(controllers::debug::get_trace_by_event_id));
427    }
428
429    let debug_router = debug_router.with_state(debug_controller.clone()).layer(auth_layer.clone());
430
431    let api_router = Router::new()
432        .merge(health_router)
433        .merge(ui_api_router)
434        .merge(session_router)
435        .merge(runtime_router)
436        .merge(artifacts_router)
437        .merge(debug_router);
438
439    let ui_router = Router::new()
440        .route("/", get(web_ui::root_redirect))
441        .route("/ui/", get(web_ui::serve_ui_index))
442        .route("/ui/assets/config/runtime-config.json", get(web_ui::serve_runtime_config))
443        .with_state(config.clone())
444        .route("/ui/{*path}", get(web_ui::serve_ui_assets));
445
446    let mut app = Router::new().nest("/api", api_router).merge(ui_router);
447
448    if let Some(base_url) = a2a_base_url {
449        let a2a_controller = A2aController::new(config.clone(), base_url);
450        // Discovery stays public — an agent card is meant to be fetched by unknown peers.
451        // The JSON-RPC routes execute agent and tool work, so they carry the same
452        // authentication as every other mutation surface. They were previously merged at the
453        // root, outside the layer applied to `/api`, so anyone who could reach the port could
454        // drive the agent and incur its costs.
455        let a2a_discovery = Router::new()
456            .route("/.well-known/agent.json", get(controllers::a2a::get_agent_card))
457            .with_state(a2a_controller.clone());
458        let a2a_rpc = Router::new()
459            .route("/a2a", post(controllers::a2a::handle_jsonrpc))
460            .route("/a2a/stream", post(controllers::a2a::handle_jsonrpc_stream))
461            .with_state(a2a_controller)
462            .layer(auth_layer.clone());
463        app = app.merge(a2a_discovery).merge(a2a_rpc);
464    }
465
466    let cors_layer = build_cors_layer(&config);
467    let trace_layer = TraceLayer::new_for_http().make_span_with(|request: &Request<Body>| {
468        let request_id =
469            request.extensions().get::<RequestId>().map(RequestId::as_str).unwrap_or("");
470        tracing::info_span!(
471            "http.request",
472            request.id = %request_id,
473            http.method = %request.method(),
474            http.path = %request.uri().path()
475        )
476    });
477
478    app.layer(
479        ServiceBuilder::new()
480            .layer(middleware::from_fn(request_id_middleware))
481            .layer(trace_layer)
482            .layer(TimeoutLayer::with_status_code(
483                StatusCode::REQUEST_TIMEOUT,
484                config.security.request_timeout,
485            ))
486            .layer(DefaultBodyLimit::max(config.security.max_body_size))
487            .layer(cors_layer)
488            .layer(SetResponseHeaderLayer::if_not_present(
489                header::X_CONTENT_TYPE_OPTIONS,
490                HeaderValue::from_static("nosniff"),
491            ))
492            .layer(SetResponseHeaderLayer::if_not_present(
493                header::X_FRAME_OPTIONS,
494                HeaderValue::from_static("DENY"),
495            ))
496            .layer(SetResponseHeaderLayer::if_not_present(
497                header::X_XSS_PROTECTION,
498                HeaderValue::from_static("1; mode=block"),
499            )),
500    )
501}
502
503// ---------------------------------------------------------------------------
504// ServerBuilder — extensible server construction with custom routes
505// ---------------------------------------------------------------------------
506
507/// Builder for constructing an ADK server with custom routes.
508///
509/// `ServerBuilder` allows registering additional Axum routers alongside the
510/// built-in REST, A2A, and UI routes. Custom routes benefit from the same
511/// middleware stack (auth, CORS, tracing, timeout, security headers) as the
512/// built-in routes.
513///
514/// # Example
515///
516/// ```rust,ignore
517/// use adk_server::{ServerBuilder, ServerConfig};
518/// use axum::{Router, routing::get};
519///
520/// let config = ServerConfig::new(agent, session_service);
521///
522/// let app = ServerBuilder::new(config)
523///     .add_api_routes(
524///         Router::new()
525///             .route("/projects", get(list_projects))
526///             .route("/projects/{id}", get(get_project))
527///     )
528///     .add_api_routes(
529///         Router::new()
530///             .route("/automations", get(list_automations))
531///     )
532///     .with_a2a("http://localhost:8080")
533///     .build();
534///
535/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
536/// axum::serve(listener, app).await?;
537/// ```
538pub struct ServerBuilder {
539    config: ServerConfig,
540    a2a_base_url: Option<String>,
541    api_routes: Vec<Router>,
542    root_routes: Vec<Router>,
543    shutdown_endpoint: bool,
544    skill_index: Option<Arc<adk_skill::SkillIndex>>,
545    #[cfg(feature = "agent-engine")]
546    agent_engine: bool,
547}
548
549impl ServerBuilder {
550    /// Create a new server builder with the given configuration.
551    pub fn new(config: ServerConfig) -> Self {
552        Self {
553            config,
554            a2a_base_url: None,
555            api_routes: Vec::new(),
556            root_routes: Vec::new(),
557            shutdown_endpoint: false,
558            skill_index: None,
559            #[cfg(feature = "agent-engine")]
560            agent_engine: false,
561        }
562    }
563
564    /// Add custom routes nested under `/api`.
565    ///
566    /// These routes are merged into the API router and benefit from the auth
567    /// middleware layer. Multiple calls accumulate routes.
568    ///
569    /// # Example
570    ///
571    /// ```rust,ignore
572    /// builder.add_api_routes(
573    ///     Router::new()
574    ///         .route("/projects", get(list_projects))
575    ///         .route("/projects/{id}", get(get_project))
576    /// )
577    /// ```
578    pub fn add_api_routes(mut self, routes: Router) -> Self {
579        self.api_routes.push(routes);
580        self
581    }
582
583    /// Add custom routes at the root level (not nested under `/api`).
584    ///
585    /// These routes are merged at the top level of the application, alongside
586    /// the UI and A2A routes. They receive the full middleware stack (CORS,
587    /// tracing, timeout, security headers) but NOT the auth middleware.
588    ///
589    /// Use this for routes that need their own auth handling or public endpoints.
590    pub fn add_root_routes(mut self, routes: Router) -> Self {
591        self.root_routes.push(routes);
592        self
593    }
594
595    /// Enable A2A protocol support at the specified base URL.
596    ///
597    /// The base URL is used to construct the agent card's endpoint URL.
598    pub fn with_a2a(mut self, base_url: impl Into<String>) -> Self {
599        self.a2a_base_url = Some(base_url.into());
600        self
601    }
602
603    /// Expose the skills in `skill_index` on the A2A agent card.
604    ///
605    /// When A2A is enabled via [`with_a2a`](Self::with_a2a), the card served at
606    /// `/.well-known/agent.json` appends one `skills[]` entry per indexed
607    /// skill, mapped by [`agent_skills_from_index`](crate::a2a::agent_skills_from_index).
608    /// Has no effect without `with_a2a`.
609    ///
610    /// # Example
611    ///
612    /// ```rust,ignore
613    /// use adk_server::{ServerBuilder, ServerConfig};
614    /// use std::sync::Arc;
615    ///
616    /// let index = adk_skill::load_skill_index(".")?;
617    /// let app = ServerBuilder::new(config)
618    ///     .with_a2a("http://localhost:8080")
619    ///     .with_skill_index(Arc::new(index))
620    ///     .build();
621    /// ```
622    pub fn with_skill_index(mut self, skill_index: Arc<adk_skill::SkillIndex>) -> Self {
623        self.skill_index = Some(skill_index);
624        self
625    }
626
627    /// Mount the Agent Engine dispatch endpoints alongside the built-in routes.
628    ///
629    /// When enabled, `POST /api/reasoning_engine` and
630    /// `POST /api/stream_reasoning_engine` serve the Gemini Enterprise Agent
631    /// Platform runtime contract for the loader's root agent, using the
632    /// configured session and artifact services. The routes carry the shared
633    /// middleware stack but **not** the auth middleware — a deployed engine is
634    /// fronted by the platform, which authenticates callers before they reach
635    /// the container. The memory class methods report `Unsupported`; use
636    /// [`serve_agent_engine`](crate::agent_engine::serve_agent_engine) with
637    /// [`AgentEngineOptions`](crate::agent_engine::AgentEngineOptions) to
638    /// configure a memory service.
639    ///
640    /// # Panics
641    ///
642    /// [`build`](Self::build) panics when the root agent's name is not a
643    /// valid app name, mirroring the builder's other misuse panics.
644    #[cfg(feature = "agent-engine")]
645    pub fn with_agent_engine(mut self, enabled: bool) -> Self {
646        self.agent_engine = enabled;
647        self
648    }
649
650    /// Enable the `POST /api/shutdown` endpoint for graceful shutdown.
651    ///
652    /// When enabled, the server exposes a shutdown endpoint that triggers
653    /// graceful shutdown: stops accepting new connections, completes in-flight
654    /// requests, and then exits. Use [`build_with_shutdown`](Self::build_with_shutdown)
655    /// to get the [`ShutdownHandle`] for wiring into `axum::serve().with_graceful_shutdown()`.
656    ///
657    /// The endpoint is protected by the auth middleware when a
658    /// `RequestContextExtractor` is configured.
659    pub fn enable_shutdown_endpoint(mut self) -> Self {
660        self.shutdown_endpoint = true;
661        self
662    }
663
664    /// Build the final Axum router with all routes and middleware applied.
665    pub fn build(self) -> Router {
666        self.build_inner().0
667    }
668
669    /// Build the final Axum router and return a [`ShutdownHandle`].
670    ///
671    /// Use this when [`enable_shutdown_endpoint()`](Self::enable_shutdown_endpoint) is set.
672    /// Pass the handle's signal to `axum::serve().with_graceful_shutdown()`.
673    ///
674    /// # Example
675    ///
676    /// ```rust,ignore
677    /// let (app, shutdown_handle) = ServerBuilder::new(config)
678    ///     .enable_shutdown_endpoint()
679    ///     .build_with_shutdown();
680    ///
681    /// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
682    /// axum::serve(listener, app)
683    ///     .with_graceful_shutdown(shutdown_handle.signal())
684    ///     .await?;
685    /// ```
686    pub fn build_with_shutdown(self) -> (Router, ShutdownHandle) {
687        let (router, handle) = self.build_inner();
688        (router, handle.expect("build_with_shutdown requires enable_shutdown_endpoint()"))
689    }
690
691    fn build_inner(self) -> (Router, Option<ShutdownHandle>) {
692        let config = &self.config;
693        let session_controller = SessionController::new(config.session_service.clone());
694        let runtime_controller = RuntimeController::new(config.clone());
695        let apps_controller = AppsController::new(config.clone());
696        let artifacts_controller = ArtifactsController::new(config.clone());
697        let debug_controller = DebugController::new(config.clone());
698        let health_controller = HealthController::new(config);
699
700        // Start YAML agent hot reload watchers if configured.
701        #[cfg(feature = "yaml-agent")]
702        {
703            let dirs = config.yaml_agent_dirs.clone();
704            if !dirs.is_empty() {
705                tokio::spawn(async move {
706                    let _watchers = start_yaml_agent_watchers(&dirs).await;
707                    std::future::pending::<()>().await;
708                });
709            }
710        }
711
712        let auth_layer = middleware::from_fn({
713            let extractor = config.request_context_extractor.clone();
714            move |request: Request<Body>, next: Next| {
715                let extractor = extractor.clone();
716                async move { auth_middleware(request, next, extractor).await }
717            }
718        });
719
720        let health_router =
721            Router::new().route("/health", get(health_check)).with_state(health_controller);
722
723        let ui_api_router = Router::new()
724            .route("/apps", get(controllers::apps::list_apps))
725            .route("/list-apps", get(controllers::apps::list_apps_compat))
726            .route("/ui/agents/{name}", get(controllers::apps::get_agent_details))
727            .with_state(apps_controller)
728            .route("/ui/capabilities", get(controllers::ui::ui_capabilities))
729            .route("/ui/initialize", post(controllers::ui::ui_initialize))
730            .route("/ui/message", post(controllers::ui::ui_message))
731            .route("/ui/update-model-context", post(controllers::ui::ui_update_model_context))
732            .route("/ui/notifications/poll", post(controllers::ui::ui_poll_notifications))
733            .route(
734                "/ui/notifications/resources-list-changed",
735                post(controllers::ui::ui_notify_resources_list_changed),
736            )
737            .route(
738                "/ui/notifications/tools-list-changed",
739                post(controllers::ui::ui_notify_tools_list_changed),
740            )
741            .route("/ui/resources", get(controllers::ui::list_ui_resources))
742            .route("/ui/resources/read", get(controllers::ui::read_ui_resource))
743            .route("/ui/resources/register", post(controllers::ui::register_ui_resource))
744            // These routes mutate and read shared bridge and resource state, so they
745            // carry the same authentication as the session, artifact, and debug routers.
746            .layer(auth_layer.clone());
747
748        let session_router = Router::new()
749            .route("/sessions", post(controllers::session::create_session))
750            .route(
751                "/sessions/{app_name}/{user_id}/{session_id}",
752                get(controllers::session::get_session).delete(controllers::session::delete_session),
753            )
754            .route(
755                "/apps/{app_name}/users/{user_id}/sessions",
756                get(controllers::session::list_sessions)
757                    .post(controllers::session::create_session_from_path),
758            )
759            .route(
760                "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
761                get(controllers::session::get_session_from_path)
762                    .post(controllers::session::create_session_from_path)
763                    .delete(controllers::session::delete_session_from_path),
764            )
765            .with_state(session_controller)
766            .layer(auth_layer.clone());
767
768        let runtime_router = Router::new()
769            .route("/run", post(controllers::runtime::run_collect))
770            .route("/run/{app_name}/{user_id}/{session_id}", post(controllers::runtime::run_sse))
771            .route("/run_sse", post(controllers::runtime::run_sse_compat))
772            .with_state(runtime_controller);
773
774        let artifacts_router = Router::new()
775            .route(
776                "/sessions/{app_name}/{user_id}/{session_id}/artifacts",
777                get(controllers::artifacts::list_artifacts),
778            )
779            .route(
780                "/sessions/{app_name}/{user_id}/{session_id}/artifacts/{artifact_name}",
781                get(controllers::artifacts::get_artifact),
782            )
783            .with_state(artifacts_controller)
784            .layer(auth_layer.clone());
785
786        let mut debug_router = Router::new()
787            .route("/debug/trace/session/{session_id}", get(controllers::debug::get_session_traces))
788            .route(
789                "/debug/graph/{app_name}/{user_id}/{session_id}/{event_id}",
790                get(controllers::debug::get_graph),
791            )
792            .route(
793                "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
794                get(controllers::debug::get_graph),
795            )
796            .route("/apps/{app_name}/eval_sets", get(controllers::debug::get_eval_sets))
797            .route(
798                "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}",
799                get(controllers::debug::get_event),
800            );
801
802        if config.request_context_extractor.is_none() || config.security.expose_admin_debug {
803            debug_router = debug_router
804                .route("/debug/trace/{event_id}", get(controllers::debug::get_trace_by_event_id));
805        }
806
807        let debug_router =
808            debug_router.with_state(debug_controller.clone()).layer(auth_layer.clone());
809
810        // Assemble the API router with built-in + custom routes
811        let mut api_router = Router::new()
812            .merge(health_router)
813            .merge(ui_api_router)
814            .merge(session_router)
815            .merge(runtime_router)
816            .merge(artifacts_router)
817            .merge(debug_router);
818
819        // Merge custom API routes — these get the same /api prefix and auth middleware
820        for custom_routes in self.api_routes {
821            api_router = api_router.merge(custom_routes.layer(auth_layer.clone()));
822        }
823
824        // Add shutdown endpoint if enabled
825        let shutdown_handle = if self.shutdown_endpoint {
826            let handle = ShutdownHandle::new();
827            let shutdown_router = Router::new()
828                .route("/shutdown", post(handle_shutdown))
829                .with_state(handle.token.clone())
830                .layer(auth_layer.clone());
831            api_router = api_router.merge(shutdown_router);
832            Some(handle)
833        } else {
834            None
835        };
836
837        let ui_router = Router::new()
838            .route("/", get(web_ui::root_redirect))
839            .route("/ui/", get(web_ui::serve_ui_index))
840            .route("/ui/assets/config/runtime-config.json", get(web_ui::serve_runtime_config))
841            .with_state(config.clone())
842            .route("/ui/{*path}", get(web_ui::serve_ui_assets));
843
844        let mut app = Router::new().nest("/api", api_router).merge(ui_router);
845
846        // Merge custom root routes
847        for custom_routes in self.root_routes {
848            app = app.merge(custom_routes);
849        }
850
851        // Agent Engine dispatch surface: root-level merge (its routes carry
852        // their own /api/... paths) so the platform host reaches it without
853        // the local auth middleware — the platform authenticates callers
854        // before they reach the container.
855        #[cfg(feature = "agent-engine")]
856        if self.agent_engine {
857            let root_agent = config.agent_loader.root_agent();
858            let mut runner_builder = adk_runner::Runner::builder()
859                .app_name(root_agent.name())
860                .agent(root_agent.clone())
861                .session_service(config.session_service.clone());
862            if let Some(artifact_service) = &config.artifact_service {
863                runner_builder = runner_builder.artifact_service(artifact_service.clone());
864            }
865            let runner = runner_builder
866                .build()
867                .expect("with_agent_engine requires the root agent's name to be a valid app name");
868            let mut state = crate::agent_engine::AgentEngineState::new(Arc::new(runner));
869            if let Some(artifact_service) = &config.artifact_service {
870                state = state.with_artifact_service(artifact_service.clone());
871            }
872            app = app.merge(crate::agent_engine::agent_engine_router(state));
873        }
874
875        if let Some(base_url) = &self.a2a_base_url {
876            let a2a_controller = match &self.skill_index {
877                Some(skill_index) => {
878                    A2aController::with_skill_index(config.clone(), base_url, skill_index.clone())
879                }
880                None => A2aController::new(config.clone(), base_url),
881            };
882            // Same split as `create_app_with_a2a`: discovery is public, RPC is authenticated.
883            let a2a_discovery = Router::new()
884                .route("/.well-known/agent.json", get(controllers::a2a::get_agent_card))
885                .with_state(a2a_controller.clone());
886            let a2a_rpc = Router::new()
887                .route("/a2a", post(controllers::a2a::handle_jsonrpc))
888                .route("/a2a/stream", post(controllers::a2a::handle_jsonrpc_stream))
889                .with_state(a2a_controller)
890                .layer(auth_layer.clone());
891            app = app.merge(a2a_discovery).merge(a2a_rpc);
892        }
893
894        let cors_layer = build_cors_layer(config);
895        let trace_layer = TraceLayer::new_for_http().make_span_with(|request: &Request<Body>| {
896            let request_id =
897                request.extensions().get::<RequestId>().map(RequestId::as_str).unwrap_or("");
898            tracing::info_span!(
899                "http.request",
900                request.id = %request_id,
901                http.method = %request.method(),
902                http.path = %request.uri().path()
903            )
904        });
905
906        (
907            app.layer(
908                ServiceBuilder::new()
909                    .layer(middleware::from_fn(request_id_middleware))
910                    .layer(trace_layer)
911                    .layer(TimeoutLayer::with_status_code(
912                        StatusCode::REQUEST_TIMEOUT,
913                        config.security.request_timeout,
914                    ))
915                    .layer(DefaultBodyLimit::max(config.security.max_body_size))
916                    .layer(cors_layer)
917                    .layer(SetResponseHeaderLayer::if_not_present(
918                        header::X_CONTENT_TYPE_OPTIONS,
919                        HeaderValue::from_static("nosniff"),
920                    ))
921                    .layer(SetResponseHeaderLayer::if_not_present(
922                        header::X_FRAME_OPTIONS,
923                        HeaderValue::from_static("DENY"),
924                    ))
925                    .layer(SetResponseHeaderLayer::if_not_present(
926                        header::X_XSS_PROTECTION,
927                        HeaderValue::from_static("1; mode=block"),
928                    )),
929            ),
930            shutdown_handle,
931        )
932    }
933}
934
935/// Wait for a process shutdown signal.
936pub async fn shutdown_signal() {
937    let ctrl_c = async {
938        let _ = tokio::signal::ctrl_c().await;
939    };
940
941    #[cfg(unix)]
942    let terminate = async {
943        if let Ok(mut signal) =
944            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
945        {
946            let _ = signal.recv().await;
947        }
948    };
949
950    #[cfg(not(unix))]
951    let terminate = std::future::pending::<()>();
952
953    tokio::select! {
954        _ = ctrl_c => {}
955        _ = terminate => {}
956    }
957}
958
959// ---------------------------------------------------------------------------
960// ShutdownHandle — programmatic graceful shutdown trigger
961// ---------------------------------------------------------------------------
962
963/// Handle for triggering graceful server shutdown.
964///
965/// Returned by [`ServerBuilder::build_with_shutdown`]. Pass the future from
966/// [`ShutdownHandle::signal()`] to `axum::serve(...).with_graceful_shutdown()`
967/// to enable both OS signal-based and HTTP endpoint-based shutdown.
968///
969/// # Example
970///
971/// ```rust,ignore
972/// use adk_server::{ServerBuilder, ServerConfig};
973///
974/// let (app, shutdown_handle) = ServerBuilder::new(config)
975///     .enable_shutdown_endpoint()
976///     .build_with_shutdown();
977///
978/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
979/// axum::serve(listener, app)
980///     .with_graceful_shutdown(shutdown_handle.signal())
981///     .await?;
982/// ```
983#[derive(Clone)]
984pub struct ShutdownHandle {
985    token: CancellationToken,
986}
987
988impl ShutdownHandle {
989    /// Create a new shutdown handle.
990    fn new() -> Self {
991        Self { token: CancellationToken::new() }
992    }
993
994    /// Trigger graceful shutdown programmatically.
995    ///
996    /// This has the same effect as calling `POST /api/shutdown` — the server
997    /// stops accepting new connections and completes in-flight requests.
998    pub fn shutdown(&self) {
999        tracing::info!("graceful shutdown triggered programmatically");
1000        self.token.cancel();
1001    }
1002
1003    /// Returns a future that resolves when shutdown is triggered.
1004    ///
1005    /// Combines OS signals (Ctrl+C, SIGTERM) with the programmatic/HTTP trigger.
1006    /// Pass this to `axum::serve(...).with_graceful_shutdown()`.
1007    pub async fn signal(self) {
1008        let token = self.token.clone();
1009
1010        let ctrl_c = async {
1011            let _ = tokio::signal::ctrl_c().await;
1012        };
1013
1014        #[cfg(unix)]
1015        let terminate = async {
1016            if let Ok(mut signal) =
1017                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
1018            {
1019                let _ = signal.recv().await;
1020            }
1021        };
1022
1023        #[cfg(not(unix))]
1024        let terminate = std::future::pending::<()>();
1025
1026        tokio::select! {
1027            _ = ctrl_c => {
1028                tracing::info!("received Ctrl+C, initiating graceful shutdown");
1029            }
1030            _ = terminate => {
1031                tracing::info!("received SIGTERM, initiating graceful shutdown");
1032            }
1033            _ = token.cancelled() => {
1034                // Shutdown triggered via POST /api/shutdown or programmatic call
1035            }
1036        }
1037    }
1038
1039    /// Returns whether shutdown has been triggered.
1040    pub fn is_shutdown(&self) -> bool {
1041        self.token.is_cancelled()
1042    }
1043}
1044
1045/// Handler for `POST /api/shutdown`.
1046///
1047/// Triggers graceful shutdown: the server stops accepting new connections,
1048/// completes in-flight requests, and then exits.
1049async fn handle_shutdown(State(token): State<CancellationToken>) -> impl IntoResponse {
1050    tracing::info!("POST /api/shutdown received, initiating graceful shutdown");
1051    token.cancel();
1052    (StatusCode::OK, Json(serde_json::json!({ "status": "shutting_down" })))
1053}