netget/server/sse/
mod.rs

1//! Server-sent events server. Each connection carries one HTTP/1.1 request; the model answers
2//! it with a stream (kept open, with keep-alive comments) or an ordinary response. Open streams
3//! are registered per server for `sse_broadcast` and as peers for injected events.
4pub mod actions;
5pub mod wire;
6
7use crate::llm::action_helper::call_llm;
8use crate::llm::actions::protocol_trait::ActionResult;
9use crate::logging::emit::Log;
10use crate::protocol::{Event, SpawnContext};
11use crate::server::accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS};
12use crate::server::connection::ConnectionId;
13use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
14use anyhow::Result;
15use serde_json::{json, Map, Value};
16use std::collections::HashMap;
17use std::future::Future;
18use std::net::SocketAddr;
19use std::pin::Pin;
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::sync::{Arc, LazyLock};
22use std::task::{Context, Poll};
23use std::time::Duration;
24use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt, WriteHalf};
25use tokio::net::{TcpListener, TcpStream};
26use tokio::sync::{Mutex, Notify};
27
28/// A request head at most.
29pub const MAX_REQUEST_HEAD_BYTES: usize = 16 * 1024;
30/// How long a request head may take to arrive.
31pub const HEAD_TIMEOUT: Duration = Duration::from_secs(10);
32pub const MAX_KEEPALIVE_SECS: u64 = 3600;
33/// How long one write to a connection may make no progress, and how long a broadcast, an
34/// injected event or a keepalive may take to reach one open stream (waiting for its writer
35/// included). A stream that misses it is dropped: removed from the broadcast table, its write
36/// side shut down, its connection ended.
37pub const SUBSCRIBER_WRITE_TIMEOUT: Duration = Duration::from_secs(5);
38
39type Writer = Arc<Mutex<StreamWriter>>;
40
41/// Marks a stream that missed `SUBSCRIBER_WRITE_TIMEOUT`. Every later write to it fails at once,
42/// and its connection task, waiting on `notify`, ends the connection.
43#[derive(Default)]
44struct Stall {
45    tripped: AtomicBool,
46    notify: Notify,
47}
48
49impl Stall {
50    fn trip(&self) {
51        if !self.tripped.swap(true, Ordering::SeqCst) {
52            self.notify.notify_one();
53        }
54    }
55
56    fn is_tripped(&self) -> bool {
57        self.tripped.load(Ordering::SeqCst)
58    }
59}
60
61fn stalled_error() -> std::io::Error {
62    std::io::Error::new(
63        std::io::ErrorKind::TimedOut,
64        format!(
65            "the peer took no bytes for {}s",
66            SUBSCRIBER_WRITE_TIMEOUT.as_secs()
67        ),
68    )
69}
70
71/// A connection's write half that fails a write making no progress for
72/// `SUBSCRIBER_WRITE_TIMEOUT` and trips the connection's `Stall`. Every writer of the
73/// connection goes through it — the request's own answer, broadcasts, keepalives and the
74/// generic peer-command task that writes injected events — so a peer that stops reading holds
75/// none of them longer than that.
76pub struct StreamWriter {
77    inner: WriteHalf<TcpStream>,
78    stall: Arc<Stall>,
79    deadline: Option<Pin<Box<tokio::time::Sleep>>>,
80}
81
82impl StreamWriter {
83    fn guard<T>(
84        &mut self,
85        cx: &mut Context<'_>,
86        poll: Poll<std::io::Result<T>>,
87    ) -> Poll<std::io::Result<T>> {
88        match poll {
89            Poll::Ready(r) => {
90                self.deadline = None;
91                Poll::Ready(r)
92            }
93            Poll::Pending => {
94                let sleep = self
95                    .deadline
96                    .get_or_insert_with(|| Box::pin(tokio::time::sleep(SUBSCRIBER_WRITE_TIMEOUT)));
97                if sleep.as_mut().poll(cx).is_ready() {
98                    self.deadline = None;
99                    self.stall.trip();
100                    return Poll::Ready(Err(stalled_error()));
101                }
102                Poll::Pending
103            }
104        }
105    }
106}
107
108impl AsyncWrite for StreamWriter {
109    fn poll_write(
110        self: Pin<&mut Self>,
111        cx: &mut Context<'_>,
112        buf: &[u8],
113    ) -> Poll<std::io::Result<usize>> {
114        let this = self.get_mut();
115        if this.stall.is_tripped() {
116            return Poll::Ready(Err(stalled_error()));
117        }
118        let poll = Pin::new(&mut this.inner).poll_write(cx, buf);
119        this.guard(cx, poll)
120    }
121
122    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
123        let this = self.get_mut();
124        if this.stall.is_tripped() {
125            return Poll::Ready(Err(stalled_error()));
126        }
127        let poll = Pin::new(&mut this.inner).poll_flush(cx);
128        this.guard(cx, poll)
129    }
130
131    /// Passes through even on a tripped stream: shutting the write side is how one is ended.
132    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
133        let this = self.get_mut();
134        let poll = Pin::new(&mut this.inner).poll_shutdown(cx);
135        this.guard(cx, poll)
136    }
137}
138
139struct Subscriber {
140    path: String,
141    writer: Writer,
142    stall: Arc<Stall>,
143}
144
145/// Writes `bytes` to one stream within `SUBSCRIBER_WRITE_TIMEOUT`, waiting for its writer
146/// included. `false` means the stream missed it or failed; the caller drops it.
147async fn deliver(writer: &Writer, bytes: &[u8]) -> bool {
148    let write = async {
149        let mut w = writer.lock().await;
150        w.write_all(bytes).await?;
151        w.flush().await
152    };
153    matches!(
154        tokio::time::timeout(SUBSCRIBER_WRITE_TIMEOUT, write).await,
155        Ok(Ok(()))
156    )
157}
158
159/// Shuts a dropped stream's write side, bounded like any other write to it.
160async fn shut(writer: &Writer) {
161    let _ = tokio::time::timeout(SUBSCRIBER_WRITE_TIMEOUT, async {
162        writer.lock().await.shutdown().await
163    })
164    .await;
165}
166
167/// Every open stream of every running SSE server, by the server's bound address (unique among
168/// running servers in a process, where server ids are not: each `AppState` counts its own) and
169/// connection.
170static HUBS: LazyLock<std::sync::Mutex<HashMap<SocketAddr, Arc<Mutex<HashMap<u32, Subscriber>>>>>> =
171    LazyLock::new(Default::default);
172
173fn hub(server: SocketAddr) -> Arc<Mutex<HashMap<u32, Subscriber>>> {
174    HUBS.lock()
175        .unwrap_or_else(|e| e.into_inner())
176        .entry(server)
177        .or_default()
178        .clone()
179}
180
181/// Removes a server's subscriber table when its accept loop ends (stop_server aborts it).
182struct HubGuard(SocketAddr);
183impl Drop for HubGuard {
184    fn drop(&mut self) {
185        HUBS.lock()
186            .unwrap_or_else(|e| e.into_inner())
187            .remove(&self.0);
188    }
189}
190
191/// Send already-encoded event bytes to every open stream of a server (optionally one path),
192/// to all of them at once. Returns how many streams took them. A stream that misses
193/// `SUBSCRIBER_WRITE_TIMEOUT` or fails is removed from the table here and its connection
194/// ended, so one peer that stops reading costs a broadcast at most that long and the others
195/// still get it.
196pub async fn broadcast(
197    state: &crate::state::app_state::AppState,
198    server_id: crate::state::ServerId,
199    path: Option<&str>,
200    bytes: &[u8],
201) -> usize {
202    let Some(addr) = state.get_server(server_id).await.and_then(|s| s.local_addr) else {
203        return 0;
204    };
205    let hub = hub(addr);
206    let targets: Vec<(u32, Writer, Arc<Stall>)> = {
207        let subs = hub.lock().await;
208        subs.iter()
209            .filter(|(_, s)| path.is_none_or(|p| s.path == p))
210            .map(|(id, s)| (*id, s.writer.clone(), s.stall.clone()))
211            .collect()
212    };
213    let results =
214        futures::future::join_all(targets.into_iter().map(|(id, writer, stall)| async move {
215            let ok = deliver(&writer, bytes).await;
216            (id, ok, stall)
217        }))
218        .await;
219    let mut delivered = 0;
220    for (id, ok, stall) in results {
221        if ok {
222            delivered += 1;
223        } else {
224            hub.lock().await.remove(&id);
225            stall.trip();
226            tracing::warn!("SSE stream {id} dropped: a broadcast did not reach it");
227        }
228    }
229    delivered
230}
231
232fn keepalive(ctx: &SpawnContext) -> Result<Duration> {
233    let secs = ctx
234        .startup_params
235        .as_ref()
236        .map(|p| p.get_optional_u64("keepalive_secs"))
237        .transpose()?
238        .flatten()
239        .unwrap_or(actions::KEEPALIVE_SECS);
240    anyhow::ensure!(
241        (1..=MAX_KEEPALIVE_SECS).contains(&secs),
242        "keepalive_secs must be 1..={MAX_KEEPALIVE_SECS}"
243    );
244    Ok(Duration::from_secs(secs))
245}
246
247pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
248    let keepalive = keepalive(&ctx)?;
249    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
250    let addr = listener.local_addr()?;
251    Log::new(Some(&ctx.status_tx)).info(format!("SSE server listening on {addr}"));
252    let state = ctx.state.clone();
253    let server_id = ctx.server_id;
254    let accept = tokio::spawn(async move {
255        let _hub = HubGuard(addr);
256        let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
257        loop {
258            let (socket, peer, permit) = match accept_bounded(
259                &listener,
260                &limiter,
261                b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
262                "SSE",
263                Some(&ctx.status_tx),
264            )
265            .await
266            {
267                Ok(v) => v,
268                Err(_) => break,
269            };
270            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
271            let now = crate::utils::clock::Instant::now();
272            ctx.state
273                .add_connection_to_server(
274                    server_id,
275                    ConnectionState {
276                        id,
277                        remote_addr: peer,
278                        local_addr: addr,
279                        bytes_sent: 0,
280                        bytes_received: 0,
281                        packets_sent: 0,
282                        packets_received: 0,
283                        last_activity: now,
284                        status: ConnectionStatus::Active,
285                        status_changed_at: now,
286                        protocol_info: ProtocolConnectionInfo::empty(),
287                    },
288                )
289                .await;
290            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
291            let child = ctx.clone();
292            ctx.state
293                .spawn_server_task(server_id, async move {
294                    let _permit = permit;
295                    if let Err(e) = connection(&child, addr, id, socket, peer, keepalive).await {
296                        Log::new(Some(&child.status_tx))
297                            .debug(format!("SSE connection {id} ended: {e}"));
298                    }
299                    hub(addr).lock().await.remove(&id.as_u32());
300                    child.state.close_connection_on_server(server_id, id).await;
301                    let _ = child.status_tx.send("__UPDATE_UI__".into());
302                })
303                .await;
304        }
305    });
306    state.register_server_task(server_id, accept).await;
307    Ok(addr)
308}
309
310struct Request {
311    method: String,
312    path: String,
313    query: Map<String, Value>,
314    headers: HashMap<String, String>,
315}
316
317fn head_end(buf: &[u8]) -> Option<usize> {
318    buf.windows(4).position(|w| w == b"\r\n\r\n").map(|i| i + 4)
319}
320
321fn parse_head(head: &[u8]) -> Option<Request> {
322    let text = std::str::from_utf8(head).ok()?;
323    let mut lines = text.split("\r\n");
324    let mut first = lines.next()?.split(' ');
325    let method = first.next()?.to_string();
326    let target = first.next()?;
327    if !first.next()?.starts_with("HTTP/1.") {
328        return None;
329    }
330    let (path, query) = target.split_once('?').unwrap_or((target, ""));
331    let query = url::form_urlencoded::parse(query.as_bytes())
332        .map(|(k, v)| (k.into_owned(), Value::String(v.into_owned())))
333        .collect();
334    let headers = lines
335        .filter_map(|l| l.split_once(':'))
336        .map(|(k, v)| (k.trim().to_ascii_lowercase(), v.trim().to_string()))
337        .collect();
338    Some(Request {
339        method,
340        path: path.to_string(),
341        query,
342        headers,
343    })
344}
345
346fn status_text(status: u16) -> &'static str {
347    match status {
348        200 => "OK",
349        201 => "Created",
350        202 => "Accepted",
351        204 => "No Content",
352        301 => "Moved Permanently",
353        302 => "Found",
354        304 => "Not Modified",
355        400 => "Bad Request",
356        401 => "Unauthorized",
357        403 => "Forbidden",
358        404 => "Not Found",
359        405 => "Method Not Allowed",
360        413 => "Content Too Large",
361        429 => "Too Many Requests",
362        500 => "Internal Server Error",
363        503 => "Service Unavailable",
364        _ => "Status",
365    }
366}
367
368fn response(status: u16, content_type: &str, body: &[u8], extra: &str) -> Vec<u8> {
369    let mut out = format!(
370        "HTTP/1.1 {status} {}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n{extra}\r\n",
371        status_text(status),
372        body.len()
373    )
374    .into_bytes();
375    out.extend_from_slice(body);
376    out
377}
378
379fn outcome(ctx: &SpawnContext, id: ConnectionId, path: &str, decision: &str) {
380    let line = format!("SSE connection {id} path={path} decision={decision}");
381    if decision.starts_with("fail_closed") {
382        tracing::warn!("{line}");
383    } else {
384        tracing::info!("{line}");
385    }
386    let _ = ctx.status_tx.send(format!("[INFO] {line}"));
387}
388
389/// What the model's answer says to do with the request.
390#[derive(Default)]
391struct Answer {
392    accept: bool,
393    retry_ms: Option<u64>,
394    respond: Option<Value>,
395    stream: Vec<u8>,
396    close: bool,
397}
398
399fn collect(results: Vec<ActionResult>, answer: &mut Answer) {
400    for r in results {
401        match r {
402            ActionResult::Output(bytes) => answer.stream.extend_from_slice(&bytes),
403            ActionResult::CloseConnection => answer.close = true,
404            ActionResult::Multiple(items) => collect(items, answer),
405            ActionResult::Custom { name, data } if name == actions::ACCEPT => {
406                answer.accept = true;
407                answer.retry_ms = data["retry_ms"].as_u64();
408            }
409            ActionResult::Custom { name, data } if name == actions::RESPOND => {
410                answer.respond.get_or_insert(data);
411            }
412            _ => {}
413        }
414    }
415}
416
417async fn connection(
418    ctx: &SpawnContext,
419    addr: SocketAddr,
420    id: ConnectionId,
421    socket: TcpStream,
422    peer: SocketAddr,
423    keepalive: Duration,
424) -> Result<()> {
425    let (mut reader, writer) = tokio::io::split(socket);
426    let stall = Arc::new(Stall::default());
427    let writer: Writer = Arc::new(Mutex::new(StreamWriter {
428        inner: writer,
429        stall: stall.clone(),
430        deadline: None,
431    }));
432    let mut buf = Vec::new();
433    let mut chunk = [0u8; 4096];
434    let deadline = tokio::time::Instant::now() + HEAD_TIMEOUT;
435    let head_len = loop {
436        if let Some(n) = head_end(&buf) {
437            break n;
438        }
439        if buf.len() > MAX_REQUEST_HEAD_BYTES {
440            outcome(ctx, id, "", "fail_closed_head_too_large");
441            let mut w = writer.lock().await;
442            w.write_all(&response(431, "text/plain", b"", "")).await?;
443            return Ok(());
444        }
445        let want = (MAX_REQUEST_HEAD_BYTES + 1 - buf.len()).min(chunk.len());
446        let n = match tokio::time::timeout_at(deadline, reader.read(&mut chunk[..want])).await {
447            Ok(r) => r?,
448            Err(_) => anyhow::bail!(
449                "request head not complete within {}s",
450                HEAD_TIMEOUT.as_secs()
451            ),
452        };
453        if n == 0 {
454            return Ok(());
455        }
456        buf.extend_from_slice(&chunk[..n]);
457    };
458    ctx.state
459        .update_connection_stats(
460            ctx.server_id,
461            id,
462            Some(buf.len() as u64),
463            None,
464            Some(1),
465            None,
466        )
467        .await;
468    let Some(req) = parse_head(&buf[..head_len]) else {
469        outcome(ctx, id, "", "fail_closed_bad_request");
470        writer
471            .lock()
472            .await
473            .write_all(&response(400, "text/plain", b"", ""))
474            .await?;
475        return Ok(());
476    };
477    if req.method != "GET" && req.method != "HEAD" {
478        outcome(ctx, id, &req.path, "fail_closed_method");
479        writer
480            .lock()
481            .await
482            .write_all(&response(405, "text/plain", b"", "Allow: GET, HEAD\r\n"))
483            .await?;
484        return Ok(());
485    }
486    let declares_body = req
487        .headers
488        .get("content-length")
489        .is_some_and(|l| l.trim() != "0")
490        || req.headers.contains_key("transfer-encoding");
491    if declares_body {
492        outcome(ctx, id, &req.path, "fail_closed_body_not_accepted");
493        writer
494            .lock()
495            .await
496            .write_all(&response(413, "text/plain", b"", ""))
497            .await?;
498        return Ok(());
499    }
500    let h = |k: &str| req.headers.get(k).cloned().unwrap_or_default();
501    let event = Event::new(
502        &actions::REQUEST_EVENT,
503        json!({
504            "method": req.method, "path": req.path, "query": req.query,
505            "last_event_id": h("last-event-id"), "accept": h("accept"),
506            "user_agent": h("user-agent"), "client_ip": peer.ip().to_string(),
507        }),
508    );
509    let protocol = Arc::new(actions::SseProtocol);
510    let mut answer = Answer::default();
511    match call_llm(
512        &ctx.llm_client,
513        &ctx.state,
514        ctx.server_id,
515        Some(id),
516        &event,
517        protocol.as_ref(),
518    )
519    .await
520    {
521        Ok(r) => collect(r.protocol_results, &mut answer),
522        Err(e) => {
523            outcome(ctx, id, &req.path, "fail_closed_llm_error");
524            let text = crate::utils::wire_failure::wire_failure_text(&e);
525            writer
526                .lock()
527                .await
528                .write_all(&response(
529                    503,
530                    "text/plain; charset=utf-8",
531                    text.as_bytes(),
532                    "Retry-After: 5\r\n",
533                ))
534                .await?;
535            return Ok(());
536        }
537    }
538    let mut w = writer.lock().await;
539    if let Some(r) = answer.respond {
540        outcome(ctx, id, &req.path, "model_respond");
541        let status = r["status"].as_u64().unwrap_or(500) as u16;
542        let body = r["body"].as_str().unwrap_or_default();
543        let ct = r["content_type"]
544            .as_str()
545            .unwrap_or("text/plain; charset=utf-8");
546        let mut bytes = response(status, ct, body.as_bytes(), "");
547        if req.method == "HEAD" {
548            bytes.truncate(bytes.len() - body.len());
549        }
550        w.write_all(&bytes).await?;
551        return Ok(());
552    }
553    if !answer.accept && answer.stream.is_empty() {
554        // An EventSource stops reconnecting on 204; nothing else says "no stream" as plainly.
555        outcome(ctx, id, &req.path, "model_silent");
556        w.write_all(&response(204, "text/plain", b"", "")).await?;
557        return Ok(());
558    }
559    outcome(ctx, id, &req.path, "model_stream");
560    let mut first = b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\nX-Accel-Buffering: no\r\n\r\n".to_vec();
561    if req.method == "HEAD" {
562        w.write_all(&first).await?;
563        return Ok(());
564    }
565    if let Some(retry) = answer.retry_ms {
566        first.extend_from_slice(&wire::encode_retry(retry));
567    }
568    first.extend_from_slice(&answer.stream);
569    if !answer.close {
570        // Registered while this stream's writer is held, so a broadcast or an injected event
571        // that races the first write waits for it instead of going missing or ahead of it.
572        hub(addr).lock().await.insert(
573            id.as_u32(),
574            Subscriber {
575                path: req.path.clone(),
576                writer: writer.clone(),
577                stall: stall.clone(),
578            },
579        );
580        let commands = crate::server::peer_support::register_peer_channel(
581            &ctx.state,
582            ctx.server_id,
583            id.as_u32(),
584        )
585        .await;
586        crate::server::peer_support::spawn_peer_command_task(
587            commands,
588            protocol,
589            ctx.state.clone(),
590            ctx.server_id,
591            id.as_u32(),
592            writer.clone(),
593            ctx.status_tx.clone(),
594        );
595    }
596    w.write_all(&first).await?;
597    w.flush().await?;
598    ctx.state
599        .update_connection_stats(
600            ctx.server_id,
601            id,
602            None,
603            Some(first.len() as u64),
604            None,
605            Some(1),
606        )
607        .await;
608    if answer.close {
609        let _ = w.shutdown().await;
610        return Ok(());
611    }
612    drop(w);
613    let _ = ctx.status_tx.send("__UPDATE_UI__".into());
614    let mut tick = tokio::time::interval(keepalive);
615    tick.tick().await;
616    loop {
617        tokio::select! {
618            read = reader.read(&mut chunk) => {
619                // A client says nothing more on an event stream; EOF or an error ends it.
620                match read {
621                    Ok(0) | Err(_) => return Ok(()),
622                    Ok(_) => {}
623                }
624            }
625            _ = stall.notify.notified() => {
626                outcome(ctx, id, &req.path, "fail_closed_subscriber_stalled");
627                hub(addr).lock().await.remove(&id.as_u32());
628                shut(&writer).await;
629                return Ok(());
630            }
631            _ = tick.tick() => {
632                if !deliver(&writer, b": keepalive\n\n").await {
633                    // Ended by the stall branch on the next turn, like any other missed write.
634                    stall.trip();
635                }
636            }
637        }
638    }
639}
netget/server/sse/mod.rs

netget/server/sse/
mod.rs
