systemprompt_api/services/server/
startup.rs

1//! Early-bind listener that answers health probes before bootstrap finishes.
2//!
3//! [`bind_and_serve`] binds the TCP listener up front and serves a minimal
4//! router — `200 {"status":"starting"}` on the health paths, `503` everywhere
5//! else — so platform health checks pass while migrations, content publish, and
6//! agent reconciliation run. Once bootstrap completes, the full router is
7//! swapped onto the same listener via [`EarlyServer::activate`]; the port is
8//! bound exactly once, so probes never hit an unbind/rebind window.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use std::convert::Infallible;
14use std::net::SocketAddr;
15use std::pin::Pin;
16use std::sync::{Arc, PoisonError, RwLock};
17use std::task::{Context, Poll};
18
19use anyhow::{Context as _, Result};
20use axum::body::Body;
21use axum::http::{Request, StatusCode};
22use axum::response::{IntoResponse, Response};
23use axum::routing::get;
24use axum::{Json, Router};
25use serde_json::json;
26use systemprompt_manifest::profile::ShutdownConfig;
27use systemprompt_models::api::ApiError;
28use systemprompt_models::modules::ApiPaths;
29use systemprompt_runtime::ShutdownRequest;
30use systemprompt_traits::{OwnedTask, StartupEvent, StartupEventExt, StartupEventSender};
31use tower::ServiceExt;
32
33#[derive(Debug)]
34pub struct EarlyServer {
35    swap: Arc<RwLock<Router>>,
36    join: OwnedTask<Result<()>>,
37    local_addr: SocketAddr,
38}
39
40impl EarlyServer {
41    pub const fn local_addr(&self) -> SocketAddr {
42        self.local_addr
43    }
44
45    pub fn activate(&self, router: Router) {
46        *self.swap.write().unwrap_or_else(PoisonError::into_inner) = router;
47        tracing::info!("Full API router activated");
48    }
49
50    pub async fn join(self) -> Result<()> {
51        self.join.join().await.context("API serve task panicked")?
52    }
53}
54
55pub async fn bind_and_serve(
56    addr: &str,
57    events: Option<StartupEventSender>,
58    shutdown: ShutdownRequest,
59    timings: ShutdownConfig,
60) -> Result<EarlyServer> {
61    if let Some(ref tx) = events
62        && tx
63            .unbounded_send(StartupEvent::ServerBinding {
64                address: addr.to_owned(),
65            })
66            .is_err()
67    {
68        tracing::debug!("Startup event receiver dropped");
69    }
70
71    let listener = tokio::net::TcpListener::bind(addr)
72        .await
73        .with_context(|| format!("Failed to bind to {addr}"))?;
74    let local_addr = listener
75        .local_addr()
76        .context("Failed to read bound address")?;
77
78    if let Some(ref tx) = events {
79        tx.server_listening(addr, std::process::id());
80    }
81
82    let swap = Arc::new(RwLock::new(starting_router()));
83    let outer = Router::new().fallback_service(SwapService {
84        swap: Arc::clone(&swap),
85    });
86
87    let join = OwnedTask::spawn("api_serve", async move {
88        axum::serve(
89            listener,
90            outer.into_make_service_with_connect_info::<SocketAddr>(),
91        )
92        .with_graceful_shutdown(super::shutdown::shutdown_signal(shutdown, timings))
93        .await
94        .map_err(Into::into)
95    });
96
97    Ok(EarlyServer {
98        swap,
99        join,
100        local_addr,
101    })
102}
103
104pub fn starting_router() -> Router {
105    Router::new()
106        .route(ApiPaths::HEALTH, get(starting_health))
107        .route("/health", get(starting_health))
108        .route(ApiPaths::LIVEZ, get(starting_livez))
109        .route(ApiPaths::READYZ, get(starting_readyz))
110        .fallback(starting_fallback)
111}
112
113async fn starting_health() -> impl IntoResponse {
114    Json(json!({ "status": "starting" }))
115}
116
117async fn starting_livez() -> impl IntoResponse {
118    Json(json!({ "status": "alive", "version": env!("CARGO_PKG_VERSION") }))
119}
120
121async fn starting_readyz() -> impl IntoResponse {
122    (
123        StatusCode::SERVICE_UNAVAILABLE,
124        Json(json!({ "status": "starting", "version": env!("CARGO_PKG_VERSION") })),
125    )
126}
127
128async fn starting_fallback() -> Response {
129    tracing::debug!("request answered with 503 while the service is starting");
130    let body = ApiError::service_unavailable("service starting").with_error_key("service_starting");
131    (StatusCode::SERVICE_UNAVAILABLE, Json(body)).into_response()
132}
133
134#[derive(Clone)]
135struct SwapService {
136    swap: Arc<RwLock<Router>>,
137}
138
139impl tower::Service<Request<Body>> for SwapService {
140    type Response = Response;
141    type Error = Infallible;
142    type Future = Pin<Box<dyn Future<Output = Result<Response, Infallible>> + Send>>;
143
144    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
145        Poll::Ready(Ok(()))
146    }
147
148    fn call(&mut self, req: Request<Body>) -> Self::Future {
149        let router = self
150            .swap
151            .read()
152            .unwrap_or_else(PoisonError::into_inner)
153            .clone();
154        Box::pin(router.oneshot(req))
155    }
156}
systemprompt_api/services/server/startup.rs

systemprompt_api/services/server/
startup.rs
