rust 74 lines · 10 steps

Proxying an SSE chat stream in Axum

An Axum handler forwards a chat request upstream and re-streams the response back to the client as Server-Sent Events.

Explained by highlit
Axum Advanced server-sent-events streaming async-generators reverse-proxy backpressure
1use axum::{
2 extract::State,
3 response::sse::{Event, KeepAlive, Sse},
4 Json,
5};
6use futures::stream::{Stream, StreamExt, TryStreamExt};
7use serde::Deserialize;
8use std::convert::Infallible;
9 
10#[derive(Clone)]
11pub struct AppState {
12 http: reqwest::Client,
13 upstream: String,
14 api_key: String,
15}
16 
17#[derive(Deserialize)]
18pub struct ChatRequest {
19 model: String,
20 messages: Vec<serde_json::Value>,
21 #[serde(default)]
22 temperature: Option<f32>,
23}
24 
25pub async fn chat_completions(
26 State(state): State<AppState>,
27 Json(req): Json<ChatRequest>,
28) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
29 let body = serde_json::json!({
30 "model": req.model,
31 "messages": req.messages,
32 "temperature": req.temperature,
33 "stream": true,
34 });
35 
36 let upstream = state
37 .http
38 .post(format!("{}/v1/chat/completions", state.upstream))
39 .bearer_auth(&state.api_key)
40 .json(&body)
41 .send()
42 .await;
43 
44 let stream = async_stream::stream! {
45 let resp = match upstream {
46 Ok(r) => r,
47 Err(e) => {
48 yield Ok(Event::default().event("error").data(e.to_string()));
49 return;
50 }
51 };
52 
53 let mut bytes = resp.bytes_stream();
54 let mut buf = Vec::new();
55 
56 while let Ok(Some(chunk)) = bytes.try_next().await {
57 buf.extend_from_slice(&chunk);
58 while let Some(pos) = buf.windows(2).position(|w| w == b"\n\n") {
59 let frame = buf.drain(..pos + 2).collect::<Vec<_>>();
60 for line in frame.split(|&b| b == b'\n') {
61 let Some(payload) = line.strip_prefix(b"data: ") else { continue };
62 let payload = String::from_utf8_lossy(payload);
63 if payload.trim() == "[DONE]" {
64 yield Ok(Event::default().data("[DONE]"));
65 return;
66 }
67 yield Ok(Event::default().data(payload.trim().to_string()));
68 }
69 }
70 }
71 };
72 
73 Sse::new(stream.boxed()).keep_alive(KeepAlive::default())
74}
01 / 01
STEP 01

Walkthrough

Space play ←→ step click any line
Three takeaways
  1. 1Server-Sent Events let a handler push tokens to the client incrementally instead of buffering the full response.
  2. 2Byte streams arrive in arbitrary chunks, so you must buffer and split on frame delimiters yourself rather than assuming message boundaries.
  3. 3async_stream::stream! turns imperative loops with yield into a Stream, keeping streaming proxy logic readable.

Related explainers

rust 
use serde::Deserialize;
 
#[derive(Debug, Deserialize)]
#[serde(untagged)]

Parsing flexible JSON shapes with serde

deserialization enums json
Intermediate 6 steps
rust 
use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;
use std::fs::File;
use std::io::{self, BufRead, BufReader, BufWriter, Lines, Write};

K-way merge of sorted logs in Rust

binary-heap k-way-merge streaming-io
Intermediate 8 steps
rust 
use axum::{
    extract::{Path, State},
    response::sse::{Event, KeepAlive, Sse},
};

Streaming import progress with SSE in Axum

server-sent-events streams watch-channel
Advanced 7 steps
rust 
use std::f64::consts::PI;
 
#[derive(Debug, Clone, Copy)]
pub struct LatLng {

Building geographic bounding boxes in Rust

geospatial value-types option
Intermediate 7 steps
python 
import random
from typing import Iterator, List
 
 

How reservoir sampling picks k items

reservoir-sampling streaming randomness
Intermediate 5 steps
rust 
use chrono::{Duration, NaiveDate};
 
#[derive(Debug)]
pub struct DateRange {

Parsing and iterating date ranges in Rust

error-handling iterators parsing
Intermediate 7 steps

Share this explainer

Here's the card — post it anywhere.

Proxying an SSE chat stream in Axum — share card
Made with highlit — turn any snippet into a walkthrough like this in about a minute.
Explain your code