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_models::modules::ApiPaths;
27use systemprompt_runtime::ShutdownRequest;
28use systemprompt_traits::{StartupEvent, StartupEventExt, StartupEventSender};
29use tokio::task::JoinHandle;
30use tower::ServiceExt;
31
32#[derive(Debug)]
33pub struct EarlyServer {
34    swap: Arc<RwLock<Router>>,
35    join: JoinHandle<Result<()>>,
36    local_addr: SocketAddr,
37}
38
39impl EarlyServer {
40    pub const fn local_addr(&self) -> SocketAddr {
41        self.local_addr
42    }
43
44    pub fn activate(&self, router: Router) {
45        *self.swap.write().unwrap_or_else(PoisonError::into_inner) = router;
46        tracing::info!("Full API router activated");
47    }
48
49    pub async fn join(self) -> Result<()> {
50        self.join.await.context("API serve task panicked")?
51    }
52}
53
54pub async fn bind_and_serve(
55    addr: &str,
56    events: Option<StartupEventSender>,
57    shutdown: ShutdownRequest,
58) -> Result<EarlyServer> {
59    if let Some(ref tx) = events
60        && tx
61            .unbounded_send(StartupEvent::ServerBinding {
62                address: addr.to_owned(),
63            })
64            .is_err()
65    {
66        tracing::debug!("Startup event receiver dropped");
67    }
68
69    let listener = tokio::net::TcpListener::bind(addr)
70        .await
71        .with_context(|| format!("Failed to bind to {addr}"))?;
72    let local_addr = listener
73        .local_addr()
74        .context("Failed to read bound address")?;
75
76    if let Some(ref tx) = events {
77        tx.server_listening(addr, std::process::id());
78    }
79
80    let swap = Arc::new(RwLock::new(starting_router()));
81    let outer = Router::new().fallback_service(SwapService {
82        swap: Arc::clone(&swap),
83    });
84
85    let join = tokio::spawn(async move {
86        axum::serve(
87            listener,
88            outer.into_make_service_with_connect_info::<SocketAddr>(),
89        )
90        .with_graceful_shutdown(super::shutdown::shutdown_signal(shutdown))
91        .await
92        .map_err(Into::into)
93    });
94
95    Ok(EarlyServer {
96        swap,
97        join,
98        local_addr,
99    })
100}
101
102pub fn starting_router() -> Router {
103    Router::new()
104        .route(ApiPaths::HEALTH, get(starting_health))
105        .route("/health", get(starting_health))
106        .route(ApiPaths::LIVEZ, get(starting_livez))
107        .route(ApiPaths::READYZ, get(starting_readyz))
108        .fallback(starting_fallback)
109}
110
111async fn starting_health() -> impl IntoResponse {
112    Json(json!({ "status": "starting" }))
113}
114
115async fn starting_livez() -> impl IntoResponse {
116    Json(json!({ "status": "alive", "version": env!("CARGO_PKG_VERSION") }))
117}
118
119async fn starting_readyz() -> impl IntoResponse {
120    (
121        StatusCode::SERVICE_UNAVAILABLE,
122        Json(json!({ "status": "starting", "version": env!("CARGO_PKG_VERSION") })),
123    )
124}
125
126async fn starting_fallback() -> impl IntoResponse {
127    (
128        StatusCode::SERVICE_UNAVAILABLE,
129        Json(json!({ "error": "service starting" })),
130    )
131}
132
133#[derive(Clone)]
134struct SwapService {
135    swap: Arc<RwLock<Router>>,
136}
137
138impl tower::Service<Request<Body>> for SwapService {
139    type Response = Response;
140    type Error = Infallible;
141    type Future = Pin<Box<dyn Future<Output = Result<Response, Infallible>> + Send>>;
142
143    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
144        Poll::Ready(Ok(()))
145    }
146
147    fn call(&mut self, req: Request<Body>) -> Self::Future {
148        let router = self
149            .swap
150            .read()
151            .unwrap_or_else(PoisonError::into_inner)
152            .clone();
153        Box::pin(router.oneshot(req))
154    }
155}
systemprompt_api/services/server/startup.rs

systemprompt_api/services/server/
startup.rs
