Skip to main content

statespace_server/
server.rs

1//! HTTP server and Axum router.
2
3use crate::content::{ContentResolver, LocalContentResolver};
4use crate::error::ErrorExt;
5use crate::templates::FAVICON_SVG;
6use axum::{
7    Json, Router,
8    extract::{Path, Query, State},
9    http::{StatusCode, header},
10    response::{IntoResponse, Response},
11    routing::get,
12};
13use statespace_tool_runtime::{
14    ActionRequest, ActionResponse, BuiltinTool, ExecutionLimits, SandboxEnv, ToolExecutor, eval,
15    expand_command_for_execution, expand_placeholders, parse_frontmatter,
16    validate_command_with_specs, validate_env_map,
17};
18use std::collections::HashMap;
19use std::path::PathBuf;
20use std::sync::Arc;
21use tokio::fs;
22use tower_http::cors::{Any, CorsLayer};
23use tower_http::trace::TraceLayer;
24use tracing::{info, warn};
25
26#[derive(Clone)]
27pub struct ServerConfig {
28    pub content_root: PathBuf,
29    pub host: String,
30    pub port: u16,
31    pub limits: ExecutionLimits,
32    pub env: HashMap<String, String>,
33    pub sandbox_env: SandboxEnv,
34}
35
36impl std::fmt::Debug for ServerConfig {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.debug_struct("ServerConfig")
39            .field("content_root", &self.content_root)
40            .field("host", &self.host)
41            .field("port", &self.port)
42            .field("limits", &self.limits)
43            .field("env_keys", &self.env.len())
44            .field("sandbox_path", &self.sandbox_env.path())
45            .finish()
46    }
47}
48
49impl ServerConfig {
50    #[must_use]
51    pub fn new(content_root: PathBuf) -> Self {
52        Self {
53            content_root,
54            host: "127.0.0.1".to_string(),
55            port: 8000,
56            limits: ExecutionLimits::default(),
57            env: HashMap::new(),
58            sandbox_env: SandboxEnv::default(),
59        }
60    }
61
62    #[must_use]
63    pub fn with_host(mut self, host: impl Into<String>) -> Self {
64        self.host = host.into();
65        self
66    }
67
68    #[must_use]
69    pub const fn with_port(mut self, port: u16) -> Self {
70        self.port = port;
71        self
72    }
73
74    #[must_use]
75    pub fn with_limits(mut self, limits: ExecutionLimits) -> Self {
76        self.limits = limits;
77        self
78    }
79
80    #[must_use]
81    pub fn with_env(mut self, env: HashMap<String, String>) -> Self {
82        self.env = env;
83        self
84    }
85
86    #[must_use]
87    pub fn with_sandbox_env(mut self, sandbox_env: SandboxEnv) -> Self {
88        self.sandbox_env = sandbox_env;
89        self
90    }
91
92    #[must_use]
93    pub fn socket_addr(&self) -> String {
94        format!("{}:{}", self.host, self.port)
95    }
96
97    #[must_use]
98    pub fn base_url(&self) -> String {
99        format!("http://{}:{}", self.host, self.port)
100    }
101}
102
103#[derive(Clone)]
104pub struct ServerState {
105    pub content_resolver: Arc<dyn ContentResolver>,
106    pub limits: ExecutionLimits,
107    pub content_root: PathBuf,
108    pub env: Arc<HashMap<String, String>>,
109    pub sandbox_env: Arc<SandboxEnv>,
110}
111
112impl std::fmt::Debug for ServerState {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct("ServerState")
115            .field("limits", &self.limits)
116            .field("content_root", &self.content_root)
117            .field("env_keys", &self.env.len())
118            .field("sandbox_path", &self.sandbox_env.path())
119            .finish_non_exhaustive()
120    }
121}
122
123impl ServerState {
124    /// # Errors
125    ///
126    /// Returns an error if the content root path cannot be canonicalized.
127    pub fn from_config(config: &ServerConfig) -> crate::error::Result<Self> {
128        Ok(Self {
129            content_resolver: Arc::new(LocalContentResolver::new(&config.content_root)?),
130            limits: config.limits.clone(),
131            content_root: config.content_root.clone(),
132            env: Arc::new(config.env.clone()),
133            sandbox_env: Arc::new(config.sandbox_env.clone()),
134        })
135    }
136}
137
138/// # Errors
139///
140/// Returns an error if the content root path cannot be canonicalized.
141pub fn build_router(config: &ServerConfig) -> crate::error::Result<Router> {
142    let state = ServerState::from_config(config)?;
143
144    let cors = CorsLayer::new()
145        .allow_origin(Any)
146        .allow_methods(Any)
147        .allow_headers(Any);
148
149    Ok(Router::new()
150        .route("/", get(index_handler).post(action_handler_root))
151        .route("/favicon.svg", get(favicon_handler))
152        .route("/favicon.ico", get(favicon_handler))
153        .route("/opengraph.png", get(opengraph_handler))
154        .route("/{*path}", get(file_handler).post(action_handler))
155        .layer(cors)
156        .layer(TraceLayer::new_for_http())
157        .with_state(state))
158}
159
160async fn index_handler(
161    Query(query_env): Query<HashMap<String, String>>,
162    State(state): State<ServerState>,
163) -> Response {
164    serve_page("AGENTS.md", &query_env, &state).await
165}
166
167async fn favicon_handler(State(state): State<ServerState>) -> Response {
168    let content = match fs::read_to_string(state.content_root.join("favicon.svg")).await {
169        Ok(custom) => custom,
170        Err(_) => FAVICON_SVG.to_string(),
171    };
172
173    (
174        StatusCode::OK,
175        [(header::CONTENT_TYPE, "image/svg+xml")],
176        content,
177    )
178        .into_response()
179}
180
181async fn opengraph_handler(State(state): State<ServerState>) -> Response {
182    match fs::read(state.content_root.join("opengraph.png")).await {
183        Ok(custom) => (
184            StatusCode::OK,
185            [(header::CONTENT_TYPE, "image/png")],
186            custom,
187        )
188            .into_response(),
189        Err(_) => StatusCode::NOT_FOUND.into_response(),
190    }
191}
192
193async fn file_handler(
194    Path(path): Path<String>,
195    Query(query_env): Query<HashMap<String, String>>,
196    State(state): State<ServerState>,
197) -> Response {
198    serve_page(&path, &query_env, &state).await
199}
200
201async fn serve_page(
202    path: &str,
203    query_env: &HashMap<String, String>,
204    state: &ServerState,
205) -> Response {
206    if let Err(e) = validate_env_map(query_env) {
207        return (StatusCode::BAD_REQUEST, e.to_string()).into_response();
208    }
209
210    let file_path = match state.content_resolver.resolve_path(path).await {
211        Ok(p) => p,
212        Err(e) => {
213            warn!("File not found: {} ({})", path, e);
214            return (e.status_code(), e.user_message()).into_response();
215        }
216    };
217
218    let content = match fs::read_to_string(&file_path).await {
219        Ok(c) => c,
220        Err(e) => {
221            warn!("Failed to read {}: {}", path, e);
222            return (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error").into_response();
223        }
224    };
225
226    let working_dir = file_path.parent().unwrap_or(&state.content_root);
227    let has_eval = !eval::parse_eval_blocks(&content).is_empty();
228    let merged_env = eval::merge_eval_env(state.env.as_ref(), query_env);
229    let rendered = eval::process_eval_blocks_with_sandbox(
230        &content,
231        working_dir,
232        &merged_env,
233        &state.sandbox_env,
234    )
235    .await;
236
237    if has_eval {
238        (
239            [
240                (header::CONTENT_TYPE, "text/markdown; charset=utf-8"),
241                (header::CACHE_CONTROL, "no-store"),
242                (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
243            ],
244            rendered,
245        )
246            .into_response()
247    } else {
248        (
249            [
250                (header::CONTENT_TYPE, "text/markdown; charset=utf-8"),
251                (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
252            ],
253            rendered,
254        )
255            .into_response()
256    }
257}
258
259async fn action_handler_root(
260    State(state): State<ServerState>,
261    Json(request): Json<ActionRequest>,
262) -> Response {
263    execute_action("", &state, request).await
264}
265
266async fn action_handler(
267    Path(path): Path<String>,
268    State(state): State<ServerState>,
269    Json(request): Json<ActionRequest>,
270) -> Response {
271    execute_action(&path, &state, request).await
272}
273
274fn error_to_action_response(e: &statespace_tool_runtime::Error) -> Response {
275    let status = e.status_code();
276    let response = ActionResponse::error(e.user_message());
277    (status, Json(response)).into_response()
278}
279
280async fn execute_action(path: &str, state: &ServerState, request: ActionRequest) -> Response {
281    if let Err(msg) = request.validate() {
282        return error_response(StatusCode::BAD_REQUEST, &msg);
283    }
284
285    let file_path = match state.content_resolver.resolve_path(path).await {
286        Ok(p) => p,
287        Err(e) => return error_to_action_response(&e),
288    };
289
290    let content = match state.content_resolver.resolve(path).await {
291        Ok(c) => c,
292        Err(e) => return error_to_action_response(&e),
293    };
294
295    let frontmatter = match parse_frontmatter(&content) {
296        Ok(fm) => fm,
297        Err(e) => return error_to_action_response(&e),
298    };
299
300    let command_with_placeholders = expand_placeholders(&request.command, &request.args);
301    let merged_env = eval::merge_eval_env(state.env.as_ref(), &request.env);
302    let expanded_command =
303        expand_command_for_execution(&command_with_placeholders, &frontmatter.specs, &merged_env);
304
305    let validation_result =
306        validate_command_with_specs(&frontmatter.specs, &command_with_placeholders)
307            .or_else(|_error| validate_command_with_specs(&frontmatter.specs, &expanded_command));
308
309    if let Err(e) = validation_result {
310        warn!(
311            "Command not allowed by frontmatter: {:?} (expanded: {:?}, file: {})",
312            command_with_placeholders, expanded_command, path
313        );
314        return error_to_action_response(&e);
315    }
316
317    let tool = match BuiltinTool::from_command(&expanded_command) {
318        Ok(t) => t,
319        Err(e) => {
320            warn!("Unknown tool: {}", e);
321            return error_to_action_response(&e);
322        }
323    };
324
325    let working_dir = file_path.parent().unwrap_or(&file_path);
326    let executor = ToolExecutor::new(working_dir.to_path_buf(), state.limits.clone())
327        .with_sandbox_env((*state.sandbox_env).clone());
328
329    info!("Executing tool: {:?}", tool);
330
331    match executor.execute(&tool).await {
332        Ok(output) => {
333            let response = ActionResponse::success(output.to_text());
334            (StatusCode::OK, Json(response)).into_response()
335        }
336        Err(e) => {
337            let status = e.status_code();
338            let response = ActionResponse::error(e.user_message());
339            (status, Json(response)).into_response()
340        }
341    }
342}
343
344fn error_response(status: StatusCode, message: &str) -> Response {
345    let response = ActionResponse::error(message.to_string());
346    (status, Json(response)).into_response()
347}
348
349#[cfg(test)]
350#[allow(clippy::unwrap_used)]
351mod tests {
352    use super::*;
353    use axum::body;
354    use std::collections::HashMap;
355
356    async fn response_text(response: Response) -> String {
357        let bytes = body::to_bytes(response.into_body(), usize::MAX)
358            .await
359            .unwrap();
360        String::from_utf8_lossy(&bytes).to_string()
361    }
362
363    #[tokio::test]
364    async fn eval_pages_set_no_store_cache_control() {
365        let dir = tempfile::tempdir().unwrap();
366        std::fs::write(
367            dir.path().join("README.md"),
368            "```component\necho hello\n```\n",
369        )
370        .unwrap();
371
372        let config = ServerConfig::new(dir.path().to_path_buf());
373        let state = ServerState::from_config(&config).unwrap();
374
375        let response = serve_page("README.md", &HashMap::new(), &state).await;
376        assert_eq!(response.status(), StatusCode::OK);
377
378        let cache_control = response
379            .headers()
380            .get(header::CACHE_CONTROL)
381            .and_then(|v| v.to_str().ok());
382        assert_eq!(cache_control, Some("no-store"));
383    }
384
385    #[tokio::test]
386    async fn query_params_injected_into_component_env() {
387        let dir = tempfile::tempdir().unwrap();
388        std::fs::write(
389            dir.path().join("README.md"),
390            "```component\nprintf '%s/%s' \"$USER_ID\" \"$PAGE\"\n```\n",
391        )
392        .unwrap();
393
394        let config = ServerConfig::new(dir.path().to_path_buf());
395        let state = ServerState::from_config(&config).unwrap();
396        let query = HashMap::from([
397            ("USER_ID".to_string(), "42".to_string()),
398            ("PAGE".to_string(), "stats".to_string()),
399        ]);
400
401        let response = serve_page("README.md", &query, &state).await;
402        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
403            .await
404            .unwrap();
405        assert_eq!(String::from_utf8_lossy(&body).trim_end(), "42/stats");
406    }
407
408    #[tokio::test]
409    async fn configured_env_overrides_query_params() {
410        let dir = tempfile::tempdir().unwrap();
411        std::fs::write(
412            dir.path().join("README.md"),
413            "```component\necho \"$USER_ID\"\n```\n",
414        )
415        .unwrap();
416
417        let config = ServerConfig::new(dir.path().to_path_buf()).with_env(HashMap::from([(
418            "USER_ID".to_string(),
419            "trusted".to_string(),
420        )]));
421        let state = ServerState::from_config(&config).unwrap();
422        let query = HashMap::from([("USER_ID".to_string(), "untrusted".to_string())]);
423
424        let response = serve_page("README.md", &query, &state).await;
425        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
426            .await
427            .unwrap();
428        assert_eq!(String::from_utf8_lossy(&body).trim_end(), "trusted");
429    }
430
431    #[tokio::test]
432    async fn invalid_query_key_returns_bad_request() {
433        let dir = tempfile::tempdir().unwrap();
434        std::fs::write(dir.path().join("README.md"), "ok\n").unwrap();
435
436        let config = ServerConfig::new(dir.path().to_path_buf());
437        let state = ServerState::from_config(&config).unwrap();
438        let query = HashMap::from([("A=B".to_string(), "1".to_string())]);
439
440        let response = serve_page("README.md", &query, &state).await;
441        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
442    }
443
444    #[tokio::test]
445    async fn invalid_query_value_returns_bad_request() {
446        let dir = tempfile::tempdir().unwrap();
447        std::fs::write(dir.path().join("README.md"), "ok\n").unwrap();
448
449        let config = ServerConfig::new(dir.path().to_path_buf());
450        let state = ServerState::from_config(&config).unwrap();
451        let query = HashMap::from([("USER_ID".to_string(), "abc\0def".to_string())]);
452
453        let response = serve_page("README.md", &query, &state).await;
454        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
455    }
456
457    #[tokio::test]
458    async fn action_expands_trusted_literal_env_segments() {
459        let dir = tempfile::tempdir().unwrap();
460        std::fs::write(
461            dir.path().join("README.md"),
462            "---\ntools:\n  - [echo, $DATABASE_URL]\n---\n",
463        )
464        .unwrap();
465
466        let config = ServerConfig::new(dir.path().to_path_buf())
467            .with_env(HashMap::from([(
468                "DATABASE_URL".to_string(),
469                "postgresql://gateway:gateway@localhost:5432/gateway_dev".to_string(),
470            )]))
471            .with_sandbox_env(SandboxEnv::from_host_process());
472        let state = ServerState::from_config(&config).unwrap();
473
474        let request = ActionRequest {
475            command: vec!["echo".to_string(), "$DATABASE_URL".to_string()],
476            args: HashMap::new(),
477            env: HashMap::new(),
478        };
479
480        let response = execute_action("README.md", &state, request).await;
481        assert_eq!(response.status(), StatusCode::OK);
482
483        let body = response_text(response).await;
484        assert!(body.contains("postgresql://gateway:gateway@localhost:5432/gateway_dev"));
485    }
486
487    #[tokio::test]
488    async fn action_does_not_expand_placeholders_into_trusted_env() {
489        let dir = tempfile::tempdir().unwrap();
490        std::fs::write(
491            dir.path().join("README.md"),
492            "---\ntools:\n  - [echo, { }]\n---\n",
493        )
494        .unwrap();
495
496        let config = ServerConfig::new(dir.path().to_path_buf())
497            .with_env(HashMap::from([(
498                "DATABASE_URL".to_string(),
499                "postgresql://gateway:gateway@localhost:5432/gateway_dev".to_string(),
500            )]))
501            .with_sandbox_env(SandboxEnv::from_host_process());
502        let state = ServerState::from_config(&config).unwrap();
503
504        let request = ActionRequest {
505            command: vec!["echo".to_string(), "$DATABASE_URL".to_string()],
506            args: HashMap::new(),
507            env: HashMap::new(),
508        };
509
510        let response = execute_action("README.md", &state, request).await;
511        assert_eq!(response.status(), StatusCode::OK);
512
513        let body = response_text(response).await;
514        assert!(body.contains("$DATABASE_URL"));
515        assert!(!body.contains("postgresql://gateway:gateway@localhost:5432/gateway_dev"));
516    }
517}