| 1 | //! A scripted OpenAI-compatible streaming endpoint that stands in for Onde Cloud. |
| 2 | //! |
| 3 | //! It answers `GET /v1/models` with [`MODEL`] and `other-model`, records every |
| 4 | //! `/chat/completions` body, and replies with the next [`Step`] of its script, chosen by how many |
| 5 | //! tool results the request already carries. Final answers report usage when the request asks |
| 6 | //! for it with `stream_options.include_usage`, as Onde Cloud does. |
| 7 | |
| 8 | use std::path::PathBuf; |
| 9 | use std::sync::atomic::{AtomicUsize, Ordering}; |
| 10 | use std::sync::{Arc, Mutex}; |
| 11 | |
| 12 | use serde_json::{Value, json}; |
| 13 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; |
| 14 | use tokio::net::TcpListener; |
| 15 | |
| 16 | /// Dummy Onde credentials; the mock asserts they arrive as the bearer token. |
| 17 | pub const TEST_KEY: &str = "test-app:test-secret"; |
| 18 | pub const MODEL: &str = "scripted-test-model"; |
| 19 | |
| 20 | #[derive(Debug, Default)] |
| 21 | pub struct MockStats { |
| 22 | pub requests: AtomicUsize, |
| 23 | pub saw_bearer: Mutex<bool>, |
| 24 | /// Every `/chat/completions` request body, in arrival order. |
| 25 | pub payloads: Mutex<Vec<Value>>, |
| 26 | } |
| 27 | |
| 28 | /// What the scripted model does on its next completion. |
| 29 | #[derive(Clone)] |
| 30 | pub enum Step { |
| 31 | /// Emit a tool call, then wait for the tool result. |
| 32 | ToolCall { |
| 33 | name: &'static str, |
| 34 | arguments: Value, |
| 35 | }, |
| 36 | /// Emit final assistant text and finish. |
| 37 | Final(&'static str), |
| 38 | /// Stop with `finish_reason: content_filter`. |
| 39 | ContentFilter, |
| 40 | } |
| 41 | |
| 42 | #[derive(Clone)] |
| 43 | pub struct Script(pub Vec<Step>); |
| 44 | |
| 45 | impl Script { |
| 46 | /// Count completed tool results to know how far into the script we are. |
| 47 | fn step_for(&self, messages: &[Value]) -> (usize, Step) { |
| 48 | let done = messages.iter().filter(|m| m["role"] == "tool").count(); |
| 49 | ( |
| 50 | done, |
| 51 | self.0 |
| 52 | .get(done) |
| 53 | .unwrap_or_else(|| self.0.last().expect("empty script")) |
| 54 | .clone(), |
| 55 | ) |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | fn sse_chunk( |
| 60 | content: Option<&str>, |
| 61 | tool_call: Option<(&str, &str, &str)>, |
| 62 | finish: Option<&str>, |
| 63 | ) -> String { |
| 64 | let mut delta = json!({}); |
| 65 | if let Some(c) = content { |
| 66 | delta["content"] = json!(c); |
| 67 | } |
| 68 | if let Some((id, name, args)) = tool_call { |
| 69 | delta["tool_calls"] = json!([{ |
| 70 | "index": 0, |
| 71 | "id": id, |
| 72 | "function": { "name": name, "arguments": args }, |
| 73 | }]); |
| 74 | } |
| 75 | let chunk = json!({ "choices": [{ "delta": delta, "finish_reason": finish }] }); |
| 76 | format!("data: {chunk}\n\n") |
| 77 | } |
| 78 | |
| 79 | async fn read_http_request(socket: &mut tokio::net::TcpStream) -> (String, String) { |
| 80 | let mut buf = Vec::new(); |
| 81 | let mut tmp = [0u8; 4096]; |
| 82 | let header_end; |
| 83 | let mut content_length = 0usize; |
| 84 | loop { |
| 85 | let n = socket.read(&mut tmp).await.unwrap(); |
| 86 | if n == 0 { |
| 87 | panic!("connection closed before headers complete"); |
| 88 | } |
| 89 | buf.extend_from_slice(&tmp[..n]); |
| 90 | if let Some(pos) = find_subslice(&buf, b"\r\n\r\n") { |
| 91 | header_end = pos + 4; |
| 92 | let headers = String::from_utf8_lossy(&buf[..pos]).to_string(); |
| 93 | for line in headers.lines() { |
| 94 | if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") { |
| 95 | content_length = v.trim().parse().unwrap(); |
| 96 | } |
| 97 | } |
| 98 | break; |
| 99 | } |
| 100 | } |
| 101 | while buf.len() < header_end + content_length { |
| 102 | let n = socket.read(&mut tmp).await.unwrap(); |
| 103 | if n == 0 { |
| 104 | break; |
| 105 | } |
| 106 | buf.extend_from_slice(&tmp[..n]); |
| 107 | } |
| 108 | let head = String::from_utf8_lossy(&buf[..header_end]).to_string(); |
| 109 | let body = String::from_utf8_lossy(&buf[header_end..header_end + content_length]).to_string(); |
| 110 | (head, body) |
| 111 | } |
| 112 | |
| 113 | fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> { |
| 114 | haystack.windows(needle.len()).position(|w| w == needle) |
| 115 | } |
| 116 | |
| 117 | async fn respond(socket: &mut tokio::net::TcpStream, content_type: &str, body: &str) { |
| 118 | let resp = format!( |
| 119 | "HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", |
| 120 | body.len() |
| 121 | ); |
| 122 | socket.write_all(resp.as_bytes()).await.unwrap(); |
| 123 | socket.shutdown().await.unwrap(); |
| 124 | } |
| 125 | |
| 126 | /// Spawn the mock server; returns its base URL (e.g. `http://127.0.0.1:PORT/v1`). |
| 127 | pub async fn start_mock_llm(script: Script, stats: Arc<MockStats>) -> String { |
| 128 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); |
| 129 | let port = listener.local_addr().unwrap().port(); |
| 130 | tokio::spawn(async move { |
| 131 | loop { |
| 132 | let Ok((mut socket, _)) = listener.accept().await else { |
| 133 | return; |
| 134 | }; |
| 135 | let stats = stats.clone(); |
| 136 | let script = script.clone(); |
| 137 | tokio::spawn(async move { |
| 138 | let (head, body) = read_http_request(&mut socket).await; |
| 139 | let first = head.lines().next().unwrap_or("").to_string(); |
| 140 | if first.starts_with("GET /v1/models") { |
| 141 | let list = |
| 142 | json!({"object": "list", "data": [{"id": MODEL}, {"id": "other-model"}]}); |
| 143 | return respond(&mut socket, "application/json", &list.to_string()).await; |
| 144 | } |
| 145 | stats.requests.fetch_add(1, Ordering::SeqCst); |
| 146 | if head |
| 147 | .to_ascii_lowercase() |
| 148 | .contains(&format!("authorization: bearer {TEST_KEY}").to_ascii_lowercase()) |
| 149 | { |
| 150 | *stats.saw_bearer.lock().unwrap() = true; |
| 151 | } |
| 152 | assert!( |
| 153 | first.starts_with("POST /v1/chat/completions"), |
| 154 | "unexpected request: {first}" |
| 155 | ); |
| 156 | let payload: Value = serde_json::from_str(&body).unwrap(); |
| 157 | assert_eq!(payload["stream"], true); |
| 158 | stats.payloads.lock().unwrap().push(payload.clone()); |
| 159 | let messages = payload["messages"].as_array().unwrap(); |
| 160 | let (done, step) = script.step_for(messages); |
| 161 | |
| 162 | let mut data = String::new(); |
| 163 | match step { |
| 164 | Step::ToolCall { name, arguments } => { |
| 165 | let id = format!("call_{}", done + 1); |
| 166 | data.push_str(&sse_chunk( |
| 167 | None, |
| 168 | Some((&id, name, &arguments.to_string())), |
| 169 | None, |
| 170 | )); |
| 171 | data.push_str(&sse_chunk(None, None, Some("tool_calls"))); |
| 172 | } |
| 173 | Step::Final(text) => { |
| 174 | data.push_str(&sse_chunk(Some(text), None, None)); |
| 175 | data.push_str(&sse_chunk(None, None, Some("stop"))); |
| 176 | // `stream_options.include_usage`: a final chunk with no choices. |
| 177 | if payload.get("stream_options").is_some() { |
| 178 | let usage = json!({"choices": [], "usage": { |
| 179 | "prompt_tokens": 120, "completion_tokens": 30, "total_tokens": 150 |
| 180 | }}); |
| 181 | data.push_str(&format!("data: {usage}\n\n")); |
| 182 | } |
| 183 | } |
| 184 | Step::ContentFilter => { |
| 185 | data.push_str(&sse_chunk(None, None, Some("content_filter"))); |
| 186 | } |
| 187 | } |
| 188 | data.push_str("data: [DONE]\n\n"); |
| 189 | respond(&mut socket, "text/event-stream", &data).await; |
| 190 | }); |
| 191 | } |
| 192 | }); |
| 193 | format!("http://127.0.0.1:{port}/v1") |
| 194 | } |
| 195 | |
| 196 | pub fn temp_dir(prefix: &str) -> PathBuf { |
| 197 | let dir = std::env::temp_dir().join(format!("{prefix}-{}", uuid::Uuid::new_v4())); |
| 198 | std::fs::create_dir_all(&dir).unwrap(); |
| 199 | dir |
| 200 | } |
| 201 | |
| 202 | /// Start a mock with `steps`; returns its base URL and the recorded requests. |
| 203 | pub async fn mock(steps: Vec<Step>) -> (String, Arc<MockStats>) { |
| 204 | let stats = Arc::new(MockStats::default()); |
| 205 | let url = start_mock_llm(Script(steps), stats.clone()).await; |
| 206 | (url, stats) |
| 207 | } |