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

systemprompt_api/services/server/
startup.rs
