trustee_api/
lib.rs

1//! Trustee API — REST + WebSocket server for the Trustee agent.
2//!
3//! Wraps a [`trustee_core::session::Session`] and exposes it over HTTP.
4//! Static frontend files are served from [`trustee_web`].
5//!
6//! Authentication is optional. When `[oidc]` or `[dev]` sections are present
7//! in the config TOML, all `/api/v1/*` endpoints require a valid JWT or dev
8//! token. Otherwise, all endpoints are open.
9
10pub mod auth;
11pub mod tls;
12mod routes;
13mod state;
14mod thq_register;
15
16use std::net::SocketAddr;
17use std::sync::Arc;
18
19use anyhow::Result;
20use axum::routing::{get, post};
21use tower_http::cors::CorsLayer;
22
23pub use auth::{AuthConfig, AuthState};
24pub use state::ServerState;
25
26/// Run the API server.
27///
28/// Creates a `Session` with the given config, starts a background task to
29/// drain workflow messages and broadcast them to WebSocket clients, then
30/// serves the REST + WebSocket + static files on `addr`.
31///
32/// If `[oidc]` or `[dev]` sections are found in the config TOML, auth is
33/// enabled — all `/api/v1/*` endpoints (except health) require a valid token.
34///
35/// By default serves over HTTPS using a self-signed certificate from
36/// `~/.trustee/certs/`. If `use_tls` is false, serves plain HTTP.
37pub async fn run(
38    config_toml: String,
39    secrets: std::collections::HashMap<String, String>,
40    build_info: trustee_core::types::BuildInfo,
41    addr: SocketAddr,
42    use_tls: bool,
43) -> Result<()> {
44    // Parse auth config from TOML (returns None if no [oidc] or [dev] sections)
45    let auth_state = AuthConfig::from_toml(&config_toml).map(|cfg| {
46        let is_dev = cfg.dev_config.local_dev_mode;
47        tracing::info!(
48            "Auth enabled: {} mode, issuer={}",
49            if is_dev { "development" } else { "production" },
50            cfg.issuer_url
51        );
52        Arc::new(AuthState::new(cfg))
53    });
54
55    // Parse THQ registration config before config_toml is moved into session
56    let thq_config = thq_register::ThqConfig::from_toml(&config_toml);
57
58    // Build the session — keep copies of secrets/build_info for per-user sessions
59    let config_toml_for_state = config_toml.clone();
60    let secrets_for_state = secrets.clone();
61    let build_info_for_state = build_info.clone();
62    let (mut session, workflow_rx) = trustee_core::session::Session::new();
63    session.config_toml = Some(config_toml);
64    session.secrets = Some(secrets);
65    session.build_info = Some(build_info);
66    session.parse_auto_handoff_config();
67
68    // Extract agent name from config TOML for stateless operation
69    if let Some(ref config_toml_str) = session.config_toml {
70        if let Ok(table) = config_toml_str.parse::<toml::Value>() {
71            if let Some(name) = table.get("agent").and_then(|a| a.get("name")).and_then(|n| n.as_str()) {
72                session.agent_name = name.to_string();
73            }
74        }
75    }
76
77    // Create the broadcast channel for WebSocket fan-out
78    let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(256);
79
80    // Wrap session in shared state (with shared config/secrets/build_info for per-user sessions)
81    // Parse max_sessions_per_user from [web] section
82    let max_sessions: usize = {
83        let config_str: &str = &config_toml_for_state;
84        match toml::from_str::<toml::Value>(config_str) {
85            Ok(v) => v
86                .get("web")
87                .and_then(|w| w.as_table())
88                .and_then(|w| w.get("max_sessions_per_user").and_then(|v| v.as_integer()))
89                .map(|v| v as usize)
90                .unwrap_or(4),
91            Err(_) => 4,
92        }
93    };
94
95    let state = ServerState::new(session, ws_tx, auth_state)
96        .with_config_toml(config_toml_for_state)
97        .with_secrets(secrets_for_state)
98        .with_build_info(build_info_for_state)
99        .with_max_sessions_per_user(max_sessions);
100
101    // Start background message drain task (owns workflow_rx directly — no deadlock)
102    state.clone().spawn_drain_task(workflow_rx);
103
104    // THQ auto-registration with Torpi (if [thq] section is present in config)
105    if let Some(cfg) = thq_config {
106        thq_register::spawn(cfg);
107    } else {
108        tracing::debug!("THQ registration not configured (no [thq] section)");
109    }
110
111    // Build router
112    //
113    // Auth middleware approach: since axum 0.8's from_fn_with_state has
114    // trait bound issues with nested routers, we apply auth checking at
115    // the handler level via a helper. Each protected route's handler
116    // calls auth::check_auth() first. This is simpler and avoids type
117    // complexity.
118    let app = axum::Router::new()
119        // Public routes
120        .route("/api/v1/health", get(routes::health))
121        .nest("/auth", auth::auth_routes())
122        // Protected API routes
123        .route("/api/v1/session", get(routes::get_session))
124        .route("/api/v1/session/command", post(routes::post_command))
125        .route("/api/v1/session/cancel", post(routes::post_cancel))
126        .route("/api/v1/session/handoff", post(routes::post_handoff))
127        .route("/api/v1/session/stream", get(routes::ws_handler))
128        // Session naming
129        .route("/api/v1/session/name", post(routes::set_session_name))
130        .route("/api/v1/session/new", post(routes::new_session))
131        .route("/api/v1/project/name", post(routes::set_project_name))
132        // Session discovery & resume
133        // Session discovery & resume (checkpoint-based, existing)
134        .route("/api/v1/sessions", get(routes::list_sessions).post(routes::create_session))
135        .route("/api/v1/sessions/live", get(routes::list_live_sessions))
136        .route("/api/v1/sessions/{id}", get(routes::get_session_detail).delete(routes::destroy_session))
137        .route("/api/v1/sessions/{id}/live", get(routes::get_live_session))
138        .route("/api/v1/sessions/{id}/resume", post(routes::resume_session))
139        .route("/api/v1/sessions/{id}/history", get(routes::get_session_history))
140        // MSU: session-scoped live routes
141        .route("/api/v1/sessions/{id}/command", post(routes::post_command_session))
142        .route("/api/v1/sessions/{id}/cancel", post(routes::post_cancel_session))
143        .route("/api/v1/sessions/{id}/handoff", post(routes::post_handoff_session))
144        .route("/api/v1/sessions/{id}/name", post(routes::set_session_name_session))
145        .route("/api/v1/sessions/{id}/stream", get(routes::ws_session_handler))
146        // Static files from trustee-web
147        .route("/", get(routes::serve_index))
148        .route("/{file}", get(routes::serve_static))
149        .layer(CorsLayer::permissive())
150        .layer(axum::extract::DefaultBodyLimit::max(10 * 1024 * 1024))
151        .with_state(state);
152
153    // Start server
154    let listener = tokio::net::TcpListener::bind(addr).await?;
155
156    if use_tls {
157        // Install ring as the process-level crypto provider (required when
158        // rustls is built with default-features=false to avoid ambiguity
159        // with aws-lc-rs pulled in transitively by other crates).
160        let _ = rustls::crypto::ring::default_provider().install_default();
161
162        // Ensure self-signed certs exist
163        let cert_dir = tls::default_cert_dir();
164        let (cert_path, key_path) = tls::ensure_certs(&cert_dir)?;
165
166        // Load TLS config
167        let tls_config = tls::load_tls_config(&cert_path, &key_path)?;
168        let acceptor = tokio_rustls::TlsAcceptor::from(std::sync::Arc::new(tls_config));
169
170        tracing::info!("Trustee API listening on https://{}", addr);
171
172        // Manual accept loop — spawn hyper-util auto connection per TLS stream
173        loop {
174            let (tcp_stream, peer_addr) = match listener.accept().await {
175                Ok(stream) => stream,
176                Err(e) => {
177                    tracing::warn!("TCP accept failed: {}", e);
178                    continue;
179                }
180            };
181
182            let acceptor = acceptor.clone();
183            let app = app.clone();
184
185            tokio::spawn(async move {
186                let tls_stream = match acceptor.accept(tcp_stream).await {
187                    Ok(s) => s,
188                    Err(e) => {
189                        tracing::debug!("TLS accept failed from {}: {}", peer_addr, e);
190                        return;
191                    }
192                };
193
194                // Use hyper-util auto builder with the tower service from axum.
195                // serve_connection_with_upgrades is required for WebSocket support.
196                let io = hyper_util::rt::TokioIo::new(tls_stream);
197                let svc = hyper_util::service::TowerToHyperService::new(app);
198
199                let _ = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new())
200                    .serve_connection_with_upgrades(io, svc)
201                    .await;
202            });
203        }
204    } else {
205        tracing::info!("Trustee API listening on http://{}", addr);
206        axum::serve(listener, app).await?;
207    }
208
209    Ok(())
210}
trustee_api/lib.rs

trustee_api/
lib.rs
