lean_ctx/proxy/
sse_keepalive.rs

1//! SSE keepalive injection for proxy-to-client streams.
2//!
3//! When the upstream (e.g. chatgpt.com during extended thinking) goes idle
4//! for longer than the keepalive interval, this wrapper injects SSE comment
5//! lines (`: keepalive\n\n`) into the downstream stream. SSE comments are
6//! ignored by all compliant clients but reset their read-idle timers,
7//! preventing "stream disconnected" errors in Codex Desktop and similar
8//! consumers.
9//!
10//! The wrapper is transparent: every upstream byte is forwarded unchanged,
11//! and keepalives are only injected during genuine idle gaps.
12
13use std::time::Duration;
14
15use axum::body::Bytes;
16use futures::{Stream, StreamExt};
17
18/// Codex Desktop's internal read timeout is ~30s; we ping well before that.
19const DEFAULT_INTERVAL_SECS: u64 = 15;
20
21/// SSE comment — ignored by all compliant clients, resets their idle timer.
22const KEEPALIVE_BYTES: &[u8] = b": keepalive\n\n";
23
24fn keepalive_interval() -> Duration {
25    let secs = std::env::var("LEAN_CTX_PROXY_SSE_KEEPALIVE_SECS")
26        .ok()
27        .and_then(|v| v.trim().parse::<u64>().ok())
28        .filter(|s| *s > 0)
29        .unwrap_or(DEFAULT_INTERVAL_SECS);
30    Duration::from_secs(secs)
31}
32
33/// Wraps an upstream SSE byte stream: forwards every upstream chunk unchanged
34/// and injects `": keepalive\n\n"` SSE comments during idle gaps so the
35/// downstream client's read-idle timer never fires.
36///
37/// Uses `tokio::time::timeout` per chunk to detect idle periods without
38/// requiring `pin_project` or manual `Pin` implementations.
39pub fn keepalive_stream<S, E>(inner: S) -> impl Stream<Item = Result<Bytes, E>> + Send + Unpin
40where
41    S: Stream<Item = Result<Bytes, E>> + Send + Unpin + 'static,
42    E: Send + 'static,
43{
44    keepalive_stream_with_interval(inner, keepalive_interval())
45}
46
47fn keepalive_stream_with_interval<S, E>(
48    inner: S,
49    interval: Duration,
50) -> impl Stream<Item = Result<Bytes, E>> + Send + Unpin
51where
52    S: Stream<Item = Result<Bytes, E>> + Send + Unpin + 'static,
53    E: Send + 'static,
54{
55    Box::pin(futures::stream::unfold(
56        (inner, interval),
57        |(mut inner, interval)| async move {
58            match tokio::time::timeout(interval, inner.next()).await {
59                Ok(Some(item)) => Some((item, (inner, interval))),
60                Ok(None) => None,
61                Err(_timeout) => Some((Ok(Bytes::from_static(KEEPALIVE_BYTES)), (inner, interval))),
62            }
63        },
64    ))
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[tokio::test]
72    async fn keepalive_injected_during_idle() {
73        let (mut tx, rx) = futures::channel::mpsc::channel::<Result<Bytes, std::io::Error>>(8);
74        let stream = rx.map(|item| item);
75
76        let mut wrapped = keepalive_stream_with_interval(Box::pin(stream), Duration::from_secs(1));
77
78        let item = tokio::time::timeout(Duration::from_secs(3), wrapped.next())
79            .await
80            .expect("should not timeout waiting for keepalive");
81
82        let bytes = item.expect("stream not ended").expect("no error");
83        assert_eq!(bytes.as_ref(), KEEPALIVE_BYTES);
84
85        use futures::SinkExt;
86        tx.send(Ok(Bytes::from_static(b"data: hello\n\n")))
87            .await
88            .unwrap();
89        let item = wrapped.next().await.expect("stream not ended");
90        assert_eq!(item.unwrap().as_ref(), b"data: hello\n\n");
91
92        drop(tx);
93    }
94
95    #[tokio::test]
96    async fn no_keepalive_when_data_flows() {
97        let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
98            Ok(Bytes::from_static(b"data: a\n\n")),
99            Ok(Bytes::from_static(b"data: b\n\n")),
100        ];
101        let stream = futures::stream::iter(chunks);
102        let mut wrapped = keepalive_stream(Box::pin(stream));
103
104        let a = wrapped.next().await.unwrap().unwrap();
105        assert_eq!(a.as_ref(), b"data: a\n\n");
106        let b = wrapped.next().await.unwrap().unwrap();
107        assert_eq!(b.as_ref(), b"data: b\n\n");
108        assert!(wrapped.next().await.is_none());
109    }
110}
lean_ctx/proxy/sse_keepalive.rs

lean_ctx/proxy/
sse_keepalive.rs
