ZacharyZhang-NY
a02b555e66
Sweep every first-party crate source (1956 .rs files) to the project comment guidelines: delete redundant restatements, decorative banners, change narration, and end-of-line comments; keep and tighten the crucial ones (invariants, bug rationale, SAFETY blocks, ported-source attribution). No functional code changed. Every edit is proven comment-only against the prior tree by a comment-stripping lexer (string/char/raw-string aware) plus a separate doctest-fence check. Where removing a comment made rustfmt or clippy want to re-lay-out adjacent code, the minimal triggering comment is restored so code tokens stay byte-identical. Gates green: cargo fmt --all --check (0 diffs), cargo check and cargo clippy --workspace --all-targets (0 warnings). Adds scripts/check_codegen_comment_guidelines.py — the enforcement gate for these guidelines (flags banners, end-of-line comments, change narration, and commented-out code).
1453 lines
55 KiB
Rust
1453 lines
55 KiB
Rust
//! Mock inference server with request logging and automatic cleanup.
|
|
//!
|
|
//! Serves `/v1/chat/completions`, `/v1/responses`, and `/v1/messages` in one
|
|
//! of two response modes: echo (default — streams `Echo: <last user message>`)
|
|
//! or a fixed text set via [`MockInferenceServer::set_response`] (streamed
|
|
//! with byte-exact reconstruction). A per-path FIFO of [`ScriptedResponse`]s
|
|
//! (see [`MockInferenceServer::enqueue_response`]) overrides the mode for
|
|
//! exact status/body/SSE control. `/v1/models` and `/v1/settings` return
|
|
//! configurable responses (settings is 404 until set). All requests are
|
|
//! logged — bodies and headers — for assertion in tests.
|
|
|
|
use std::collections::{HashMap, VecDeque};
|
|
use std::convert::Infallible;
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
|
use std::time::Duration;
|
|
|
|
use anyhow::Context as _;
|
|
use axum::http::{HeaderMap, StatusCode};
|
|
use axum::response::sse::{KeepAlive, Sse};
|
|
use axum::response::{IntoResponse, Response};
|
|
use axum::routing::{get, post};
|
|
use axum::{Json, Router};
|
|
use futures_util::stream;
|
|
use serde_json::{Value, json};
|
|
use tokio::net::TcpListener;
|
|
use tokio::sync::oneshot;
|
|
|
|
pub use crate::scripted::{ScriptedBody, ScriptedResponse, SseEvent};
|
|
use crate::sse;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct LogEntry {
|
|
pub method: String,
|
|
pub path: String,
|
|
pub body: Option<Value>,
|
|
pub authorization: Option<String>,
|
|
/// Request headers (lowercase names, arrival order), captured on the
|
|
/// inference POST endpoints; the GET endpoints log an empty list.
|
|
pub headers: Vec<(String, String)>,
|
|
}
|
|
|
|
impl LogEntry {
|
|
/// First value of `name` (case-insensitive), if the request carried it.
|
|
pub fn header(&self, name: &str) -> Option<&str> {
|
|
let name = name.to_ascii_lowercase();
|
|
self.headers
|
|
.iter()
|
|
.find(|(k, _)| *k == name)
|
|
.map(|(_, v)| v.as_str())
|
|
}
|
|
}
|
|
|
|
pub struct RequestLog {
|
|
count: AtomicU32,
|
|
entries: std::sync::Mutex<Vec<LogEntry>>,
|
|
}
|
|
|
|
impl RequestLog {
|
|
fn new() -> Self {
|
|
Self {
|
|
count: AtomicU32::new(0),
|
|
entries: std::sync::Mutex::new(Vec::new()),
|
|
}
|
|
}
|
|
|
|
fn record(
|
|
&self,
|
|
method: &str,
|
|
path: &str,
|
|
body: Option<&Value>,
|
|
authorization: Option<&str>,
|
|
headers: Vec<(String, String)>,
|
|
) {
|
|
self.count.fetch_add(1, Ordering::SeqCst);
|
|
self.entries.lock().unwrap().push(LogEntry {
|
|
method: method.to_string(),
|
|
path: path.to_string(),
|
|
body: body.cloned(),
|
|
authorization: authorization.map(String::from),
|
|
headers,
|
|
});
|
|
}
|
|
}
|
|
|
|
type ScriptQueues = Arc<std::sync::Mutex<HashMap<String, VecDeque<ScriptedResponse>>>>;
|
|
|
|
/// A model entry for the mock `/v1/models` endpoint.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MockModelEntry {
|
|
pub id: String,
|
|
/// Emitted as `agentType` inside `_meta` when set (e.g. `"cursor"`).
|
|
pub agent_type: Option<String>,
|
|
/// Emitted as `apiBackend` when set (e.g. `"messages"`); absent means the
|
|
/// shell's default backend.
|
|
pub api_backend: Option<String>,
|
|
/// Emitted as `supportsBackendSearch` when true.
|
|
pub supports_backend_search: bool,
|
|
/// Emitted as `supportsReasoningEffort` (top-level) when true.
|
|
pub supports_reasoning_effort: bool,
|
|
/// Emitted as `reasoningEffort` (top-level) when set.
|
|
pub reasoning_effort: Option<String>,
|
|
/// Emitted as `reasoningEfforts` (top-level) when non-empty. Each entry is a
|
|
/// raw JSON option (a table `{ "value": ..., "id"?, "label"?, ... }` or a
|
|
/// bare value string), matching what `parse_remote_model_value` reads.
|
|
pub reasoning_efforts: Vec<Value>,
|
|
}
|
|
|
|
impl MockModelEntry {
|
|
pub fn new(id: impl Into<String>) -> Self {
|
|
Self {
|
|
id: id.into(),
|
|
agent_type: None,
|
|
api_backend: None,
|
|
supports_backend_search: false,
|
|
supports_reasoning_effort: false,
|
|
reasoning_effort: None,
|
|
reasoning_efforts: Vec::new(),
|
|
}
|
|
}
|
|
|
|
pub fn with_agent_type(id: impl Into<String>, agent_type: impl Into<String>) -> Self {
|
|
Self {
|
|
agent_type: Some(agent_type.into()),
|
|
..Self::new(id)
|
|
}
|
|
}
|
|
|
|
pub fn with_api_backend(mut self, api_backend: impl Into<String>) -> Self {
|
|
self.api_backend = Some(api_backend.into());
|
|
self
|
|
}
|
|
|
|
pub fn with_supports_backend_search(mut self, supports: bool) -> Self {
|
|
self.supports_backend_search = supports;
|
|
self
|
|
}
|
|
|
|
pub fn with_supports_reasoning_effort(mut self, supports: bool) -> Self {
|
|
self.supports_reasoning_effort = supports;
|
|
self
|
|
}
|
|
|
|
pub fn with_reasoning_effort(mut self, effort: impl Into<String>) -> Self {
|
|
self.reasoning_effort = Some(effort.into());
|
|
self
|
|
}
|
|
|
|
pub fn with_reasoning_efforts(mut self, efforts: Vec<Value>) -> Self {
|
|
self.reasoning_efforts = efforts;
|
|
self
|
|
}
|
|
|
|
fn to_json(&self) -> Value {
|
|
let mut obj = json!({
|
|
"id": self.id,
|
|
"object": "model",
|
|
"created": 1234567890,
|
|
"owned_by": "test"
|
|
});
|
|
if let Some(ref at) = self.agent_type {
|
|
obj["_meta"] = json!({ "agentType": at });
|
|
}
|
|
if let Some(ref backend) = self.api_backend {
|
|
obj["apiBackend"] = json!(backend);
|
|
}
|
|
if self.supports_backend_search {
|
|
obj["supportsBackendSearch"] = json!(true);
|
|
}
|
|
if self.supports_reasoning_effort {
|
|
obj["supportsReasoningEffort"] = json!(true);
|
|
}
|
|
if let Some(ref effort) = self.reasoning_effort {
|
|
obj["reasoningEffort"] = json!(effort);
|
|
}
|
|
if !self.reasoning_efforts.is_empty() {
|
|
obj["reasoningEfforts"] = json!(self.reasoning_efforts);
|
|
}
|
|
obj
|
|
}
|
|
}
|
|
|
|
/// What the inference endpoints stream back.
|
|
enum ResponseMode {
|
|
/// Echo the last user message as `Echo: <msg>` (whitespace-collapsing).
|
|
Echo,
|
|
/// Stream a fixed text whose deltas reconstruct it byte-for-byte
|
|
/// (newlines preserved — required for fenced code blocks).
|
|
Fixed(String),
|
|
}
|
|
|
|
/// Opt-in barrier that holds an **agent turn's terminal SSE event** until the
|
|
/// test releases it, so the turn stays deterministically "running" while the
|
|
/// test interacts with it (queue edits/removals) — eliminating turn-end races.
|
|
///
|
|
/// Inert by default (`held == false`): [`wait_if_held`] returns immediately, so
|
|
/// every test that never calls [`MockInferenceServer::hold_agent_completions`]
|
|
/// is completely unaffected.
|
|
///
|
|
/// [`wait_if_held`]: CompletionGate::wait_if_held
|
|
#[derive(Default)]
|
|
struct CompletionGate {
|
|
held: AtomicBool,
|
|
notify: tokio::sync::Notify,
|
|
}
|
|
|
|
impl CompletionGate {
|
|
fn hold(&self) {
|
|
self.held.store(true, Ordering::SeqCst);
|
|
}
|
|
|
|
fn release(&self) {
|
|
self.held.store(false, Ordering::SeqCst);
|
|
self.notify.notify_waiters();
|
|
}
|
|
|
|
/// Block while the gate is held. Registers the wake-up interest *before*
|
|
/// re-checking `held` so a concurrent `release` can never be missed.
|
|
async fn wait_if_held(&self) {
|
|
loop {
|
|
let notified = self.notify.notified();
|
|
if !self.held.load(Ordering::SeqCst) {
|
|
return;
|
|
}
|
|
notified.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Wrap SSE `events` in a stream that emits each one after `delay`. `None`
|
|
/// keeps instant emission (the default fast path); `Some(d)` paces the stream
|
|
/// so tests can interact with a turn while it is visibly in flight.
|
|
///
|
|
/// When `gate` is `Some`, the stream additionally blocks on the gate right
|
|
/// before emitting the **final** event (the SSE terminator), so a held gate
|
|
/// keeps the turn streaming-but-not-complete until released.
|
|
fn paced_events(
|
|
events: Vec<axum::response::sse::Event>,
|
|
delay: Option<Duration>,
|
|
gate: Option<Arc<CompletionGate>>,
|
|
) -> impl futures_util::Stream<Item = Result<axum::response::sse::Event, Infallible>> {
|
|
use futures_util::StreamExt as _;
|
|
let last_idx = events.len().saturating_sub(1);
|
|
stream::iter(events.into_iter().enumerate()).then(move |(idx, event)| {
|
|
let gate = gate.clone();
|
|
async move {
|
|
if let Some(d) = delay {
|
|
tokio::time::sleep(d).await;
|
|
}
|
|
if idx == last_idx
|
|
&& let Some(gate) = gate.as_deref()
|
|
{
|
|
gate.wait_if_held().await;
|
|
}
|
|
Ok::<_, Infallible>(event)
|
|
}
|
|
})
|
|
}
|
|
|
|
/// Max body bytes retained on each accepted [`StorageUpload`] (keeps large
|
|
/// e2e artifacts from ballooning test memory; meta/small dumps stay intact).
|
|
const STORAGE_BODY_CAPTURE_CAP: usize = 256 * 1024;
|
|
|
|
/// One accepted (HTTP 200) mock `/v1/storage` upload.
|
|
#[derive(Debug, Clone)]
|
|
pub struct StorageUpload {
|
|
pub path: String,
|
|
pub size: usize,
|
|
/// Request body when `size <= 256 KiB`; empty for larger payloads.
|
|
pub body: Vec<u8>,
|
|
pub authorization: Option<String>,
|
|
}
|
|
|
|
/// Mock `/v1/storage` state: a flippable 401 gate plus a record of accepted
|
|
/// uploads, so e2e tests can simulate an auth outage window and assert the
|
|
/// trace upload queue parks, then drains after the gate heals.
|
|
#[derive(Default)]
|
|
struct StorageState {
|
|
unauthorized: AtomicBool,
|
|
request_count: AtomicU32,
|
|
uploads: std::sync::Mutex<Vec<StorageUpload>>,
|
|
}
|
|
|
|
/// Mock `/v1/chat/completions` + `/v1/responses` + `/v1/messages` +
|
|
/// `/v1/models` + `/v1/settings` + `/v1/storage` server.
|
|
/// Logs all requests. Shuts down on drop.
|
|
pub struct MockInferenceServer {
|
|
addr: SocketAddr,
|
|
shutdown_tx: Option<oneshot::Sender<()>>,
|
|
log: Arc<RequestLog>,
|
|
models: Arc<std::sync::RwLock<Vec<Value>>>,
|
|
settings: Arc<std::sync::RwLock<Option<Value>>>,
|
|
response_mode: Arc<std::sync::RwLock<ResponseMode>>,
|
|
scripted: ScriptQueues,
|
|
/// Per-agent-turn assistant texts (see [`set_agent_turns`]).
|
|
///
|
|
/// [`set_agent_turns`]: Self::set_agent_turns
|
|
agent_turns: Arc<std::sync::Mutex<VecDeque<String>>>,
|
|
/// `stop_reason` emitted by the `/v1/messages` terminal `message_delta`.
|
|
messages_stop_reason: Arc<std::sync::RwLock<String>>,
|
|
/// Optional per-SSE-event delay on all inference endpoints. `None`
|
|
/// (default) streams instantly; `Some(d)` holds the turn "streaming" long
|
|
/// enough for tests to interact with it mid-flight (e.g. Esc-cancel).
|
|
chunk_delay: Arc<std::sync::RwLock<Option<Duration>>>,
|
|
/// Mock `/v1/storage` 401 gate + accepted-upload record.
|
|
storage: Arc<StorageState>,
|
|
/// Opt-in barrier holding agent turns' terminal event (see
|
|
/// [`Self::hold_agent_completions`]). Inert until a test holds it.
|
|
completion_gate: Arc<CompletionGate>,
|
|
/// See [`Self::set_user_subscription_tier`].
|
|
user_tier: Arc<std::sync::RwLock<Option<String>>>,
|
|
}
|
|
|
|
impl MockInferenceServer {
|
|
/// Start with a single default `test-model` (no agent_type).
|
|
pub async fn start() -> anyhow::Result<Self> {
|
|
Self::start_with_models(vec![MockModelEntry::new("test-model")]).await
|
|
}
|
|
|
|
/// Start with custom models. Use [`MockModelEntry::with_agent_type`] to
|
|
/// configure models with specific harness types for agent-type tests.
|
|
pub async fn start_with_models(models: Vec<MockModelEntry>) -> anyhow::Result<Self> {
|
|
Self::start_inner(models, None).await
|
|
}
|
|
|
|
/// Start a mock that returns 401 on inference requests missing
|
|
/// `Authorization: Bearer <required_token>`.
|
|
pub async fn start_with_required_auth(
|
|
models: Vec<MockModelEntry>,
|
|
required_token: impl Into<String>,
|
|
) -> anyhow::Result<Self> {
|
|
Self::start_inner(models, Some(required_token.into())).await
|
|
}
|
|
|
|
async fn start_inner(
|
|
models: Vec<MockModelEntry>,
|
|
required_token: Option<String>,
|
|
) -> anyhow::Result<Self> {
|
|
let log = Arc::new(RequestLog::new());
|
|
let models_json: Vec<Value> = models.iter().map(MockModelEntry::to_json).collect();
|
|
let shared_models = Arc::new(std::sync::RwLock::new(models_json));
|
|
let shared_settings = Arc::new(std::sync::RwLock::new(None::<Value>));
|
|
let response_mode = Arc::new(std::sync::RwLock::new(ResponseMode::Echo));
|
|
let scripted: ScriptQueues = Arc::new(std::sync::Mutex::new(HashMap::new()));
|
|
let agent_turns = Arc::new(std::sync::Mutex::new(VecDeque::new()));
|
|
let messages_stop_reason = Arc::new(std::sync::RwLock::new("end_turn".to_string()));
|
|
let chunk_delay = Arc::new(std::sync::RwLock::new(None::<Duration>));
|
|
let storage = Arc::new(StorageState::default());
|
|
let completion_gate = Arc::new(CompletionGate::default());
|
|
let user_tier = Arc::new(std::sync::RwLock::new(None::<String>));
|
|
let app = Self::build_router(
|
|
log.clone(),
|
|
shared_models.clone(),
|
|
shared_settings.clone(),
|
|
response_mode.clone(),
|
|
scripted.clone(),
|
|
agent_turns.clone(),
|
|
messages_stop_reason.clone(),
|
|
chunk_delay.clone(),
|
|
storage.clone(),
|
|
completion_gate.clone(),
|
|
user_tier.clone(),
|
|
required_token,
|
|
);
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.context("bind mock inference server")?;
|
|
let addr = listener.local_addr().context("local_addr")?;
|
|
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
|
|
|
tokio::spawn(async move {
|
|
axum::serve(listener, app)
|
|
.with_graceful_shutdown(async {
|
|
let _ = shutdown_rx.await;
|
|
})
|
|
.await
|
|
.unwrap();
|
|
});
|
|
|
|
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
|
|
while tokio::net::TcpStream::connect(addr).await.is_err() {
|
|
if tokio::time::Instant::now() >= deadline {
|
|
anyhow::bail!("mock server not ready within 5s");
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
}
|
|
|
|
Ok(Self {
|
|
addr,
|
|
shutdown_tx: Some(shutdown_tx),
|
|
log,
|
|
models: shared_models,
|
|
settings: shared_settings,
|
|
response_mode,
|
|
scripted,
|
|
agent_turns,
|
|
messages_stop_reason,
|
|
chunk_delay,
|
|
storage,
|
|
completion_gate,
|
|
user_tier,
|
|
})
|
|
}
|
|
|
|
/// Replace the model list at runtime. The next `/v1/models` request
|
|
/// (e.g. during session resume) will return the new list.
|
|
pub fn set_models(&self, models: Vec<MockModelEntry>) {
|
|
let mut guard = self.models.write().unwrap();
|
|
*guard = models.iter().map(MockModelEntry::to_json).collect();
|
|
}
|
|
|
|
/// Stream this fixed text from all inference endpoints instead of echoing
|
|
/// the user message. Deltas reconstruct the text byte-for-byte (newlines
|
|
/// preserved). Subsequent calls replace the text.
|
|
pub fn set_response(&self, text: impl Into<String>) {
|
|
*self.response_mode.write().unwrap() = ResponseMode::Fixed(text.into());
|
|
}
|
|
|
|
/// Queue a [`ScriptedResponse`] for the next request on `path` (e.g.
|
|
/// `"/v1/chat/completions"`). Scripts are consumed FIFO per path by the
|
|
/// three inference endpoints; when a path's queue is empty, requests fall
|
|
/// back to the active response mode (echo/fixed).
|
|
pub fn enqueue_response(&self, path: impl Into<String>, response: ScriptedResponse) {
|
|
// Fail at the call site, not at serve time.
|
|
response.validate();
|
|
self.scripted
|
|
.lock()
|
|
.unwrap()
|
|
.entry(path.into())
|
|
.or_default()
|
|
.push_back(response);
|
|
}
|
|
|
|
/// Queue one byte-exact response per agent turn, consumed FIFO. Only
|
|
/// requests carrying 2+ tools count as agent turns, so aux requests
|
|
/// (title/classifier) never steal a turn; an empty queue falls back to
|
|
/// the active response mode.
|
|
pub fn set_agent_turns(&self, turns: impl IntoIterator<Item = String>) {
|
|
*self.agent_turns.lock().unwrap() = turns.into_iter().collect();
|
|
}
|
|
|
|
/// Replace the settings at runtime. The next `GET /v1/settings` request
|
|
/// will return the new value as JSON. Until set, `/v1/settings` returns 404.
|
|
pub fn set_settings(&self, settings: impl serde::Serialize) {
|
|
let value = serde_json::to_value(settings).expect("serialize settings");
|
|
let mut guard = self.settings.write().unwrap();
|
|
*guard = Some(value);
|
|
}
|
|
|
|
/// Preset `/v1/settings` to the minimal `{"allow_access": true}` payload
|
|
/// that opens the subscription gate (clients treat a missing field as
|
|
/// `false` and would sit on the upsell screen).
|
|
pub fn preset_allow_access(&self) {
|
|
self.set_settings(json!({ "allow_access": true }));
|
|
}
|
|
|
|
/// Set the `subscriptionTier` served by `GET /v1/user`. `None`
|
|
/// (default) omits the field, which the shell treats as "no qualifying
|
|
/// subscription" (free tier).
|
|
pub fn set_user_subscription_tier(&self, tier: Option<&str>) {
|
|
*self.user_tier.write().unwrap() = tier.map(str::to_owned);
|
|
}
|
|
|
|
/// Set the `stop_reason` emitted by the `/v1/messages` terminal
|
|
/// `message_delta` (default `"end_turn"`).
|
|
pub fn set_messages_stop_reason(&self, stop_reason: impl Into<String>) {
|
|
*self.messages_stop_reason.write().unwrap() = stop_reason.into();
|
|
}
|
|
|
|
/// Pace all inference SSE streams: each event is emitted after `delay`.
|
|
/// `None` (default) restores instant streaming. Lets PTY e2e tests hold a
|
|
/// turn visibly "streaming" long enough to interact with it mid-flight
|
|
/// (e.g. Esc-cancel). Applies to requests started after the call.
|
|
pub fn set_chunk_delay(&self, delay: Option<Duration>) {
|
|
*self.chunk_delay.write().unwrap() = delay;
|
|
}
|
|
|
|
/// Hold every agent turn's terminal SSE event until
|
|
/// [`release_agent_completions`] is called, keeping the turn
|
|
/// deterministically "streaming-but-not-complete". Lets a test interact
|
|
/// with a running turn (e.g. queue edits/removals) without racing turn
|
|
/// end. Content deltas still stream normally; only completion is gated.
|
|
/// Inert for tests that never call this.
|
|
///
|
|
/// [`release_agent_completions`]: Self::release_agent_completions
|
|
pub fn hold_agent_completions(&self) {
|
|
self.completion_gate.hold();
|
|
}
|
|
|
|
/// Release a hold set by [`hold_agent_completions`], letting held (and
|
|
/// future) agent turns emit their terminal event and complete.
|
|
///
|
|
/// [`hold_agent_completions`]: Self::hold_agent_completions
|
|
pub fn release_agent_completions(&self) {
|
|
self.completion_gate.release();
|
|
}
|
|
|
|
/// e.g. `http://127.0.0.1:12345/v1`
|
|
pub fn url(&self) -> String {
|
|
format!("http://{}/v1", self.addr)
|
|
}
|
|
|
|
pub fn request_count(&self) -> u32 {
|
|
self.log.count.load(Ordering::SeqCst)
|
|
}
|
|
|
|
pub fn requests(&self) -> Vec<LogEntry> {
|
|
self.log.entries.lock().unwrap().clone()
|
|
}
|
|
|
|
/// Bodies of all received requests, in arrival order (body-less requests
|
|
/// such as `GET /v1/models` are skipped).
|
|
pub fn request_bodies(&self) -> Vec<Value> {
|
|
self.log
|
|
.entries
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.filter_map(|e| e.body.clone())
|
|
.collect()
|
|
}
|
|
|
|
pub fn has_chat_completion_request(&self) -> bool {
|
|
self.log
|
|
.entries
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.any(|e| e.path.contains("chat/completions"))
|
|
}
|
|
|
|
pub fn has_responses_request(&self) -> bool {
|
|
self.log
|
|
.entries
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.any(|e| e.path.contains("responses"))
|
|
}
|
|
|
|
pub fn messages_request_count(&self) -> usize {
|
|
self.log
|
|
.entries
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.filter(|e| e.path == "/v1/messages")
|
|
.count()
|
|
}
|
|
|
|
pub fn request_log_summary(&self) -> String {
|
|
let entries = self.log.entries.lock().unwrap();
|
|
if entries.is_empty() {
|
|
return "(no requests received)".to_string();
|
|
}
|
|
entries
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, e)| format!(" [{}] {} {}", i, e.method, e.path))
|
|
.collect::<Vec<_>>()
|
|
.join("\n")
|
|
}
|
|
|
|
pub fn last_system_prompt(&self) -> Option<String> {
|
|
let entries = self.log.entries.lock().unwrap();
|
|
entries
|
|
.iter()
|
|
.rev()
|
|
.find(|e| e.path.contains("chat/completions") || e.path.contains("responses"))
|
|
.and_then(|e| e.body.as_ref())
|
|
// Chat completions carries the system prompt as the first message;
|
|
// the Responses API carries it in `instructions` instead.
|
|
.and_then(|body| {
|
|
body.get("messages")
|
|
.and_then(|m| m.as_array())
|
|
.and_then(|msgs| msgs.first())
|
|
.and_then(|msg| msg.get("content"))
|
|
.and_then(|c| c.as_str())
|
|
.map(String::from)
|
|
.or_else(|| {
|
|
body.get("instructions")
|
|
.and_then(|s| s.as_str())
|
|
.map(String::from)
|
|
})
|
|
})
|
|
}
|
|
|
|
/// Flip the mock `/v1/storage` 401 gate. While `true`, every upload is
|
|
/// rejected with 401 (the auth-outage window the park-on-401 e2e drives).
|
|
pub fn set_storage_unauthorized(&self, unauthorized: bool) {
|
|
self.storage
|
|
.unauthorized
|
|
.store(unauthorized, Ordering::SeqCst);
|
|
}
|
|
|
|
/// Total `/v1/storage` upload attempts seen, including 401-rejected ones.
|
|
pub fn storage_request_count(&self) -> u32 {
|
|
self.storage.request_count.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Snapshot of accepted (HTTP 200) `/v1/storage` uploads.
|
|
pub fn storage_uploads(&self) -> Vec<StorageUpload> {
|
|
self.storage.uploads.lock().unwrap().clone()
|
|
}
|
|
|
|
/// Mock `/v1/storage` upload: count the attempt, reject with 401 while the
|
|
/// gate is closed, else record the upload and mirror the proxy's
|
|
/// `UploadResponse` JSON shape.
|
|
fn storage_upload_handler(
|
|
storage: &StorageState,
|
|
headers: &HeaderMap,
|
|
body: &axum::body::Bytes,
|
|
) -> Response {
|
|
storage.request_count.fetch_add(1, Ordering::SeqCst);
|
|
if storage.unauthorized.load(Ordering::SeqCst) {
|
|
return (
|
|
StatusCode::UNAUTHORIZED,
|
|
r#"{"error":"Invalid or expired credentials (mock)"}"#,
|
|
)
|
|
.into_response();
|
|
}
|
|
|
|
let path = headers
|
|
.get("X-Storage-Path")
|
|
.and_then(|v| v.to_str().ok())
|
|
.unwrap_or("")
|
|
.to_owned();
|
|
let size = body.len();
|
|
let captured_body = if size <= STORAGE_BODY_CAPTURE_CAP {
|
|
body.to_vec()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
let authorization = Self::extract_auth(headers);
|
|
let response = (
|
|
StatusCode::OK,
|
|
[(axum::http::header::CONTENT_TYPE, "application/json")],
|
|
json!({
|
|
"bucket": "mock-bucket",
|
|
"path": path,
|
|
"size": size,
|
|
"content_type": "application/octet-stream",
|
|
"generation": 1,
|
|
})
|
|
.to_string(),
|
|
);
|
|
storage.uploads.lock().unwrap().push(StorageUpload {
|
|
path,
|
|
size,
|
|
body: captured_body,
|
|
authorization,
|
|
});
|
|
response.into_response()
|
|
}
|
|
|
|
fn extract_auth(headers: &HeaderMap) -> Option<String> {
|
|
headers
|
|
.get("authorization")
|
|
.and_then(|v| v.to_str().ok())
|
|
.map(String::from)
|
|
}
|
|
|
|
fn headers_vec(headers: &HeaderMap) -> Vec<(String, String)> {
|
|
headers
|
|
.iter()
|
|
.map(|(k, v)| {
|
|
(
|
|
k.as_str().to_string(),
|
|
String::from_utf8_lossy(v.as_bytes()).into_owned(),
|
|
)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn pop_scripted(scripted: &ScriptQueues, path: &str) -> Option<ScriptedResponse> {
|
|
scripted
|
|
.lock()
|
|
.unwrap()
|
|
.get_mut(path)
|
|
.and_then(VecDeque::pop_front)
|
|
}
|
|
|
|
/// Pop the next scripted turn, gated to agent turns (2+ tools) so aux
|
|
/// requests don't consume one.
|
|
fn pop_agent_turn(
|
|
agent_turns: &Arc<std::sync::Mutex<VecDeque<String>>>,
|
|
body: &Value,
|
|
) -> Option<String> {
|
|
let tool_count = body
|
|
.get("tools")
|
|
.and_then(Value::as_array)
|
|
.map_or(0, Vec::len);
|
|
if tool_count < 2 {
|
|
return None;
|
|
}
|
|
agent_turns.lock().unwrap().pop_front()
|
|
}
|
|
|
|
/// Returns `Some(401)` if auth is required and the Bearer token doesn't match.
|
|
fn check_auth(auth: Option<&str>, required_token: Option<&str>) -> Option<Response> {
|
|
let expected = required_token?;
|
|
let valid = auth.is_some_and(|v| {
|
|
v.strip_prefix("Bearer ")
|
|
.or_else(|| v.strip_prefix("bearer "))
|
|
.is_some_and(|token| token == expected)
|
|
});
|
|
if valid {
|
|
return None;
|
|
}
|
|
Some((
|
|
StatusCode::UNAUTHORIZED,
|
|
Json(json!({
|
|
"error": "missing API key; set the x-api-key header or Authorization: Bearer header"
|
|
})),
|
|
).into_response())
|
|
}
|
|
|
|
fn build_router(
|
|
log: Arc<RequestLog>,
|
|
models: Arc<std::sync::RwLock<Vec<Value>>>,
|
|
settings: Arc<std::sync::RwLock<Option<Value>>>,
|
|
response_mode: Arc<std::sync::RwLock<ResponseMode>>,
|
|
scripted: ScriptQueues,
|
|
agent_turns: Arc<std::sync::Mutex<VecDeque<String>>>,
|
|
messages_stop_reason: Arc<std::sync::RwLock<String>>,
|
|
chunk_delay: Arc<std::sync::RwLock<Option<Duration>>>,
|
|
storage: Arc<StorageState>,
|
|
completion_gate: Arc<CompletionGate>,
|
|
user_tier: Arc<std::sync::RwLock<Option<String>>>,
|
|
required_token: Option<String>,
|
|
) -> Router {
|
|
let log_cc = log.clone();
|
|
let log_rs = log.clone();
|
|
let log_msg = log.clone();
|
|
let token_cc = required_token.clone();
|
|
let token_msg = required_token.clone();
|
|
let token_rs = required_token;
|
|
let mode_cc = response_mode.clone();
|
|
let mode_rs = response_mode.clone();
|
|
let mode_msg = response_mode;
|
|
let scripted_cc = scripted.clone();
|
|
let scripted_rs = scripted.clone();
|
|
let scripted_settings = scripted.clone();
|
|
let scripted_msg = scripted;
|
|
let delay_cc = chunk_delay.clone();
|
|
let delay_rs = chunk_delay.clone();
|
|
let delay_msg = chunk_delay;
|
|
|
|
Router::new()
|
|
.route(
|
|
"/v1/chat/completions",
|
|
post(move |headers: HeaderMap, Json(body): Json<Value>| {
|
|
let log = log_cc.clone();
|
|
let required = token_cc.clone();
|
|
let mode = mode_cc.clone();
|
|
let scripted = scripted_cc.clone();
|
|
let agent_turns = agent_turns.clone();
|
|
let delay = delay_cc.clone();
|
|
let completion_gate = completion_gate.clone();
|
|
async move {
|
|
let auth = Self::extract_auth(&headers);
|
|
log.record(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
Some(&body),
|
|
auth.as_deref(),
|
|
Self::headers_vec(&headers),
|
|
);
|
|
|
|
if let Some(s) = Self::pop_scripted(&scripted, "/v1/chat/completions") {
|
|
return s.into_response_paced(*delay.read().unwrap());
|
|
}
|
|
|
|
if let Some(rejection) =
|
|
Self::check_auth(auth.as_deref(), required.as_deref())
|
|
{
|
|
return rejection;
|
|
}
|
|
|
|
let user_msg = body
|
|
.get("messages")
|
|
.and_then(|m| m.as_array())
|
|
.and_then(|msgs| {
|
|
msgs.iter()
|
|
.rev()
|
|
.find(|m| m.get("role").and_then(Value::as_str) == Some("user"))
|
|
})
|
|
.and_then(|m| m.get("content"))
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("hello");
|
|
|
|
let model = body
|
|
.get("model")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("test-model");
|
|
|
|
// Only agent turns are gate-eligible: aux requests
|
|
// (title/classifier) must never block session startup.
|
|
let (events, gate) = match Self::pop_agent_turn(&agent_turns, &body) {
|
|
Some(text) => (
|
|
sse::chat_completion_events_exact(&text, model),
|
|
Some(completion_gate.clone()),
|
|
),
|
|
None => {
|
|
let events = match &*mode.read().unwrap() {
|
|
ResponseMode::Echo => sse::chat_completion_events(
|
|
&format!("Echo: {user_msg}"),
|
|
model,
|
|
),
|
|
ResponseMode::Fixed(text) => {
|
|
sse::chat_completion_events_exact(text, model)
|
|
}
|
|
};
|
|
(events, None)
|
|
}
|
|
};
|
|
let stream = paced_events(events, *delay.read().unwrap(), gate);
|
|
Sse::new(stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response()
|
|
}
|
|
}),
|
|
)
|
|
.route(
|
|
"/v1/responses",
|
|
post(move |headers: HeaderMap, Json(body): Json<Value>| {
|
|
let log = log_rs.clone();
|
|
let required = token_rs.clone();
|
|
let mode = mode_rs.clone();
|
|
let scripted = scripted_rs.clone();
|
|
let delay = delay_rs.clone();
|
|
async move {
|
|
let auth = Self::extract_auth(&headers);
|
|
log.record(
|
|
"POST",
|
|
"/v1/responses",
|
|
Some(&body),
|
|
auth.as_deref(),
|
|
Self::headers_vec(&headers),
|
|
);
|
|
|
|
if let Some(s) = Self::pop_scripted(&scripted, "/v1/responses") {
|
|
return s.into_response_paced(*delay.read().unwrap());
|
|
}
|
|
|
|
if let Some(rejection) =
|
|
Self::check_auth(auth.as_deref(), required.as_deref())
|
|
{
|
|
return rejection;
|
|
}
|
|
|
|
let user_msg = body
|
|
.get("input")
|
|
.and_then(|i| i.as_array())
|
|
.and_then(|items| {
|
|
items.iter().rev().find(|item| {
|
|
item.get("role").and_then(Value::as_str) == Some("user")
|
|
})
|
|
})
|
|
.and_then(|item| {
|
|
item.get("content").and_then(|c| {
|
|
c.as_str().map(String::from).or_else(|| {
|
|
c.as_array().and_then(|parts| {
|
|
parts.iter().find_map(|p| {
|
|
if p.get("type").and_then(Value::as_str)
|
|
== Some("input_text")
|
|
{
|
|
p.get("text")
|
|
.and_then(Value::as_str)
|
|
.map(String::from)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
})
|
|
})
|
|
})
|
|
})
|
|
.unwrap_or_else(|| "hello".to_string());
|
|
|
|
let model = body
|
|
.get("model")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("test-model");
|
|
|
|
let events = match &*mode.read().unwrap() {
|
|
ResponseMode::Echo => {
|
|
sse::responses_api_events(&format!("Echo: {user_msg}"), model)
|
|
}
|
|
ResponseMode::Fixed(text) => {
|
|
sse::responses_api_events_exact(text, model)
|
|
}
|
|
};
|
|
let stream = paced_events(events, *delay.read().unwrap(), None);
|
|
Sse::new(stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response()
|
|
}
|
|
}),
|
|
)
|
|
.route(
|
|
"/v1/messages",
|
|
post(move |headers: HeaderMap, Json(body): Json<Value>| {
|
|
let log = log_msg.clone();
|
|
let required = token_msg.clone();
|
|
let mode = mode_msg.clone();
|
|
let scripted = scripted_msg.clone();
|
|
let stop_reason = messages_stop_reason.clone();
|
|
let delay = delay_msg.clone();
|
|
async move {
|
|
let auth = Self::extract_auth(&headers);
|
|
log.record(
|
|
"POST",
|
|
"/v1/messages",
|
|
Some(&body),
|
|
auth.as_deref(),
|
|
Self::headers_vec(&headers),
|
|
);
|
|
|
|
if let Some(s) = Self::pop_scripted(&scripted, "/v1/messages") {
|
|
return s.into_response_paced(*delay.read().unwrap());
|
|
}
|
|
|
|
if let Some(rejection) =
|
|
Self::check_auth(auth.as_deref(), required.as_deref())
|
|
{
|
|
return rejection;
|
|
}
|
|
|
|
// Anthropic content is either a plain string or an
|
|
// array of typed blocks; extract the last user text.
|
|
let user_msg = body
|
|
.get("messages")
|
|
.and_then(|m| m.as_array())
|
|
.and_then(|msgs| {
|
|
msgs.iter()
|
|
.rev()
|
|
.find(|m| m.get("role").and_then(Value::as_str) == Some("user"))
|
|
})
|
|
.and_then(|m| m.get("content"))
|
|
.and_then(|c| {
|
|
c.as_str().map(String::from).or_else(|| {
|
|
c.as_array().and_then(|blocks| {
|
|
blocks.iter().find_map(|b| {
|
|
if b.get("type").and_then(Value::as_str) == Some("text")
|
|
{
|
|
b.get("text")
|
|
.and_then(Value::as_str)
|
|
.map(String::from)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
})
|
|
})
|
|
})
|
|
.unwrap_or_else(|| "hello".to_string());
|
|
|
|
let model = body
|
|
.get("model")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("test-model");
|
|
|
|
let stop = stop_reason.read().unwrap().clone();
|
|
// Messages streams its text as a single delta, so the
|
|
// fixed text is byte-exact by construction.
|
|
let events = match &*mode.read().unwrap() {
|
|
ResponseMode::Echo => {
|
|
sse::messages_api_events(&format!("Echo: {user_msg}"), model, &stop)
|
|
}
|
|
ResponseMode::Fixed(text) => {
|
|
sse::messages_api_events(text, model, &stop)
|
|
}
|
|
};
|
|
let stream = paced_events(events, *delay.read().unwrap(), None);
|
|
Sse::new(stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response()
|
|
}
|
|
}),
|
|
)
|
|
.route(
|
|
"/v1/models",
|
|
get({
|
|
let log = log.clone();
|
|
move || {
|
|
let log = log.clone();
|
|
let models = models.clone();
|
|
async move {
|
|
log.record("GET", "/v1/models", None, None, Vec::new());
|
|
let models_json = models.read().unwrap().clone();
|
|
Json(json!({
|
|
"object": "list",
|
|
"data": models_json,
|
|
}))
|
|
}
|
|
}
|
|
}),
|
|
)
|
|
.route(
|
|
"/v1/settings",
|
|
get({
|
|
let log = log.clone();
|
|
move || {
|
|
let log = log.clone();
|
|
let settings = settings.clone();
|
|
let scripted = scripted_settings.clone();
|
|
async move {
|
|
log.record("GET", "/v1/settings", None, None, Vec::new());
|
|
// Scripted one-shots take precedence (FIFO), so a
|
|
// test can serve a transient payload (e.g. one
|
|
// stale gated snapshot) and fall back to the
|
|
// steady-state `set_settings` value afterwards.
|
|
if let Some(s) = Self::pop_scripted(&scripted, "/v1/settings") {
|
|
return s.into_response_paced(None);
|
|
}
|
|
let maybe = settings.read().unwrap().clone();
|
|
match maybe {
|
|
Some(s) => Json(s).into_response(),
|
|
None => StatusCode::NOT_FOUND.into_response(),
|
|
}
|
|
}
|
|
}
|
|
}),
|
|
)
|
|
.route(
|
|
"/v1/user",
|
|
get(
|
|
move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
|
|
let log = log.clone();
|
|
let user_tier = user_tier.clone();
|
|
async move {
|
|
// Keep the query string in the log so tests can
|
|
// count `?include=subscription` checks separately
|
|
// from plain enrichment fetches.
|
|
let path = match query {
|
|
Some(q) if !q.is_empty() => format!("/v1/user?{q}"),
|
|
_ => "/v1/user".to_owned(),
|
|
};
|
|
log.record("GET", &path, None, None, Vec::new());
|
|
let tier = user_tier.read().unwrap().clone();
|
|
let mut body = json!({
|
|
"userId": "mock-user",
|
|
"email": "mock-user@test.invalid",
|
|
});
|
|
if let Some(t) = tier {
|
|
body["subscriptionTier"] = json!(t);
|
|
}
|
|
Json(body).into_response()
|
|
}
|
|
},
|
|
),
|
|
)
|
|
.route(
|
|
"/v1/storage",
|
|
post({
|
|
let storage = storage.clone();
|
|
move |headers: HeaderMap, body: axum::body::Bytes| {
|
|
let storage = storage.clone();
|
|
async move { Self::storage_upload_handler(&storage, &headers, &body) }
|
|
}
|
|
}),
|
|
)
|
|
// The shell probes these before/alongside per-file uploads. Answer
|
|
// 404 ("old proxy") so it falls back to plain `POST /v1/storage`,
|
|
// which is the path the park-on-401 e2e exercises.
|
|
.route(
|
|
"/v1/storage/exists",
|
|
get(|| async { StatusCode::NOT_FOUND }),
|
|
)
|
|
.route(
|
|
"/v1/storage/batch_exists",
|
|
post(|| async { StatusCode::NOT_FOUND }),
|
|
)
|
|
.route(
|
|
"/v1/storage/batch_upload_json",
|
|
post(|| async { StatusCode::NOT_FOUND }),
|
|
)
|
|
.route(
|
|
"/v1/storage/batch_upload",
|
|
post(|| async { StatusCode::NOT_FOUND }),
|
|
)
|
|
.route(
|
|
"/v1/storage/limits",
|
|
get(|| async { StatusCode::NOT_FOUND }),
|
|
)
|
|
// Body limit: repo-context archives can exceed axum's 2 MB default.
|
|
.layer(axum::extract::DefaultBodyLimit::max(256 * 1024 * 1024))
|
|
}
|
|
}
|
|
|
|
impl Drop for MockInferenceServer {
|
|
fn drop(&mut self) {
|
|
if let Some(tx) = self.shutdown_tx.take() {
|
|
let _ = tx.send(());
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
const MERMAID_TEXT: &str =
|
|
"Here is a flow:\n\n```mermaid\nflowchart TD\n A --> B\n```\n\nDone.\n";
|
|
|
|
fn sse_data_payloads(body: &str) -> Vec<String> {
|
|
body.lines()
|
|
.filter_map(|l| l.strip_prefix("data:"))
|
|
.map(|d| d.trim_start().to_owned())
|
|
.filter(|d| d != "[DONE]")
|
|
.collect()
|
|
}
|
|
|
|
fn chat_stream_text(body: &str) -> String {
|
|
sse_data_payloads(body)
|
|
.iter()
|
|
.filter_map(|d| serde_json::from_str::<Value>(d).ok())
|
|
.filter_map(|v| {
|
|
v.get("choices")
|
|
.and_then(|c| c.get(0))
|
|
.and_then(|c| c.get("delta"))
|
|
.and_then(|d| d.get("content"))
|
|
.and_then(Value::as_str)
|
|
.map(String::from)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn responses_stream_text(body: &str) -> String {
|
|
sse_data_payloads(body)
|
|
.iter()
|
|
.filter_map(|d| serde_json::from_str::<Value>(d).ok())
|
|
.filter(|v| v.get("type").and_then(Value::as_str) == Some("response.output_text.delta"))
|
|
.filter_map(|v| v.get("delta").and_then(Value::as_str).map(String::from))
|
|
.collect()
|
|
}
|
|
|
|
fn messages_stream_text(body: &str) -> String {
|
|
sse_data_payloads(body)
|
|
.iter()
|
|
.filter_map(|d| serde_json::from_str::<Value>(d).ok())
|
|
.filter(|v| v.get("type").and_then(Value::as_str) == Some("content_block_delta"))
|
|
.filter_map(|v| {
|
|
v.get("delta")
|
|
.and_then(|d| d.get("text"))
|
|
.and_then(Value::as_str)
|
|
.map(String::from)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
async fn post_chat(server: &MockInferenceServer, content: &str) -> reqwest::Response {
|
|
reqwest::Client::new()
|
|
.post(format!("{}/chat/completions", server.url()))
|
|
.json(&json!({
|
|
"model": "test-model",
|
|
"messages": [{ "role": "user", "content": content }]
|
|
}))
|
|
.send()
|
|
.await
|
|
.expect("POST /v1/chat/completions")
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn echo_mode_echoes_last_user_message() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
|
|
let body = post_chat(&server, "ping pong").await.text().await.unwrap();
|
|
assert_eq!(chat_stream_text(&body), "Echo: ping pong");
|
|
|
|
// Echo mode collapses whitespace by design; only fixed mode is exact.
|
|
let body = post_chat(&server, "a b\nc").await.text().await.unwrap();
|
|
assert_eq!(chat_stream_text(&body), "Echo: a b c");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn fixed_mode_reconstructs_byte_exact_over_http() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response(MERMAID_TEXT);
|
|
|
|
let body = post_chat(&server, "ignored").await.text().await.unwrap();
|
|
assert_eq!(chat_stream_text(&body), MERMAID_TEXT);
|
|
|
|
let body = reqwest::Client::new()
|
|
.post(format!("{}/responses", server.url()))
|
|
.json(&json!({
|
|
"model": "test-model",
|
|
"input": [{ "role": "user", "content": "ignored" }]
|
|
}))
|
|
.send()
|
|
.await
|
|
.expect("POST /v1/responses")
|
|
.text()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(responses_stream_text(&body), MERMAID_TEXT);
|
|
|
|
let body = reqwest::Client::new()
|
|
.post(format!("{}/messages", server.url()))
|
|
.json(&json!({
|
|
"model": "test-model",
|
|
"messages": [{ "role": "user", "content": "ignored" }]
|
|
}))
|
|
.send()
|
|
.await
|
|
.expect("POST /v1/messages")
|
|
.text()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(messages_stream_text(&body), MERMAID_TEXT);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn settings_404_until_set_then_200() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
let url = format!("{}/settings", server.url());
|
|
|
|
let resp = reqwest::get(&url).await.unwrap();
|
|
assert_eq!(resp.status(), 404);
|
|
|
|
server.set_settings(json!({ "tips": ["t1"] }));
|
|
let resp = reqwest::get(&url).await.unwrap();
|
|
assert_eq!(resp.status(), 200);
|
|
let body: Value = resp.json().await.unwrap();
|
|
assert_eq!(body, json!({ "tips": ["t1"] }));
|
|
|
|
server.preset_allow_access();
|
|
let resp = reqwest::get(&url).await.unwrap();
|
|
assert_eq!(resp.status(), 200);
|
|
let body: Value = resp.json().await.unwrap();
|
|
assert_eq!(body, json!({ "allow_access": true }));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn request_bodies_returns_bodies_in_arrival_order() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
|
|
post_chat(&server, "first").await.text().await.unwrap();
|
|
// Body-less request in between must be skipped, not break ordering.
|
|
reqwest::get(format!("{}/models", server.url()))
|
|
.await
|
|
.unwrap();
|
|
post_chat(&server, "second").await.text().await.unwrap();
|
|
|
|
let bodies = server.request_bodies();
|
|
assert_eq!(bodies.len(), 2);
|
|
assert_eq!(
|
|
bodies[0]["messages"][0]["content"],
|
|
json!("first"),
|
|
"bodies must be in arrival order"
|
|
);
|
|
assert_eq!(bodies[1]["messages"][0]["content"], json!("second"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn scripted_responses_serve_fifo_per_path_then_fall_back() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::text(401, "Unauthorized"),
|
|
);
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::json(500, json!({ "error": { "message": "boom" } })),
|
|
);
|
|
|
|
// FIFO: first the 401 text, then the 500 json.
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 401);
|
|
assert_eq!(resp.text().await.unwrap(), "Unauthorized");
|
|
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 500);
|
|
let body: Value = resp.json().await.unwrap();
|
|
assert_eq!(body, json!({ "error": { "message": "boom" } }));
|
|
|
|
// Queue drained: falls back to the active mode (echo).
|
|
let body = post_chat(&server, "ping pong").await.text().await.unwrap();
|
|
assert_eq!(chat_stream_text(&body), "Echo: ping pong");
|
|
|
|
// Queues are per path: an unrelated endpoint is unaffected.
|
|
server.enqueue_response("/v1/chat/completions", ScriptedResponse::text(503, "later"));
|
|
let resp = reqwest::Client::new()
|
|
.post(format!("{}/responses", server.url()))
|
|
.json(&json!({
|
|
"model": "test-model",
|
|
"input": [{ "role": "user", "content": "hi there" }]
|
|
}))
|
|
.send()
|
|
.await
|
|
.expect("POST /v1/responses");
|
|
assert_eq!(resp.status(), 200);
|
|
assert_eq!(
|
|
responses_stream_text(&resp.text().await.unwrap()),
|
|
"Echo: hi there "
|
|
);
|
|
}
|
|
|
|
/// Pins the documented precedence: a script bypasses the required-auth
|
|
/// gate; once the queue empties, the gate is back.
|
|
#[tokio::test]
|
|
async fn scripted_response_takes_precedence_over_required_auth() {
|
|
let server = MockInferenceServer::start_with_required_auth(
|
|
vec![MockModelEntry::new("test-model")],
|
|
"secret-token",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::text(200, "scripted"),
|
|
);
|
|
|
|
// No token: the script still serves.
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 200);
|
|
assert_eq!(resp.text().await.unwrap(), "scripted");
|
|
|
|
// Queue drained: the auth gate applies again.
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 401);
|
|
}
|
|
|
|
/// Scripted headers must reach the client: 429 + `Retry-After` injection
|
|
/// depends on them.
|
|
#[tokio::test]
|
|
async fn scripted_response_headers_reach_the_client() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
let mut rate_limited = ScriptedResponse::text(429, "slow down");
|
|
rate_limited
|
|
.headers
|
|
.push(("retry-after".to_string(), "7".to_string()));
|
|
server.enqueue_response("/v1/chat/completions", rate_limited);
|
|
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 429);
|
|
assert_eq!(
|
|
resp.headers()
|
|
.get("retry-after")
|
|
.and_then(|v| v.to_str().ok()),
|
|
Some("7")
|
|
);
|
|
assert_eq!(resp.text().await.unwrap(), "slow down");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn scripted_raw_body_served_byte_exact() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
let raw = "data: {\"choices\":[]}\n\ndata: not-json-at-all\n\ndata: [DONE]\n\n";
|
|
server.enqueue_response("/v1/chat/completions", ScriptedResponse::text(200, raw));
|
|
|
|
let resp = post_chat(&server, "hi").await;
|
|
assert_eq!(resp.status(), 200);
|
|
assert_eq!(resp.text().await.unwrap(), raw);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn scripted_sse_preserves_event_names_and_order() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::sse(vec![
|
|
SseEvent::with_event("custom.kind", "{\"a\":1}"),
|
|
SseEvent::data("{\"b\":2}"),
|
|
]),
|
|
);
|
|
|
|
let body = post_chat(&server, "hi").await.text().await.unwrap();
|
|
let named_then_plain = body
|
|
.find("event: custom.kind")
|
|
.zip(body.find("data: {\"b\":2}"))
|
|
.is_some_and(|(named, plain)| named < plain);
|
|
assert!(
|
|
body.contains("event: custom.kind") && body.contains("data: {\"a\":1}"),
|
|
"named event must carry both fields, got:\n{body}"
|
|
);
|
|
assert!(named_then_plain, "events must be served in order:\n{body}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn request_log_captures_arbitrary_headers() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
|
|
reqwest::Client::new()
|
|
.post(format!("{}/chat/completions", server.url()))
|
|
.header("authorization", "Bearer log-me")
|
|
.header("x-test-marker", "zap")
|
|
.json(&json!({
|
|
"model": "test-model",
|
|
"messages": [{ "role": "user", "content": "hi" }]
|
|
}))
|
|
.send()
|
|
.await
|
|
.expect("POST /v1/chat/completions");
|
|
|
|
let entry = server.requests().pop().expect("one logged request");
|
|
assert_eq!(entry.header("x-test-marker"), Some("zap"));
|
|
assert_eq!(entry.header("X-Test-Marker"), Some("zap"));
|
|
assert_eq!(entry.header("authorization"), Some("Bearer log-me"));
|
|
assert_eq!(entry.authorization.as_deref(), Some("Bearer log-me"));
|
|
assert_eq!(entry.header("x-absent"), None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn required_auth_enforced_in_both_response_modes() {
|
|
let server = MockInferenceServer::start_with_required_auth(
|
|
vec![MockModelEntry::new("test-model")],
|
|
"secret-token",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let client = reqwest::Client::new();
|
|
let url = format!("{}/chat/completions", server.url());
|
|
let req_body = json!({
|
|
"model": "test-model",
|
|
"messages": [{ "role": "user", "content": "hi there" }]
|
|
});
|
|
|
|
// Echo mode: missing auth rejected, valid auth streams the echo.
|
|
let resp = client.post(&url).json(&req_body).send().await.unwrap();
|
|
assert_eq!(resp.status(), 401);
|
|
let resp = client
|
|
.post(&url)
|
|
.header("authorization", "Bearer secret-token")
|
|
.json(&req_body)
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(resp.status(), 200);
|
|
assert_eq!(
|
|
chat_stream_text(&resp.text().await.unwrap()),
|
|
"Echo: hi there"
|
|
);
|
|
|
|
// Fixed mode: same auth gate, fixed text streamed byte-exact.
|
|
server.set_response(MERMAID_TEXT);
|
|
let resp = client.post(&url).json(&req_body).send().await.unwrap();
|
|
assert_eq!(resp.status(), 401);
|
|
let resp = client
|
|
.post(&url)
|
|
.header("authorization", "Bearer secret-token")
|
|
.json(&req_body)
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(resp.status(), 200);
|
|
assert_eq!(chat_stream_text(&resp.text().await.unwrap()), MERMAID_TEXT);
|
|
}
|
|
}
|