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
59    let (mut session, workflow_rx) = trustee_core::session::Session::new();
60    session.config_toml = Some(config_toml);
61    session.secrets = Some(secrets);
62    session.build_info = Some(build_info);
63    session.parse_auto_handoff_config();
64
65    // Extract agent name from config TOML for stateless operation
66    if let Some(ref config_toml_str) = session.config_toml {
67        if let Ok(table) = config_toml_str.parse::<toml::Value>() {
68            if let Some(name) = table.get("agent").and_then(|a| a.get("name")).and_then(|n| n.as_str()) {
69                session.agent_name = name.to_string();
70            }
71        }
72    }
73
74    // Create the broadcast channel for WebSocket fan-out
75    let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(256);
76
77    // Wrap session in shared state
78    let state = ServerState::new(session, ws_tx, auth_state);
79
80    // Start background message drain task (owns workflow_rx directly — no deadlock)
81    state.clone().spawn_drain_task(workflow_rx);
82
83    // THQ auto-registration with Torpi (if [thq] section is present in config)
84    if let Some(cfg) = thq_config {
85        thq_register::spawn(cfg);
86    } else {
87        tracing::debug!("THQ registration not configured (no [thq] section)");
88    }
89
90    // Build router
91    //
92    // Auth middleware approach: since axum 0.8's from_fn_with_state has
93    // trait bound issues with nested routers, we apply auth checking at
94    // the handler level via a helper. Each protected route's handler
95    // calls auth::check_auth() first. This is simpler and avoids type
96    // complexity.
97    let app = axum::Router::new()
98        // Public routes
99        .route("/api/v1/health", get(routes::health))
100        .nest("/auth", auth::auth_routes())
101        // Protected API routes
102        .route("/api/v1/session", get(routes::get_session))
103        .route("/api/v1/session/command", post(routes::post_command))
104        .route("/api/v1/session/cancel", post(routes::post_cancel))
105        .route("/api/v1/session/handoff", post(routes::post_handoff))
106        .route("/api/v1/session/stream", get(routes::ws_handler))
107        // Session naming
108        .route("/api/v1/session/name", post(routes::set_session_name))
109        .route("/api/v1/session/new", post(routes::new_session))
110        .route("/api/v1/project/name", post(routes::set_project_name))
111        // Session discovery & resume
112        .route("/api/v1/sessions", get(routes::list_sessions))
113        .route("/api/v1/sessions/{id}", get(routes::get_session_detail))
114        .route("/api/v1/sessions/{id}/resume", post(routes::resume_session))
115        .route("/api/v1/sessions/{id}/history", get(routes::get_session_history))
116        // Static files from trustee-web
117        .route("/", get(routes::serve_index))
118        .route("/{file}", get(routes::serve_static))
119        .layer(CorsLayer::permissive())
120        .layer(axum::extract::DefaultBodyLimit::max(10 * 1024 * 1024))
121        .with_state(state);
122
123    // Start server
124    let listener = tokio::net::TcpListener::bind(addr).await?;
125
126    if use_tls {
127        // Install ring as the process-level crypto provider (required when
128        // rustls is built with default-features=false to avoid ambiguity
129        // with aws-lc-rs pulled in transitively by other crates).
130        let _ = rustls::crypto::ring::default_provider().install_default();
131
132        // Ensure self-signed certs exist
133        let cert_dir = tls::default_cert_dir();
134        let (cert_path, key_path) = tls::ensure_certs(&cert_dir)?;
135
136        // Load TLS config
137        let tls_config = tls::load_tls_config(&cert_path, &key_path)?;
138        let acceptor = tokio_rustls::TlsAcceptor::from(std::sync::Arc::new(tls_config));
139
140        tracing::info!("Trustee API listening on https://{}", addr);
141
142        // Manual accept loop — spawn hyper-util auto connection per TLS stream
143        loop {
144            let (tcp_stream, peer_addr) = match listener.accept().await {
145                Ok(stream) => stream,
146                Err(e) => {
147                    tracing::warn!("TCP accept failed: {}", e);
148                    continue;
149                }
150            };
151
152            let acceptor = acceptor.clone();
153            let app = app.clone();
154
155            tokio::spawn(async move {
156                let tls_stream = match acceptor.accept(tcp_stream).await {
157                    Ok(s) => s,
158                    Err(e) => {
159                        tracing::debug!("TLS accept failed from {}: {}", peer_addr, e);
160                        return;
161                    }
162                };
163
164                // Use hyper-util auto builder with the tower service from axum.
165                // serve_connection_with_upgrades is required for WebSocket support.
166                let io = hyper_util::rt::TokioIo::new(tls_stream);
167                let svc = hyper_util::service::TowerToHyperService::new(app);
168
169                let _ = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new())
170                    .serve_connection_with_upgrades(io, svc)
171                    .await;
172            });
173        }
174    } else {
175        tracing::info!("Trustee API listening on http://{}", addr);
176        axum::serve(listener, app).await?;
177    }
178
179    Ok(())
180}
trustee_api/lib.rs

trustee_api/
lib.rs
