Skip to main content

forge_gateway/
splice.rs

1//! The byte-level `model` splice.
2//!
3//! The request body is never deserialized: the splice finds the
4//! top-level `"model"` key, optionally replaces its string value, and
5//! fixes `Content-Length`. Everything else - tool schemas, betas,
6//! `context_management` - passes through byte for byte, which is the
7//! whole reason the gateway forwards instead of re-encoding.
8
9/// Why a body could not be spliced. Both are loud by design: a splice
10/// that silently forwarded the wrong model would be worse than a
11/// failure naming the session.
12#[derive(Debug, thiserror::Error, PartialEq, Eq)]
13pub enum SpliceError {
14    #[error("request body has no top-level \"model\" field")]
15    ModelMissing,
16    #[error("request body's top-level \"model\" is not a string")]
17    ModelNotAString,
18}
19
20/// Extract the top-level `model` string and, when `replacement` is
21/// `Some`, rewrite it in place. Returns the (possibly new) body and
22/// the extracted model. Nested `"model"` keys are never touched: only
23/// a string key at brace depth 1 counts.
24pub fn splice_model(
25    body: &[u8],
26    replacement: Option<&str>,
27) -> Result<(Vec<u8>, String), SpliceError> {
28    let Some((value_span, model)) = find_top_level_model(body)? else {
29        return Err(SpliceError::ModelMissing);
30    };
31    let (start, end) = value_span;
32    let rewritten = match replacement {
33        None => body.to_vec(),
34        Some(new_value) => {
35            // The span excludes the quotes; body[..start] ends on the
36            // opening quote and body[end..] starts on the closing one,
37            // so the replacement lands between them untouched.
38            let mut out = Vec::with_capacity(body.len() - (end - start) + new_value.len());
39            out.extend_from_slice(&body[..start]);
40            out.extend_from_slice(new_value.as_bytes());
41            out.extend_from_slice(&body[end..]);
42            out
43        }
44    };
45    Ok((rewritten, model))
46}
47
48/// The located top-level model: its byte span (excluding the quotes)
49/// and its content.
50type ModelSpan = ((usize, usize), String);
51
52/// Locate the top-level `"model"` string value: its byte span
53/// (excluding the quotes) and its content. Scans with a string/escape
54/// state machine so a `"model"` that is only a nested key or a value
55/// elsewhere is skipped.
56fn find_top_level_model(body: &[u8]) -> Result<Option<ModelSpan>, SpliceError> {
57    let mut depth: usize = 0;
58    let mut in_string = false;
59    let mut escaped = false;
60    let mut i = 0;
61    while i < body.len() {
62        let byte = body[i];
63        if in_string {
64            if escaped {
65                escaped = false;
66            } else if byte == b'\\' {
67                escaped = true;
68            } else if byte == b'"' {
69                in_string = false;
70            }
71            i += 1;
72            continue;
73        }
74        match byte {
75            b'"' => {
76                let (content, next) = read_string(body, i);
77                in_string = false;
78                // A string at depth 1 in an object is a KEY; the value
79                // follows the colon. In an array it is a value and is
80                // skipped - "model" only counts as a key.
81                if depth == 1 && content == b"model" {
82                    let colon = skip_ws(body, next);
83                    if body.get(colon) == Some(&b':') {
84                        let value_start = skip_ws(body, colon + 1);
85                        return match body.get(value_start) {
86                            Some(b'"') => {
87                                let (value, after) = read_string(body, value_start);
88                                let extracted =
89                                    unescape(value).ok_or(SpliceError::ModelNotAString)?;
90                                Ok(Some(((value_start + 1, after - 1), extracted)))
91                            }
92                            Some(_) => Err(SpliceError::ModelNotAString),
93                            None => Err(SpliceError::ModelMissing),
94                        };
95                    }
96                }
97                i = next;
98            }
99            b'{' | b'[' => {
100                depth += 1;
101                i += 1;
102            }
103            b'}' | b']' => {
104                depth = depth.saturating_sub(1);
105                i += 1;
106            }
107            _ => i += 1,
108        }
109    }
110    Ok(None)
111}
112
113/// Read a JSON string starting at `start` (the opening quote). Returns
114/// the raw content bytes and the index just past the closing quote.
115fn read_string(body: &[u8], start: usize) -> (&[u8], usize) {
116    let mut i = start + 1;
117    let mut escaped = false;
118    while let Some(&byte) = body.get(i) {
119        if escaped {
120            escaped = false;
121        } else if byte == b'\\' {
122            escaped = true;
123        } else if byte == b'"' {
124            return (&body[start + 1..i], i + 1);
125        }
126        i += 1;
127    }
128    (&body[start + 1..], body.len())
129}
130
131fn skip_ws(body: &[u8], mut i: usize) -> usize {
132    while let Some(&byte) = body.get(i) {
133        if byte.is_ascii_whitespace() {
134            i += 1;
135        } else {
136            break;
137        }
138    }
139    i
140}
141
142/// Model names are plain ASCII slugs in practice; unescape the two
143/// escapes that could appear around them and reject exotic ones rather
144/// than misreporting the name.
145fn unescape(raw: &[u8]) -> Option<String> {
146    let text = std::str::from_utf8(raw).ok()?;
147    if !text.contains('\\') {
148        return Some(text.to_owned());
149    }
150    let mut out = String::with_capacity(text.len());
151    let mut chars = text.chars();
152    while let Some(c) = chars.next() {
153        if c != '\\' {
154            out.push(c);
155            continue;
156        }
157        match chars.next()? {
158            '"' => out.push('"'),
159            '\\' => out.push('\\'),
160            '/' => out.push('/'),
161            _ => return None,
162        }
163    }
164    Some(out)
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    /// A real CLI-emitted request body, captured through the gateway
172    /// against a local 401 stub (nothing billed) and redacted of
173    /// session-local identifiers. Committed so the splice parser and its
174    /// pin share no blind spots: every field the splice touches, and the
175    /// body's structure throughout, are the real capture's. One block was
176    /// replaced with a placeholder, so that text is the exception.
177    fn live_capture_body() -> Vec<u8> {
178        let path =
179            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("fixtures/messages_request.json");
180        std::fs::read(path).expect("the committed live capture is readable")
181    }
182
183    /// Small synthetic bodies for the failure arms, where the shape
184    /// must be exact.
185    fn tiny_body(model_json: Option<&str>) -> Vec<u8> {
186        match model_json {
187            Some(model) => format!(r#"{{"model":{model},"max_tokens":8}}"#).into_bytes(),
188            None => br#"{"max_tokens":8}"#.to_vec(),
189        }
190    }
191
192    #[test]
193    fn extracts_the_top_level_model_from_a_live_capture() {
194        let body = live_capture_body();
195        let (_, model) = splice_model(&body, None).expect("splice");
196        assert!(!model.is_empty(), "the real body's model is extracted");
197    }
198
199    /// The control that makes the next redactor mistake visible: the
200    /// byte-level splice tests ran green against a corrupt fixture
201    /// because they never parse, so this parses independently.
202    #[test]
203    fn the_committed_capture_parses_as_json() {
204        let body = live_capture_body();
205        serde_json::from_str::<serde_json::Value>(std::str::from_utf8(&body).expect("utf8"))
206            .expect("the committed capture must be valid JSON");
207    }
208
209    #[test]
210    fn a_rewrite_changes_only_the_model_and_fixes_the_length() {
211        let body = live_capture_body();
212        let (rewritten, model) =
213            splice_model(&body, Some("deepseek/deepseek-v4.1-flash")).expect("splice");
214        assert_eq!(model, "glm-5.3-flash", "the extracted name is the old one");
215        // The parse is the assertion: a splice that corrupted the body
216        // (double-quoted value, stray bytes) fails here, where a
217        // substring check against the corrupted text would pass.
218        let parsed: serde_json::Value =
219            serde_json::from_slice(&rewritten).expect("the rewritten body parses as JSON");
220        assert_eq!(
221            parsed.get("model").and_then(|m| m.as_str()),
222            Some("deepseek/deepseek-v4.1-flash"),
223            "the top-level model is exactly the replacement",
224        );
225        let text = String::from_utf8(rewritten.clone()).expect("utf8");
226        // The untouched bulk is byte-identical: same tool count, same
227        // system prompt bytes.
228        let original = live_capture_body();
229        let original_text = String::from_utf8(original).expect("utf8");
230        for kept in ["context_management", "output_config", "stream"] {
231            assert!(
232                text.contains(kept) && original_text.contains(kept),
233                "{kept} survives the splice",
234            );
235        }
236    }
237
238    /// The rewrite lands the value BETWEEN the existing quotes: a
239    /// double-quoted value (`""name""`) is invalid JSON that a
240    /// substring assertion cannot see.
241    #[test]
242    fn a_rewrite_never_doubles_the_quotes() {
243        let body = br#"{"model":"glm-5.3-flash","max_tokens":8}"#;
244        let (rewritten, _) = splice_model(body, Some("z-ai/glm-5.3-flash")).expect("splice");
245        let text = String::from_utf8(rewritten).expect("utf8");
246        assert_eq!(
247            text, r#"{"model":"z-ai/glm-5.3-flash","max_tokens":8}"#,
248            "the rewritten body is byte-exact",
249        );
250    }
251
252    #[test]
253    fn a_nested_model_key_is_never_touching() {
254        let body = br#"{"messages":[{"role":"user","content":"model"}],"metadata":{"model":"nested-model"},"model":"real-model"}"#;
255        let (rewritten, model) = splice_model(body, Some("new")).expect("splice");
256        assert_eq!(model, "real-model", "only the top-level key counts");
257        let text = String::from_utf8(rewritten).expect("utf8");
258        assert!(text.contains("nested-model"), "the nested key survives");
259        assert!(text.contains("\"new\""), "the top-level value is rewritten");
260    }
261
262    #[test]
263    fn a_non_string_model_is_a_loud_failure() {
264        let body = tiny_body(Some("42"));
265        assert_eq!(splice_model(&body, None), Err(SpliceError::ModelNotAString));
266    }
267
268    #[test]
269    fn a_body_with_no_model_is_a_loud_failure() {
270        assert_eq!(splice_model(&tiny_body(None), None), Err(SpliceError::ModelMissing));
271        assert_eq!(splice_model(b"[]", None), Err(SpliceError::ModelMissing));
272        assert_eq!(splice_model(b"", None), Err(SpliceError::ModelMissing));
273    }
274
275    #[test]
276    fn whitespace_between_key_and_value_is_tolerated() {
277        let body = br#"{"model" : "claude-opus-5"}"#;
278        let (rewritten, model) = splice_model(body, Some("x")).expect("splice");
279        assert_eq!(model, "claude-opus-5");
280        assert!(String::from_utf8(rewritten).expect("utf8").contains("\"x\""));
281    }
282
283    #[test]
284    fn a_multi_byte_utf8_model_rewrites_without_corrupting_the_body() {
285        // The replacement is inserted as raw UTF-8 bytes; the byte-span
286        // arithmetic must not split a multi-byte sequence either in the
287        // replacement or in the keys around it. Values use unicode
288        // escapes so the source stays ASCII.
289        let model_value = "caf\u{e9}-latte";
290        let content = "union \u{2014} \u{7d75}\u{6587}\u{5b57}";
291        let body =
292            format!("{{\"model\":\"{model_value}\",\"messages\":[{{\"content\":\"{content}\"}}]}}");
293        let (rewritten, model) =
294            splice_model(body.as_bytes(), Some("glm-5.3-flash")).expect("splice");
295        assert_eq!(model, "caf\u{e9}-latte");
296        // Parsed, not a substring: a re-introduced double-quote around
297        // the value would satisfy a contains check and still be
298        // invalid JSON.
299        let parsed: serde_json::Value =
300            serde_json::from_slice(&rewritten).expect("the rewritten body parses as JSON");
301        assert_eq!(
302            parsed.get("model").and_then(|m| m.as_str()),
303            Some("glm-5.3-flash"),
304            "the top-level model is exactly the replacement",
305        );
306        let text = String::from_utf8(rewritten).expect("the rewrite keeps the body valid UTF-8");
307        assert!(text.contains(content), "the multi-byte sibling content survives");
308    }
309}