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

systemprompt_api/services/server/
startup.rs
