onwards 0.26.0

A flexible LLM proxy library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266

//! Shared utilities for response ID override.
//!
//! When [`AppState::response_id_header`] is configured, the caller can supply a
//! response ID via a request header. These helpers extract the override value
//! and patch it into the HTTP response body, handling gzip/brotli decompression
//! and returning the body uncompressed (with `Content-Encoding` stripped).

use axum::body::Body;
use axum::http::{HeaderMap, HeaderValue, Response, header};
use std::io::Read as _;
use tracing::debug;

/// Extract a response ID override from request headers.
///
/// Looks up the header named `header_name` and normalises the value with a
/// `resp_` prefix (added if not already present).
pub fn extract_override_id(headers: &HeaderMap, header_name: &str) -> Option<String> {
    headers
        .get(header_name)
        .and_then(|v| v.to_str().ok())
        .map(|id| {
            if id.starts_with("resp_") {
                id.to_string()
            } else {
                format!("resp_{id}")
            }
        })
}

/// Patch the `id` field in a JSON response body with the given override.
///
/// Handles `Content-Encoding: gzip` and `Content-Encoding: br` transparently:
/// the body is decompressed, the `id` field is overwritten, and the response is
/// returned **uncompressed** (with `Content-Encoding` / `Transfer-Encoding`
/// stripped and `Content-Length` updated).
///
/// If the body is not JSON, cannot be decompressed, or does not contain an `id`
/// field, the response is returned unchanged.
pub async fn patch_response_body_id(response: &mut Response<Body>, override_id: String) {
    let is_json = response
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .is_some_and(|ct| ct.contains("application/json"));

    if !is_json {
        return;
    }

    let content_encoding = response
        .headers()
        .get("content-encoding")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_lowercase());

    let bytes = match axum::body::to_bytes(std::mem::take(response.body_mut()), usize::MAX).await {
        Ok(b) => b,
        Err(_) => return,
    };

    // Decompress if needed.
    let decompressed = match content_encoding.as_deref() {
        Some("gzip") => {
            let mut decoder = flate2::read::GzDecoder::new(&bytes[..]);
            let mut buf = Vec::new();
            if decoder.read_to_end(&mut buf).is_ok() {
                buf
            } else {
                debug!("Failed to gzip-decompress response for ID patching, passing through");
                *response.body_mut() = Body::from(bytes);
                return;
            }
        }
        Some("br") | Some("brotli") => {
            let mut buf = Vec::new();
            if brotli::Decompressor::new(&bytes[..], 4096)
                .read_to_end(&mut buf)
                .is_ok()
            {
                buf
            } else {
                debug!("Failed to brotli-decompress response for ID patching, passing through");
                *response.body_mut() = Body::from(bytes);
                return;
            }
        }
        _ => bytes.to_vec(),
    };

    // Parse, patch, rewrite.
    if let Ok(mut json) = serde_json::from_slice::<serde_json::Value>(&decompressed) {
        if json.get("id").is_some() {
            json["id"] = serde_json::Value::String(override_id);
        }
        let patched = serde_json::to_vec(&json).unwrap_or(decompressed);
        let content_length = patched.len();
        *response.body_mut() = Body::from(patched);

        // Return uncompressed — the body size changed so the original encoding
        // is no longer valid. Strip encoding headers and set Content-Length.
        response.headers_mut().remove(header::CONTENT_ENCODING);
        response.headers_mut().remove(header::TRANSFER_ENCODING);
        response
            .headers_mut()
            .insert(header::CONTENT_LENGTH, HeaderValue::from(content_length));
    } else {
        // Not valid JSON — restore the original (possibly compressed) body.
        *response.body_mut() = Body::from(bytes);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::StatusCode;
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use std::io::Write as _;

    // --- extract_override_id ---

    #[test]
    fn extract_override_id_adds_prefix() {
        let mut headers = HeaderMap::new();
        headers.insert("x-custom-id", HeaderValue::from_static("abc-123"));
        assert_eq!(
            extract_override_id(&headers, "x-custom-id"),
            Some("resp_abc-123".to_string())
        );
    }

    #[test]
    fn extract_override_id_preserves_existing_prefix() {
        let mut headers = HeaderMap::new();
        headers.insert("x-custom-id", HeaderValue::from_static("resp_abc-123"));
        assert_eq!(
            extract_override_id(&headers, "x-custom-id"),
            Some("resp_abc-123".to_string())
        );
    }

    #[test]
    fn extract_override_id_missing_header() {
        let headers = HeaderMap::new();
        assert_eq!(extract_override_id(&headers, "x-custom-id"), None);
    }

    #[test]
    fn extract_override_id_wrong_header_name() {
        let mut headers = HeaderMap::new();
        headers.insert("x-other", HeaderValue::from_static("abc"));
        assert_eq!(extract_override_id(&headers, "x-custom-id"), None);
    }

    // --- patch_response_body_id ---

    fn build_json_response(body: &[u8], content_type: &str) -> Response<Body> {
        Response::builder()
            .status(StatusCode::OK)
            .header("content-type", content_type)
            .body(Body::from(body.to_vec()))
            .unwrap()
    }

    fn gzip_compress(data: &[u8]) -> Vec<u8> {
        let mut encoder = GzEncoder::new(Vec::new(), Compression::fast());
        encoder.write_all(data).unwrap();
        encoder.finish().unwrap()
    }

    fn brotli_compress(data: &[u8]) -> Vec<u8> {
        let mut buf = Vec::new();
        {
            let mut writer = brotli::CompressorWriter::new(&mut buf, 4096, 4, 22);
            writer.write_all(data).unwrap();
        }
        buf
    }

    #[tokio::test]
    async fn patch_uncompressed_json() {
        let body = br#"{"id":"original","model":"gpt-4","status":"completed"}"#;
        let mut response = build_json_response(body, "application/json");

        patch_response_body_id(&mut response, "resp_override".to_string()).await;

        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let json: serde_json::Value = serde_json::from_slice(&result).unwrap();
        assert_eq!(json["id"], "resp_override");
        assert_eq!(json["model"], "gpt-4");
    }

    #[tokio::test]
    async fn patch_gzip_compressed_json() {
        let body = br#"{"id":"original","model":"gpt-4","status":"completed"}"#;
        let compressed = gzip_compress(body);
        let mut response = Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "application/json")
            .header("content-encoding", "gzip")
            .body(Body::from(compressed))
            .unwrap();

        patch_response_body_id(&mut response, "resp_patched".to_string()).await;

        // Should be decompressed after patching
        assert!(response.headers().get("content-encoding").is_none());
        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let json: serde_json::Value = serde_json::from_slice(&result).unwrap();
        assert_eq!(json["id"], "resp_patched");
    }

    #[tokio::test]
    async fn patch_brotli_compressed_json() {
        let body = br#"{"id":"original","model":"gpt-4","status":"completed"}"#;
        let compressed = brotli_compress(body);
        let mut response = Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "application/json")
            .header("content-encoding", "br")
            .body(Body::from(compressed))
            .unwrap();

        patch_response_body_id(&mut response, "resp_br_patched".to_string()).await;

        assert!(response.headers().get("content-encoding").is_none());
        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let json: serde_json::Value = serde_json::from_slice(&result).unwrap();
        assert_eq!(json["id"], "resp_br_patched");
    }

    #[tokio::test]
    async fn patch_skips_non_json() {
        let body = b"not json";
        let mut response = build_json_response(body, "text/event-stream");

        patch_response_body_id(&mut response, "resp_ignored".to_string()).await;

        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        assert_eq!(result.as_ref(), b"not json");
    }

    #[tokio::test]
    async fn patch_preserves_body_without_id_field() {
        let body = br#"{"model":"gpt-4","status":"completed"}"#;
        let mut response = build_json_response(body, "application/json");

        patch_response_body_id(&mut response, "resp_noop".to_string()).await;

        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let json: serde_json::Value = serde_json::from_slice(&result).unwrap();
        assert!(json.get("id").is_none());
    }

    #[tokio::test]
    async fn patch_sets_content_length() {
        let body = br#"{"id":"short"}"#;
        let mut response = build_json_response(body, "application/json");

        patch_response_body_id(&mut response, "resp_much-longer-id-value".to_string()).await;

        let cl: usize = response.headers().get("content-length").unwrap().to_str().unwrap().parse().unwrap();
        let result = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
        assert_eq!(cl, result.len());
    }
}