compute/src/gateway.rs

//! Streaming-safe TCP gateway for OpenAI-compatible model servers.
//!
//! The gateway deliberately stays below HTTP: it forwards bytes unchanged, so
//! chunked responses and server-sent-event token streams retain their timing and
//! semantics while clients keep one stable Lumbridge Compute address.

use anyhow::{Context, Result};
use std::io;
use std::net::{Shutdown, TcpListener, TcpStream};
use std::thread;

pub fn run(listen: &str, upstream: &str) -> Result<()> {
    let listener = TcpListener::bind(listen)
        .with_context(|| format!("binding Lumbridge gateway at {listen}"))?;
    println!("Lumbridge gateway {listen} -> {upstream}");
    serve_listener(listener, upstream, None)
}

fn serve_listener(
    listener: TcpListener,
    upstream: &str,
    max_connections: Option<usize>,
) -> Result<()> {
    let mut accepted = 0usize;
    for incoming in listener.incoming() {
        let client = match incoming {
            Ok(client) => client,
            Err(error) => {
                eprintln!("gateway accept failed: {error}");
                continue;
            }
        };
        let upstream = upstream.to_string();
        thread::spawn(move || {
            if let Err(error) = proxy(client, &upstream) {
                eprintln!("gateway request failed: {error:#}");
            }
        });
        accepted += 1;
        if max_connections.is_some_and(|limit| accepted >= limit) {
            break;
        }
    }
    Ok(())
}

fn proxy(mut client: TcpStream, upstream_addr: &str) -> Result<()> {
    client.set_nodelay(true).ok();
    let mut upstream = TcpStream::connect(upstream_addr)
        .with_context(|| format!("connecting gateway upstream {upstream_addr}"))?;
    upstream.set_nodelay(true).ok();

    let mut client_reader = client.try_clone()?;
    let mut upstream_writer = upstream.try_clone()?;
    let request = thread::spawn(move || -> io::Result<u64> {
        let copied = io::copy(&mut client_reader, &mut upstream_writer)?;
        upstream_writer.shutdown(Shutdown::Write).ok();
        Ok(copied)
    });

    io::copy(&mut upstream, &mut client)?;
    client.shutdown(Shutdown::Write).ok();
    request
        .join()
        .map_err(|_| anyhow::anyhow!("gateway request-copy thread panicked"))??;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};

    #[test]
    fn gateway_forwards_bidirectional_bytes_without_buffering_protocols() {
        let upstream = TcpListener::bind("127.0.0.1:0").unwrap();
        let upstream_addr = upstream.local_addr().unwrap();
        let upstream_thread = thread::spawn(move || {
            let (mut socket, _) = upstream.accept().unwrap();
            let mut request = [0u8; 4];
            socket.read_exact(&mut request).unwrap();
            assert_eq!(&request, b"ping");
            socket.write_all(b"pong").unwrap();
        });

        let gateway = TcpListener::bind("127.0.0.1:0").unwrap();
        let gateway_addr = gateway.local_addr().unwrap();
        let gateway_thread = thread::spawn(move || {
            serve_listener(gateway, &upstream_addr.to_string(), Some(1)).unwrap();
        });

        let mut client = TcpStream::connect(gateway_addr).unwrap();
        client.write_all(b"ping").unwrap();
        client.shutdown(Shutdown::Write).unwrap();
        let mut response = Vec::new();
        client.read_to_end(&mut response).unwrap();
        assert_eq!(response, b"pong");

        upstream_thread.join().unwrap();
        gateway_thread.join().unwrap();
    }
}