superfluid_daemon/
apiweb.rs

1//! api-web.
2
3use std::sync::Arc;
4
5use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
6use axum::extract::State;
7use axum::http::{HeaderMap, StatusCode};
8use axum::response::sse::{Event, KeepAlive, Sse};
9use axum::response::IntoResponse;
10use axum::routing::{get, post};
11use axum::{Json, Router};
12use tokio_stream::wrappers::ReceiverStream;
13use tokio_stream::StreamExt;
14
15use crate::api::{handle_web, serve_requests, EventMsg, Request, Response, Sink};
16use crate::{Daemon, DaemonError, EventBody};
17
18#[derive(Clone)]
19struct WebState {
20    daemon: Arc<Daemon>,
21    token: Arc<String>,
22    origins: Arc<Vec<String>>,
23    fim_bucket: Arc<std::sync::Mutex<crate::completion_bucket::TokenBucket>>,
24}
25
26fn ct_eq(a: &str, b: &str) -> bool {
27    let (a, b) = (a.as_bytes(), b.as_bytes());
28    if a.len() != b.len() {
29        return false;
30    }
31    let mut diff = 0u8;
32    for i in 0..a.len() {
33        diff |= a[i] ^ b[i];
34    }
35    diff == 0
36}
37
38fn header_values<'a>(headers: &'a HeaderMap, name: &str) -> Vec<&'a str> {
39    headers
40        .get_all(name)
41        .iter()
42        .map(|v| v.to_str().unwrap_or(""))
43        .collect()
44}
45
46fn bearer_tokens(headers: &HeaderMap) -> Vec<&str> {
47    header_values(headers, "authorization")
48        .into_iter()
49        .map(|v| v.strip_prefix("Bearer ").unwrap_or(""))
50        .collect()
51}
52
53fn all_match(presented: &[&str], token: &str) -> bool {
54    !presented.is_empty() && presented.iter().all(|t| ct_eq(t, token))
55}
56
57fn origins_allowed(state: &WebState, headers: &HeaderMap) -> bool {
58    header_values(headers, "origin")
59        .iter()
60        .all(|origin| state.origins.iter().any(|o| o == origin))
61}
62
63fn authorize(state: &WebState, headers: &HeaderMap) -> Result<(), (StatusCode, &'static str)> {
64    if !all_match(&bearer_tokens(headers), &state.token) {
65        return Err((StatusCode::UNAUTHORIZED, "bad or missing bearer token"));
66    }
67    if !origins_allowed(state, headers) {
68        return Err((StatusCode::FORBIDDEN, "origin not allowed"));
69    }
70    if !all_match(&header_values(headers, "x-superfluid-csrf"), &state.token) {
71        return Err((StatusCode::FORBIDDEN, "missing or bad CSRF token"));
72    }
73    Ok(())
74}
75
76async fn rpc(
77    State(state): State<WebState>,
78    headers: HeaderMap,
79    Json(req): Json<Request>,
80) -> axum::response::Response {
81    if let Err((code, msg)) = authorize(&state, &headers) {
82        return (code, Json(Response::Err { message: msg.into() })).into_response();
83    }
84    if matches!(req, Request::Complete { .. }) {
85        let cfg = state.daemon.completion_bucket();
86        let taken = state
87            .fim_bucket
88            .lock()
89            .expect("api-web fim bucket")
90            .try_take(cfg, std::time::Instant::now());
91        if let Err(wait) = taken {
92            state.daemon.note_completion_throttled();
93            let secs = crate::completion_bucket::retry_after_secs(wait);
94            let body = Response::Throttled {
95                retry_after_ms: crate::completion_bucket::retry_after_ms(wait),
96            };
97            return (
98                StatusCode::TOO_MANY_REQUESTS,
99                [(axum::http::header::RETRY_AFTER, secs.to_string())],
100                Json(body),
101            )
102                .into_response();
103        }
104    }
105    let resp = handle_web(&state.daemon, req);
106    (StatusCode::OK, Json(resp)).into_response()
107}
108
109async fn health() -> impl IntoResponse {
110    StatusCode::OK
111}
112
113#[derive(serde::Deserialize)]
114struct StreamReq {
115    session: u64,
116    max_tokens: u32,
117    #[serde(default)]
118    provisional: bool,
119}
120
121async fn stream(
122    State(state): State<WebState>,
123    headers: HeaderMap,
124    Json(req): Json<StreamReq>,
125) -> axum::response::Response {
126    if let Err((code, msg)) = authorize(&state, &headers) {
127        return (code, Json(Response::Err { message: msg.into() })).into_response();
128    }
129    let (tx, rx) = tokio::sync::mpsc::channel::<serde_json::Value>(256);
130    let daemon = Arc::clone(&state.daemon);
131    let max_tokens = req.max_tokens;
132    let session = req.session;
133    let bus = daemon.bus();
134    let sub_id = if req.provisional {
135        let mut sub = bus.subscribe(session, true);
136        let rx_bus = sub.take_receiver();
137        let tx_p = tx.clone();
138        tokio::task::spawn_blocking(move || {
139            while let Ok(msg) = rx_bus.recv() {
140                if let crate::bus::BusMsg::Provisional { channel, text } = msg {
141                    if tx_p
142                        .blocking_send(serde_json::json!({"provisional": {"channel": channel, "text": text}}))
143                        .is_err()
144                    {
145                        break;
146                    }
147                }
148            }
149        });
150        Some((sub.session, sub.id))
151    } else {
152        None
153    };
154    let bus_end = Arc::clone(&daemon);
155    tokio::task::spawn_blocking(move || {
156        let _closed = crate::keepalive::cancel_when_closed(&daemon, session, &tx);
157        let tx2 = tx.clone();
158        let out = daemon.generate_streaming(session, max_tokens, move |ev| {
159            let msg: EventMsg = ev.clone().into();
160            let carry = matches!(
161                &ev.body,
162                EventBody::Generated { .. }
163                    | EventBody::ToolUse { .. }
164                    | EventBody::ToolResult { .. }
165                    | EventBody::ToolParseFailure { .. }
166            );
167            if carry {
168                tx2.blocking_send(serde_json::json!({ "committed": msg }))
169                    .map_err(|_| DaemonError::Protocol("generation consumer gone"))?;
170            }
171            Ok(())
172        });
173        if let Some((s, id)) = sub_id {
174            bus_end.bus().unsubscribe(s, id);
175        }
176        match out {
177            Ok(o) => {
178                let _ = tx.blocking_send(serde_json::json!({
179                    "done": {"tokens_generated": o.tokens_generated, "finish": o.finish, "warm_prefix": o.warm_prefix}
180                }));
181            }
182            Err(e) => {
183                let _ = tx.blocking_send(serde_json::json!({"error": e.to_string()}));
184            }
185        }
186    });
187    let body = ReceiverStream::new(rx)
188        .map(|v| Event::default().data(v.to_string()))
189        .chain(tokio_stream::once(Event::default().data("[DONE]")))
190        .map(Ok::<_, std::convert::Infallible>);
191    Sse::new(body).keep_alive(KeepAlive::default()).into_response()
192}
193
194const WS_INBOUND_QUEUE: usize = 32;
195const WS_INBOUND_BYTES: usize = 64 << 20;
196
197type QueuedRequest = (Result<Request, String>, tokio::sync::OwnedSemaphorePermit);
198
199const WS_OUTBOUND_BYTES: usize = 16 << 20;
200
201#[derive(Default)]
202struct OutBudget {
203    state: std::sync::Mutex<(usize, bool)>,
204    freed: std::sync::Condvar,
205    turn: std::sync::Mutex<()>,
206}
207
208impl OutBudget {
209    fn reserve(&self, len: usize) -> bool {
210        let mut st = self.state.lock().expect("ws budget");
211        while !st.1 && st.0 > 0 && st.0 + len > WS_OUTBOUND_BYTES {
212            st = self.freed.wait(st).expect("ws budget");
213        }
214        if st.1 {
215            return false;
216        }
217        st.0 += len;
218        true
219    }
220
221    fn release(&self, len: usize) {
222        let mut st = self.state.lock().expect("ws budget");
223        st.0 = st.0.saturating_sub(len);
224        self.freed.notify_all();
225    }
226
227    fn close(&self) {
228        self.state.lock().expect("ws budget").1 = true;
229        self.freed.notify_all();
230    }
231}
232
233struct CloseOnDrop(Arc<OutBudget>);
234
235impl Drop for CloseOnDrop {
236    fn drop(&mut self) {
237        self.0.close();
238    }
239}
240
241const WS_PROTOCOL: &str = "superfluid.v1";
242const WS_TOKEN_PREFIX: &str = "superfluid.token.";
243
244fn offered_protocols(headers: &HeaderMap) -> Vec<String> {
245    headers
246        .get_all(axum::http::header::SEC_WEBSOCKET_PROTOCOL)
247        .iter()
248        .filter_map(|v| v.to_str().ok())
249        .flat_map(|v| v.split(','))
250        .map(|p| p.trim().to_string())
251        .filter(|p| !p.is_empty())
252        .collect()
253}
254
255fn authorize_ws(state: &WebState, headers: &HeaderMap) -> Result<(), (StatusCode, &'static str)> {
256    let offered = offered_protocols(headers);
257    let presented: Vec<&str> = bearer_tokens(headers)
258        .into_iter()
259        .chain(
260            offered
261                .iter()
262                .filter_map(|p| p.strip_prefix(WS_TOKEN_PREFIX)),
263        )
264        .collect();
265    if !all_match(&presented, &state.token) {
266        return Err((StatusCode::UNAUTHORIZED, "bad or missing bearer token"));
267    }
268    if header_values(headers, "origin").is_empty() {
269        return Err((
270            StatusCode::FORBIDDEN,
271            "missing Origin (required on /web/ws)",
272        ));
273    }
274    if !origins_allowed(state, headers) {
275        return Err((StatusCode::FORBIDDEN, "origin not allowed"));
276    }
277    if !offered.is_empty() && !offered.iter().any(|p| p == WS_PROTOCOL) {
278        return Err((StatusCode::BAD_REQUEST, "offer the superfluid.v1 subprotocol"));
279    }
280    Ok(())
281}
282
283async fn ws(
284    State(state): State<WebState>,
285    headers: HeaderMap,
286    upgrade: WebSocketUpgrade,
287) -> axum::response::Response {
288    if let Err((code, msg)) = authorize_ws(&state, &headers) {
289        return (code, Json(Response::Err { message: msg.into() })).into_response();
290    }
291    let daemon = Arc::clone(&state.daemon);
292    upgrade
293        .protocols([WS_PROTOCOL])
294        .on_upgrade(move |socket| ws_session(socket, daemon))
295}
296
297async fn ws_session(mut socket: WebSocket, daemon: Arc<Daemon>) {
298    let (out_tx, mut out_rx) = tokio::sync::mpsc::channel::<String>(256);
299    let budget = Arc::new(OutBudget::default());
300    let _closer = CloseOnDrop(Arc::clone(&budget));
301    let producer_budget = Arc::clone(&budget);
302    let (req_tx, mut req_rx) = tokio::sync::mpsc::channel::<QueuedRequest>(WS_INBOUND_QUEUE);
303    let inbound = Arc::new(tokio::sync::Semaphore::new(WS_INBOUND_BYTES));
304    let sink: Sink = Arc::new(move |msg: &Response| {
305        let _turn = producer_budget.turn.lock().expect("ws producer turn");
306        let text = serde_json::to_string(msg)
307            .map_err(|_| DaemonError::Protocol("unserializable api-web response"))?;
308        let len = text.len();
309        if !producer_budget.reserve(len) {
310            return Err(DaemonError::Protocol("websocket closed"));
311        }
312        out_tx.blocking_send(text).map_err(|_| {
313            producer_budget.release(len);
314            DaemonError::Protocol("websocket closed")
315        })
316    });
317    let loop_sink = Arc::clone(&sink);
318    let spawned = std::thread::Builder::new()
319        .name("superfluid-web-ws".into())
320        .spawn(move || {
321            let cancels = daemon.cancel_registry();
322            let active = daemon.active_registry();
323            let err_sink = Arc::clone(&loop_sink);
324            let requests = std::iter::from_fn(move || loop {
325                match req_rx.blocking_recv().map(|(item, _permit)| item) {
326                    Some(Ok(req)) => return Some(Ok(req)),
327                    Some(Err(message)) => {
328                        if err_sink(&Response::Err { message }).is_err() {
329                            return None;
330                        }
331                    }
332                    None => return None,
333                }
334            });
335            let _ = serve_requests(requests, loop_sink, daemon, &cancels, &active);
336        });
337    drop(sink);
338    if spawned.is_err() {
339        return;
340    }
341    let mut slot: Option<tokio::sync::mpsc::Permit<'_, QueuedRequest>> = None;
342    let mut pending: Option<(Result<Request, String>, u32)> = None;
343    loop {
344        tokio::select! {
345            permit = inbound.clone().acquire_many_owned(pending.as_ref().map_or(0, |p| p.1)),
346                if pending.is_some() => {
347                let Ok(permit) = permit else { break };
348                let (item, _) = pending.take().expect("guarded");
349                slot.take().expect("slot precedes the read").send((item, permit));
350            }
351            reserved = req_tx.reserve(), if slot.is_none() && pending.is_none() => match reserved {
352                Ok(p) => slot = Some(p),
353                Err(_) => break,
354            },
355            frame = socket.recv(), if slot.is_some() && pending.is_none() => {
356                let (parsed, len) = match frame {
357                    Some(Ok(Message::Text(t))) => {
358                        (serde_json::from_str::<Request>(t.as_str()), t.len())
359                    }
360                    Some(Ok(Message::Binary(b))) => {
361                        (serde_json::from_slice::<Request>(&b), b.len())
362                    }
363                    Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue,
364                    Some(Ok(Message::Close(_)) | Err(_)) | None => break,
365                };
366                let item = parsed.map_err(|e| format!("invalid request frame: {e}"));
367                pending = Some((item, len.clamp(1, WS_INBOUND_BYTES) as u32));
368            }
369            out = out_rx.recv() => {
370                let Some(text) = out else { break };
371                let len = text.len();
372                let sent = socket.send(Message::Text(text.into())).await;
373                budget.release(len);
374                if sent.is_err() {
375                    break;
376                }
377            }
378        }
379    }
380    let _ = futures_util::SinkExt::close(&mut socket).await;
381    drop(slot);
382    drop(req_tx);
383}
384
385pub fn serve_blocking(
386    listener: std::net::TcpListener,
387    daemon: Arc<Daemon>,
388    token: String,
389    origins: Vec<String>,
390) -> std::io::Result<()> {
391    let rt = tokio::runtime::Runtime::new()?;
392    rt.block_on(async move {
393        listener.set_nonblocking(true)?;
394        let listener = tokio::net::TcpListener::from_std(listener)?;
395        let app = Router::new()
396            .route("/web/rpc", post(rpc))
397            .route("/web/stream", post(stream))
398            .route("/web/ws", get(ws))
399            .route("/web/health", get(health))
400            .with_state(WebState {
401                fim_bucket: Arc::new(std::sync::Mutex::new(
402                    crate::completion_bucket::TokenBucket::new(
403                        daemon.completion_bucket(),
404                        std::time::Instant::now(),
405                    ),
406                )),
407                daemon,
408                token: Arc::new(token),
409                origins: Arc::new(origins),
410            });
411        axum::serve(listener, app).await
412    })
413}
414
415pub fn mint_token(path: &std::path::Path) -> std::io::Result<String> {
416    let mut buf = [0u8; 32];
417    {
418        use std::io::Read;
419        std::fs::File::open("/dev/urandom")?.read_exact(&mut buf)?;
420    }
421    let token: String = buf.iter().map(|b| format!("{b:02x}")).collect();
422    let mut o = std::fs::OpenOptions::new();
423    o.create(true).write(true).truncate(true);
424    #[cfg(unix)]
425    {
426        use std::os::unix::fs::OpenOptionsExt;
427        o.mode(0o600);
428    }
429    use std::io::Write;
430    let mut f = o.open(path)?;
431    #[cfg(unix)]
432    {
433        use std::os::unix::fs::PermissionsExt;
434        f.set_permissions(std::fs::Permissions::from_mode(0o600))?;
435    }
436    f.write_all(token.as_bytes())?;
437    Ok(token)
438}
439
440#[cfg(test)]
441mod token_tests {
442    #[cfg(unix)]
443    #[test]
444    fn mint_token_forces_owner_only_even_over_a_pre_existing_file() {
445        use std::os::unix::fs::PermissionsExt;
446        let dir = std::env::temp_dir().join(format!("superfluid-tok-{}", std::process::id()));
447        std::fs::create_dir_all(&dir).unwrap();
448        let path = dir.join("web-token");
449        std::fs::write(&path, b"old").unwrap();
450        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
451        let _tok = super::mint_token(&path).unwrap();
452        let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
453        assert_eq!(mode, 0o600, "token file must be owner-only after mint");
454    }
455}
456
457#[cfg(test)]
458mod budget_tests {
459    use super::*;
460
461    #[test]
462    fn outbound_budget_backpressures_and_closes() {
463        let b = Arc::new(OutBudget::default());
464        assert!(
465            b.reserve(WS_OUTBOUND_BYTES * 2),
466            "an oversized frame alone passes"
467        );
468        b.release(WS_OUTBOUND_BYTES * 2);
469        assert!(b.reserve(WS_OUTBOUND_BYTES - 1));
470        let waiter = {
471            let b = Arc::clone(&b);
472            std::thread::spawn(move || b.reserve(2))
473        };
474        std::thread::sleep(std::time::Duration::from_millis(50));
475        assert!(!waiter.is_finished(), "over budget: the producer waits");
476        b.release(WS_OUTBOUND_BYTES - 1);
477        assert!(waiter.join().unwrap(), "released bytes let it through");
478        let waiter = {
479            let b = Arc::clone(&b);
480            std::thread::spawn(move || b.reserve(WS_OUTBOUND_BYTES))
481        };
482        std::thread::sleep(std::time::Duration::from_millis(50));
483        drop(CloseOnDrop(Arc::clone(&b)));
484        assert!(!waiter.join().unwrap(), "a closed socket fails the wait");
485    }
486}
superfluid_daemon/apiweb.rs

superfluid_daemon/
apiweb.rs
