loadmaster_probe/
fake.rs

1//! A scriptable fake OpenAI-compatible server for tests (and `loadmaster __fake-runtime`).
2//!
3//! By default it behaves like a correct runtime: tools in the request produce `tool_calls`
4//! whose arguments are built from the tool's parameter schema, `response_format` produces
5//! JSON built from the schema, a `tool` message produces a final answer, and streaming
6//! splits everything across deltas. Each [`Misbehavior`] breaks one of those the way real
7//! runtimes do, so every failure class the probe knows can be exercised.
8
9use std::collections::HashSet;
10use std::sync::Arc;
11use std::sync::atomic::{AtomicU64, Ordering};
12
13use axum::Router;
14use axum::body::Body;
15use axum::extract::State;
16use axum::http::{HeaderValue, StatusCode, header};
17use axum::response::{IntoResponse, Response};
18use axum::routing::{get, post};
19use serde_json::{Value, json};
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub enum Misbehavior {
23    /// The tool call is written as JSON text in `content`, `tool_calls` stays empty.
24    CallInContent,
25    /// The tool call appears in `reasoning_content` (or a `<think>` block), not `tool_calls`.
26    CallInReasoning,
27    /// `arguments` is not valid JSON.
28    ArgsInvalidJson,
29    /// Answers in prose instead of calling a tool.
30    NoCall,
31    /// Several requested tools: only the first is called.
32    ParallelSerialised,
33    /// Several requested tools: all calls written into `content`.
34    ParallelInContent,
35    /// Nested object/array arguments flattened to strings.
36    SchemaFlattened,
37    /// After a tool result, calls the same tool again instead of answering.
38    LoopsOnTool,
39    /// After a tool result, answers without using it.
40    IgnoresResult,
41    /// `response_format` ignored: HTTP 200 with prose.
42    SchemaDroppedSilent,
43    /// `response_format` with a schema answered with HTTP 400.
44    SchemaRejected,
45    /// `enum` constraints ignored.
46    EnumIgnored,
47    /// `pattern` constraints ignored.
48    PatternIgnored,
49    /// Stream ends mid tool call: no finish_reason, no `[DONE]`.
50    StreamTruncated,
51    /// `/v1/embeddings` returns a different vector for the same text each time.
52    EmbedNondeterministic,
53    /// Streamed argument fragments do not concatenate to valid JSON.
54    StreamArgsSplitInvalid,
55    /// The runtime's tool-call parser rejects the model's output, as llama-server does
56    /// with "does not match the expected peg-native format": HTTP 500, or an `error`
57    /// event in the middle of a stream.
58    RuntimeParseError,
59}
60
61impl Misbehavior {
62    pub const ALL: &[(&str, Misbehavior)] = &[
63        ("call_in_content", Misbehavior::CallInContent),
64        ("call_in_reasoning", Misbehavior::CallInReasoning),
65        ("args_invalid_json", Misbehavior::ArgsInvalidJson),
66        ("no_call", Misbehavior::NoCall),
67        ("parallel_serialised", Misbehavior::ParallelSerialised),
68        ("parallel_in_content", Misbehavior::ParallelInContent),
69        ("schema_flattened", Misbehavior::SchemaFlattened),
70        ("loops_on_tool", Misbehavior::LoopsOnTool),
71        ("ignores_result", Misbehavior::IgnoresResult),
72        ("schema_dropped_silent", Misbehavior::SchemaDroppedSilent),
73        ("schema_rejected", Misbehavior::SchemaRejected),
74        ("enum_ignored", Misbehavior::EnumIgnored),
75        ("pattern_ignored", Misbehavior::PatternIgnored),
76        ("stream_truncated", Misbehavior::StreamTruncated),
77        ("embed_nondeterministic", Misbehavior::EmbedNondeterministic),
78        (
79            "stream_args_split_invalid",
80            Misbehavior::StreamArgsSplitInvalid,
81        ),
82        ("runtime_parse_error", Misbehavior::RuntimeParseError),
83    ];
84
85    pub fn parse(name: &str) -> Option<Self> {
86        Self::ALL.iter().find(|(n, _)| *n == name).map(|(_, m)| *m)
87    }
88}
89
90/// Which runtime's identifying endpoints and headers to imitate.
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum Impersonate {
93    LlamaCpp,
94    Vllm,
95    Ollama,
96    LmStudio,
97    SgLang,
98    TabbyApi,
99}
100
101impl Impersonate {
102    pub fn parse(s: &str) -> Option<Self> {
103        Some(match s {
104            "llama.cpp" => Impersonate::LlamaCpp,
105            "vllm" => Impersonate::Vllm,
106            "ollama" => Impersonate::Ollama,
107            "lmstudio" => Impersonate::LmStudio,
108            "sglang" => Impersonate::SgLang,
109            "tabbyapi" => Impersonate::TabbyApi,
110            _ => return None,
111        })
112    }
113}
114
115#[derive(Debug, Clone)]
116pub struct Config {
117    pub model: String,
118    pub misbehave: HashSet<Misbehavior>,
119    pub impersonate: Impersonate,
120    /// Chat template text served on llama.cpp's `/props`.
121    pub chat_template: String,
122    /// Prompt and generation speed (tokens/s) reported in llama.cpp `/completion` timings.
123    pub speed: (f64, f64),
124}
125
126impl Default for Config {
127    fn default() -> Self {
128        Config {
129            model: "fake-model".into(),
130            misbehave: HashSet::new(),
131            impersonate: Impersonate::LlamaCpp,
132            chat_template: "{% for m in messages %}{{ m.content }}{% endfor %}".into(),
133            speed: (1000.0, 50.0),
134        }
135    }
136}
137
138struct Fake {
139    cfg: Config,
140    requests: AtomicU64,
141}
142
143pub fn router(cfg: Config) -> Router {
144    let state = Arc::new(Fake {
145        cfg,
146        requests: AtomicU64::new(0),
147    });
148    let mut r = Router::new()
149        .route(
150            "/health",
151            get(|| async { axum::Json(json!({ "status": "ok" })) }),
152        )
153        .route("/v1/models", get(models))
154        .route("/v1/chat/completions", post(chat))
155        .route("/v1/embeddings", post(embeddings));
156    r = match state.cfg.impersonate {
157        Impersonate::LlamaCpp => r
158            .route("/props", get(props))
159            .route("/completion", post(native_completion)),
160        Impersonate::Vllm => r.route(
161            "/version",
162            get(|| async { axum::Json(json!({ "version": "0.11.0" })) }),
163        ),
164        Impersonate::Ollama => r.route(
165            "/api/version",
166            get(|| async { axum::Json(json!({ "version": "0.12.5" })) }),
167        ),
168        Impersonate::LmStudio => r.route("/api/v0/models", get(lmstudio_models)),
169        Impersonate::SgLang => r.route("/get_model_info", get(sglang_info)),
170        Impersonate::TabbyApi => r.route("/v1/model", get(tabby_model)),
171    };
172    r.with_state(state)
173}
174
175/// Serve on `listener` until the process ends.
176pub async fn serve(listener: tokio::net::TcpListener, cfg: Config) -> std::io::Result<()> {
177    axum::serve(listener, router(cfg)).await
178}
179
180type St = State<Arc<Fake>>;
181
182async fn models(State(f): St) -> Response {
183    let mut m = json!({ "id": f.cfg.model, "object": "model", "created": 0, "owned_by": "fake" });
184    match f.cfg.impersonate {
185        Impersonate::LlamaCpp => m["owned_by"] = json!("llamacpp"),
186        Impersonate::Vllm => {
187            m["owned_by"] = json!("vllm");
188            m["max_model_len"] = json!(32768);
189        }
190        Impersonate::Ollama => m["owned_by"] = json!("library"),
191        Impersonate::LmStudio => m["owned_by"] = json!("organization_owner"),
192        Impersonate::SgLang => m["owned_by"] = json!("sglang"),
193        Impersonate::TabbyApi => m["owned_by"] = json!("tabbyAPI"),
194    }
195    let mut resp = axum::Json(json!({ "object": "list", "data": [m] })).into_response();
196    if f.cfg.impersonate == Impersonate::Ollama {
197        resp.headers_mut()
198            .insert(header::SERVER, HeaderValue::from_static("ollama"));
199    }
200    resp
201}
202
203async fn props(State(f): St) -> Response {
204    axum::Json(json!({
205        "chat_template": f.cfg.chat_template,
206        "build_info": "b6700-abc1234",
207        "model_path": format!("/models/{}.gguf", f.cfg.model),
208    }))
209    .into_response()
210}
211
212/// `/v1/embeddings`: a deterministic 384-dimensional vector per input, derived from the
213/// text, so equal inputs give equal vectors and different inputs differ.
214async fn embeddings(State(f): St, axum::Json(body): axum::Json<Value>) -> Response {
215    let inputs: Vec<String> = match &body["input"] {
216        Value::String(s) => vec![s.clone()],
217        Value::Array(a) => a
218            .iter()
219            .filter_map(|v| v.as_str().map(str::to_owned))
220            .collect(),
221        _ => Vec::new(),
222    };
223    let nondeterministic = f
224        .cfg
225        .misbehave
226        .contains(&Misbehavior::EmbedNondeterministic);
227    let data: Vec<Value> = inputs
228        .iter()
229        .enumerate()
230        .map(|(i, text)| {
231            let mut seed = blake3::hash(text.as_bytes()).as_bytes().to_vec();
232            if nondeterministic {
233                seed[0] ^= f.requests.fetch_add(1, Ordering::Relaxed) as u8;
234            }
235            let v: Vec<f64> = (0..384)
236                .map(|j| f64::from(seed[j % 32]) / 255.0 - 0.5 + j as f64 * 1e-4)
237                .collect();
238            json!({ "object": "embedding", "index": i, "embedding": v })
239        })
240        .collect();
241    axum::Json(json!({ "object": "list", "model": f.cfg.model, "data": data })).into_response()
242}
243
244/// llama.cpp's native `/completion`: text plus `timings` at the configured speeds. A
245/// string prompt counts one token per whitespace-separated word.
246async fn native_completion(State(f): St, axum::Json(body): axum::Json<Value>) -> Response {
247    let prompt_n = match &body["prompt"] {
248        Value::Array(t) => t.len(),
249        Value::String(s) => s.split_whitespace().count(),
250        _ => 0,
251    } as f64;
252    let predicted_n = body["n_predict"].as_u64().unwrap_or(16) as f64;
253    let (pp, tg) = f.cfg.speed;
254    let (prompt_ms, predicted_ms) = (prompt_n / pp * 1000.0, predicted_n / tg * 1000.0);
255    axum::Json(json!({
256        "content": "fake ".repeat(predicted_n as usize),
257        "tokens_predicted": predicted_n,
258        "tokens_evaluated": prompt_n,
259        "timings": {
260            "cache_n": 0,
261            "prompt_n": prompt_n,
262            "prompt_ms": prompt_ms,
263            "prompt_per_token_ms": if prompt_n > 0.0 { prompt_ms / prompt_n } else { 0.0 },
264            "prompt_per_second": pp,
265            "predicted_n": predicted_n,
266            "predicted_ms": predicted_ms,
267            "predicted_per_token_ms": predicted_ms / predicted_n.max(1.0),
268            "predicted_per_second": tg,
269        }
270    }))
271    .into_response()
272}
273
274async fn lmstudio_models(State(f): St) -> Response {
275    axum::Json(json!({ "data": [{ "id": f.cfg.model, "type": "llm", "arch": "qwen3", "quantization": "Q4_K_M", "state": "loaded" }] })).into_response()
276}
277
278async fn sglang_info(State(f): St) -> Response {
279    axum::Json(json!({ "model_path": f.cfg.model, "is_generation": true, "version": "0.5.3" }))
280        .into_response()
281}
282
283async fn tabby_model(State(f): St) -> Response {
284    axum::Json(json!({ "id": f.cfg.model, "object": "model", "owned_by": "tabbyAPI", "parameters": { "max_seq_len": 32768 } })).into_response()
285}
286
287/// A minimal instance of a JSON Schema: required properties, first enum value, first
288/// `examples` entry (which is how fixtures satisfy `pattern`), first `anyOf` branch.
289pub fn instance(schema: &Value) -> Value {
290    if let Some(v) = schema.get("const") {
291        return v.clone();
292    }
293    if let Some(v) = schema.get("examples").and_then(|e| e.get(0)) {
294        return v.clone();
295    }
296    if let Some(v) = schema.get("default") {
297        return v.clone();
298    }
299    if let Some(v) = schema.get("enum").and_then(|e| e.get(0)) {
300        return v.clone();
301    }
302    for key in ["anyOf", "oneOf"] {
303        if let Some(first) = schema.get(key).and_then(|b| b.get(0)) {
304            return instance(first);
305        }
306    }
307    let ty = match schema.get("type") {
308        Some(Value::String(t)) => t.as_str(),
309        Some(Value::Array(ts)) => ts
310            .iter()
311            .filter_map(Value::as_str)
312            .find(|t| *t != "null")
313            .unwrap_or("null"),
314        _ if schema.get("properties").is_some() => "object",
315        _ => "string",
316    };
317    match ty {
318        "object" => {
319            let props = schema.get("properties").and_then(Value::as_object);
320            let mut out = serde_json::Map::new();
321            for (k, v) in props.into_iter().flatten() {
322                out.insert(k.clone(), instance(v));
323            }
324            Value::Object(out)
325        }
326        "array" => {
327            let n = schema
328                .get("minItems")
329                .and_then(Value::as_u64)
330                .unwrap_or(1)
331                .max(1);
332            let item = schema.get("items").map_or(json!("item"), instance);
333            Value::Array(vec![item; n as usize])
334        }
335        "integer" => json!(schema.get("minimum").and_then(Value::as_i64).unwrap_or(1)),
336        "number" => json!(schema.get("minimum").and_then(Value::as_f64).unwrap_or(1.5)),
337        "boolean" => json!(true),
338        "null" => Value::Null,
339        _ => json!("example"),
340    }
341}
342
343/// Break constraints the way runtimes that ignore them do.
344fn violate(v: &mut Value, schema: &Value, bad: &HashSet<Misbehavior>) {
345    if bad.contains(&Misbehavior::EnumIgnored) && schema.get("enum").is_some() {
346        *v = json!("not-in-enum");
347    }
348    if bad.contains(&Misbehavior::PatternIgnored) && schema.get("pattern").is_some() {
349        *v = json!("does not match");
350    }
351    if let (Some(obj), Some(props)) = (
352        v.as_object_mut(),
353        schema.get("properties").and_then(Value::as_object),
354    ) {
355        for (k, s) in props {
356            if let Some(child) = obj.get_mut(k) {
357                violate(child, s, bad);
358            }
359        }
360    }
361}
362
363/// llama-server's message when its tool-call parser rejects the model's output.
364const PARSE_ERROR: &str =
365    "The model produced output that does not match the expected peg-native format";
366
367struct Call {
368    name: String,
369    arguments: String,
370}
371
372enum Reply {
373    Text {
374        content: String,
375        reasoning: Option<String>,
376    },
377    Calls(Vec<Call>),
378    Error(StatusCode, String),
379}
380
381fn decide(f: &Fake, req: &Value) -> Reply {
382    let bad = &f.cfg.misbehave;
383    let messages = req["messages"].as_array().cloned().unwrap_or_default();
384    let last = messages.last().cloned().unwrap_or(Value::Null);
385
386    // A tool result: answer with it, unless told to misbehave.
387    if last["role"] == "tool" {
388        let result = last["content"].as_str().unwrap_or("").to_owned();
389        if bad.contains(&Misbehavior::LoopsOnTool) {
390            if let Some(tool) = req["tools"].get(0) {
391                return Reply::Calls(vec![call_for(tool, bad)]);
392            }
393        }
394        let content = if bad.contains(&Misbehavior::IgnoresResult) {
395            "I could not find that information.".to_owned()
396        } else {
397            format!("Based on the tool result: {result}")
398        };
399        return Reply::Text {
400            content,
401            reasoning: None,
402        };
403    }
404
405    if let Some(tools) = req["tools"].as_array().filter(|t| !t.is_empty()) {
406        if bad.contains(&Misbehavior::RuntimeParseError) {
407            return Reply::Error(StatusCode::INTERNAL_SERVER_ERROR, PARSE_ERROR.into());
408        }
409        if bad.contains(&Misbehavior::NoCall) {
410            return Reply::Text {
411                content: "I would need to look that up.".into(),
412                reasoning: None,
413            };
414        }
415        let mut calls: Vec<Call> = tools.iter().map(|t| call_for(t, bad)).collect();
416        if bad.contains(&Misbehavior::ParallelSerialised) {
417            calls.truncate(1);
418        }
419        let as_text = |calls: &[Call]| {
420            calls
421                .iter()
422                .map(|c| {
423                    format!(
424                        "{{\"name\": \"{}\", \"arguments\": {}}}",
425                        c.name, c.arguments
426                    )
427                })
428                .collect::<Vec<_>>()
429                .join("\n")
430        };
431        if bad.contains(&Misbehavior::ParallelInContent) && calls.len() > 1
432            || bad.contains(&Misbehavior::CallInContent)
433        {
434            return Reply::Text {
435                content: format!("<tool_call>\n{}\n</tool_call>", as_text(&calls)),
436                reasoning: None,
437            };
438        }
439        if bad.contains(&Misbehavior::CallInReasoning) {
440            return Reply::Text {
441                content: String::new(),
442                reasoning: Some(format!(
443                    "I should call the tool.\n<tool_call>{}</tool_call>",
444                    as_text(&calls)
445                )),
446            };
447        }
448        return Reply::Calls(calls);
449    }
450
451    if let Some(rf) = req
452        .get("response_format")
453        .filter(|r| r["type"] == "json_schema")
454    {
455        if bad.contains(&Misbehavior::SchemaRejected) {
456            return Reply::Error(
457                StatusCode::BAD_REQUEST,
458                "response_format json_schema is not supported".into(),
459            );
460        }
461        if bad.contains(&Misbehavior::SchemaDroppedSilent) {
462            return Reply::Text {
463                content: "Sure! Here is the information you asked for.".into(),
464                reasoning: None,
465            };
466        }
467        let schema = &rf["json_schema"]["schema"];
468        let mut v = instance(schema);
469        violate(&mut v, schema, bad);
470        return Reply::Text {
471            content: v.to_string(),
472            reasoning: None,
473        };
474    }
475
476    Reply::Text {
477        content: "Hello from the fake runtime.".into(),
478        reasoning: None,
479    }
480}
481
482fn call_for(tool: &Value, bad: &HashSet<Misbehavior>) -> Call {
483    let f = &tool["function"];
484    let schema = &f["parameters"];
485    let mut args = instance(schema);
486    if bad.contains(&Misbehavior::SchemaFlattened)
487        && let Some(obj) = args.as_object_mut()
488    {
489        for v in obj.values_mut() {
490            if v.is_object() || v.is_array() {
491                *v = Value::String(v.to_string());
492            }
493        }
494    }
495    let mut arguments = args.to_string();
496    if bad.contains(&Misbehavior::ArgsInvalidJson) {
497        arguments = arguments
498            .replacen('"', "'", 2)
499            .trim_end_matches('}')
500            .to_owned();
501    }
502    Call {
503        name: f["name"].as_str().unwrap_or("tool").to_owned(),
504        arguments,
505    }
506}
507
508async fn chat(State(f): St, axum::Json(req): axum::Json<Value>) -> Response {
509    let n = f.requests.fetch_add(1, Ordering::SeqCst);
510    let id = format!("chatcmpl-fake{n}");
511    let stream = req["stream"].as_bool().unwrap_or(false);
512    match decide(&f, &req) {
513        // A server error after streaming began arrives as an SSE `error` event.
514        reply @ Reply::Error(status, _) if stream && status.is_server_error() => {
515            sse(&f, &id, reply)
516        }
517        Reply::Error(status, msg) => {
518            let kind = if status.is_server_error() {
519                "server_error"
520            } else {
521                "invalid_request_error"
522            };
523            (
524                status,
525                axum::Json(
526                    json!({ "error": { "code": status.as_u16(), "message": msg, "type": kind } }),
527                ),
528            )
529                .into_response()
530        }
531        reply if stream => sse(&f, &id, reply),
532        Reply::Text { content, reasoning } => {
533            let mut msg = json!({ "role": "assistant", "content": content });
534            if let Some(r) = reasoning {
535                msg["reasoning_content"] = json!(r);
536            }
537            axum::Json(completion(&f, &id, msg, "stop")).into_response()
538        }
539        Reply::Calls(calls) => {
540            let tool_calls: Vec<Value> = calls
541                .iter()
542                .enumerate()
543                .map(|(i, c)| json!({ "id": format!("call_{i}"), "type": "function", "function": { "name": c.name, "arguments": c.arguments } }))
544                .collect();
545            let msg = json!({ "role": "assistant", "content": null, "tool_calls": tool_calls });
546            axum::Json(completion(&f, &id, msg, "tool_calls")).into_response()
547        }
548    }
549}
550
551fn completion(f: &Fake, id: &str, message: Value, finish: &str) -> Value {
552    json!({
553        "id": id, "object": "chat.completion", "created": 0, "model": f.cfg.model,
554        "choices": [{ "index": 0, "message": message, "finish_reason": finish }],
555        "usage": { "prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20 },
556    })
557}
558
559fn sse(f: &Fake, id: &str, reply: Reply) -> Response {
560    let bad = &f.cfg.misbehave;
561    let chunk = |delta: Value, finish: Option<&str>| {
562        let c = json!({ "id": id, "object": "chat.completion.chunk", "created": 0, "model": f.cfg.model,
563            "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }] });
564        format!("data: {c}\n\n")
565    };
566    let mut events = vec![chunk(json!({ "role": "assistant" }), None)];
567    let mut truncated = false;
568    match reply {
569        Reply::Text { content, reasoning } => {
570            if let Some(r) = reasoning {
571                events.push(chunk(json!({ "reasoning_content": r }), None));
572            }
573            for piece in split_in(&content, 3) {
574                events.push(chunk(json!({ "content": piece }), None));
575            }
576            events.push(chunk(json!({}), Some("stop")));
577        }
578        Reply::Calls(calls) => {
579            for (i, c) in calls.iter().enumerate() {
580                events.push(chunk(json!({ "tool_calls": [{ "index": i, "id": format!("call_{i}"), "type": "function", "function": { "name": c.name, "arguments": "" } }] }), None));
581                let mut pieces = split_in(&c.arguments, 3);
582                if bad.contains(&Misbehavior::StreamArgsSplitInvalid) && pieces.len() > 1 {
583                    pieces.remove(1);
584                }
585                for (j, p) in pieces.iter().enumerate() {
586                    if bad.contains(&Misbehavior::StreamTruncated) && j == pieces.len() - 1 {
587                        truncated = true;
588                        break;
589                    }
590                    events.push(chunk(
591                        json!({ "tool_calls": [{ "index": i, "function": { "arguments": p } }] }),
592                        None,
593                    ));
594                }
595                if truncated {
596                    break;
597                }
598            }
599            if !truncated {
600                events.push(chunk(json!({}), Some("tool_calls")));
601            }
602        }
603        Reply::Error(status, msg) => {
604            // llama-server: the role chunk, then the error, then nothing (no `[DONE]`).
605            let e = json!({ "error": { "code": status.as_u16(), "message": msg, "type": "server_error" } });
606            events.push(format!("data: {e}\n\n"));
607            truncated = true;
608        }
609    }
610    if !truncated {
611        events.push("data: [DONE]\n\n".into());
612    }
613    let body = Body::from_stream(futures_util::stream::iter(
614        events.into_iter().map(Ok::<_, std::io::Error>),
615    ));
616    Response::builder()
617        .header(header::CONTENT_TYPE, "text/event-stream")
618        .header(header::CACHE_CONTROL, "no-cache")
619        .body(body)
620        .expect("static response")
621}
622
623/// `s` in up to `n` pieces on char boundaries.
624fn split_in(s: &str, n: usize) -> Vec<String> {
625    let chars: Vec<char> = s.chars().collect();
626    if chars.is_empty() {
627        return vec![];
628    }
629    let size = chars.len().div_ceil(n);
630    chars.chunks(size).map(|c| c.iter().collect()).collect()
631}
loadmaster_probe/fake.rs

loadmaster_probe/
fake.rs
