Skip to main content

systemprompt_api/services/server/
builder.rs

1use anyhow::Result;
2use axum::Router;
3use axum::extract::DefaultBodyLimit;
4use systemprompt_runtime::AppContext;
5use systemprompt_traits::{StartupEvent, StartupEventExt, StartupEventSender};
6
7use super::routes::configure_routes;
8use crate::models::ServerConfig;
9use crate::services::middleware::{
10    AnalyticsMiddleware, ContextMiddleware, CorsMiddleware, JwtContextExtractor, SessionMiddleware,
11    inject_security_headers, inject_trace_header, remove_trailing_slash,
12};
13
14pub use super::discovery::*;
15pub use super::health::handle_health;
16
17#[derive(Debug)]
18pub struct ApiServer {
19    router: Router,
20    _config: ServerConfig,
21    events: Option<StartupEventSender>,
22}
23
24impl ApiServer {
25    pub fn new(router: Router, events: Option<StartupEventSender>) -> Self {
26        Self::with_config(router, ServerConfig::default(), events)
27    }
28
29    pub const fn with_config(
30        router: Router,
31        config: ServerConfig,
32        events: Option<StartupEventSender>,
33    ) -> Self {
34        Self {
35            router,
36            _config: config,
37            events,
38        }
39    }
40
41    pub async fn serve(self, addr: &str) -> Result<()> {
42        if let Some(ref tx) = self.events {
43            if tx
44                .unbounded_send(StartupEvent::ServerBinding {
45                    address: addr.to_string(),
46                })
47                .is_err()
48            {
49                tracing::debug!("Startup event receiver dropped");
50            }
51        }
52
53        let listener = self.create_listener(addr).await?;
54
55        if let Some(ref tx) = self.events {
56            tx.server_listening(addr, std::process::id());
57        }
58
59        axum::serve(
60            listener,
61            self.router
62                .into_make_service_with_connect_info::<std::net::SocketAddr>(),
63        )
64        .await?;
65        Ok(())
66    }
67
68    async fn create_listener(&self, addr: &str) -> Result<tokio::net::TcpListener> {
69        tokio::net::TcpListener::bind(addr)
70            .await
71            .map_err(|e| anyhow::anyhow!("Failed to bind to {addr}: {e}"))
72    }
73}
74
75pub fn setup_api_server(ctx: &AppContext, events: Option<StartupEventSender>) -> Result<ApiServer> {
76    let rate_config = &ctx.config().rate_limits;
77
78    if rate_config.disabled {
79        if let Some(ref tx) = events {
80            tx.warning("Rate limiting disabled - development mode only");
81        }
82    }
83
84    let router = configure_routes(ctx, events.as_ref())?;
85    let router = apply_global_middleware(router, ctx)?;
86
87    Ok(ApiServer::new(router, events))
88}
89
90fn apply_global_middleware(router: Router, ctx: &AppContext) -> Result<Router> {
91    let mut router = router;
92
93    router = router.layer(DefaultBodyLimit::max(100 * 1024 * 1024));
94
95    let analytics_middleware = AnalyticsMiddleware::new(ctx)?;
96    router = router.layer(axum::middleware::from_fn({
97        let middleware = analytics_middleware;
98        move |req, next| {
99            let middleware = middleware.clone();
100            async move { middleware.track_request(req, next).await }
101        }
102    }));
103
104    let jwt_extractor = JwtContextExtractor::new(
105        systemprompt_models::SecretsBootstrap::jwt_secret()?,
106        ctx.db_pool(),
107    );
108    let global_context_middleware = ContextMiddleware::public(jwt_extractor);
109    router = router.layer(axum::middleware::from_fn({
110        let middleware = global_context_middleware;
111        move |req, next| {
112            let middleware = middleware.clone();
113            async move { middleware.handle(req, next).await }
114        }
115    }));
116
117    let session_middleware = SessionMiddleware::new(ctx)?;
118    router = router.layer(axum::middleware::from_fn({
119        let middleware = session_middleware;
120        move |req, next| {
121            let middleware = middleware.clone();
122            async move { middleware.handle(req, next).await }
123        }
124    }));
125
126    let cors = CorsMiddleware::build_layer(ctx.config())?;
127    router = router.layer(cors);
128
129    router = router.layer(axum::middleware::from_fn(remove_trailing_slash));
130
131    router = router.layer(axum::middleware::from_fn(inject_trace_header));
132
133    if ctx.config().content_negotiation.enabled {
134        router = router.layer(axum::middleware::from_fn(
135            crate::services::middleware::content_negotiation_middleware,
136        ));
137    }
138
139    if ctx.config().security_headers.enabled {
140        let security_config = ctx.config().security_headers.clone();
141        router = router.layer(axum::middleware::from_fn(move |req, next| {
142            let config = security_config.clone();
143            inject_security_headers(config, req, next)
144        }));
145    }
146
147    Ok(router)
148}