#![cfg(feature = "net")]
use std::collections::VecDeque;
use std::io;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::client::conn::http1::SendRequest;
use hyper::{Request, Response};
use hyper_util::rt::TokioIo;
use llmp::proxy::{self, BodyMode, JsonlSink, ProxyConfig, Sink};
use llmp::{ExchangeError, Record, StreamKind, read_jsonl};
use serde_json::{Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
enum Step {
Write(Vec<u8>),
Sleep(u64),
Close,
}
struct Upstream {
addr: SocketAddr,
requests: Arc<Mutex<Vec<Vec<u8>>>>,
connections: Arc<AtomicUsize>,
}
impl Upstream {
fn request(&self, index: usize) -> String {
String::from_utf8_lossy(&self.requests.lock().unwrap()[index]).into_owned()
}
}
async fn upstream(replies: Vec<Vec<Step>>) -> Upstream {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let requests = Arc::new(Mutex::new(Vec::new()));
let connections = Arc::new(AtomicUsize::new(0));
let replies = Arc::new(Mutex::new(VecDeque::from(replies)));
let (seen, counter) = (Arc::clone(&requests), Arc::clone(&connections));
tokio::spawn(async move {
loop {
let Ok((mut stream, _)) = listener.accept().await else {
return;
};
counter.fetch_add(1, Ordering::SeqCst);
let (seen, replies) = (Arc::clone(&seen), Arc::clone(&replies));
tokio::spawn(async move {
'requests: while let Some(raw) = read_request(&mut stream).await {
seen.lock().unwrap().push(raw);
let Some(steps) = replies.lock().unwrap().pop_front() else {
return;
};
for step in steps {
match step {
Step::Write(bytes) => {
if stream.write_all(&bytes).await.is_err() {
break 'requests;
}
}
Step::Sleep(ms) => tokio::time::sleep(Duration::from_millis(ms)).await,
Step::Close => break 'requests,
}
}
}
});
}
});
Upstream {
addr,
requests,
connections,
}
}
async fn read_request(stream: &mut TcpStream) -> Option<Vec<u8>> {
let mut buf = Vec::new();
let mut tmp = [0u8; 8192];
loop {
if let Some(end) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
let head = String::from_utf8_lossy(&buf[..end]).to_ascii_lowercase();
let length = head
.lines()
.find_map(|l| l.strip_prefix("content-length:"))
.and_then(|v| v.trim().parse::<usize>().ok())
.unwrap_or(0);
while buf.len() < end + 4 + length {
let n = stream.read(&mut tmp).await.ok()?;
if n == 0 {
return None;
}
buf.extend_from_slice(&tmp[..n]);
}
return Some(buf[..end + 4 + length].to_vec());
}
let n = stream.read(&mut tmp).await.ok()?;
if n == 0 {
return None;
}
buf.extend_from_slice(&tmp[..n]);
}
}
fn json_reply(body: &str) -> Vec<Step> {
vec![Step::Write(
format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}",
body.len()
)
.into_bytes(),
)]
}
fn chunked_head(content_type: &str) -> Step {
Step::Write(
format!(
"HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nTransfer-Encoding: chunked\r\n\r\n"
)
.into_bytes(),
)
}
fn chunk(data: &str) -> Step {
Step::Write(format!("{:x}\r\n{data}\r\n", data.len()).into_bytes())
}
fn last_chunk() -> Step {
Step::Write(b"0\r\n\r\n".to_vec())
}
struct Proxy {
addr: SocketAddr,
records: mpsc::UnboundedReceiver<Record>,
stop: Option<oneshot::Sender<()>>,
task: JoinHandle<io::Result<()>>,
}
impl Proxy {
async fn record(&mut self) -> Record {
tokio::time::timeout(Duration::from_secs(5), self.records.recv())
.await
.expect("no record within 5 s")
.expect("sink closed")
}
async fn stop(mut self) -> Vec<Record> {
if let Some(stop) = self.stop.take() {
let _ = stop.send(());
}
tokio::time::timeout(Duration::from_secs(5), self.task)
.await
.expect("serve did not return")
.unwrap()
.unwrap();
let mut rest = Vec::new();
while let Ok(record) = self.records.try_recv() {
rest.push(record);
}
rest
}
}
fn free_addr() -> SocketAddr {
std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
}
async fn start(upstream: SocketAddr, body_mode: BodyMode, max_body_bytes: usize) -> Proxy {
let addr = free_addr();
let (tx, records) = mpsc::unbounded_channel();
let (stop_tx, stop_rx) = oneshot::channel::<()>();
let config = ProxyConfig {
listen: addr,
upstream,
body_mode,
max_body_bytes,
};
let sink = move |record: Record| {
let _ = tx.send(record);
};
let task = tokio::spawn(proxy::serve(config, sink, async {
let _ = stop_rx.await;
}));
let deadline = Instant::now() + Duration::from_secs(5);
while TcpStream::connect(addr).await.is_err() {
assert!(Instant::now() < deadline, "proxy did not start");
tokio::time::sleep(Duration::from_millis(10)).await;
}
Proxy {
addr,
records,
stop: Some(stop_tx),
task,
}
}
async fn client(addr: SocketAddr) -> (SendRequest<Full<Bytes>>, JoinHandle<()>) {
let stream = TcpStream::connect(addr).await.unwrap();
let (sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(stream))
.await
.unwrap();
let task = tokio::spawn(async move {
let _ = conn.await;
});
(sender, task)
}
fn post(path: &str, body: &str) -> Request<Full<Bytes>> {
Request::post(path)
.header("host", "llmp.test")
.header("content-type", "application/json")
.header("user-agent", "test/1")
.header("authorization", "Bearer s3cret")
.body(Full::new(Bytes::from(body.to_owned())))
.unwrap()
}
async fn send(
sender: &mut SendRequest<Full<Bytes>>,
request: Request<Full<Bytes>>,
) -> Response<Incoming> {
sender.ready().await.unwrap();
sender.send_request(request).await.unwrap()
}
async fn body_text(response: Response<Incoming>) -> String {
let bytes = response.into_body().collect().await.unwrap().to_bytes();
String::from_utf8(bytes.to_vec()).unwrap()
}
const CHAT: &str = r#"{"model":"m","messages":[{"role":"user","content":"hello"}]}"#;
#[tokio::test(flavor = "multi_thread")]
async fn non_streamed_exchange_is_forwarded_and_recorded() {
let reply = r#"{"choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":5}}"#;
let upstream = upstream(vec![json_reply(reply)]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions?x=1", CHAT)).await;
assert_eq!(response.status(), 200);
assert_eq!(response.headers()["content-type"], "application/json");
assert_eq!(body_text(response).await, reply);
let forwarded = upstream.request(0);
assert!(forwarded.starts_with("POST /v1/chat/completions?x=1 HTTP/1.1\r\n"));
assert!(forwarded.ends_with(CHAT));
assert!(
forwarded
.to_ascii_lowercase()
.contains("authorization: bearer s3cret")
);
let record = proxy.record().await;
assert_eq!(record.id, 1);
assert_eq!(record.error, None);
assert_eq!(record.request.method, "POST");
assert_eq!(record.request.path, "/v1/chat/completions?x=1");
assert_eq!(record.request.body_bytes, CHAT.len() as u64);
assert_eq!(
record.request.body,
Some(serde_json::from_str::<Value>(CHAT).unwrap())
);
assert!(!record.request.headers.contains_key("authorization"));
assert_eq!(record.request.headers["user-agent"], "test/1");
let response = record.response.unwrap();
assert_eq!(response.status, 200);
assert_eq!(response.stream, None);
assert_eq!((response.chunks, response.bad_chunks), (None, None));
assert_eq!(response.body_bytes, reply.len() as u64);
assert_eq!(
response.body,
Some(serde_json::from_str::<Value>(reply).unwrap())
);
assert!(!response.truncated);
let first_byte = record.t_first_byte_ms.unwrap();
assert!(record.t_start_ms <= first_byte && first_byte <= record.t_end_ms);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn sse_stream_is_passed_through_as_it_arrives_and_merged() {
let events = [
r#"{"object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"}}]}"#,
r#"{"object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"lo"},"finish_reason":"stop"}]}"#,
r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":7},"timings":{"cache_n":3}}"#,
];
let mut steps = vec![chunked_head("text/event-stream")];
steps.push(chunk(&format!("data: {}\n\n", events[0])));
steps.push(Step::Sleep(400));
steps.push(chunk(&format!(
"data: {}\n\ndata: {}\n\n",
events[1], events[2]
)));
steps.push(chunk("data: [DONE]\n\n"));
steps.push(last_chunk());
let upstream = upstream(vec![steps]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let started = Instant::now();
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
let mut body = response.into_body();
let first = body.frame().await.unwrap().unwrap().into_data().unwrap();
assert!(
started.elapsed() < Duration::from_millis(300),
"first chunk was held back"
);
assert!(String::from_utf8_lossy(&first).contains("Hel"));
let rest = body.collect().await.unwrap().to_bytes();
let received = [first.to_vec(), rest.to_vec()].concat();
let expected = format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
events[0], events[1], events[2]
);
assert_eq!(String::from_utf8(received).unwrap(), expected);
let record = proxy.record().await;
assert_eq!(record.error, None);
assert!(record.t_end_ms - record.t_first_byte_ms.unwrap() >= 300);
let response = record.response.unwrap();
assert_eq!(response.stream, Some(StreamKind::Sse));
assert_eq!((response.chunks, response.bad_chunks), (Some(3), Some(0)));
assert_eq!(response.body_bytes, expected.len() as u64);
assert_eq!(response.headers["transfer-encoding"], "chunked");
let body = response.body.unwrap();
assert_eq!(body["choices"][0]["message"]["content"], "Hello");
assert_eq!(body["usage"], json!({"prompt_tokens": 7}));
assert_eq!(body["timings"], json!({"cache_n": 3}));
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn ollama_ndjson_stream_is_merged() {
let lines = [
r#"{"model":"m","message":{"role":"assistant","content":"Hi"},"done":false}"#,
r#"{"model":"m","message":{"role":"assistant","content":" there"},"done":false}"#,
r#"{"model":"m","message":{"role":"assistant","content":""},"done":true,"prompt_eval_count":12,"eval_count":2}"#,
];
let steps = vec![
chunked_head("application/x-ndjson"),
chunk(&format!("{}\n", lines[0])),
chunk(&format!("{}\n{}\n", lines[1], lines[2])),
last_chunk(),
];
let upstream = upstream(vec![steps]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/api/chat", CHAT)).await;
assert_eq!(
body_text(response).await,
lines.map(|l| format!("{l}\n")).concat()
);
let record = proxy.record().await;
let response = record.response.unwrap();
assert_eq!(response.stream, Some(StreamKind::Ndjson));
assert_eq!(response.chunks, Some(3));
let body = response.body.unwrap();
assert_eq!(body["message"]["content"], "Hi there");
assert_eq!(body["prompt_eval_count"], 12);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn keep_alive_reuses_one_upstream_connection() {
let upstream = upstream(
(0..3)
.map(|i| json_reply(&format!("{{\"n\":{i}}}")))
.collect(),
)
.await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
for i in 0..3 {
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
assert_eq!(body_text(response).await, format!("{{\"n\":{i}}}"));
}
let mut records = Vec::new();
for _ in 0..3 {
records.push(proxy.record().await);
}
records.sort_by_key(|r| r.id);
assert_eq!(
records.iter().map(|r| r.id).collect::<Vec<_>>(),
vec![1, 2, 3]
);
assert!(
records
.iter()
.all(|r| r.conn == records[0].conn && r.error.is_none())
);
assert_eq!(upstream.connections.load(Ordering::SeqCst), 1);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn upstream_that_closes_after_each_reply_is_reconnected() {
let reply = |i: usize| {
let body = format!("{{\"n\":{i}}}");
vec![
Step::Write(
format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
)
.into_bytes(),
),
Step::Close,
]
};
let upstream = upstream((0..3).map(reply).collect()).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
for i in 0..3 {
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
assert_eq!(response.status(), 200);
assert!(response.headers().get("connection").is_none());
assert_eq!(body_text(response).await, format!("{{\"n\":{i}}}"));
assert_eq!(proxy.record().await.error, None);
}
assert_eq!(upstream.connections.load(Ordering::SeqCst), 3);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn upstream_down_gives_502_and_a_record() {
let mut proxy = start(free_addr(), BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
assert_eq!(response.status(), 502);
let record = proxy.record().await;
assert_eq!(record.error, Some(ExchangeError::UpstreamUnreachable));
assert_eq!(record.response, None);
assert_eq!(record.t_first_byte_ms, None);
assert!(record.request.body.is_some());
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn upstream_reset_mid_stream_is_recorded() {
let steps = vec![
chunked_head("text/event-stream"),
chunk("data: {\"choices\":[{\"delta\":{\"content\":\"par\"}}]}\n\n"),
Step::Sleep(50),
Step::Close,
];
let upstream = upstream(vec![steps]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
assert!(response.into_body().collect().await.is_err());
let record = proxy.record().await;
assert_eq!(record.error, Some(ExchangeError::UpstreamReset));
let response = record.response.unwrap();
assert_eq!(
response.body.unwrap()["choices"][0]["message"]["content"],
"par"
);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn client_disconnect_mid_stream_is_recorded() {
let steps = vec![
chunked_head("text/event-stream"),
chunk("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"),
Step::Sleep(300),
chunk("data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n"),
Step::Sleep(300),
chunk("data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n"),
last_chunk(),
];
let upstream = upstream(vec![steps]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
let mut body = response.into_body();
body.frame().await.unwrap().unwrap();
drop(body);
drop(sender);
conn.abort();
let record = proxy.record().await;
assert_eq!(record.error, Some(ExchangeError::ClientDisconnected));
let response = record.response.unwrap();
assert!(response.body_bytes > 0);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn hop_by_hop_headers_are_not_forwarded() {
let steps = vec![Step::Write(
b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 2\r\nConnection: X-Drop\r\nX-Drop: 1\r\nX-Keep: 2\r\n\r\n{}"
.to_vec(),
)];
let upstream = upstream(vec![steps]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let mut request = post("/v1/chat/completions", CHAT);
let headers = request.headers_mut();
headers.insert("connection", "keep-alive, x-private".parse().unwrap());
headers.insert("x-private", "1".parse().unwrap());
headers.insert("keep-alive", "timeout=5".parse().unwrap());
headers.insert("x-public", "1".parse().unwrap());
let response = send(&mut sender, request).await;
assert!(response.headers().get("x-drop").is_none());
assert_eq!(response.headers()["x-keep"], "2");
assert_eq!(body_text(response).await, "{}");
let forwarded = upstream.request(0).to_ascii_lowercase();
assert!(!forwarded.contains("x-private"));
assert!(!forwarded.contains("keep-alive"));
assert!(forwarded.contains("x-public: 1"));
proxy.record().await;
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn digest_mode_keeps_no_text() {
let reply = r#"{"choices":[{"message":{"role":"assistant","content":"secret reply"}}],"usage":{"prompt_tokens":5}}"#;
let upstream = upstream(vec![json_reply(reply)]).await;
let mut proxy = start(upstream.addr, BodyMode::Digest, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let body = r#"{"model":"m","messages":[{"role":"user","content":"secret prompt"}]}"#;
let response = send(&mut sender, post("/v1/chat/completions", body)).await;
assert_eq!(body_text(response).await, reply);
let record = proxy.record().await;
let line = serde_json::to_string(&record).unwrap();
assert!(!line.contains("secret"));
let digest = record.request.digest.unwrap();
assert_eq!(digest.messages[0].chars, 13);
let response = record.response.unwrap();
assert_eq!(response.digest.unwrap().output.content_chars, 12);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn large_bodies_are_forwarded_whole_and_recorded_within_the_limit() {
let reply = format!("{{\"pad\":\"{}\"}}", "x".repeat(100));
let upstream = upstream(vec![json_reply(&reply), json_reply("{}")]).await;
let mut proxy = start(upstream.addr, BodyMode::Full, 64).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
assert_eq!(body_text(response).await, reply);
let record = proxy.record().await;
assert_eq!(record.error, None);
let response = record.response.unwrap();
assert!(response.truncated);
assert_eq!(response.body_bytes, reply.len() as u64);
assert_eq!(response.body, None);
assert_eq!(response.body_text.unwrap(), reply[..64]);
let big = format!("{{\"messages\":[],\"pad\":\"{}\"}}", "y".repeat(100));
let response = send(&mut sender, post("/v1/chat/completions", &big)).await;
assert_eq!(response.status(), 200);
assert!(upstream.request(1).ends_with(&big));
let record = proxy.record().await;
assert_eq!(record.error, Some(ExchangeError::BodyTooLarge));
assert_eq!(record.request.body_bytes, big.len() as u64);
assert_eq!(
(record.request.body, record.request.body_text),
(None, None)
);
proxy.stop().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn shutdown_waits_for_the_exchange_under_way() {
let steps = vec![
chunked_head("text/event-stream"),
chunk("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"),
Step::Sleep(300),
chunk("data: [DONE]\n\n"),
last_chunk(),
];
let upstream = upstream(vec![steps]).await;
let proxy = start(upstream.addr, BodyMode::Full, 1 << 20).await;
let (mut sender, _conn) = client(proxy.addr).await;
let response = send(&mut sender, post("/v1/chat/completions", CHAT)).await;
let reader = tokio::spawn(body_text(response));
let records = proxy.stop().await;
assert!(reader.await.unwrap().ends_with("data: [DONE]\n\n"));
assert_eq!(records.len(), 1);
assert_eq!(records[0].error, None);
}
#[test]
fn jsonl_sink_writes_every_record() {
let path = std::path::Path::new(env!("CARGO_TARGET_TMPDIR")).join("jsonl_sink.jsonl");
let _ = std::fs::remove_file(&path);
let sink = JsonlSink::create(&path).unwrap();
let line = r#"{"v":1,"id":0,"conn":1,"t_start_ms":0,"t_first_byte_ms":null,"t_end_ms":1,
"request":{"method":"GET","path":"/","headers":{},"body_bytes":0,"body":null,"body_text":null,"digest":null},
"response":null,"error":null,"tokens":null}"#;
let template: Record = serde_json::from_str(line).unwrap();
for id in (1..=50).rev() {
let mut record = template.clone();
record.id = id;
sink.record(record);
}
assert_eq!(sink.close().unwrap(), 50);
let read = read_jsonl(std::io::BufReader::new(std::fs::File::open(&path).unwrap()));
assert!(read.bad_lines.is_empty());
assert_eq!(
read.records.iter().map(|r| r.id).collect::<Vec<_>>(),
(1..=50).collect::<Vec<_>>()
);
}