1pub mod auth;
57pub(crate) mod constrained;
58pub(crate) mod darklane;
62pub(crate) mod health;
66pub(crate) mod lanes {
67 pub use memra_lanes::*;
68}
69mod anthropic;
75mod dsv4_serve;
76mod embed_api;
77pub mod metering;
84mod responses_api;
85mod surfaces;
86mod toolcall;
87mod ttft;
88mod worker;
89
90use std::collections::HashMap;
91use std::net::{SocketAddr, ToSocketAddrs};
92use std::sync::Arc;
93use std::sync::mpsc::Sender;
94
95use axum::{
96 Extension, Json, Router,
97 body::Body,
98 extract::{DefaultBodyLimit, Query, Request as AxumRequest, State},
99 http::{
100 HeaderMap, StatusCode,
101 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
102 },
103 middleware::{self, Next},
104 response::{
105 IntoResponse, Response,
106 sse::{Event as SseEvent, Sse},
107 },
108 routing::{get, post},
109};
110use futures_core::Stream as _;
111use serde::{Deserialize, Serialize};
112use serde_json::json;
113
114use memra_engine::decode::GenParams;
115use memra_engine::sampler::SamplerConfig;
116use memra_tokenizer::{
117 Tokenizer,
118 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
119};
120use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
121use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
122
123const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
140const MAX_BODY_ADMISSIONS: usize = 4;
141const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
142const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
146const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
147const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
148const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
149const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
150const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
151
152fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
153 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
154 SEMAPHORE
155 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
156 .clone()
157}
158
159fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
160 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
161 SEMAPHORE
162 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
163 .clone()
164}
165
166fn declared_body_length(req: &AxumRequest) -> Option<usize> {
167 req.headers()
168 .get(CONTENT_LENGTH)
169 .and_then(|value| value.to_str().ok())
170 .and_then(|value| value.parse().ok())
171}
172
173fn body_requires_admission(req: &AxumRequest) -> bool {
174 if req.headers().contains_key(TRANSFER_ENCODING) {
178 return true;
179 }
180 declared_body_length(req).map_or(true, |length| length > BODY_ADMISSION_BYPASS_BYTES)
181}
182
183fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
187 let Some(length) = declared_body_length(req) else {
188 return BODY_READ_TIMEOUT;
189 };
190 let bytes = length as u64;
191 let extra_seconds =
192 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
193 let seconds = BODY_READ_TIMEOUT
194 .as_secs()
195 .saturating_add(extra_seconds)
196 .min(BODY_READ_TIMEOUT_MAX.as_secs());
197 std::time::Duration::from_secs(seconds)
198}
199
200async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
205 let resp = next.run(req).await;
206 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
207 return resp;
208 }
209 error_response_coded(
210 StatusCode::PAYLOAD_TOO_LARGE,
211 &format!(
212 "request body exceeds the {} MiB limit",
213 MAX_BODY_BYTES / (1024 * 1024)
214 ),
215 "invalid_request_error",
216 None,
217 Some("request_too_large"),
218 )
219}
220
221fn apply_body_limit(app: Router) -> Router {
224 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
225 .layer(middleware::from_fn(shape_payload_too_large))
226}
227
228fn protected_inference_path(path: &str) -> bool {
229 matches!(
230 path,
231 "/v1/auth/check"
232 | "/v1/completions"
233 | "/v1/chat/completions"
234 | "/v1/messages"
235 | "/v1/responses"
236 | "/v1/embeddings"
237 | "/v1/rerank"
238 )
239}
240
241async fn shape_inference_early_response(path: &str, response: Response) -> Response {
245 let request_id = Envelope::new(path != "/v1/completions");
246 if path == "/v1/messages" {
247 anthropic::with_anthropic_request_id(
248 &request_id.id,
249 anthropic::reshape_error(response, &request_id.id).await,
250 )
251 } else {
252 with_request_id(&request_id.id, response)
253 }
254}
255
256async fn authenticate_inference_before_body(
261 State(st): State<AppState>,
262 mut req: AxumRequest,
263 next: Next,
264) -> Response {
265 if !protected_inference_path(req.uri().path()) {
266 return next.run(req).await;
267 }
268 let path = req.uri().path().to_string();
269 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
273 return shape_inference_early_response(
274 &path,
275 error_response_coded(
276 StatusCode::PAYLOAD_TOO_LARGE,
277 &format!(
278 "request body exceeds the {} MiB limit",
279 MAX_BODY_BYTES / (1024 * 1024)
280 ),
281 "invalid_request_error",
282 None,
283 Some("request_too_large"),
284 ),
285 )
286 .await;
287 }
288 let headers = req.headers();
289 let bearer = bearer_token(headers);
290 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
291 let api_key = headers
292 .get("x-api-key")
293 .and_then(|value| value.to_str().ok());
294 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
295 } else {
296 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
297 };
298 if let Err(why) = auth {
299 return shape_inference_early_response(&path, authentication_error(why)).await;
300 }
301 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
309 let body_admission = if body_requires_admission(&req) {
310 body_admission_semaphore()
311 } else {
312 small_body_admission_semaphore()
313 };
314 let body_permit = match body_admission.try_acquire_owned() {
315 Ok(permit) => Some(permit),
316 Err(tokio::sync::TryAcquireError::Closed) => {
317 let response = retry_contract_response(
318 error_response_coded(
319 StatusCode::SERVICE_UNAVAILABLE,
320 "request body admission is unavailable",
321 "server_error",
322 None,
323 Some("body_admission_unavailable"),
324 ),
325 Some(BODY_ADMISSION_RETRY_AFTER_S),
326 );
327 return shape_inference_early_response(&path, response).await;
328 }
329 Err(tokio::sync::TryAcquireError::NoPermits) => {
330 let response = retry_contract_response(
331 error_response_coded(
332 StatusCode::TOO_MANY_REQUESTS,
333 "request body admission is busy",
334 "rate_limit_error",
335 None,
336 Some("body_admission_busy"),
337 ),
338 Some(BODY_ADMISSION_RETRY_AFTER_S),
339 );
340 return shape_inference_early_response(&path, response).await;
341 }
342 };
343 let body = std::mem::replace(req.body_mut(), Body::empty());
347 let mut body = Box::pin(body.into_data_stream());
348 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
349 let body_timed_out_flag = body_timed_out.clone();
350 let guarded_body = async_stream::stream! {
351 loop {
352 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
353 if remaining.is_zero() {
354 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
355 yield Err(std::io::Error::new(
356 std::io::ErrorKind::TimedOut,
357 "request body read deadline exceeded",
358 ));
359 break;
360 }
361 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
362 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
363 Ok(frame) => frame,
364 Err(_) => {
365 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
366 yield Err(std::io::Error::new(
367 std::io::ErrorKind::TimedOut,
368 "request body idle timeout exceeded",
369 ));
370 break;
371 }
372 };
373 match frame {
374 Some(Ok(bytes)) => yield Ok(bytes),
375 Some(Err(error)) => {
376 yield Err(std::io::Error::other(error.to_string()));
377 break;
378 }
379 None => break,
380 }
381 }
382 drop(body_permit);
383 };
384 *req.body_mut() = Body::from_stream(guarded_body);
385 let response = next.run(req).await;
386 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
387 let request_id = Envelope::new(path != "/v1/completions");
388 let timeout = error_response_coded(
389 StatusCode::REQUEST_TIMEOUT,
390 "request body read timed out",
391 "invalid_request_error",
392 None,
393 Some("request_body_timeout"),
394 );
395 return if path == "/v1/messages" {
396 anthropic::with_anthropic_request_id(
397 &request_id.id,
398 anthropic::reshape_error(timeout, &request_id.id).await,
399 )
400 } else {
401 with_request_id(&request_id.id, timeout)
402 };
403 }
404 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
405 let request_id = Envelope::new(true);
406 return anthropic::with_anthropic_request_id(
407 &request_id.id,
408 anthropic::reshape_error(response, &request_id.id).await,
409 );
410 }
411 response
412}
413
414#[cfg(test)]
415mod body_limit_tests {
416 use super::*;
417 use tower::ServiceExt as _;
418
419 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
420
421 fn test_app() -> Router {
425 let app = Router::new()
426 .route(
427 "/bytes",
428 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
429 )
430 .route(
431 "/json",
432 post(|Json(v): Json<serde_json::Value>| async move {
433 v["pad"].as_str().unwrap_or("").len().to_string()
434 }),
435 );
436 apply_body_limit(app)
437 }
438
439 fn streamed_body(chunks: usize) -> Body {
440 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
443 Body::from_stream(async_stream::stream! {
444 for _ in 0..chunks {
445 yield Ok::<_, std::io::Error>(chunk.clone());
446 }
447 })
448 }
449
450 #[tokio::test]
451 async fn bodies_past_the_old_2mib_default_are_accepted() {
452 for (path, body) in [
455 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
456 (
457 "/json",
458 Body::from(
459 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
460 ),
461 ),
462 ] {
463 let resp = test_app()
464 .oneshot(
465 axum::http::Request::post(path)
466 .header(CONTENT_TYPE, "application/json")
467 .body(body)
468 .unwrap(),
469 )
470 .await
471 .unwrap();
472 assert_eq!(resp.status(), StatusCode::OK, "{path}");
473 }
474 }
475
476 #[tokio::test]
477 async fn body_at_exactly_the_limit_is_accepted() {
478 let resp = test_app()
479 .oneshot(
480 axum::http::Request::post("/bytes")
481 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
482 .unwrap(),
483 )
484 .await
485 .unwrap();
486 assert_eq!(resp.status(), StatusCode::OK);
487 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
488 .await
489 .unwrap();
490 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
491 }
492
493 #[tokio::test]
494 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
495 for path in ["/bytes", "/json"] {
499 let resp = test_app()
500 .oneshot(
501 axum::http::Request::post(path)
502 .header(CONTENT_TYPE, "application/json")
503 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
504 .unwrap(),
505 )
506 .await
507 .unwrap();
508 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
509 assert_eq!(
510 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
511 Some(b"false".as_ref()),
512 "{path}: retrying identical bytes cannot fix a 413"
513 );
514 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
515 .await
516 .unwrap();
517 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
518 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
519 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
520 assert!(
521 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
522 "{path}: message names the limit"
523 );
524 }
525 }
526
527 #[tokio::test]
528 async fn authenticated_body_admission_is_finite() {
529 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
530 let semaphore = body_admission_semaphore();
531 let mut permits = Vec::new();
532 for _ in 0..MAX_BODY_ADMISSIONS {
533 permits.push(semaphore.clone().acquire_owned().await.unwrap());
534 }
535 assert!(
536 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
537 .await
538 .is_err(),
539 "body parser admission must not be unbounded"
540 );
541 drop(permits);
542 assert!(semaphore.acquire().await.is_ok());
543 }
544
545 #[tokio::test]
546 async fn small_body_admission_is_finite_and_separate() {
547 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
548 let large = body_admission_semaphore();
549 let small = small_body_admission_semaphore();
550 let mut small_permits = Vec::new();
551 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
552 small_permits.push(small.clone().acquire_owned().await.unwrap());
553 }
554 assert!(
555 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
556 .await
557 .is_err(),
558 "small body parser admission must be bounded"
559 );
560 assert!(
561 large.clone().try_acquire().is_ok(),
562 "small uploads must not consume large-upload permits"
563 );
564 drop(small_permits);
565 assert!(small.acquire().await.is_ok());
566 }
567
568 #[test]
569 fn small_declared_bodies_bypass_large_upload_admission() {
570 let request = axum::http::Request::post("/v1/chat/completions")
571 .header(CONTENT_LENGTH, "2048")
572 .body(Body::empty())
573 .unwrap();
574 assert!(!body_requires_admission(&request));
575
576 let request = axum::http::Request::post("/v1/chat/completions")
577 .header(
578 CONTENT_LENGTH,
579 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
580 )
581 .body(Body::empty())
582 .unwrap();
583 assert!(body_requires_admission(&request));
584
585 let request = axum::http::Request::post("/v1/chat/completions")
586 .header(CONTENT_LENGTH, "2048")
587 .header(TRANSFER_ENCODING, "chunked")
588 .body(Body::empty())
589 .unwrap();
590 assert!(body_requires_admission(&request));
591 }
592
593 #[test]
594 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
595 let unknown = axum::http::Request::post("/v1/chat/completions")
596 .body(Body::empty())
597 .unwrap();
598 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
599
600 let large = axum::http::Request::post("/v1/chat/completions")
601 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
602 .body(Body::empty())
603 .unwrap();
604 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
605 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
606
607 let absurd = axum::http::Request::post("/v1/chat/completions")
608 .header(CONTENT_LENGTH, u64::MAX.to_string())
609 .body(Body::empty())
610 .unwrap();
611 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
612 }
613
614 #[tokio::test]
615 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
616 let too_large = shape_inference_early_response(
617 "/v1/messages",
618 error_response_coded(
619 StatusCode::PAYLOAD_TOO_LARGE,
620 "request body exceeds the 192 MiB limit",
621 "invalid_request_error",
622 None,
623 Some("request_too_large"),
624 ),
625 )
626 .await;
627 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
628 let house_id = too_large.headers()["x-request-id"].clone();
629 assert_eq!(too_large.headers()["request-id"], house_id);
630 assert_eq!(too_large.headers()["x-should-retry"], "false");
631 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
632 .await
633 .unwrap();
634 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
635 assert_eq!(payload["type"], "error");
636 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
637
638 let busy = shape_inference_early_response(
639 "/v1/chat/completions",
640 retry_contract_response(
641 error_response_coded(
642 StatusCode::TOO_MANY_REQUESTS,
643 "request body admission is busy",
644 "rate_limit_error",
645 None,
646 Some("body_admission_busy"),
647 ),
648 Some(BODY_ADMISSION_RETRY_AFTER_S),
649 ),
650 )
651 .await;
652 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
653 assert!(!busy.headers()["x-request-id"].is_empty());
654 assert_eq!(busy.headers()["retry-after"], "1");
655 assert_eq!(busy.headers()["retry-after-ms"], "1000");
656 assert!(busy.headers().get("x-should-retry").is_none());
657 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
658 .await
659 .unwrap();
660 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
661 assert_eq!(payload["error"]["code"], "body_admission_busy");
662 }
663}
664
665#[derive(Clone, Default)]
666struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
667
668fn is_sse_data_frame(bytes: &[u8]) -> bool {
669 bytes
670 .windows(b"data:".len())
671 .any(|window| window == b"data:")
672}
673
674async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
675 let trace = ttft::start(req.uri().path());
676 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
677 let response = next.run(req).await;
678 let Some(trace) = trace else {
679 return response;
680 };
681 let is_sse = response
682 .headers()
683 .get(CONTENT_TYPE)
684 .and_then(|value| value.to_str().ok())
685 .is_some_and(|value| value.starts_with("text/event-stream"));
686 if !is_sse {
687 return response;
688 }
689
690 let (parts, body) = response.into_parts();
693 let mut body = Box::pin(body.into_data_stream());
694 let stream = async_stream::stream! {
695 while let Some(frame) =
696 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
697 {
698 if frame
699 .as_ref()
700 .is_ok_and(|bytes| is_sse_data_frame(bytes))
701 {
702 trace.mark_first_sse_byte();
703 }
704 yield frame;
705 }
706 };
707 Response::from_parts(parts, Body::from_stream(stream))
708}
709
710const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
711const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
712
713#[derive(Debug, Clone, Default, Deserialize)]
714#[serde(deny_unknown_fields)]
715struct OpenRouterMetadataFile {
716 #[serde(default)]
717 models: HashMap<String, OpenRouterModelMetadata>,
718 #[serde(default)]
721 planned_models: HashMap<String, OpenRouterModelMetadata>,
722 #[serde(default)]
725 provider: Option<ProviderMetadata>,
726}
727
728#[derive(Debug, Clone, Deserialize)]
732#[serde(deny_unknown_fields)]
733struct ProviderMetadata {
734 id: String,
735 #[serde(default)]
736 status_url: Option<String>,
737 #[serde(default)]
738 support_contact: Option<String>,
739 #[serde(default)]
740 incident_contact: Option<String>,
741 #[serde(default)]
742 regions: Vec<String>,
743}
744
745#[derive(Debug, Clone, Default, Deserialize)]
747#[serde(deny_unknown_fields)]
748struct LifecycleMetadata {
749 #[serde(default)]
750 status: Option<String>,
751 #[serde(default)]
752 deprecation_at: Option<String>,
753 #[serde(default)]
754 retirement_at: Option<String>,
755 #[serde(default)]
756 replacement_model_id: Option<String>,
757}
758
759#[derive(Debug, Clone, Default, Deserialize)]
761#[serde(deny_unknown_fields)]
762struct ReliabilityMetadata {
763 #[serde(default)]
764 first_token_timeout_seconds: Option<u64>,
765 #[serde(default)]
766 completion_timeout_seconds: Option<u64>,
767 #[serde(default)]
768 stream_idle_timeout_seconds: Option<u64>,
769 #[serde(default)]
770 capacity_scope: Option<String>,
771}
772
773#[derive(Debug, Clone, Default, Deserialize)]
774#[serde(deny_unknown_fields)]
775struct OpenRouterModelMetadata {
776 #[serde(default)]
778 owned_by: Option<String>,
779 #[serde(default)]
780 lifecycle: Option<LifecycleMetadata>,
781 #[serde(default)]
782 reliability: Option<ReliabilityMetadata>,
783 #[serde(default)]
784 hugging_face_id: Option<String>,
785 #[serde(default)]
786 created: Option<u64>,
787 #[serde(default)]
788 quantization: Option<String>,
789 #[serde(default)]
790 description: Option<String>,
791 #[serde(default)]
792 max_prompt_length: Option<u64>,
793 #[serde(default)]
794 max_output_length: Option<u64>,
795 #[serde(default)]
798 default_output_length: Option<u64>,
799 #[serde(default)]
800 pricing: OpenRouterPricing,
801 #[serde(default)]
802 capacity: OpenRouterCapacity,
803 #[serde(default)]
804 is_ready: Option<bool>,
805 #[serde(default)]
806 is_free: Option<bool>,
807 #[serde(default)]
808 discount_to_user: Option<f64>,
809 #[serde(default)]
810 openrouter_slug: Option<String>,
811 #[serde(default)]
812 datacenters: Vec<OpenRouterDatacenter>,
813 #[serde(default)]
817 input_modalities: Vec<String>,
818 #[serde(default)]
833 surface: Option<String>,
834 #[serde(default)]
835 zdr: Option<bool>,
836 #[serde(default)]
837 hipaa: Option<bool>,
838 #[serde(default)]
848 default_reasoning_effort: Option<String>,
849 #[serde(default)]
867 default_temperature: Option<f32>,
868 #[serde(default)]
869 default_top_p: Option<f32>,
870 #[serde(default)]
872 default_top_k: Option<usize>,
873 #[serde(default)]
874 default_min_p: Option<f32>,
875 #[serde(default)]
876 default_presence_penalty: Option<f32>,
877 #[serde(default)]
878 default_frequency_penalty: Option<f32>,
879 #[serde(default)]
881 default_repetition_penalty: Option<f32>,
882 #[serde(default)]
903 non_thinking_sampling: Option<SamplingArmMetadata>,
904}
905
906#[derive(Debug, Clone, Default, Deserialize)]
912#[serde(deny_unknown_fields)]
913struct SamplingArmMetadata {
914 #[serde(default)]
915 temperature: Option<f32>,
916 #[serde(default)]
917 top_p: Option<f32>,
918 #[serde(default)]
919 top_k: Option<usize>,
920 #[serde(default)]
921 min_p: Option<f32>,
922 #[serde(default)]
923 presence_penalty: Option<f32>,
924 #[serde(default)]
925 frequency_penalty: Option<f32>,
926 #[serde(default)]
927 repetition_penalty: Option<f32>,
928}
929
930impl SamplingArmMetadata {
931 fn is_empty(&self) -> bool {
932 self.temperature.is_none()
933 && self.top_p.is_none()
934 && self.top_k.is_none()
935 && self.min_p.is_none()
936 && self.presence_penalty.is_none()
937 && self.frequency_penalty.is_none()
938 && self.repetition_penalty.is_none()
939 }
940}
941
942#[derive(Debug, Clone, Default, Deserialize)]
943#[serde(deny_unknown_fields)]
944struct OpenRouterPricing {
945 #[serde(default)]
946 prompt: Option<String>,
947 #[serde(default)]
948 cached_prompt: Option<String>,
949 #[serde(default)]
950 cache_write: Option<String>,
951 #[serde(default)]
952 completion: Option<String>,
953 #[serde(default)]
954 internal_reasoning: Option<String>,
955 #[serde(default)]
956 request: Option<String>,
957}
958
959#[derive(Debug, Clone, Default, Deserialize)]
960#[serde(deny_unknown_fields)]
961struct OpenRouterCapacity {
962 #[serde(default)]
963 prompt_tpm: Option<u64>,
964 #[serde(default)]
965 cached_prompt_tpm: Option<u64>,
966 #[serde(default)]
967 completion_tpm: Option<u64>,
968 #[serde(default)]
969 request_rpm: Option<u64>,
970 #[serde(default)]
971 concurrency: Option<u64>,
972}
973
974#[derive(Debug, Clone, Deserialize, Serialize)]
975#[serde(deny_unknown_fields)]
976struct OpenRouterDatacenter {
977 country_code: String,
978 #[serde(default, skip_serializing_if = "Option::is_none")]
979 region: Option<String>,
980}
981
982impl OpenRouterMetadataFile {
983 fn parse(
984 text: &str,
985 ) -> Result<
986 (
987 HashMap<String, OpenRouterModelMetadata>,
988 Option<ProviderMetadata>,
989 ),
990 String,
991 > {
992 let file: Self =
993 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
994 for (alias, metadata) in &file.models {
995 validate_openrouter_metadata(alias, metadata)?;
996 }
997 for (alias, metadata) in &file.planned_models {
998 validate_openrouter_metadata(alias, metadata)?;
999 if file.models.contains_key(alias) {
1000 return Err(format!(
1001 "model alias {alias:?} appears in both models and planned_models"
1002 ));
1003 }
1004 }
1005 if let Some(provider) = &file.provider {
1006 if provider.id.is_empty() {
1007 return Err("provider.id must be a non-empty slug".into());
1008 }
1009 for (field, value) in [
1011 ("provider.support_contact", &provider.support_contact),
1012 ("provider.incident_contact", &provider.incident_contact),
1013 ] {
1014 if let Some(value) = value {
1015 if !value.contains(':') {
1016 return Err(format!(
1017 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1018 ));
1019 }
1020 }
1021 }
1022 }
1023 Ok((file.models, file.provider))
1024 }
1025
1026 #[cfg(test)]
1027 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1028 Self::parse(text).map(|(models, _)| models)
1029 }
1030}
1031
1032fn per_million_price(per_token: &str) -> Option<String> {
1036 if !valid_price_string(per_token) {
1037 return None;
1038 }
1039 let (whole, frac) = match per_token.split_once('.') {
1040 Some((whole, frac)) => (whole, frac),
1041 None => (per_token, ""),
1042 };
1043 let mut digits = format!("{whole}{frac}");
1044 let point = whole.len() + 6;
1045 while digits.len() < point {
1046 digits.push('0');
1047 }
1048 let (int_part, frac_part) = digits.split_at(point);
1049 let int_part = int_part.trim_start_matches('0');
1050 let int_part = if int_part.is_empty() { "0" } else { int_part };
1051 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1052 while frac_out.len() < 2 {
1053 frac_out.push('0');
1054 }
1055 Some(format!("{int_part}.{frac_out}"))
1056}
1057
1058fn valid_price_string(value: &str) -> bool {
1059 let mut parts = value.split('.');
1060 let whole = parts.next().unwrap_or_default();
1061 let fraction = parts.next();
1062 !whole.is_empty()
1063 && whole.bytes().all(|b| b.is_ascii_digit())
1064 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1065 && parts.next().is_none()
1066}
1067
1068fn validate_openrouter_metadata(
1069 alias: &str,
1070 metadata: &OpenRouterModelMetadata,
1071) -> Result<(), String> {
1072 if alias.is_empty() {
1073 return Err("models metadata contains an empty model alias".into());
1074 }
1075 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1078 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1079 {
1080 return Err(format!(
1081 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1082 reasoning_effort level (none|minimal|low|medium|high)"
1083 ));
1084 }
1085 validate_sampling_defaults(alias, metadata)?;
1086 for m in &metadata.input_modalities {
1087 if m != "image" && m != "video" {
1088 return Err(format!(
1089 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1090 ));
1091 }
1092 }
1093 if let Some(sfc) = metadata.surface.as_deref()
1094 && !matches!(sfc, "chat" | "embedding" | "rerank")
1095 {
1096 return Err(format!(
1097 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1098 ));
1099 }
1100 if let Some(q) = metadata.quantization.as_deref()
1101 && !matches!(
1102 q,
1103 "int4"
1104 | "int8"
1105 | "fp4"
1106 | "mxfp4"
1107 | "nvfp4"
1108 | "fp6"
1109 | "fp8"
1110 | "mxfp8"
1111 | "fp16"
1112 | "bf16"
1113 | "fp32"
1114 )
1115 {
1116 return Err(format!(
1117 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1118 ));
1119 }
1120 for (field, value) in [
1121 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1122 (
1123 "pricing.cached_prompt",
1124 metadata.pricing.cached_prompt.as_deref(),
1125 ),
1126 (
1127 "pricing.cache_write",
1128 metadata.pricing.cache_write.as_deref(),
1129 ),
1130 ("pricing.completion", metadata.pricing.completion.as_deref()),
1131 (
1132 "pricing.internal_reasoning",
1133 metadata.pricing.internal_reasoning.as_deref(),
1134 ),
1135 ("pricing.request", metadata.pricing.request.as_deref()),
1136 ] {
1137 if let Some(value) = value
1138 && !valid_price_string(value)
1139 {
1140 return Err(format!(
1141 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1142 ));
1143 }
1144 }
1145 for (field, value) in [
1146 ("created", metadata.created),
1147 ("max_prompt_length", metadata.max_prompt_length),
1148 ("max_output_length", metadata.max_output_length),
1149 ("default_output_length", metadata.default_output_length),
1150 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1151 (
1152 "capacity.cached_prompt_tpm",
1153 metadata.capacity.cached_prompt_tpm,
1154 ),
1155 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1156 ("capacity.request_rpm", metadata.capacity.request_rpm),
1157 ("capacity.concurrency", metadata.capacity.concurrency),
1158 ] {
1159 if let Some(value) = value
1160 && value > JSON_SAFE_INTEGER_MAX
1161 {
1162 return Err(format!(
1163 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1164 ));
1165 }
1166 }
1167 for (field, value) in [
1168 ("max_prompt_length", metadata.max_prompt_length),
1169 ("max_output_length", metadata.max_output_length),
1170 ("default_output_length", metadata.default_output_length),
1171 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1172 (
1173 "capacity.cached_prompt_tpm",
1174 metadata.capacity.cached_prompt_tpm,
1175 ),
1176 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1177 ("capacity.request_rpm", metadata.capacity.request_rpm),
1178 ("capacity.concurrency", metadata.capacity.concurrency),
1179 ] {
1180 if value == Some(0) {
1181 return Err(format!(
1182 "model {alias:?}: {field} must be greater than zero when declared"
1183 ));
1184 }
1185 }
1186 if let (Some(default), Some(maximum)) =
1187 (metadata.default_output_length, metadata.max_output_length)
1188 && default > maximum
1189 {
1190 return Err(format!(
1191 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1192 ));
1193 }
1194 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1195 return Err(format!(
1196 "model {alias:?}: default_output_length requires max_output_length"
1197 ));
1198 }
1199 if let Some(discount) = metadata.discount_to_user
1200 && (!discount.is_finite() || discount >= 1.0)
1201 {
1202 return Err(format!(
1203 "model {alias:?}: discount_to_user must be finite and less than 1"
1204 ));
1205 }
1206 if metadata
1207 .openrouter_slug
1208 .as_deref()
1209 .is_some_and(str::is_empty)
1210 {
1211 return Err(format!(
1212 "model {alias:?}: openrouter_slug must not be empty when declared"
1213 ));
1214 }
1215 for dc in &metadata.datacenters {
1216 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1217 return Err(format!(
1218 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1219 dc.country_code
1220 ));
1221 }
1222 }
1223 Ok(())
1224}
1225
1226fn validate_sampling_defaults(
1241 alias: &str,
1242 metadata: &OpenRouterModelMetadata,
1243) -> Result<(), String> {
1244 validate_sampling_arm(
1245 alias,
1246 &[
1247 "default_temperature",
1248 "default_top_p",
1249 "default_min_p",
1250 "default_presence_penalty",
1251 "default_frequency_penalty",
1252 "default_repetition_penalty",
1253 ],
1254 metadata.default_temperature,
1255 metadata.default_top_p,
1256 metadata.default_min_p,
1257 metadata.default_presence_penalty,
1258 metadata.default_frequency_penalty,
1259 metadata.default_repetition_penalty,
1260 )?;
1261 if let Some(arm) = &metadata.non_thinking_sampling {
1262 if arm.is_empty() {
1266 return Err(format!(
1267 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1268 least one vendor recommendation or delete the table"
1269 ));
1270 }
1271 validate_sampling_arm(
1272 alias,
1273 &[
1274 "non_thinking_sampling.temperature",
1275 "non_thinking_sampling.top_p",
1276 "non_thinking_sampling.min_p",
1277 "non_thinking_sampling.presence_penalty",
1278 "non_thinking_sampling.frequency_penalty",
1279 "non_thinking_sampling.repetition_penalty",
1280 ],
1281 arm.temperature,
1282 arm.top_p,
1283 arm.min_p,
1284 arm.presence_penalty,
1285 arm.frequency_penalty,
1286 arm.repetition_penalty,
1287 )?;
1288 }
1289 Ok(())
1290}
1291
1292#[allow(clippy::too_many_arguments)]
1298fn validate_sampling_arm(
1299 alias: &str,
1300 keys: &[&str; 6],
1301 temperature: Option<f32>,
1302 top_p: Option<f32>,
1303 min_p: Option<f32>,
1304 presence_penalty: Option<f32>,
1305 frequency_penalty: Option<f32>,
1306 repetition_penalty: Option<f32>,
1307) -> Result<(), String> {
1308 if let Some(t) = temperature {
1309 if !t.is_finite() || t <= 0.0 || t > 2.0 {
1310 return Err(format!(
1311 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1312 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1313 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1314 vendor recommendation, not greedy); clients reach greedy by sending an \
1315 explicit temperature 0.",
1316 keys[0]
1317 ));
1318 }
1319 }
1320 if let Some(p) = top_p
1321 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1322 {
1323 return Err(format!(
1324 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1325 keys[1]
1326 ));
1327 }
1328 if let Some(m) = min_p
1329 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1330 {
1331 return Err(format!(
1332 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1333 keys[2]
1334 ));
1335 }
1336 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1337 if let Some(v) = value
1338 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1339 {
1340 return Err(format!(
1341 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1342 ));
1343 }
1344 }
1345 if let Some(r) = repetition_penalty
1346 && (!r.is_finite() || r <= 0.0)
1347 {
1348 return Err(format!(
1349 "model {alias:?}: {} {r} must be finite and \
1350 greater than zero (1.0 = off)",
1351 keys[5]
1352 ));
1353 }
1354 Ok(())
1355}
1356
1357fn load_openrouter_metadata(
1358 models: &[(String, String, Option<String>)],
1359) -> Result<
1360 (
1361 HashMap<String, OpenRouterModelMetadata>,
1362 Option<ProviderMetadata>,
1363 ),
1364 String,
1365> {
1366 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1367 Ok(path) => path,
1368 Err(_) => return Ok((HashMap::new(), None)),
1369 };
1370 let p = std::path::Path::new(&path);
1371 if !p.is_file() {
1372 return Err(format!(
1373 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1374 ));
1375 }
1376 let text =
1377 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1378 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1379 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1380 for alias in metadata.keys() {
1381 if !models.iter().any(|(name, _, _)| name == alias) {
1382 return Err(format!(
1383 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1384 ));
1385 }
1386 }
1387 eprintln!(
1388 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1389 metadata.len()
1390 );
1391 Ok((metadata, provider))
1392}
1393
1394#[derive(Clone)]
1395struct AppState {
1396 cmd_tx: Sender<Cmd>,
1397 models: Arc<Vec<String>>,
1398 caps: Arc<HashMap<String, ModelCaps>>,
1399 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1400 provider_metadata: Arc<Option<ProviderMetadata>>,
1402 metering: Option<Arc<dyn metering::Metering>>,
1408 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1411 api_auth: ApiAuth,
1414 metrics_auth: MetricsAuth,
1416 metrics: SharedMetrics,
1417 started: u64,
1420 inflight: InflightCounts,
1424 tenant_inflight: TenantGauge,
1427 health: health::SharedHealth,
1431 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1435}
1436
1437impl AppState {
1438 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1452 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1453 }
1454}
1455
1456#[derive(Clone, Default)]
1457struct ApiAuth {
1458 keyring: Option<&'static auth::KeyStore>,
1459 single_key: Option<Arc<str>>,
1460}
1461
1462impl ApiAuth {
1463 fn from_env() -> Result<ApiAuth, String> {
1464 let single_key = match std::env::var("MEMRA_API_KEY") {
1465 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1466 Ok(key) => Some(Arc::from(key)),
1467 Err(std::env::VarError::NotPresent) => None,
1468 Err(std::env::VarError::NotUnicode(_)) => {
1469 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1470 }
1471 };
1472 Ok(ApiAuth {
1473 keyring: auth::global(),
1474 single_key,
1475 })
1476 }
1477
1478 fn configured(&self) -> bool {
1479 self.keyring.is_some() || self.single_key.is_some()
1480 }
1481}
1482
1483#[derive(Clone, Default)]
1484struct MetricsAuth {
1485 required: bool,
1486 token: Option<Arc<str>>,
1487}
1488
1489impl MetricsAuth {
1490 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1491 let token = token.map(Arc::from);
1492 MetricsAuth {
1493 required: !bind_loopback || api_auth_configured || token.is_some(),
1494 token,
1495 }
1496 }
1497}
1498
1499fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1500 let mut resolved = addr
1501 .to_socket_addrs()
1502 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1503 let first = resolved
1504 .next()
1505 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1506 let mut loopback = first.ip().to_canonical().is_loopback();
1507 for socket in resolved {
1508 loopback &= socket.ip().to_canonical().is_loopback();
1509 }
1510 Ok((first, loopback))
1511}
1512
1513fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1514 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1515}
1516
1517fn validate_bind_security(
1518 addr: &str,
1519 api_auth_configured: bool,
1520 allow_open_bind: bool,
1521) -> Result<bool, String> {
1522 let loopback = bind_is_loopback(addr)?;
1523 if !loopback && !api_auth_configured && !allow_open_bind {
1524 return Err(format!(
1525 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1526 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1527 ));
1528 }
1529 Ok(loopback)
1530}
1531
1532type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1550
1551type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1554
1555struct InflightGuard {
1559 counts: InflightCounts,
1560 idx: usize,
1561 tenants: TenantGauge,
1562 tenant: String,
1563}
1564
1565impl InflightGuard {
1566 fn try_acquire(
1570 counts: InflightCounts,
1571 lane: lanes::Lane,
1572 tenants: TenantGauge,
1573 tenant: &str,
1574 tenant_cap: Option<usize>,
1575 ) -> Result<(Self, usize, usize), usize> {
1576 let idx = lane.idx();
1577 let nt = {
1578 let mut m = tenants.lock().unwrap();
1579 let e = m.entry(tenant.to_string()).or_insert(0);
1580 if tenant_cap.is_some_and(|cap| *e >= cap) {
1581 return Err(*e);
1582 }
1583 *e += 1;
1584 *e
1585 };
1586 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1587 Ok((
1588 InflightGuard {
1589 counts,
1590 idx,
1591 tenants,
1592 tenant: tenant.to_string(),
1593 },
1594 n,
1595 nt,
1596 ))
1597 }
1598}
1599
1600impl Drop for InflightGuard {
1601 fn drop(&mut self) {
1602 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1603 let mut m = self.tenants.lock().unwrap();
1604 if let Some(e) = m.get_mut(&self.tenant) {
1605 *e -= 1;
1606 if *e == 0 {
1607 m.remove(&self.tenant);
1608 }
1609 }
1610 }
1611}
1612
1613fn lane_cap(lane: lanes::Lane) -> usize {
1617 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1618 CAPS.get_or_init(|| {
1619 let batching = std::env::var("MEMRA_SERVE_BATCH")
1620 .map(|v| v != "0")
1621 .unwrap_or(true);
1622 let interactive = if batching {
1623 std::env::var("MEMRA_MAX_SESSIONS")
1624 .ok()
1625 .and_then(|v| v.parse().ok())
1626 .unwrap_or(64)
1627 } else {
1628 worker::MAX_ACTIVE
1629 };
1630 let p = lanes::LanePolicy::from_env();
1631 [interactive, p.max_sessions[1], p.max_sessions[2]]
1632 })[lane.idx()]
1633}
1634
1635fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1638 if m.completed > 0 && m.step_p50_ms > 0.0 {
1639 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1640 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1641 }
1642 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1643 *D.get_or_init(|| {
1644 std::env::var("MEMRA_RL_RESET_S")
1645 .ok()
1646 .and_then(|v| v.parse().ok())
1647 .unwrap_or(2)
1648 })
1649}
1650
1651pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1668pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1669pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1670
1671pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1675 let Some(v) = v.filter(|v| !v.is_null()) else {
1676 return Ok(TIMEOUT_MS_DEFAULT);
1677 };
1678 let Some(ms) = v.as_u64() else {
1679 return Err(format!(
1680 "timeout_ms must be an integer number of milliseconds in \
1681 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX}, got {v}; for work longer than \
1682 {TIMEOUT_MS_MAX} ms use \"stream\": true — the deadline then bounds only the \
1683 time to first token and the stream may run as long as it needs"
1684 ));
1685 };
1686 if !(TIMEOUT_MS_MIN..=TIMEOUT_MS_MAX).contains(&ms) {
1687 return Err(format!(
1688 "timeout_ms {ms} is outside the accepted range \
1689 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX} (milliseconds). {TIMEOUT_MS_MAX} is a \
1690 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1691 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1692 lie. For work longer than {TIMEOUT_MS_MAX} ms use \"stream\": true — the \
1693 deadline then bounds only the time to first token and the stream may run as \
1694 long as it needs"
1695 ));
1696 }
1697 Ok(ms)
1698}
1699
1700#[derive(Clone, Copy)]
1703pub(crate) struct RequestDeadline {
1704 pub(crate) at: tokio::time::Instant,
1705 pub(crate) ms: u64,
1706}
1707
1708impl RequestDeadline {
1709 pub(crate) fn starting_now(ms: u64) -> Self {
1710 Self {
1711 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
1712 ms,
1713 }
1714 }
1715
1716 pub(crate) fn remaining(&self) -> std::time::Duration {
1717 self.at
1718 .saturating_duration_since(tokio::time::Instant::now())
1719 }
1720}
1721
1722pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
1728 let what = if stream {
1729 "the first token was produced"
1730 } else {
1731 "the response completed"
1732 };
1733 let msg = format!(
1734 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
1735 {what}; generation was cancelled and this request is not billed"
1736 );
1737 error_response_coded(
1738 StatusCode::REQUEST_TIMEOUT,
1739 &msg,
1740 "timeout",
1741 Some("timeout_ms"),
1742 Some("deadline_exceeded"),
1743 )
1744}
1745
1746pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
1786
1787pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
1792
1793pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
1797
1798fn env_flag_on(name: &'static str, default_on: bool) -> bool {
1804 match std::env::var(name) {
1805 Ok(v) => !matches!(
1806 v.trim().to_ascii_lowercase().as_str(),
1807 "0" | "off" | "false"
1808 ),
1809 Err(_) => default_on,
1810 }
1811}
1812
1813fn env_u64(name: &'static str, default: u64) -> u64 {
1816 std::env::var(name)
1817 .ok()
1818 .and_then(|v| v.parse::<u64>().ok())
1819 .filter(|v| *v > 0)
1820 .unwrap_or(default)
1821}
1822
1823const CHARS_PER_TOKEN_FLOOR: usize = 6;
1834
1835pub(crate) fn prompt_tokens_estimate(
1836 request: &worker::Request,
1837 tokenizer: Option<&Tokenizer>,
1838) -> u64 {
1839 if !request.prompt_ids.is_empty() {
1840 return request.prompt_ids.len() as u64;
1841 }
1842 let mut text = String::new();
1843 text.push_str(&request.prompt_text);
1844 for turn in &request.chat_turns {
1845 text.push_str(&turn.content);
1846 }
1847 for tool in &request.tools_json {
1848 text.push_str(tool);
1849 }
1850 if let Some(tokenizer) = tokenizer {
1851 return tokenizer.encode(text.as_str(), false).len() as u64;
1852 }
1853 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
1854}
1855
1856pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
1860 let prefill_ms = prompt_tokens
1861 .saturating_mul(1_000)
1862 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
1863 .unwrap_or(u64::MAX);
1864 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
1865 if decode_ms == 0 {
1866 return None;
1867 }
1868 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
1869}
1870
1871pub(crate) fn nonstream_deadline_gate(
1881 request: &worker::Request,
1882 stream: bool,
1883 deadline: RequestDeadline,
1884 caller_declared_max_tokens: bool,
1885 tokenizer: Option<&Tokenizer>,
1886) -> Result<(), String> {
1887 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
1888 return Ok(());
1889 }
1890 let max_new = request.params.max_new as u64;
1891 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
1898 {
1899 return Ok(());
1900 }
1901 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
1902 let remaining_ms = deadline.remaining().as_millis() as u64;
1903 let prefill_ms = prompt_tokens.saturating_mul(1_000)
1904 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
1905 let decode_ms = max_new.saturating_mul(1_000)
1906 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
1907 let est_ms = prefill_ms.saturating_add(decode_ms);
1908 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
1909 if est_ms <= bound_ms {
1910 return Ok(());
1911 }
1912 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
1913 let advice = match fits {
1914 Some(fits) if fits > 0 => format!(
1915 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
1916 stream's deadline bounds only the time to first token, so it may run as long \
1917 as it needs"
1918 ),
1919 _ => format!(
1920 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
1921 token, so no max_tokens fits: set \"stream\": true"
1922 ),
1923 };
1924 Err(format!(
1925 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
1926 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
1927 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
1928 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
1929 rather than after the deadline: {advice}",
1930 est_ms / 1_000,
1931 ))
1932}
1933
1934fn max_queue_depth(cap: usize) -> usize {
1938 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1939 D.get_or_init(|| {
1940 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
1941 .ok()
1942 .and_then(|v| v.parse().ok())
1943 })
1944 .unwrap_or(cap.saturating_mul(4))
1945}
1946
1947pub(crate) fn admission_backpressure(
1960 st: &AppState,
1961 lane: lanes::Lane,
1962 rl: &RateLimit,
1963 deadline: RequestDeadline,
1964) -> Result<(), (Response, &'static str)> {
1965 if lane != lanes::Lane::Interactive || rl.remaining > 0 {
1966 return Ok(());
1967 }
1968 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
1969 let backlog = m.queued_requests as usize
1970 + worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire);
1971 let cap = lane_cap(lane).max(1);
1972 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
1973 let bound = max_queue_depth(cap);
1974 if backlog >= bound {
1975 let msg = format!(
1976 "interactive queue is at its bound ({backlog} queued, bound {bound}); this \
1977 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
1978 coarse estimate, not a promise)"
1979 );
1980 let resp = retry_contract_response(
1981 (
1982 StatusCode::TOO_MANY_REQUESTS,
1983 Json(error_body(
1984 &msg,
1985 "rate_limit_error",
1986 None,
1987 Some("shed_queue"),
1988 )),
1989 )
1990 .into_response(),
1991 Some(est_wait_s),
1992 );
1993 return Err((resp, "shed_queue"));
1994 }
1995 let remaining_ms = deadline.remaining().as_millis() as u64;
1996 if est_wait_s.saturating_mul(1_000) > remaining_ms {
1997 let msg = format!(
1998 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
1999 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2000 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2001 estimate, not a promise)"
2002 );
2003 let resp = retry_contract_response(
2004 (
2005 StatusCode::TOO_MANY_REQUESTS,
2006 Json(error_body(
2007 &msg,
2008 "rate_limit_error",
2009 None,
2010 Some("shed_deadline"),
2011 )),
2012 )
2013 .into_response(),
2014 Some(est_wait_s),
2015 );
2016 return Err((resp, "shed_deadline"));
2017 }
2018 Ok(())
2019}
2020
2021pub(crate) struct PendingAdmissionGuard {
2026 reserved: bool,
2027 lane: lanes::Lane,
2028}
2029
2030impl PendingAdmissionGuard {
2031 pub(crate) fn commit(mut self) {
2035 self.reserved = false;
2036 std::mem::forget(self);
2037 }
2038}
2039
2040impl Drop for PendingAdmissionGuard {
2041 fn drop(&mut self) {
2042 if self.reserved {
2043 worker::release_pending_admit();
2044 worker::release_admission_reservation(self.lane);
2045 }
2046 }
2047}
2048
2049pub(crate) fn reserve_pending_admit(
2050 st: &AppState,
2051 lane: lanes::Lane,
2052 rl: &RateLimit,
2053 deadline: RequestDeadline,
2054) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2055 let cap = lane_cap(lane).max(1);
2060 let bound = max_queue_depth(cap);
2061 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2062 loop {
2063 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2064 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2065 let backlog = reservations;
2069 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2070 if backlog >= bound {
2071 let msg = format!(
2072 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2073 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2074 coarse estimate, not a promise)",
2075 lane.as_str()
2076 );
2077 let resp = retry_contract_response(
2078 (
2079 StatusCode::TOO_MANY_REQUESTS,
2080 Json(error_body(
2081 &msg,
2082 "rate_limit_error",
2083 None,
2084 Some("shed_queue"),
2085 )),
2086 )
2087 .into_response(),
2088 Some(est_wait_s),
2089 );
2090 return Err((resp, "shed_queue"));
2091 }
2092 let remaining_ms = deadline.remaining().as_millis() as u64;
2093 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2097 if lane == lanes::Lane::Interactive
2098 && waits_for_capacity
2099 && est_wait_s.saturating_mul(1_000) > remaining_ms
2100 {
2101 let msg = format!(
2102 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2103 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2104 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2105 estimate, not a promise)"
2106 );
2107 let resp = retry_contract_response(
2108 (
2109 StatusCode::TOO_MANY_REQUESTS,
2110 Json(error_body(
2111 &msg,
2112 "rate_limit_error",
2113 None,
2114 Some("shed_deadline"),
2115 )),
2116 )
2117 .into_response(),
2118 Some(est_wait_s),
2119 );
2120 return Err((resp, "shed_deadline"));
2121 }
2122 if reservations_for_lane
2123 .compare_exchange(
2124 reservations,
2125 reservations.saturating_add(1),
2126 std::sync::atomic::Ordering::AcqRel,
2127 std::sync::atomic::Ordering::Acquire,
2128 )
2129 .is_ok()
2130 {
2131 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2135 return Ok(PendingAdmissionGuard {
2136 reserved: true,
2137 lane,
2138 });
2139 }
2140 }
2141}
2142
2143static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2154
2155fn draining() -> bool {
2156 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2157}
2158
2159fn drain_deadline_s() -> u64 {
2161 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2162 *D.get_or_init(|| {
2163 std::env::var("MEMRA_DRAIN_S")
2164 .ok()
2165 .and_then(|v| v.parse().ok())
2166 .unwrap_or(30)
2167 })
2168}
2169
2170fn drain_response() -> Response {
2180 let resp = (
2181 StatusCode::SERVICE_UNAVAILABLE,
2182 Json(error_body(
2183 "server is draining (shutdown in progress); retry",
2184 "server_error",
2185 None,
2186 Some("draining"),
2187 )),
2188 )
2189 .into_response();
2190 retry_contract_response(resp, Some(drain_deadline_s()))
2191}
2192
2193struct RateLimit {
2195 limit: usize,
2196 remaining: usize,
2197 reset_s: u64,
2198}
2199
2200impl RateLimit {
2201 fn at_admit(
2206 lane: lanes::Lane,
2207 n_inflight: usize,
2208 metrics: &SharedMetrics,
2209 tenant: &auth::TenantCtx,
2210 n_tenant: usize,
2211 ) -> Self {
2212 let global = lane_cap(lane);
2213 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2214 return Self::compute(global, n_inflight, metrics);
2215 };
2216 let headroom = t
2217 .saturating_sub(n_tenant)
2218 .min(global.saturating_sub(n_inflight));
2219 Self::compute(t, t - headroom, metrics)
2221 }
2222
2223 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2224 let remaining = limit.saturating_sub(n_inflight);
2225 let reset_s = if remaining > 0 {
2226 0
2227 } else {
2228 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2229 reset_estimate_s(&m)
2230 };
2231 RateLimit {
2232 limit,
2233 remaining,
2234 reset_s,
2235 }
2236 }
2237
2238 fn attach(&self, mut resp: Response) -> Response {
2240 let h = resp.headers_mut();
2241 for (k, v) in [
2242 ("x-ratelimit-limit", self.limit as u64),
2243 ("x-ratelimit-remaining", self.remaining as u64),
2244 ("x-ratelimit-reset", self.reset_s),
2245 ] {
2246 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2247 h.insert(axum::http::HeaderName::from_static(k), v);
2248 }
2249 }
2250 resp
2251 }
2252}
2253
2254fn acquire_request_slot(
2258 st: &AppState,
2259 lane: lanes::Lane,
2260 tenant: &auth::TenantCtx,
2261 env: &Envelope,
2262) -> Result<(InflightGuard, RateLimit), Response> {
2263 let global = lane_cap(lane);
2264 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2265 match InflightGuard::try_acquire(
2266 st.inflight.clone(),
2267 lane,
2268 st.tenant_inflight.clone(),
2269 &tenant.tenant,
2270 tenant_cap,
2271 ) {
2272 Ok((guard, n_inflight, n_tenant)) => {
2273 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2274 Ok((guard, rl))
2275 }
2276 Err(n_tenant) => {
2277 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2278 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2279 let error =
2280 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2281 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2282 }
2283 }
2284}
2285
2286#[derive(Deserialize)]
2288struct CompletionReq {
2289 model: String,
2290 #[serde(default)]
2291 prompt: String,
2292 #[serde(default)]
2294 prompt_ids: Vec<u32>,
2295 #[serde(default)]
2298 max_tokens: Option<usize>,
2299 #[serde(default)]
2312 temperature: Option<f32>,
2313 #[serde(default)]
2314 top_p: Option<f32>,
2315 #[serde(default)]
2317 top_k: Option<usize>,
2318 #[serde(default)]
2320 min_p: Option<f32>,
2321 #[serde(default)]
2323 frequency_penalty: Option<f32>,
2324 #[serde(default)]
2325 presence_penalty: Option<f32>,
2326 #[serde(default)]
2328 repetition_penalty: Option<f32>,
2329 #[serde(default)]
2335 seed: Option<u64>,
2336 #[serde(default)]
2337 stop: StopSequences,
2338 #[serde(default)]
2341 logit_bias: Option<serde_json::Value>,
2342 #[serde(default)]
2343 logprobs: Option<serde_json::Value>,
2344 #[serde(default)]
2345 n: Option<usize>,
2346 #[serde(default)]
2347 best_of: Option<usize>,
2348 #[serde(default)]
2350 chat: bool,
2351 #[serde(default)]
2353 stream: bool,
2354 #[serde(default)]
2356 max_ctx: Option<usize>,
2357 #[serde(default)]
2359 trace_id: Option<String>,
2360 #[serde(default)]
2364 cache_salt: Option<String>,
2365 #[serde(default)]
2369 session_id: Option<String>,
2370 #[serde(default)]
2371 user: Option<String>,
2372 #[serde(default)]
2376 timeout_ms: Option<serde_json::Value>,
2377}
2378
2379#[derive(Deserialize)]
2380struct ChatMessage {
2381 role: String,
2382 #[serde(default)]
2384 content: serde_json::Value,
2385 #[serde(default)]
2387 tool_calls: Vec<ReqToolCall>,
2388 #[serde(default)]
2391 tool_call_id: Option<String>,
2392 #[serde(default)]
2395 name: Option<String>,
2396 #[serde(default, alias = "reasoning_content")]
2406 reasoning: Option<String>,
2407}
2408
2409#[derive(Deserialize)]
2410struct ReqToolCall {
2411 #[serde(default)]
2412 #[allow(dead_code)]
2413 id: Option<String>,
2414 function: ReqToolFunction,
2415}
2416
2417#[derive(Deserialize)]
2418struct ReqToolFunction {
2419 name: String,
2420 #[serde(default)]
2422 arguments: serde_json::Value,
2423}
2424
2425#[derive(Clone, Default, Deserialize)]
2426#[serde(untagged)]
2427enum StopSequences {
2428 One(String),
2429 Many(Vec<String>),
2430 #[default]
2431 None,
2432}
2433
2434impl StopSequences {
2435 fn into_vec(self) -> Vec<String> {
2441 let stops = match self {
2442 Self::One(stop) => vec![stop],
2443 Self::Many(stops) => stops,
2444 Self::None => Vec::new(),
2445 };
2446 stops.into_iter().filter(|s| !s.is_empty()).collect()
2447 }
2448}
2449
2450#[derive(Deserialize)]
2456struct ChatCompletionReq {
2457 model: String,
2458 messages: Vec<ChatMessage>,
2459 #[serde(default, alias = "max_completion_tokens")]
2462 max_tokens: Option<usize>,
2463 #[serde(default)]
2466 temperature: Option<f32>,
2467 #[serde(default)]
2468 top_p: Option<f32>,
2469 #[serde(default)]
2473 top_k: Option<usize>,
2474 #[serde(default)]
2476 min_p: Option<f32>,
2477 #[serde(default)]
2479 frequency_penalty: Option<f32>,
2480 #[serde(default)]
2481 presence_penalty: Option<f32>,
2482 #[serde(default)]
2484 repetition_penalty: Option<f32>,
2485 #[serde(default)]
2487 seed: Option<u64>,
2488 #[serde(default)]
2489 stop: StopSequences,
2490 #[serde(default)]
2491 stream: bool,
2492 #[serde(default)]
2493 max_ctx: Option<usize>,
2494 #[serde(default)]
2499 response_format: Option<serde_json::Value>,
2500 #[serde(default)]
2501 logit_bias: Option<serde_json::Value>,
2502 #[serde(default)]
2503 logprobs: Option<serde_json::Value>,
2504 #[serde(default)]
2505 top_logprobs: Option<usize>,
2506 #[serde(default)]
2507 n: Option<usize>,
2508 #[serde(default)]
2510 tools: Vec<serde_json::Value>,
2511 #[serde(default)]
2513 tool_choice: Option<serde_json::Value>,
2514 #[serde(default)]
2520 reasoning_effort: Option<String>,
2521 #[serde(default)]
2530 reasoning: Option<serde_json::Value>,
2531 #[serde(default)]
2546 include_reasoning: Option<bool>,
2547 #[serde(default)]
2556 enable_thinking: Option<bool>,
2557 #[serde(default)]
2563 chat_template_kwargs: Option<serde_json::Value>,
2564 #[serde(default)]
2568 cache_salt: Option<String>,
2569 #[serde(default)]
2571 session_id: Option<String>,
2572 #[serde(default)]
2573 user: Option<String>,
2574 #[serde(default)]
2579 timeout_ms: Option<serde_json::Value>,
2580}
2581fn one() -> f32 {
2582 1.0
2583}
2584fn default_temperature() -> f32 {
2590 1.0
2591}
2592
2593#[derive(Debug, Clone, Copy, Default, PartialEq)]
2613struct SamplingDefaults {
2614 temperature: Option<f32>,
2615 top_p: Option<f32>,
2616 top_k: Option<usize>,
2617 min_p: Option<f32>,
2618 frequency_penalty: Option<f32>,
2619 presence_penalty: Option<f32>,
2620 repetition_penalty: Option<f32>,
2621}
2622
2623impl SamplingDefaults {
2624 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2627 SamplingDefaults {
2628 temperature: metadata
2629 .and_then(|m| m.default_temperature)
2630 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2631 top_p: metadata
2632 .and_then(|m| m.default_top_p)
2633 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2634 top_k: metadata.and_then(|m| m.default_top_k),
2635 min_p: metadata.and_then(|m| m.default_min_p),
2636 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2637 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2638 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2639 }
2640 }
2641}
2642
2643#[derive(Debug, Clone, Copy, Default, PartialEq)]
2656struct ModelSamplingDefaults {
2657 thinking: SamplingDefaults,
2658 non_thinking: Option<SamplingDefaults>,
2659}
2660
2661impl ModelSamplingDefaults {
2662 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2663 ModelSamplingDefaults {
2664 thinking: SamplingDefaults::resolve(metadata, caps),
2665 non_thinking: metadata
2670 .and_then(|m| m.non_thinking_sampling.as_ref())
2671 .map(|arm| SamplingDefaults {
2672 temperature: arm.temperature,
2673 top_p: arm.top_p,
2674 top_k: arm.top_k,
2675 min_p: arm.min_p,
2676 frequency_penalty: arm.frequency_penalty,
2677 presence_penalty: arm.presence_penalty,
2678 repetition_penalty: arm.repetition_penalty,
2679 }),
2680 }
2681 }
2682
2683 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
2696 match (think, &self.non_thinking) {
2697 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
2698 _ => &self.thinking,
2699 }
2700 }
2701
2702 fn single(thinking: SamplingDefaults) -> Self {
2705 ModelSamplingDefaults {
2706 thinking,
2707 non_thinking: None,
2708 }
2709 }
2710}
2711
2712#[derive(Debug, Clone, Copy, Default)]
2718struct ClientSampling {
2719 temperature: Option<f32>,
2720 top_p: Option<f32>,
2721 top_k: Option<usize>,
2722 min_p: Option<f32>,
2723 frequency_penalty: Option<f32>,
2724 presence_penalty: Option<f32>,
2725 repetition_penalty: Option<f32>,
2726 seed: Option<u64>,
2727}
2728
2729impl From<&CompletionReq> for ClientSampling {
2730 fn from(r: &CompletionReq) -> Self {
2731 ClientSampling {
2732 temperature: r.temperature,
2733 top_p: r.top_p,
2734 top_k: r.top_k,
2735 min_p: r.min_p,
2736 frequency_penalty: r.frequency_penalty,
2737 presence_penalty: r.presence_penalty,
2738 repetition_penalty: r.repetition_penalty,
2739 seed: r.seed,
2740 }
2741 }
2742}
2743
2744impl From<&ChatCompletionReq> for ClientSampling {
2745 fn from(r: &ChatCompletionReq) -> Self {
2746 ClientSampling {
2747 temperature: r.temperature,
2748 top_p: r.top_p,
2749 top_k: r.top_k,
2750 min_p: r.min_p,
2751 frequency_penalty: r.frequency_penalty,
2752 presence_penalty: r.presence_penalty,
2753 repetition_penalty: r.repetition_penalty,
2754 seed: r.seed,
2755 }
2756 }
2757}
2758
2759fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
2765 sampler_config(
2766 client
2767 .temperature
2768 .or(defaults.temperature)
2769 .unwrap_or_else(default_temperature),
2770 client.top_k.or(defaults.top_k).unwrap_or(0),
2771 client.top_p.or(defaults.top_p).unwrap_or_else(one),
2772 client.min_p.or(defaults.min_p).unwrap_or(0.0),
2773 client
2774 .frequency_penalty
2775 .or(defaults.frequency_penalty)
2776 .unwrap_or(0.0),
2777 client
2778 .presence_penalty
2779 .or(defaults.presence_penalty)
2780 .unwrap_or(0.0),
2781 client
2782 .repetition_penalty
2783 .or(defaults.repetition_penalty)
2784 .unwrap_or_else(one),
2785 client.seed,
2786 )
2787}
2788
2789#[derive(Serialize)]
2790struct CompletionResp {
2791 model: String,
2792 text: String,
2793 tokens: Vec<u32>,
2794 stop_reason: String,
2798 #[serde(default, skip_serializing_if = "Option::is_none")]
2803 error: Option<serde_json::Value>,
2804 n_tokens: usize,
2805 prompt_tokens: usize,
2808 cached_tokens: usize,
2809 elapsed_s: f64,
2810}
2811
2812fn usage_json(
2820 n_prompt: usize,
2821 n_tokens: usize,
2822 n_cached: usize,
2823 elapsed_s: f64,
2824 spec: Option<worker::SpecUsage>,
2825) -> serde_json::Value {
2826 let mut u = json!({
2827 "prompt_tokens": n_prompt,
2828 "completion_tokens": n_tokens,
2829 "total_tokens": n_prompt + n_tokens,
2830 "prompt_tokens_details": { "cached_tokens": n_cached },
2831 "elapsed_s": elapsed_s,
2832 });
2833 if let Some(sp) = spec {
2834 u["spec"] = json!({
2835 "rounds": sp.rounds,
2836 "drafted": sp.drafted,
2837 "accepted": sp.accepted,
2838 "acceptance_rate": if sp.drafted > 0 {
2839 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
2840 });
2841 }
2842 u
2843}
2844
2845const SYSTEM_FINGERPRINT: &str = concat!("memra-", env!("MEMRA_BUILD_SHA"));
2859
2860fn gen_hex128() -> String {
2863 use std::hash::{BuildHasher, Hasher};
2864 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2865 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2866 let t = std::time::SystemTime::now()
2867 .duration_since(std::time::UNIX_EPOCH)
2868 .map(|d| d.as_nanos() as u64)
2869 .unwrap_or(0);
2870 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
2871 h1.write_u64(n);
2872 h1.write_u64(t);
2873 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
2874 h2.write_u64(t.rotate_left(17));
2875 h2.write_u64(n);
2876 format!("{:016x}{:016x}", h1.finish(), h2.finish())
2877}
2878
2879#[derive(Clone)]
2882struct Envelope {
2883 id: String,
2884 created: u64,
2885}
2886
2887impl Envelope {
2888 fn new(chat: bool) -> Self {
2889 Envelope {
2890 id: format!(
2891 "{}-{}",
2892 if chat { "chatcmpl" } else { "cmpl" },
2893 gen_hex128()
2894 ),
2895 created: std::time::SystemTime::now()
2896 .duration_since(std::time::UNIX_EPOCH)
2897 .map(|d| d.as_secs())
2898 .unwrap_or(0),
2899 }
2900 }
2901
2902 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
2904 v["id"] = json!(self.id);
2905 v["created"] = json!(self.created);
2906 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
2907 v
2908 }
2909}
2910
2911fn with_request_id(id: &str, mut resp: Response) -> Response {
2913 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
2914 resp.headers_mut()
2915 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
2916 }
2917 resp
2918}
2919
2920fn openai_compat() -> bool {
2928 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2929 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
2930 Ok("openai") => true,
2931 Ok(_) => false,
2932 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
2933 })
2934}
2935
2936fn cache_namespace(cache_salt: &Option<String>) -> String {
2948 cache_salt.clone().unwrap_or_default()
2949}
2950
2951const CACHE_SALT_MAX_BYTES: usize = 64;
2952
2953fn validate_cache_namespace(
2954 cache_salt: &Option<String>,
2955 keyring_configured: bool,
2956) -> Result<String, &'static str> {
2957 let raw = cache_namespace(cache_salt);
2958 if raw.len() > CACHE_SALT_MAX_BYTES {
2959 return Err("cache_salt must be at most 64 bytes");
2960 }
2961 if !keyring_configured && raw.starts_with("t:") {
2962 return Err("cache_salt must not use the reserved t: prefix without a keyring");
2963 }
2964 if !raw
2965 .bytes()
2966 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
2967 {
2968 return Err("cache_salt contains unsupported characters");
2969 }
2970 Ok(raw)
2971}
2972
2973fn affinity_key(
2989 session_id: &Option<String>,
2990 user: &Option<String>,
2991 headers: &axum::http::HeaderMap,
2992) -> Option<String> {
2993 let clean = |s: &str| -> Option<String> {
2994 let t = s.trim();
2995 if t.is_empty() {
2996 None
2997 } else {
2998 Some(t.to_string())
2999 }
3000 };
3001 session_id
3002 .as_deref()
3003 .and_then(clean)
3004 .or_else(|| user.as_deref().and_then(clean))
3005 .or_else(|| {
3006 headers
3007 .get("x-session-id")
3008 .and_then(|v| v.to_str().ok())
3009 .and_then(clean)
3010 })
3011}
3012
3013fn error_body(
3018 message: &str,
3019 etype: &str,
3020 param: Option<&str>,
3021 code: Option<&str>,
3022) -> serde_json::Value {
3023 json!({ "error": {
3024 "message": message,
3025 "type": etype,
3026 "param": param,
3027 "code": code,
3028 } })
3029}
3030
3031fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3032 error_response_coded(status, message, etype, param, None)
3033}
3034
3035fn error_response_coded(
3040 status: StatusCode,
3041 message: &str,
3042 etype: &str,
3043 param: Option<&str>,
3044 code: Option<&str>,
3045) -> Response {
3046 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3047 if status.is_client_error()
3048 && status != StatusCode::TOO_MANY_REQUESTS
3049 && status != StatusCode::REQUEST_TIMEOUT
3050 && status != StatusCode::CONFLICT
3051 {
3052 resp.headers_mut().insert(
3053 "x-should-retry",
3054 axum::http::HeaderValue::from_static("false"),
3055 );
3056 }
3057 resp
3058}
3059
3060fn bad_request(message: &str, param: Option<&str>) -> Response {
3061 error_response(
3062 StatusCode::BAD_REQUEST,
3063 message,
3064 "invalid_request_error",
3065 param,
3066 )
3067}
3068
3069const RETRY_AFTER_S_RATE_LIMIT: u64 = 2; const RETRY_AFTER_S_OVERLOADED: u64 = 5; fn class_http(class: worker::ErrClass) -> (StatusCode, &'static str, Option<&'static str>) {
3098 use worker::ErrClass as C;
3099 match class {
3100 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3101 C::ContextLength => (
3102 StatusCode::BAD_REQUEST,
3103 "invalid_request_error",
3104 Some("context_length_exceeded"),
3105 ),
3106 C::ModelNotFound => (
3107 StatusCode::BAD_REQUEST,
3108 "invalid_request_error",
3109 Some("model_not_found"),
3110 ),
3111 C::RateLimit => (
3112 StatusCode::TOO_MANY_REQUESTS,
3113 "rate_limit_error",
3114 Some("rate_limit_exceeded"),
3115 ),
3116 C::Overloaded => (
3117 StatusCode::SERVICE_UNAVAILABLE,
3118 "server_error",
3119 Some("overloaded"),
3120 ),
3121 C::Engine => (
3122 StatusCode::INTERNAL_SERVER_ERROR,
3123 "server_error",
3124 Some("engine_error"),
3125 ),
3126 }
3127}
3128
3129fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3131 use worker::ErrClass as C;
3132 match class {
3133 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3134 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3135 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3139 }
3140}
3141
3142fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3145 let (_, etype, code) = class_http(e.class);
3146 error_body(&e.message, etype, e.param, code)
3147}
3148
3149fn engine_error_response(e: &worker::EngineError) -> Response {
3151 engine_error_response_with_retry_after(e, class_retry_after_s(e.class))
3152}
3153
3154fn engine_error_response_with_retry_after(
3155 e: &worker::EngineError,
3156 retry_after_s: Option<u64>,
3157) -> Response {
3158 let (status, _, _) = class_http(e.class);
3159 let resp = (status, Json(engine_error_body(e))).into_response();
3160 retry_contract_response(resp, retry_after_s)
3161}
3162
3163fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3165 let status = resp.status();
3166 let h = resp.headers_mut();
3167 match retry_after_s {
3168 Some(secs) => {
3169 let secs = secs.clamp(1, 60);
3171 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3172 h.insert(axum::http::header::RETRY_AFTER, v);
3173 }
3174 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3175 h.insert("retry-after-ms", v);
3176 }
3177 }
3178 None if status.is_client_error() => {
3179 h.insert(
3182 "x-should-retry",
3183 axum::http::HeaderValue::from_static("false"),
3184 );
3185 }
3186 None => {}
3187 }
3188 resp
3189}
3190
3191fn worker_unavailable_response() -> Response {
3192 engine_error_response_with_retry_after(
3193 &worker::EngineError::overloaded("worker unavailable"),
3194 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3195 )
3196}
3197
3198fn stop_reason_to_finish(r: &str) -> &'static str {
3199 match r {
3200 "Eos" | "Callback" => "stop",
3201 "MaxNew" | "ContextFull" => "length",
3202 _ => "stop",
3203 }
3204}
3205
3206fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3210 match v {
3211 serde_json::Value::Null => Ok(String::new()),
3212 serde_json::Value::String(s) => Ok(s.clone()),
3213 serde_json::Value::Array(parts) => {
3214 let mut out = String::new();
3215 for p in parts {
3216 match p.get("type").and_then(|t| t.as_str()) {
3217 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3218 Some(t) => out.push_str(t),
3219 None => return Err("content part has no text field".into()),
3220 },
3221 Some(other) => {
3222 return Err(format!(
3223 "unsupported content part type {other:?} (text only)"
3224 ));
3225 }
3226 }
3227 }
3228 Ok(out)
3229 }
3230 _ => Err("content must be a string, null, or an array of text parts".into()),
3231 }
3232}
3233
3234fn vision_enabled() -> bool {
3237 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3238 *ON.get_or_init(|| {
3239 std::env::var("MEMRA_VISION_DIR").is_ok()
3240 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3241 })
3242}
3243
3244fn gemma_vision_enabled() -> bool {
3249 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3250 *ON.get_or_init(|| {
3251 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3252 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3253 })
3254}
3255
3256const VISION_MAX_IMAGES: usize = 8;
3258
3259pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3265 std::sync::atomic::AtomicUsize::new(0);
3266pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3270 tokio::sync::Semaphore::const_new(1);
3271
3272pub(crate) struct VisionMemoryPermit {
3273 bytes: usize,
3274}
3275
3276#[derive(Debug)]
3277pub(crate) enum VisionMemoryError {
3278 Request(String),
3279 Capacity(String),
3280}
3281
3282impl Drop for VisionMemoryPermit {
3283 fn drop(&mut self) {
3284 if self.bytes != 0 {
3285 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3286 }
3287 }
3288}
3289
3290fn try_reserve_vision_memory(
3291 bytes: usize,
3292) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3293 if bytes == 0 {
3294 return Ok(None);
3295 }
3296 if bytes > MAX_VISION_PATCH_BYTES {
3297 return Err(VisionMemoryError::Request(format!(
3298 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3299 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3300 )));
3301 }
3302 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3303 loop {
3304 let Some(next) = in_use.checked_add(bytes) else {
3305 return Err(VisionMemoryError::Capacity(
3306 "vision patch memory reservation overflowed".into(),
3307 ));
3308 };
3309 if next > MAX_VISION_PATCH_BYTES {
3310 return Err(VisionMemoryError::Capacity(format!(
3311 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3312 in_use / (1024 * 1024),
3313 bytes / (1024 * 1024)
3314 )));
3315 }
3316 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3317 in_use,
3318 next,
3319 std::sync::atomic::Ordering::AcqRel,
3320 std::sync::atomic::Ordering::Acquire,
3321 ) {
3322 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3323 Err(actual) => in_use = actual,
3324 }
3325 }
3326}
3327
3328pub(crate) fn vision_memory_error_response(
3329 error: VisionMemoryError,
3330 param: Option<&str>,
3331) -> Response {
3332 match error {
3333 VisionMemoryError::Request(message) => bad_request(&message, param),
3334 VisionMemoryError::Capacity(message) => retry_contract_response(
3335 error_response_coded(
3336 StatusCode::SERVICE_UNAVAILABLE,
3337 &message,
3338 "server_error",
3339 None,
3340 Some("vision_memory_busy"),
3341 ),
3342 Some(RETRY_AFTER_S_OVERLOADED),
3343 ),
3344 }
3345}
3346
3347enum PendingVisionUnit {
3353 Still {
3354 bytes: Vec<u8>,
3355 gh: usize,
3356 gw: usize,
3357 },
3358 Video {
3359 bytes: Vec<u8>,
3360 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3361 video: usize,
3362 },
3363}
3364
3365struct PendingGemmaImage {
3367 bytes: Vec<u8>,
3368 gw: usize,
3369 gh: usize,
3370}
3371
3372fn content_to_text_vision(
3381 v: &serde_json::Value,
3382 images: &mut Vec<PendingVisionUnit>,
3383 gemma_images: &mut Vec<PendingGemmaImage>,
3384 next_video: &mut usize,
3385) -> Result<String, String> {
3386 let parts = match v {
3387 serde_json::Value::Array(parts) => parts,
3388 _ => return content_to_text(v),
3389 };
3390 let mut out = String::new();
3391 for p in parts {
3392 match p.get("type").and_then(|t| t.as_str()) {
3393 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3394 Some(t) => out.push_str(t),
3395 None => return Err("content part has no text field".into()),
3396 },
3397 Some("image_url") if gemma_vision_enabled() => {
3398 let url = p
3399 .get("image_url")
3400 .and_then(|u| {
3401 if u.is_string() {
3402 u.as_str()
3403 } else {
3404 u.get("url").and_then(|x| x.as_str())
3405 }
3406 })
3407 .ok_or("image_url part has no url")?;
3408 if !url.starts_with("data:") {
3409 return Err(
3410 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3411 );
3412 }
3413 if gemma_images.len() >= VISION_MAX_IMAGES {
3414 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3415 }
3416 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
3420 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3421 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
3422 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3423 out.push_str("<|image>");
3425 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
3426 out.push_str("<|image|>");
3427 }
3428 out.push_str("<image|>");
3429 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
3430 }
3431 Some("image_url") => {
3432 if !vision_enabled() {
3433 return Err("image input is not enabled on this deployment".into());
3434 }
3435 let url = p
3436 .get("image_url")
3437 .and_then(|u| {
3438 if u.is_string() {
3439 u.as_str()
3440 } else {
3441 u.get("url").and_then(|x| x.as_str())
3442 }
3443 })
3444 .ok_or("image_url part has no url")?;
3445 if !url.starts_with("data:") {
3446 return Err(
3447 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3448 );
3449 }
3450 if images
3451 .iter()
3452 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
3453 .count()
3454 >= VISION_MAX_IMAGES
3455 {
3456 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3457 }
3458 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3463 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3464 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
3465 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3466 out.push_str("<|vision_start|>");
3467 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
3468 out.push_str("<|image_pad|>");
3469 }
3470 out.push_str("<|vision_end|>");
3471 images.push(PendingVisionUnit::Still { bytes, gh, gw });
3472 }
3473 Some("video_url") if gemma_vision_enabled() => {
3474 return Err("gemma-4 has no video input (image-only projector)".into());
3475 }
3476 Some("video_url") => {
3477 if !vision_enabled() {
3478 return Err("video input is not enabled on this deployment".into());
3479 }
3480 let url = p
3481 .get("video_url")
3482 .and_then(|u| {
3483 if u.is_string() {
3484 u.as_str()
3485 } else {
3486 u.get("url").and_then(|x| x.as_str())
3487 }
3488 })
3489 .ok_or("video_url part has no url")?;
3490 if !url.starts_with("data:") {
3491 return Err(
3492 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3493 );
3494 }
3495 if *next_video >= 2 {
3496 return Err("too many videos (max 2)".into());
3497 }
3498 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
3501 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
3502 .map_err(|e| format!("video: {e}"))?;
3503 let vidx = *next_video;
3504 *next_video += 1;
3505 for group in &vid.groups {
3507 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
3508 out.push_str("<|vision_start|>");
3509 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
3510 out.push_str("<|video_pad|>");
3511 }
3512 out.push_str("<|vision_end|>");
3513 }
3514 images.push(PendingVisionUnit::Video {
3517 bytes,
3518 groups: vid.groups,
3519 video: vidx,
3520 });
3521 }
3522 Some(other) => {
3523 return Err(format!("unsupported content part type {other:?}"));
3524 }
3525 }
3526 }
3527 Ok(out)
3528}
3529
3530fn pyjson(v: &serde_json::Value, out: &mut String) {
3534 match v {
3535 serde_json::Value::Object(m) => {
3536 out.push('{');
3537 for (i, (k, val)) in m.iter().enumerate() {
3538 if i > 0 {
3539 out.push_str(", ");
3540 }
3541 out.push_str(&serde_json::Value::String(k.clone()).to_string());
3542 out.push_str(": ");
3543 pyjson(val, out);
3544 }
3545 out.push('}');
3546 }
3547 serde_json::Value::Array(a) => {
3548 out.push('[');
3549 for (i, val) in a.iter().enumerate() {
3550 if i > 0 {
3551 out.push_str(", ");
3552 }
3553 pyjson(val, out);
3554 }
3555 out.push(']');
3556 }
3557 scalar => out.push_str(&scalar.to_string()),
3558 }
3559}
3560
3561fn pyjson_str(v: &serde_json::Value) -> String {
3562 let mut s = String::new();
3563 pyjson(v, &mut s);
3564 s
3565}
3566
3567fn sampler_config(
3574 temperature: f32,
3575 top_k: usize,
3576 top_p: f32,
3577 min_p: f32,
3578 frequency_penalty: f32,
3579 presence_penalty: f32,
3580 repetition_penalty: f32,
3581 seed: Option<u64>,
3582) -> SamplerConfig {
3583 let penalties_on =
3584 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
3585 SamplerConfig {
3586 temperature,
3587 top_k,
3588 top_p,
3589 min_p,
3590 penalty_last_n: if penalties_on {
3591 memra_engine::spec::PEN_WINDOW_MAX
3592 } else {
3593 0
3594 },
3595 penalty_repeat: repetition_penalty,
3596 penalty_freq: frequency_penalty,
3597 penalty_present: presence_penalty,
3598 seed: seed.unwrap_or_else(fresh_seed),
3601 }
3602}
3603
3604fn fresh_seed() -> u64 {
3609 use std::sync::atomic::{AtomicU64, Ordering};
3610 static COUNTER: AtomicU64 = AtomicU64::new(0);
3611 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
3612 let nanos = std::time::SystemTime::now()
3613 .duration_since(std::time::UNIX_EPOCH)
3614 .map(|d| d.as_nanos() as u64)
3615 .unwrap_or(0);
3616 let mut z = nanos
3617 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
3618 .wrapping_add(0x9E3779B97F4A7C15);
3619 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
3620 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
3621 z ^= z >> 31;
3622 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
3625}
3626
3627fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
3632 for (param, present, why) in fields {
3633 if *present {
3634 return Err((format!("{param} is not supported{why}"), param.to_string()));
3635 }
3636 }
3637 Ok(())
3638}
3639
3640#[derive(PartialEq)]
3641enum ToolChoice {
3642 Auto,
3643 None,
3644}
3645
3646fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
3647 match v {
3648 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
3649 Some(serde_json::Value::String(s)) => match s.as_str() {
3650 "auto" => Ok(ToolChoice::Auto),
3651 "none" => Ok(ToolChoice::None),
3652 "required" => Err("tool_choice \"required\" is not supported (no constrained \
3653 decoding); use \"auto\""
3654 .into()),
3655 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
3656 },
3657 Some(serde_json::Value::Object(_)) => {
3658 Err("named-function tool_choice is not supported; use \"auto\"".into())
3659 }
3660 Some(other) => Err(format!("bad tool_choice: {other}")),
3661 }
3662}
3663
3664fn parse_think(
3714 reasoning_effort: &Option<String>,
3715 reasoning: &Option<serde_json::Value>,
3716 vllm_switch: Option<bool>,
3717 suppress_switch: Option<bool>,
3718 default_effort: Option<&str>,
3719 dsv4: bool,
3720) -> Result<(ThinkMode, Option<String>, bool), String> {
3721 let mut effort = reasoning_effort.clone();
3722 let ReasoningObject {
3723 mut enabled,
3724 effort: object_effort,
3725 exclude,
3726 } = parse_reasoning_object(reasoning)?;
3727 if let Some(e) = object_effort {
3728 effort = Some(e);
3729 }
3730 match (enabled, vllm_switch) {
3735 (Some(a), Some(b)) if a != b => {
3736 return Err(format!(
3737 "contradictory reasoning switches: reasoning.enabled={a} and \
3738 enable_thinking={b} — send one"
3739 ));
3740 }
3741 (None, Some(b)) => enabled = Some(b),
3742 _ => {}
3743 }
3744 let suppress = match (exclude, suppress_switch) {
3756 (Some(true), _) | (_, Some(false)) => Some(false),
3757 _ => None,
3758 };
3759 match (enabled, suppress) {
3760 (Some(true), Some(false)) => {
3761 return Err(
3762 "contradictory reasoning switches: reasoning is enabled but \
3763 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
3764 on this server not delivering reasoning means not generating it, so send one"
3765 .into(),
3766 );
3767 }
3768 (None, Some(b)) => enabled = Some(b),
3769 _ => {}
3770 }
3771 let client_explicit = effort.is_some() || enabled.is_some();
3775 if effort.is_none() && enabled.is_none() {
3780 effort = default_effort.map(str::to_string);
3781 }
3782 let effort_arm = match effort.as_deref() {
3787 None => None,
3788 Some(raw) => {
3789 let level = canonical_effort_for(raw, dsv4).ok_or_else(|| {
3790 format!(
3791 "bad reasoning_effort {raw:?} \
3792 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
3793 highest level this model's template distinguishes)"
3794 )
3795 })?;
3796 Some(match level {
3797 "none" | "minimal" => (ThinkMode::NoThink, "low"),
3798 "low" => (ThinkMode::Think, "low"),
3799 "medium" => (ThinkMode::Think, "medium"),
3800 "max" => (ThinkMode::Think, "max"),
3801 _ => (ThinkMode::Think, "high"),
3802 })
3803 }
3804 };
3805 let (think, level) = match (enabled, effort_arm) {
3806 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
3809 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
3810 (None, Some((think, level))) => (think, Some(level.to_string())),
3811 (None, None) => (ThinkMode::Default, None),
3812 };
3813 Ok((think, level, client_explicit))
3814}
3815
3816struct ReasoningObject {
3818 enabled: Option<bool>,
3819 effort: Option<String>,
3820 exclude: Option<bool>,
3821}
3822
3823fn parse_reasoning_object(
3842 reasoning: &Option<serde_json::Value>,
3843) -> Result<ReasoningObject, String> {
3844 let mut out = ReasoningObject {
3845 enabled: None,
3846 effort: None,
3847 exclude: None,
3848 };
3849 let Some(v) = reasoning else { return Ok(out) };
3850 let obj = match v {
3851 serde_json::Value::Null => return Ok(out),
3852 serde_json::Value::Object(obj) => obj,
3853 _ => return Err("reasoning must be an object".into()),
3854 };
3855 for (key, value) in obj {
3856 match key.as_str() {
3864 "enabled" => {
3865 if !value.is_null() {
3866 out.enabled = Some(
3867 value
3868 .as_bool()
3869 .ok_or("reasoning.enabled must be true or false")?,
3870 );
3871 }
3872 }
3873 "exclude" => {
3874 if !value.is_null() {
3875 out.exclude = Some(
3876 value
3877 .as_bool()
3878 .ok_or("reasoning.exclude must be true or false")?,
3879 );
3880 }
3881 }
3882 "effort" => {
3883 if !value.is_null() {
3884 out.effort = Some(
3885 value
3886 .as_str()
3887 .ok_or("reasoning.effort must be a string")?
3888 .to_string(),
3889 );
3890 }
3891 }
3892 "max_tokens" => {
3893 return Err(
3894 "reasoning.max_tokens is not supported by this server: reasoning tokens \
3895 are output tokens here, and max_tokens is the ONE output budget covering \
3896 reasoning and content together — there is no separate reasoning budget to \
3897 spend against, so honouring this field is impossible rather than merely \
3898 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
3899 reasoning.enabled:false) to spend less of it on reasoning"
3900 .into(),
3901 );
3902 }
3903 other => {
3904 return Err(format!(
3905 "reasoning.{other} is not a field this server implements (it would change \
3906 nothing about the request); the supported keys are enabled, effort and \
3907 exclude"
3908 ));
3909 }
3910 }
3911 }
3912 Ok(out)
3913}
3914
3915fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
3936 let Some(v) = kwargs else { return Ok(None) };
3937 let obj = match v {
3938 serde_json::Value::Null => return Ok(None),
3939 serde_json::Value::Object(obj) => obj,
3940 _ => return Err("chat_template_kwargs must be an object".into()),
3941 };
3942 let mut switch = None;
3943 for (key, value) in obj {
3944 match key.as_str() {
3945 "enable_thinking" => {
3946 switch = Some(
3947 value
3948 .as_bool()
3949 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
3950 );
3951 }
3952 "preserve_thinking" => {
3953 let preserve = value
3954 .as_bool()
3955 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
3956 if !preserve {
3957 return Err(
3958 "chat_template_kwargs.preserve_thinking:false is not supported by this \
3959 server: the renderer implements the vendor DEFAULT (replay every prior \
3960 assistant turn's <think> block, empty when no reasoning was sent) but \
3961 not the strip arm — serving replay bytes under a strip request would \
3962 misdescribe the prompt. Omit the flag or send true"
3963 .into(),
3964 );
3965 }
3966 }
3968 other => {
3969 return Err(format!(
3970 "chat_template_kwargs.{other} is not supported by this server's \
3971 template renderer (it would change nothing about the prompt); the only \
3972 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
3973 refuses in both directions — see its own message)"
3974 ));
3975 }
3976 }
3977 }
3978 Ok(switch)
3979}
3980
3981fn resolve_vllm_think_switch(
3985 enable_thinking: Option<bool>,
3986 kwargs: &Option<serde_json::Value>,
3987) -> Result<Option<bool>, String> {
3988 let from_kwargs = parse_template_kwargs(kwargs)?;
3989 match (enable_thinking, from_kwargs) {
3990 (Some(a), Some(b)) if a != b => Err(format!(
3991 "contradictory reasoning switches: enable_thinking={a} and \
3992 chat_template_kwargs.enable_thinking={b} — send one"
3993 )),
3994 (Some(a), _) => Ok(Some(a)),
3995 (None, b) => Ok(b),
3996 }
3997}
3998
3999pub(crate) fn canonical_effort_for(value: &str, dsv4_max: bool) -> Option<&'static str> {
4019 match value {
4020 "none" => Some("none"),
4021 "minimal" => Some("minimal"),
4022 "low" => Some("low"),
4023 "medium" => Some("medium"),
4024 "high" => Some("high"),
4025 "xhigh" | "max" | "ultra" => Some(if dsv4_max { "max" } else { "high" }),
4026 _ => None,
4027 }
4028}
4029
4030pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4033 canonical_effort_for(value, false)
4034}
4035
4036fn json_to_val(v: &serde_json::Value) -> chat::Val {
4039 match v {
4040 serde_json::Value::Null => chat::Val::Null,
4041 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4042 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4043 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4044 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4045 serde_json::Value::Object(o) => chat::Val::Obj(
4048 o.iter()
4049 .map(|(k, val)| (k.clone(), json_to_val(val)))
4050 .collect(),
4051 ),
4052 }
4053}
4054
4055#[allow(clippy::type_complexity)]
4059fn prepare_tools(
4060 tools: &[serde_json::Value],
4061) -> Result<
4062 (
4063 Vec<String>,
4064 Vec<chat::Val>,
4065 HashMap<String, HashMap<String, String>>,
4066 ),
4067 String,
4068> {
4069 let mut tools_json = Vec::with_capacity(tools.len());
4070 let mut tools_struct = Vec::with_capacity(tools.len());
4071 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4072 for t in tools {
4073 let f = t
4074 .get("function")
4075 .ok_or("each tool needs a function object")?;
4076 let name = f
4077 .get("name")
4078 .and_then(|n| n.as_str())
4079 .ok_or("each tool needs function.name")?;
4080 let mut params: HashMap<String, String> = HashMap::new();
4081 if let Some(props) = f
4082 .get("parameters")
4083 .and_then(|p| p.get("properties"))
4084 .and_then(|p| p.as_object())
4085 {
4086 for (p, def) in props {
4087 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4088 params.insert(p.clone(), ty.to_string());
4089 }
4090 }
4091 }
4092 schemas.insert(name.to_string(), params);
4093 tools_json.push(pyjson_str(t));
4094 tools_struct.push(json_to_val(f));
4096 }
4097 Ok((tools_json, tools_struct, schemas))
4098}
4099
4100fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4105 let parsed: serde_json::Value = match &tc.function.arguments {
4106 serde_json::Value::Null => json!({}),
4107 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4108 serde_json::Value::String(s) => serde_json::from_str(s)
4109 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4110 v @ serde_json::Value::Object(_) => v.clone(),
4111 _ => return Err("tool_calls arguments must be a JSON object".into()),
4112 };
4113 let obj = parsed
4114 .as_object()
4115 .ok_or("tool_calls arguments must decode to a JSON object")?;
4116 let params = obj
4117 .iter()
4118 .map(|(k, v)| {
4119 let rendered = match v {
4120 serde_json::Value::String(s) => s.clone(),
4121 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4122 scalar => scalar.to_string(),
4123 };
4124 (k.clone(), rendered)
4125 })
4126 .collect();
4127 let args = obj
4130 .iter()
4131 .map(|(k, v)| (k.clone(), json_to_val(v)))
4132 .collect();
4133 Ok(TmplToolCall {
4134 name: tc.function.name.clone(),
4135 params,
4136 args,
4137 id: tc.id.clone(),
4138 })
4139}
4140
4141fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4143 json!({ "id": c.id, "type": "function",
4144 "function": { "name": c.name, "arguments": c.arguments } })
4145}
4146
4147#[tokio::main]
4151pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4152 serve_with(ServerWiring::stock()).await
4153}
4154
4155enum MeteringWiring {
4157 Stock,
4161 Custom(metering::MeteringFactory),
4166}
4167
4168pub struct ServerWiring {
4172 metering: MeteringWiring,
4173 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
4176 claimed_env: Vec<&'static str>,
4180}
4181
4182impl ServerWiring {
4183 pub fn stock() -> Self {
4185 ServerWiring {
4186 metering: MeteringWiring::Stock,
4187 on_ready: None,
4188 claimed_env: Vec::new(),
4189 }
4190 }
4191
4192 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
4195 ServerWiring {
4196 metering: MeteringWiring::Custom(factory),
4197 on_ready: None,
4198 claimed_env: Vec::new(),
4199 }
4200 }
4201
4202 pub fn claiming(mut self, var: &'static str) -> Self {
4206 self.claimed_env.push(var);
4207 self
4208 }
4209
4210 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
4211 self.on_ready = Some(Box::new(hook));
4212 self
4213 }
4214}
4215
4216pub struct RuntimeHandles {
4219 pub trim: TrimHandle,
4220 pub shutdown: tokio::sync::watch::Receiver<bool>,
4225}
4226
4227#[derive(Clone)]
4230pub struct TrimHandle {
4231 cmd_tx: Sender<Cmd>,
4232}
4233
4234impl TrimHandle {
4235 pub async fn trim(&self) -> Result<serde_json::Value, String> {
4237 let (tx, rx) = tokio::sync::oneshot::channel();
4238 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
4239 return Err("worker is down".into());
4240 }
4241 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4242 Ok(Ok(report)) => Ok(json!(report)),
4243 _ => Err("worker did not answer the trim within 30s".into()),
4244 }
4245 }
4246}
4247
4248pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
4249 let args: Vec<String> = std::env::args().skip(1).collect();
4252 if let Some(code) = auth::run_cli(&args) {
4253 std::process::exit(code);
4254 }
4255 auth::init_from_env();
4258 let api_auth = match ApiAuth::from_env() {
4259 Ok(auth) => auth,
4260 Err(err) => {
4261 eprintln!("[server] FATAL: {err}");
4262 std::process::exit(1);
4263 }
4264 };
4265 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
4266 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
4267 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
4268 Ok(resolved) => resolved,
4269 Err(err) => {
4270 eprintln!("[server] FATAL: {err}");
4271 std::process::exit(1);
4272 }
4273 };
4274 if !bind_loopback && !api_auth.configured() && !allow_open_bind {
4275 let message = format!(
4276 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
4277 );
4278 eprintln!("[server] FATAL: {message}");
4279 std::process::exit(1);
4280 }
4281 if !bind_loopback && !api_auth.configured() {
4282 eprintln!(
4283 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
4284 metrics remain bearer-protected"
4285 );
4286 }
4287 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
4288 Ok(token) if token.is_empty() => {
4289 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
4290 std::process::exit(1);
4291 }
4292 Ok(token) => Some(token),
4293 Err(std::env::VarError::NotPresent) => None,
4294 Err(std::env::VarError::NotUnicode(_)) => {
4295 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
4296 std::process::exit(1);
4297 }
4298 };
4299 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
4300
4301 let models = parse_models_config();
4302 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
4303 Ok(loaded) => loaded,
4304 Err(err) => {
4305 eprintln!("[server] FATAL: {err}");
4306 std::process::exit(1);
4307 }
4308 };
4309 let metering_obj: Option<Arc<dyn metering::Metering>> = {
4315 let factory = match wiring.metering {
4316 MeteringWiring::Stock => None,
4317 MeteringWiring::Custom(factory) => Some(factory),
4318 };
4319 for deployment_only in [
4320 "MEMRA_REQUEST_LEDGER",
4321 "MEMRA_TENANT_BUDGETS",
4322 "MEMRA_ADMIN_ADDR",
4323 "MEMRA_ADMIN_TOKEN_FILE",
4324 "MEMRA_CAPTURE_DIR",
4325 ] {
4326 if std::env::var_os(deployment_only).is_some()
4327 && !wiring.claimed_env.contains(&deployment_only)
4328 {
4329 eprintln!(
4330 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
4331 build ships no accounting/admin/capture. Wire a Metering implementation \
4332 through ServerWiring and claim the vars it consumes."
4333 );
4334 std::process::exit(1);
4335 }
4336 }
4337 match factory {
4338 None => None,
4339 Some(factory) => {
4340 let model_ids: Vec<String> =
4341 models.iter().map(|(name, _, _)| name.clone()).collect();
4342 match factory(&metering::MeteringInit { models: &model_ids }) {
4343 Ok(metering_obj) => metering_obj,
4344 Err(err) => {
4345 eprintln!("[server] FATAL: metering wiring: {err}");
4346 std::process::exit(1);
4347 }
4348 }
4349 }
4350 }
4351 };
4352 let budget_tokenizers = if metering_obj
4353 .as_ref()
4354 .is_some_and(|manager| manager.enforces_limits())
4355 {
4356 match load_budget_tokenizers(&models) {
4357 Ok(tokenizers) => Some(tokenizers),
4358 Err(err) => {
4359 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
4360 std::process::exit(1);
4361 }
4362 }
4363 } else {
4364 None
4365 };
4366 eprintln!("[server] starting; models config = {models:?}");
4367
4368 let health_state = health::WorkerHealth::new();
4373 health::spawn_gpu_watch(health_state.clone());
4377 health::spawn_sd_watchdog(health_state.clone());
4378
4379 let (cmd_tx, model_names, caps, metrics, worker_thread) =
4381 match worker::spawn(models, health_state.clone()) {
4382 Ok(v) => v,
4383 Err(err) => {
4384 eprintln!("[server] FATAL: worker init failed: {err}");
4385 health_state.mark_dead(format!("worker init failed: {err}"));
4386 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
4387 std::process::exit(1);
4388 }
4389 };
4390 eprintln!("[server] worker ready; serving models: {model_names:?}");
4391
4392 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
4399 if let Some(on_ready) = wiring.on_ready {
4400 on_ready(RuntimeHandles {
4401 trim: TrimHandle {
4402 cmd_tx: cmd_tx.clone(),
4403 },
4404 shutdown: drain_shutdown_rx.clone(),
4405 });
4406 }
4407
4408 let bg_handle = darklane::spawn_from_env(health_state.clone());
4411 let bg_state = bg_handle.as_ref().map(|h| {
4412 let mode = darklane::BgConfig::from_env()
4413 .map(|c| c.yield_mode.as_str())
4414 .unwrap_or("stop");
4415 (h.state.clone(), mode)
4416 });
4417
4418 let state = AppState {
4419 cmd_tx,
4420 models: model_names,
4421 caps,
4422 openrouter_metadata: Arc::new(openrouter_metadata),
4423 provider_metadata: Arc::new(provider_metadata),
4424 metering: metering_obj,
4425 budget_tokenizers,
4426 api_auth,
4427 metrics_auth,
4428 metrics,
4429 started: std::time::SystemTime::now()
4430 .duration_since(std::time::UNIX_EPOCH)
4431 .map(|d| d.as_secs())
4432 .unwrap_or(0),
4433 inflight: Arc::new(Default::default()),
4434 tenant_inflight: Arc::new(Default::default()),
4435 health: health_state.clone(),
4436 bg: bg_state,
4437 };
4438 let inflight_handle = state.inflight.clone();
4439 let drain_metering = state.metering.clone();
4442 worker::register_http_inflight(state.inflight.clone());
4447 let app = Router::new()
4448 .route("/health", get(health_live))
4453 .route("/livez", get(health_live))
4454 .route("/readyz", get(health_ready))
4455 .route("/models", get(list_models))
4456 .route("/v1/models", get(list_models_v1))
4457 .route("/v1/auth/check", get(auth_check))
4458 .route("/v1/completions", post(completions))
4459 .route("/v1/embeddings", post(embed_api::embeddings))
4460 .route("/v1/rerank", post(embed_api::rerank))
4461 .route("/v1/chat/completions", post(chat_completions))
4462 .route("/v1/messages", post(anthropic::messages))
4466 .route("/v1/responses", post(responses_api::responses))
4467 .route("/metrics", get(get_metrics))
4468 .route("/yield/metrics", get(yield_metrics))
4469 .with_state(state.clone());
4470 let app = apply_body_limit(app);
4473 let app = app.layer(middleware::from_fn_with_state(
4476 state,
4477 authenticate_inference_before_body,
4478 ));
4479 let app = if ttft::enabled() {
4480 app.layer(middleware::from_fn(ttft_request_start))
4481 } else {
4482 app
4483 };
4484
4485 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
4486 eprintln!("[server] listening on http://{bind_addr}");
4487 drop(drain_shutdown_rx);
4488 health::sd_notify("READY=1\nSTATUS=serving");
4492 let inflight = inflight_handle;
4499 let signal_admin_shutdown = drain_shutdown_tx.clone();
4500 let serve_result = axum::serve(listener, app)
4501 .with_graceful_shutdown(async move {
4502 let mut sigterm =
4503 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
4504 Ok(s) => s,
4505 Err(err) => {
4506 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
4507 std::future::pending::<()>().await;
4508 unreachable!()
4509 }
4510 };
4511 sigterm.recv().await;
4512 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
4513 let _ = signal_admin_shutdown.send(true);
4514 health::sd_notify(&format!(
4518 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
4519 (drain_deadline_s() + 5) * 1_000_000
4520 ));
4521 let n: usize = inflight
4522 .iter()
4523 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4524 .sum();
4525 eprintln!(
4526 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
4527 drain_deadline_s()
4528 );
4529 let deadline = std::time::Duration::from_secs(drain_deadline_s());
4530 let t0 = std::time::Instant::now();
4531 loop {
4532 let n: usize = inflight
4533 .iter()
4534 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4535 .sum();
4536 if n == 0 {
4537 eprintln!(
4538 "[server] drain complete in {:.1}s; exiting",
4539 t0.elapsed().as_secs_f64()
4540 );
4541 break;
4542 }
4543 if t0.elapsed() >= deadline {
4544 eprintln!(
4545 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
4546 drain_deadline_s()
4547 );
4548 if let Some(metering) = drain_metering.as_ref() {
4555 metering.drain_kill();
4556 }
4557 break;
4558 }
4559 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4560 }
4561 })
4562 .await;
4563 let _ = drain_shutdown_tx.send(true);
4566 serve_result?;
4567 if let Some(h) = bg_handle {
4571 h.shutdown();
4572 }
4573 worker_thread.join().map_err(|_| {
4582 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
4583 })?;
4584 eprintln!("[server] GPU worker shutdown complete");
4585 Ok(())
4586}
4587
4588fn validate_model_path(path: &str) -> Result<(), String> {
4594 let p = std::path::Path::new(path);
4595 if !p.exists() {
4596 return Err(format!("model path {path:?} does not exist"));
4597 }
4598 if p.is_file() {
4599 return Ok(()); }
4601 if p.join("manifest.json").exists() {
4602 return Ok(()); }
4604 let has_st =
4605 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
4606 if !has_st {
4607 return Err(format!(
4608 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
4609 model.safetensors.index.json + config.json (HF safetensors dir), or \
4610 manifest.json (memra repack dir)"
4611 ));
4612 }
4613 if !p.join("config.json").exists() {
4614 return Err(format!(
4615 "model dir {path:?} has safetensors weights but no config.json"
4616 ));
4617 }
4618 Ok(())
4619}
4620
4621fn parse_models_config() -> Vec<(String, String, Option<String>)> {
4629 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
4630 let mut out = Vec::new();
4631 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
4632 if let Some((name, path)) = entry.split_once('=') {
4633 let (mpath, dpath) = match path.trim().split_once('+') {
4636 Some((m, d)) => (m.trim(), Some(d.trim())),
4637 None => (path.trim(), None),
4638 };
4639 let resolve = |p: &str| {
4640 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
4641 eprintln!("[server] FATAL: model {name:?}: {err}");
4642 std::process::exit(1);
4643 })
4644 };
4645 let mpath = resolve(mpath);
4646 if let Err(err) = validate_model_path(&mpath) {
4647 eprintln!("[server] FATAL: model {name:?}: {err}");
4648 std::process::exit(1);
4649 }
4650 let dpath = dpath.map(|d| {
4660 let d = resolve(d);
4661 let p = std::path::Path::new(&d);
4662 if !p.exists() {
4663 eprintln!(
4664 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
4665 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
4666 rather than serving plain decode under a config that asked \
4667 for speculative decoding."
4668 );
4669 std::process::exit(1);
4670 }
4671 if !p.is_file() {
4672 eprintln!(
4673 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
4674 file — a '+draft' attach must be a NextN/MTP GGUF file."
4675 );
4676 std::process::exit(1);
4677 }
4678 d
4679 });
4680 out.push((name.trim().to_string(), mpath, dpath));
4681 } else {
4682 eprintln!(
4683 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
4684 );
4685 }
4686 }
4687 if !out.is_empty() {
4688 return out;
4689 }
4690 }
4691 vec![
4693 (
4694 "main".into(),
4695 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
4696 None,
4697 ),
4698 (
4699 "judge".into(),
4700 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
4701 None,
4702 ),
4703 ]
4704}
4705
4706fn load_budget_tokenizers(
4707 models: &[(String, String, Option<String>)],
4708) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
4709 let mut tokenizers = HashMap::new();
4710 for (alias, path, _) in models {
4711 let path = std::path::Path::new(path);
4712 let tokenizer = if path.is_dir() {
4713 let tokenizer_dir = if path.join("manifest.json").exists() {
4714 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
4715 format!("model {alias:?}: open repack tokenizer source: {err}")
4716 })?;
4717 repack
4718 .source_dir()
4719 .filter(|source| source.join("tokenizer.json").exists())
4720 .unwrap_or(path)
4721 .to_path_buf()
4722 } else {
4723 path.to_path_buf()
4724 };
4725 Tokenizer::from_hf_dir(&tokenizer_dir)
4726 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4727 } else {
4728 let gguf = memra_gguf::GgufFile::open(path)
4729 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
4730 Tokenizer::from_gguf(&gguf)
4731 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4732 };
4733 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
4734 }
4735 Ok(Arc::new(tokenizers))
4736}
4737
4738fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4740 let s = st.health.snapshot();
4741 let mut v = json!({
4742 "status": status,
4743 "models": *st.models,
4744 "worker": {
4745 "phase": health::phase_name(s.phase),
4746 "beat_age_ms": s.beat_age_ms,
4747 "tick_max_ms": s.tick_max_ms,
4748 "stall_threshold_ms": s.stall_threshold_ms,
4749 "generation": s.generation,
4750 "xid_warnings": s.xid_warns,
4751 },
4752 });
4753 if let Some(d) = detail {
4754 v["detail"] = json!(d);
4755 }
4756 v
4757}
4758
4759fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4763 let mut v = health_payload(st, status, detail);
4764 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
4765 v
4766}
4767
4768async fn auth_check() -> impl IntoResponse {
4772 StatusCode::NO_CONTENT
4773}
4774
4775async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
4791 if draining() {
4792 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
4795 }
4796 match st.health.live() {
4797 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
4798 Err(why) => retry_contract_response(
4799 (
4800 StatusCode::SERVICE_UNAVAILABLE,
4801 Json(health_payload(&st, "unhealthy", Some(&why))),
4802 )
4803 .into_response(),
4804 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
4805 ),
4806 }
4807}
4808
4809async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
4820 let is_draining = draining();
4821 match st.health.ready(is_draining) {
4822 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
4823 Err(why) => retry_contract_response(
4824 (
4825 StatusCode::SERVICE_UNAVAILABLE,
4826 Json(readiness_payload(&st, "not_ready", Some(&why))),
4827 )
4828 .into_response(),
4829 Some(if is_draining {
4830 drain_deadline_s()
4831 } else {
4832 worker::WORKER_RESPAWN_BACKOFF_BASE_S
4833 }),
4834 ),
4835 }
4836}
4837
4838#[derive(Clone, Copy)]
4839struct DualPpMetricsSnapshot {
4840 stage_ns: [u64; 4],
4841 stage_samples: [usize; 4],
4842 dropped_timing_samples: usize,
4843 overlaps: usize,
4844 slot_pairs: usize,
4845 slot_uses: [usize; 2],
4846 slot_collisions: usize,
4847}
4848
4849impl DualPpMetricsSnapshot {
4850 fn current() -> Self {
4851 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
4852 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
4853 Self {
4854 stage_ns,
4855 stage_samples,
4856 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
4857 overlaps: memra_engine::pp::dual_pp_overlaps(),
4858 slot_pairs,
4859 slot_uses,
4860 slot_collisions,
4861 }
4862 }
4863
4864 fn populated(self) -> bool {
4865 self.stage_samples.iter().any(|&n| n > 0)
4866 || self.dropped_timing_samples > 0
4867 || self.slot_pairs > 0
4868 || self.slot_collisions > 0
4869 }
4870}
4871
4872fn insert_dual_pp_metrics(
4873 body: &mut serde_json::Value,
4874 metrics_scope: &MetricsScope,
4875 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
4876) {
4877 if !metrics_scope.operator() {
4880 return;
4881 }
4882 let snapshot = snapshot();
4883 if !snapshot.populated() {
4884 return;
4885 }
4886 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
4887 .iter()
4888 .enumerate()
4889 .map(|(i, name)| {
4890 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
4891 (
4892 name.to_string(),
4893 json!({
4894 "samples": snapshot.stage_samples[i],
4895 "total_ms": total_ms,
4896 "mean_ms": if snapshot.stage_samples[i] > 0 {
4897 total_ms / snapshot.stage_samples[i] as f64
4898 } else { 0.0 },
4899 }),
4900 )
4901 })
4902 .collect();
4903 body["dual_pp"] = json!({
4904 "overlaps": snapshot.overlaps,
4905 "slot_pairs": snapshot.slot_pairs,
4906 "slot_uses": snapshot.slot_uses,
4907 "slot_collisions": snapshot.slot_collisions,
4908 "cuda_event_spans": timings,
4909 "dropped_timing_samples": snapshot.dropped_timing_samples,
4910 });
4911}
4912
4913fn insert_spec_acceptance_metrics(
4914 body: &mut serde_json::Value,
4915 metrics_scope: &MetricsScope,
4916 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
4917) {
4918 if !metrics_scope.operator() {
4921 return;
4922 }
4923 let snapshot = snapshot();
4924 if snapshot.is_empty() {
4925 return;
4926 }
4927
4928 let mut tau = serde_json::Map::new();
4929 let mut by_position = serde_json::Map::new();
4930 for (model, telemetry) in snapshot {
4931 if telemetry.rounds == 0 {
4932 continue;
4933 }
4934 let n_pos = telemetry
4935 .pos_drafted
4936 .iter()
4937 .rposition(|&n| n > 0)
4938 .map_or(0, |position| position + 1);
4939 tau.insert(model.clone(), json!(telemetry.tau()));
4940 by_position.insert(
4941 model,
4942 json!({
4943 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
4944 "rounds": telemetry.rounds,
4945 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
4946 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
4947 "accept_rate": (0..n_pos).map(|position| {
4948 let offered = telemetry.pos_drafted[position];
4949 if offered > 0 {
4950 telemetry.pos_accepted[position] as f64 / offered as f64
4951 } else {
4952 0.0
4953 }
4954 }).collect::<Vec<f64>>(),
4955 }),
4956 );
4957 }
4958 if !tau.is_empty() {
4959 body["spec_tau"] = serde_json::Value::Object(tau);
4960 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
4961 }
4962}
4963
4964fn insert_peer_probe_metrics(
4965 body: &mut serde_json::Value,
4966 metrics_scope: &MetricsScope,
4967 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
4968) {
4969 if !metrics_scope.operator() {
4972 return;
4973 }
4974 let snapshot = snapshot();
4975 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
4976 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
4977 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
4978 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
4979 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
4980 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
4981 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
4982}
4983
4984async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
4986 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
4987 Ok(scope) => scope,
4988 Err(response) => return response,
4989 };
4990 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
4991 let mut body = if metrics_scope.process_wide() {
4995 json!({
4996 "admitted": m.admitted,
4997 "completed": m.completed,
4998 "tokens_out": m.tokens_out,
4999 "step_p50_ms": m.step_p50_ms,
5000 "step_p99_ms": m.step_p99_ms,
5001 "prompt_tokens_in": m.prompt_tokens_in,
5003 "cached_tokens_in": m.cached_tokens_in,
5004 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5007 "admission_session_defers": m.admission_session_defers,
5010 "admission_vram_defers": m.admission_vram_defers,
5011 "step_oom_parks": m.step_oom_parks,
5012 "continuation_pool_hits": m.continuation_pool_hits,
5013 "continuation_pool_evictions": m.continuation_pool_evictions,
5014 "plain_affinity_rewinds": m.plain_affinity_rewinds,
5015 "served_dspark": m.served_dspark,
5016 "served_spec": m.served_spec,
5017 "served_plain": m.served_plain,
5018 "spec_pool_hits": m.spec_pool_hits,
5019 "spec_pool_misses": m.spec_pool_misses,
5020 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
5021 "spec_pool_evictions": m.spec_pool_evictions,
5022 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
5025 })
5026 } else {
5027 json!({})
5028 };
5029 if metrics_scope.operator() {
5033 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
5034 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
5035 body["budget_source_reload_consecutive"] =
5036 json!(budget_health.source_reload_consecutive);
5037 body["budget_source_available"] = json!(budget_health.source_available);
5038 }
5039 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
5041 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
5042 } else {
5043 0.0
5044 });
5045 body["prefix_cache_hits"] = json!(m.prefix_hits);
5046 body["prefix_cache_misses"] = json!(m.prefix_misses);
5047 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
5048 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
5049 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
5050 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
5051 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
5052 body["lcp_histogram"] = json!({
5055 "edges": worker::LCP_HIST_EDGES.to_vec(),
5056 "counts": m.lcp_hist.to_vec(),
5057 });
5058 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
5062 body["prefix_cache_entries"] = json!(m.prefix_entries);
5063 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
5064 body["active_sessions"] = json!(m.active_sessions);
5065 body["queued_requests"] = json!(m.queued_requests);
5066 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
5067 body["spec_pool_entries"] = json!(m.spec_pool_entries);
5068 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
5069 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
5070 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
5071 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
5072 if !m.constraint_compiler_fail_closed.is_empty() {
5073 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
5074 m.constraint_compiler_fail_closed
5075 .iter()
5076 .map(|(model, gauge)| {
5077 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
5078 (model.clone(), json!(value))
5079 })
5080 .collect(),
5081 );
5082 }
5083 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
5084 }
5085 if !m.ns_tokens.is_empty() {
5090 let tenants: serde_json::Map<String, serde_json::Value> = m
5091 .ns_tokens
5092 .iter()
5093 .filter(|(ns, _)| metrics_scope.includes(ns))
5094 .map(|(ns, [p, c])| {
5095 (
5096 ns.clone(),
5097 json!({
5098 "prompt_tokens_in": p,
5099 "cached_tokens_in": c,
5100 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
5101 }),
5102 )
5103 })
5104 .collect();
5105 if !tenants.is_empty() {
5106 body["tenants"] = serde_json::Value::Object(tenants);
5107 }
5108 }
5109 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
5110 .adsd_suspect_total
5111 .iter()
5112 .filter(|(tenant, _)| metrics_scope.includes(tenant))
5113 .map(|(tenant, total)| (tenant.clone(), json!(total)))
5114 .collect();
5115 if !adsd_suspect_total.is_empty() {
5116 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
5117 }
5118 if metrics_scope.operator() {
5120 if let Some((bg, mode)) = &st.bg {
5121 body["bg"] = bg.to_json(mode);
5122 }
5123 }
5124 if metrics_scope.operator() {
5131 let spec: serde_json::Map<String, serde_json::Value> = m
5132 .spec
5133 .iter()
5134 .map(|(model, t)| {
5135 let n_pos = t
5136 .pos_drafted
5137 .iter()
5138 .rposition(|&d| d > 0)
5139 .map_or(0, |p| p + 1);
5140 (
5141 model.clone(),
5142 json!({
5143 "rounds": t.rounds,
5144 "drafted": t.drafted,
5145 "accepted": t.accepted,
5146 "acceptance_rate": if t.drafted > 0 {
5147 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
5148 "tokens_per_round": if t.rounds > 0 {
5149 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
5150 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
5151 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
5152 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
5153 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
5154 .collect::<Vec<f64>>(),
5155 }),
5156 )
5157 })
5158 .collect();
5159 if !spec.is_empty() {
5160 body["spec"] = serde_json::Value::Object(spec);
5161 }
5162 }
5163 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
5164 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
5165 insert_peer_probe_metrics(
5166 &mut body,
5167 &metrics_scope,
5168 memra_engine::pp::peer_probe_metrics,
5169 );
5170 Json(body).into_response()
5171}
5172
5173#[derive(Debug, Default, Deserialize)]
5174struct ModelsQuery {
5175 #[serde(default)]
5176 schema: Option<String>,
5177}
5178
5179fn models_openai_body(models: &[String]) -> serde_json::Value {
5180 let data: Vec<_> = models
5181 .iter()
5182 .map(|m| json!({ "id": m, "object": "model" }))
5183 .collect();
5184 json!({ "object": "list", "data": data })
5185}
5186
5187fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
5192 match metadata.and_then(|m| m.surface.as_deref()) {
5193 Some("embedding") => "embedding",
5194 Some("rerank") => "rerank",
5195 _ => "chat",
5196 }
5197}
5198
5199fn openrouter_supported_parameters(
5200 caps: Option<&ModelCaps>,
5201 max_output_length: Option<u64>,
5202 is_chat: bool,
5203) -> serde_json::Value {
5204 let mut parameters = serde_json::Map::new();
5205 if !is_chat {
5213 return serde_json::Value::Object(parameters);
5214 }
5215 for name in [
5216 "temperature",
5217 "top_p",
5218 "min_p",
5219 "frequency_penalty",
5220 "presence_penalty",
5221 "repetition_penalty",
5222 "stop",
5223 ] {
5224 parameters.insert(name.into(), json!({ "type": "unknown" }));
5225 }
5226 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
5227 parameters.insert(
5228 "seed".into(),
5229 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
5230 );
5231 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
5232 if let Some(max) = max_output_length {
5233 max_tokens["max"] = json!(max);
5234 }
5235 parameters.insert("max_tokens".into(), max_tokens);
5236 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
5237 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
5238 if is_chat && caps.is_some_and(|c| c.tools_branch) {
5239 parameters.insert("tools".into(), json!({ "type": "boolean" }));
5240 parameters.insert(
5241 "tool_choice".into(),
5242 json!({ "type": "enum", "values": ["auto", "none"] }),
5243 );
5244 }
5245 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5246 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
5247 }
5248 serde_json::Value::Object(parameters)
5249}
5250
5251fn model_entry_openrouter(
5252 name: &str,
5253 caps: Option<&ModelCaps>,
5254 metadata: Option<&OpenRouterModelMetadata>,
5255) -> serde_json::Value {
5256 let empty = OpenRouterModelMetadata::default();
5257 let metadata = metadata.unwrap_or(&empty);
5258 let context_length = caps
5259 .map(|c| c.context_length as u64)
5260 .filter(|&v| v > 0 && v <= JSON_SAFE_INTEGER_MAX);
5261 let tokenizer = caps
5262 .map(|c| c.tokenizer.as_str())
5263 .filter(|tokenizer| !tokenizer.is_empty());
5264
5265 let mut input = serde_json::Map::new();
5266 input.insert("type".into(), json!("text"));
5267 let mut supported_inputs = serde_json::Map::new();
5268 if let Some(value) = context_length {
5269 supported_inputs.insert(
5270 "max_context_length".into(),
5271 json!({ "value": value, "unit": "token" }),
5272 );
5273 }
5274 if let Some(value) = metadata.max_prompt_length {
5275 supported_inputs.insert(
5276 "max_prompt_length".into(),
5277 json!({ "value": value, "unit": "token" }),
5278 );
5279 }
5280 if !supported_inputs.is_empty() {
5281 input.insert(
5282 "supported_inputs".into(),
5283 serde_json::Value::Object(supported_inputs),
5284 );
5285 }
5286 let mut input_pricing = Vec::new();
5287 for (kind, cost) in [
5288 ("prompt", metadata.pricing.prompt.as_deref()),
5289 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
5290 ("cache_write", metadata.pricing.cache_write.as_deref()),
5291 ] {
5292 if let Some(cost) = cost {
5293 input_pricing.push(json!({
5294 "type": kind,
5295 "unit": "token",
5296 "cost_usd": cost,
5297 }));
5298 }
5299 }
5300 if !input_pricing.is_empty() {
5301 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
5302 }
5303 let mut input_capacity = Vec::new();
5304 for (kind, value) in [
5305 ("prompt", metadata.capacity.prompt_tpm),
5306 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
5307 ] {
5308 if let Some(value) = value {
5309 input_capacity.push(json!({
5310 "type": kind,
5311 "unit": "token",
5312 "per": "minute",
5313 "value": value,
5314 }));
5315 }
5316 }
5317 if !input_capacity.is_empty() {
5318 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
5319 }
5320
5321 let or_surface = declared_surface(Some(metadata));
5322 let or_is_chat = or_surface == "chat";
5323 let mut output = serde_json::Map::new();
5324 output.insert(
5332 "type".into(),
5333 json!(match or_surface {
5334 "embedding" => "embeddings",
5335 "rerank" => "rerank",
5336 _ => "text",
5337 }),
5338 );
5339 output.insert(
5340 "supported_parameters".into(),
5341 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
5342 );
5343 if or_is_chat {
5347 output.insert("streaming".into(), json!(true));
5348 }
5349 if let Some(value) = metadata.max_output_length
5352 && or_is_chat
5353 {
5354 output.insert(
5355 "max_length".into(),
5356 json!({ "value": value, "unit": "token" }),
5357 );
5358 }
5359 let mut output_pricing = Vec::new();
5360 for (kind, cost) in [
5361 ("completion", metadata.pricing.completion.as_deref()),
5362 (
5363 "internal_reasoning",
5364 metadata.pricing.internal_reasoning.as_deref(),
5365 ),
5366 ] {
5367 if let Some(cost) = cost {
5368 output_pricing.push(json!({
5369 "type": kind,
5370 "unit": "token",
5371 "cost_usd": cost,
5372 }));
5373 }
5374 }
5375 if !output_pricing.is_empty() {
5376 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
5377 }
5378 let mut output_capacity = Vec::new();
5379 if let Some(value) = metadata.capacity.completion_tpm {
5380 output_capacity.push(json!({
5381 "type": "completion",
5382 "unit": "token",
5383 "per": "minute",
5384 "value": value,
5385 }));
5386 }
5387 if let Some(value) = metadata.capacity.concurrency {
5388 output_capacity.push(json!({
5389 "type": "concurrency",
5390 "unit": "request",
5391 "value": value,
5392 }));
5393 }
5394 if !output_capacity.is_empty() {
5395 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
5396 }
5397
5398 let mut entry = serde_json::Map::new();
5399 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
5400 entry.insert("id".into(), json!(name));
5401 entry.insert("name".into(), json!(name));
5402 if let Some(value) = metadata.hugging_face_id.as_deref() {
5403 entry.insert("hugging_face_id".into(), json!(value));
5404 }
5405 if let Some(value) = metadata.created {
5406 entry.insert("created".into(), json!(value));
5407 }
5408 if let Some(value) = metadata.quantization.as_deref() {
5409 entry.insert("quantization".into(), json!(value));
5410 }
5411 if let Some(value) = tokenizer {
5412 entry.insert("tokenizer".into(), json!(value));
5413 }
5414 if let Some(value) = metadata.description.as_deref() {
5415 entry.insert("description".into(), json!(value));
5416 }
5417 let mut input_modalities = vec![serde_json::Value::Object(input)];
5418 for m in &metadata.input_modalities {
5419 let mut extra = serde_json::Map::new();
5420 extra.insert("type".into(), json!(m));
5421 if let Some(cost) = metadata.pricing.prompt.as_deref() {
5422 extra.insert(
5424 "pricing".into(),
5425 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
5426 );
5427 }
5428 input_modalities.push(serde_json::Value::Object(extra));
5429 }
5430 entry.insert(
5431 "input_modalities".into(),
5432 serde_json::Value::Array(input_modalities),
5433 );
5434 entry.insert(
5435 "output_modalities".into(),
5436 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
5437 );
5438 if let Some(cost) = metadata.pricing.request.as_deref() {
5439 entry.insert(
5440 "pricing".into(),
5441 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
5442 );
5443 }
5444 if let Some(value) = metadata.capacity.request_rpm {
5445 entry.insert(
5446 "capacity".into(),
5447 json!([{
5448 "type": "request",
5449 "unit": "request",
5450 "per": "minute",
5451 "value": value,
5452 }]),
5453 );
5454 }
5455 if let Some(value) = metadata.is_ready {
5456 entry.insert("is_ready".into(), json!(value));
5457 }
5458 if let Some(value) = metadata.is_free {
5459 entry.insert("is_free".into(), json!(value));
5460 }
5461 if let Some(value) = metadata.discount_to_user {
5462 entry.insert("discount_to_user".into(), json!(value));
5463 }
5464 if let Some(value) = metadata.openrouter_slug.as_deref() {
5465 entry.insert("openrouter".into(), json!({ "slug": value }));
5466 }
5467 if !metadata.datacenters.is_empty() {
5468 entry.insert("datacenters".into(), json!(metadata.datacenters));
5469 }
5470 let mut compliance = serde_json::Map::new();
5471 if let Some(value) = metadata.zdr {
5472 compliance.insert("zdr".into(), json!(value));
5473 }
5474 if let Some(value) = metadata.hipaa {
5475 compliance.insert("hipaa".into(), json!(value));
5476 }
5477 if !compliance.is_empty() {
5478 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
5479 }
5480 serde_json::Value::Object(entry)
5481}
5482
5483fn models_openrouter_body(st: &AppState) -> serde_json::Value {
5484 let data: Vec<_> = st
5485 .models
5486 .iter()
5487 .map(|model| {
5488 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
5489 })
5490 .collect();
5491 json!({ "data": data })
5492}
5493
5494fn model_entry_openmodels(
5495 name: &str,
5496 caps: Option<&ModelCaps>,
5497 metadata: Option<&OpenRouterModelMetadata>,
5498) -> Result<serde_json::Value, String> {
5499 let metadata = metadata.ok_or_else(|| {
5500 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
5501 })?;
5502 let context_length = caps
5503 .map(|c| c.context_length as u64)
5504 .filter(|&value| value > 0 && value <= JSON_SAFE_INTEGER_MAX)
5505 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
5506 let created = metadata
5507 .created
5508 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
5509 let max_output_length = metadata
5510 .max_output_length
5511 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
5512 let prompt = metadata
5513 .pricing
5514 .prompt
5515 .as_deref()
5516 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
5517 let completion =
5518 metadata.pricing.completion.as_deref().ok_or_else(|| {
5519 format!("OpenModels feed requires pricing.completion for model {name:?}")
5520 })?;
5521 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
5522 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
5523 })?;
5524 let is_ready = metadata
5525 .is_ready
5526 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
5527 let is_free = metadata
5528 .is_free
5529 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
5530 let discount_to_user = metadata
5531 .discount_to_user
5532 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
5533
5534 let mut pricing = serde_json::Map::new();
5535 pricing.insert("prompt".into(), json!(prompt));
5536 pricing.insert("completion".into(), json!(completion));
5537 pricing.insert("input_cache_read".into(), json!(input_cache_read));
5538 if let Some(value) = metadata.pricing.request.as_deref() {
5539 pricing.insert("request".into(), json!(value));
5540 }
5541
5542 let om_surface = declared_surface(Some(metadata));
5543 let om_is_chat = om_surface == "chat";
5544 let mut supported_features = Vec::new();
5545 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
5546 supported_features.push("tool_calling");
5547 }
5548 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5549 supported_features.push("reasoning");
5550 }
5551
5552 let mut entry = serde_json::Map::new();
5553 entry.insert("id".into(), json!(name));
5554 entry.insert("name".into(), json!(name));
5555 entry.insert("created".into(), json!(created));
5556 entry.insert("input_modalities".into(), json!(["text"]));
5557 entry.insert(
5558 "output_modalities".into(),
5559 json!(match om_surface {
5560 "embedding" => ["embeddings"],
5561 "rerank" => ["rerank"],
5562 _ => ["text"],
5563 }),
5564 );
5565 entry.insert("context_length".into(), json!(context_length));
5566 entry.insert("max_output_length".into(), json!(max_output_length));
5567 entry.insert("currency".into(), json!("USD"));
5570 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
5571 entry.insert("supported_features".into(), json!(supported_features));
5572 entry.insert("is_ready".into(), json!(is_ready));
5573 entry.insert("is_free".into(), json!(is_free));
5574 entry.insert("discount_to_user".into(), json!(discount_to_user));
5575 Ok(serde_json::Value::Object(entry))
5576}
5577
5578fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
5579 let data: Result<Vec<_>, _> = st
5580 .models
5581 .iter()
5582 .map(|model| {
5583 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
5584 })
5585 .collect();
5586 Ok(json!({ "data": data? }))
5587}
5588
5589async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
5590 match query.schema.as_deref() {
5591 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
5592 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
5593 Some("openmodels") => match models_openmodels_body(&st) {
5594 Ok(body) => Json(body).into_response(),
5595 Err(error) => bad_request(&error, Some("schema")),
5596 },
5597 Some(schema) => bad_request(
5598 &format!(
5599 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
5600 ),
5601 Some("schema"),
5602 ),
5603 }
5604}
5605
5606fn model_entry_v1(
5614 name: &str,
5615 caps: Option<&ModelCaps>,
5616 metadata: Option<&OpenRouterModelMetadata>,
5617) -> serde_json::Value {
5618 let ctx = caps.map(|c| c.context_length).filter(|&c| c > 0);
5619 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
5623 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
5626 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
5627 Some(p) => json!(p),
5628 None => serde_json::Value::Null,
5629 };
5630 let owned_by = metadata
5631 .and_then(|m| m.owned_by.as_deref())
5632 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
5633 let mut input_modalities = vec!["text"];
5634 if let Some(meta) = metadata {
5635 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
5636 }
5637 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
5638 let reliability = metadata.and_then(|m| m.reliability.as_ref());
5639 let surface = declared_surface(metadata);
5645 let (model_type, endpoints, output_modalities) = match surface {
5646 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
5650 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
5651 _ => ("chat", vec!["chat/completions"], vec!["text"]),
5652 };
5653 let is_chat = surface == "chat";
5654 json!({
5655 "id": name,
5656 "name": name,
5657 "object": "model",
5658 "owned_by": owned_by,
5659 "type": model_type,
5660 "context_length": ctx,
5661 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
5664 "endpoints": endpoints,
5665 "input_modalities": input_modalities,
5666 "output_modalities": output_modalities,
5667 "capabilities": {
5668 "streaming": is_chat,
5671 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
5672 "structured_output": is_chat && !is_dsv4,
5673 "reasoning": is_chat && thinking,
5674 "prompt_caching": is_chat && !is_dsv4,
5675 },
5676 "pricing": {
5677 "currency": "USD",
5678 "unit": "per_1m_tokens",
5679 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
5680 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
5681 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
5682 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
5683 "minimum_request": metadata
5685 .and_then(|m| m.pricing.request.as_deref())
5686 .unwrap_or("0"),
5687 },
5688 "lifecycle": {
5689 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
5690 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
5691 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
5692 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
5693 },
5694 "reliability": {
5695 "first_token_timeout_seconds":
5696 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
5697 "completion_timeout_seconds":
5698 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
5699 "stream_idle_timeout_seconds":
5700 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
5701 "capacity_scope":
5702 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
5703 },
5704 })
5705}
5706
5707async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
5710 let data: Vec<_> = st
5711 .models
5712 .iter()
5713 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
5714 .collect();
5715 let mut body = json!({
5716 "object": "list",
5717 "contract_version": "2.0",
5718 "data": data,
5719 });
5720 if let Some(provider) = st.provider_metadata.as_ref() {
5725 body["provider"] = json!({
5726 "id": provider.id,
5727 "status_url": provider.status_url,
5728 "support_contact": provider.support_contact,
5729 "incident_contact": provider.incident_contact,
5730 "regions": provider.regions,
5731 "request_id_header": "x-request-id",
5732 "error_contract": {
5733 "rate_limit_status": 429,
5734 "overload_status": 503,
5735 "retry_after_header": "Retry-After",
5736 "account_quota_error_codes": ["insufficient_balance"],
5737 },
5738 });
5739 }
5740 Json(body)
5741}
5742
5743async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5746 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5747 Ok(scope) => scope,
5748 Err(response) => return response,
5749 };
5750 if !metrics_scope.process_wide() {
5751 return error_response(
5752 StatusCode::FORBIDDEN,
5753 "completion api keys do not authorize process-wide yield metrics; configure \
5754 MEMRA_METRICS_TOKEN",
5755 "authentication_error",
5756 None,
5757 );
5758 }
5759 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5760 let lane = |i: usize| {
5761 json!({
5762 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
5763 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
5764 })
5765 };
5766 let mut body = json!({
5767 "lanes": {
5768 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
5769 },
5770 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
5771 });
5772 if metrics_scope.operator() {
5773 body["batch_size_last"] = json!(m.batch_size_last);
5774 }
5775 Json(body).into_response()
5776}
5777
5778async fn peek_admission(
5791 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5792) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, (Response, &'static str)> {
5793 match rx.recv().await {
5794 Some(Event::Error(e)) => {
5799 let error_code = engine_error_code(e.class);
5800 Err((engine_error_response(&e), error_code))
5801 }
5802 first => {
5803 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
5804 if let Some(ev) = first {
5805 let _ = tx2.send(ev);
5806 }
5807 tokio::spawn(forward_events(rx, tx2));
5808 Ok(rx2)
5809 }
5810 }
5811}
5812
5813async fn forward_events(
5824 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5825 tx2: tokio::sync::mpsc::UnboundedSender<Event>,
5826) {
5827 loop {
5828 tokio::select! {
5829 biased;
5830 () = tx2.closed() => break,
5831 ev = rx.recv() => match ev {
5832 Some(ev) => {
5833 if tx2.send(ev).is_err() {
5834 break;
5835 }
5836 }
5837 None => break,
5838 },
5839 }
5840 }
5841}
5842
5843async fn peek_first_token(
5858 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5859 deadline: RequestDeadline,
5860) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, ()> {
5861 let mut buffered: Vec<Event> = Vec::new();
5862 loop {
5863 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
5864 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
5867 let first_delivery = matches!(
5868 ev,
5869 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
5870 );
5871 buffered.push(ev);
5872 if first_delivery {
5873 break;
5874 }
5875 }
5876 }
5877 }
5878 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
5879 for ev in buffered {
5880 let _ = tx2.send(ev);
5881 }
5882 tokio::spawn(forward_events(rx, tx2));
5883 Ok(rx2)
5884}
5885
5886#[cfg(test)]
5888fn build_request(
5892 req: &CompletionReq,
5893 tx: tokio::sync::mpsc::UnboundedSender<Event>,
5894 lane: lanes::Lane,
5895 affinity: Option<String>,
5896) -> Request {
5897 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
5898}
5899
5900fn build_request_with_trace(
5901 req: &CompletionReq,
5902 tx: tokio::sync::mpsc::UnboundedSender<Event>,
5903 lane: lanes::Lane,
5904 affinity: Option<String>,
5905 ttft: Option<Arc<ttft::Trace>>,
5906 sampling_defaults: &SamplingDefaults,
5907) -> Request {
5908 let params = GenParams {
5909 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
5910 max_ctx: req.max_ctx,
5911 eos: Vec::new(), };
5913 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
5918 Request {
5919 model: req.model.clone(),
5920 prompt_ids: req.prompt_ids.clone(),
5921 prompt_text: req.prompt.clone(),
5922 chat: req.chat,
5923 chat_turns: Vec::new(),
5924 tools_json: Vec::new(),
5925 tools_struct: Vec::new(),
5926 think: ThinkMode::Default,
5927 reasoning_effort: None, params,
5929 sampler_cfg,
5930 stop_strings: req.stop.clone().into_vec(),
5931 trace_id: req.trace_id.clone(),
5932 max_prompt_tokens: None,
5933 cache_ns: cache_namespace(&req.cache_salt),
5934 affinity,
5935 lane,
5936 grammar: None, prepared_constraint: None,
5938 constraint_ready: None,
5939 oom_retries: 0, spec_k_replay: None,
5941 prepared_prompt: None,
5942 capture: None, images: Vec::new(), gemma_images: Vec::new(),
5945 vision_memory: None,
5946 ttft,
5947 tx,
5948 }
5949}
5950
5951struct ChatPlan {
5954 request: Request,
5955 parser: Option<ToolStreamParser>,
5958 pending_images: Vec<PendingVisionUnit>,
5961 pending_gemma: Vec<PendingGemmaImage>,
5962 vision_memory: Option<VisionMemoryPermit>,
5966}
5967
5968pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
5969 req.messages.iter().any(|message| {
5970 message.content.as_array().is_some_and(|parts| {
5971 parts.iter().any(|part| {
5972 matches!(
5973 part.get("type").and_then(serde_json::Value::as_str),
5974 Some("image_url" | "video_url")
5975 )
5976 })
5977 })
5978 })
5979}
5980
5981fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
5982 let mut total = 0usize;
5983 let mut add = |bytes: usize| {
5984 total = total.checked_add(bytes).ok_or_else(|| {
5985 "vision patch memory reservation overflowed while planning".to_string()
5986 })?;
5987 Ok::<(), String>(())
5988 };
5989 for unit in &plan.pending_images {
5990 let bytes = match unit {
5991 PendingVisionUnit::Still { gh, gw, .. } => gh
5992 .checked_mul(*gw)
5993 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
5994 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
5995 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
5996 PendingVisionUnit::Video { groups, .. } => {
5997 groups.iter().try_fold(0usize, |total, group| {
5998 let bytes = group
5999 .gh
6000 .checked_mul(group.gw)
6001 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6002 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6003 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6004 total.checked_add(bytes).ok_or_else(|| {
6005 "vision patch memory reservation overflowed while planning".to_string()
6006 })
6007 })?
6008 }
6009 };
6010 add(bytes)?;
6011 }
6012 for unit in &plan.pending_gemma {
6013 let bytes = unit
6014 .gw
6015 .checked_mul(unit.gh)
6016 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
6017 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6018 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6019 add(bytes)?;
6020 }
6021 Ok(total)
6022}
6023
6024pub(crate) fn reserve_vision_memory(
6025 plan: &ChatPlan,
6026) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
6027 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
6028 try_reserve_vision_memory(bytes)
6029}
6030
6031#[cfg(test)]
6032fn build_chat_request(
6033 req: ChatCompletionReq,
6034 caps: Option<&ModelCaps>,
6035 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6036 lane: lanes::Lane,
6037 affinity: Option<String>,
6038) -> Result<ChatPlan, String> {
6039 let defaults = ModelSamplingDefaults::resolve(None, caps);
6042 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
6043}
6044
6045#[allow(clippy::too_many_arguments)]
6059fn build_chat_request_with_trace(
6060 req: ChatCompletionReq,
6061 caps: Option<&ModelCaps>,
6062 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6063 lane: lanes::Lane,
6064 affinity: Option<String>,
6065 ttft: Option<Arc<ttft::Trace>>,
6066 default_effort: Option<&str>,
6067 sampling_defaults: &ModelSamplingDefaults,
6068) -> Result<ChatPlan, String> {
6069 let client_sampling: ClientSampling = (&req).into();
6072 let tool_choice = parse_tool_choice(&req.tool_choice)?;
6073 if let Some(c) = caps {
6078 if !c.chat_ok {
6079 return Err(format!(
6080 "model {:?} has no chat template (checkpoint carries neither \
6081 tokenizer_config.json chat_template nor chat_template.jinja) — \
6082 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
6083 req.model
6084 ));
6085 }
6086 }
6087 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
6088 let (mut think, effort_level, think_client_explicit) = parse_think(
6089 &req.reasoning_effort,
6090 &req.reasoning,
6091 vllm_switch,
6092 req.include_reasoning,
6093 default_effort,
6094 caps.is_some_and(|c| c.dsv4),
6095 )?;
6096 let level_template = caps
6100 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort)
6101 .unwrap_or(false);
6102 if think_client_explicit && think == ThinkMode::NoThink {
6123 if let Some(c) = caps {
6124 if c.qwen_think && !c.think_switch && !c.dsv4 {
6125 return Err(format!(
6126 "model {:?} cannot disable reasoning: its chat template opens a think \
6127 tail unconditionally and carries no enable_thinking switch, so \
6128 reasoning_effort/enable_thinking cannot turn it off on this model",
6129 req.model
6130 ));
6131 }
6132 }
6133 }
6134 let reasoning_effort = if level_template { effort_level } else { None };
6167 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
6170 if grammar.is_some() {
6176 if let Some(c) = caps {
6177 if c.qwen_think && think != ThinkMode::NoThink {
6178 if c.think_switch {
6179 think = ThinkMode::NoThink;
6180 } else {
6181 return Err(
6182 "response_format requires disabling the model's think tail, \
6183 but this chat template has no enable_thinking switch"
6184 .into(),
6185 );
6186 }
6187 }
6188 }
6189 }
6190
6191 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
6198
6199 let (tools_json, tools_struct, schemas) =
6202 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
6203 prepare_tools(&req.tools)?
6204 } else {
6205 (Vec::new(), Vec::new(), HashMap::new())
6206 };
6207
6208 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
6209 let mut images: Vec<PendingVisionUnit> = Vec::new();
6210 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
6211 let mut next_video = 0usize;
6212 for msg in &req.messages {
6213 let content = content_to_text_vision(
6214 &msg.content,
6215 &mut images,
6216 &mut gemma_images,
6217 &mut next_video,
6218 )
6219 .map_err(|e| format!("{} message: {e}", msg.role))?;
6220 let tool_calls = msg
6221 .tool_calls
6222 .iter()
6223 .map(render_req_tool_call)
6224 .collect::<Result<Vec<_>, _>>()?;
6225 if !tool_calls.is_empty() && msg.role != "assistant" {
6226 return Err("tool_calls are only valid on assistant messages".into());
6227 }
6228 let role = if msg.role == "developer" {
6231 "system".to_string()
6232 } else {
6233 msg.role.clone()
6234 };
6235 turns.push(TmplTurn {
6236 role,
6237 content,
6238 tool_calls,
6239 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
6241 tool_call_id: msg.tool_call_id.clone(),
6242 tool_name: msg.name.clone(),
6243 tool_responses: Vec::new(),
6244 task: None,
6248 tools: Vec::new(),
6249 });
6250 }
6251
6252 let has_tool_features = !tools_json.is_empty()
6255 || turns
6256 .iter()
6257 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
6258 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
6259 return Err(format!(
6260 "model {:?} chat template has no tools branch",
6261 req.model
6262 ));
6263 }
6264
6265 let think_open = caps
6268 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
6269 .unwrap_or(false);
6270 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
6288 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
6294 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
6295 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
6296 let parser = if is_dsv4 && (dsv4_tools || dsv4_think_open) {
6297 Some(ToolStreamParser::dsv4(dsv4_think_open))
6298 } else if gemma_tools {
6299 Some(ToolStreamParser::gemma_tools())
6300 } else if !tools_json.is_empty() {
6301 Some(ToolStreamParser::new(schemas, think_open))
6302 } else if think_open {
6303 Some(ToolStreamParser::reasoning_only())
6304 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
6305 Some(ToolStreamParser::gemma_thought())
6313 } else {
6314 None
6315 };
6316
6317 Ok(ChatPlan {
6318 request: Request {
6319 model: req.model,
6320 prompt_ids: Vec::new(),
6321 prompt_text: String::new(),
6322 chat: false,
6323 chat_turns: turns,
6324 tools_json,
6325 tools_struct,
6326 think,
6327 reasoning_effort,
6328 params: GenParams {
6329 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6330 max_ctx: req.max_ctx,
6331 eos: Vec::new(),
6332 },
6333 sampler_cfg,
6334 stop_strings: {
6335 let mut stops = req.stop.into_vec();
6340 if gemma_tools {
6341 stops.push("<tool_call|>".to_string());
6342 }
6343 if dsv4_tools {
6348 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
6349 }
6350 stops
6351 },
6352 trace_id: None,
6353 max_prompt_tokens: None,
6354 cache_ns: cache_namespace(&req.cache_salt),
6355 affinity,
6356 lane,
6357 grammar,
6358 prepared_constraint: None,
6359 constraint_ready: None,
6360 oom_retries: 0, spec_k_replay: None,
6362 prepared_prompt: None,
6363 images: Vec::new(),
6368 gemma_images: Vec::new(),
6369 capture: None, vision_memory: None,
6371 ttft,
6372 tx,
6373 },
6374 parser,
6375 pending_images: images,
6376 pending_gemma: gemma_images,
6377 vision_memory: None,
6378 })
6379}
6380
6381fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
6387 for (i, unit) in plan.pending_images.drain(..).enumerate() {
6388 match unit {
6389 PendingVisionUnit::Still { bytes, gh, gw } => {
6390 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
6391 .map_err(|e| format!("image {}: {e}", i + 1))?;
6392 if (prep.gh, prep.gw) != (gh, gw) {
6393 return Err(format!(
6394 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
6395 i + 1,
6396 prep.gh,
6397 prep.gw
6398 ));
6399 }
6400 plan.request
6401 .images
6402 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
6403 }
6404 PendingVisionUnit::Video {
6405 bytes,
6406 groups,
6407 video,
6408 } => {
6409 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
6410 .map_err(|e| format!("video {}: {e}", i + 1))?;
6411 if prepared.groups.len() != groups.len() {
6412 return Err(format!(
6413 "video {}: decoded {} groups differ from its header-planned {} groups",
6414 i + 1,
6415 prepared.groups.len(),
6416 groups.len()
6417 ));
6418 }
6419 for ((group, prep), timestamp) in
6420 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
6421 {
6422 if (prep.gh, prep.gw) != (group.gh, group.gw) {
6423 return Err(format!(
6424 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
6425 i + 1,
6426 prep.gh,
6427 prep.gw,
6428 group.gh,
6429 group.gw
6430 ));
6431 }
6432 if (timestamp - group.timestamp).abs() > 0.001 {
6433 return Err(format!(
6434 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
6435 i + 1,
6436 group.timestamp
6437 ));
6438 }
6439 plan.request
6440 .images
6441 .push(memra_engine::vision_pre::VisionUnit {
6442 prep,
6443 video: Some(video),
6444 });
6445 }
6446 }
6447 }
6448 }
6449 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
6450 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
6451 .map_err(|e| format!("image {}: {e}", i + 1))?;
6452 if (gw, gh) != (unit.gw, unit.gh) {
6453 return Err(format!(
6454 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
6455 i + 1,
6456 unit.gw,
6457 unit.gh
6458 ));
6459 }
6460 plan.request
6461 .gemma_images
6462 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
6463 }
6464 Ok(())
6465}
6466
6467fn bearer_token(headers: &HeaderMap) -> Option<&str> {
6475 headers
6476 .get("authorization")
6477 .and_then(|value| value.to_str().ok())
6478 .and_then(|value| value.strip_prefix("Bearer "))
6479}
6480
6481fn authentication_error(why: auth::AuthDenied) -> Response {
6482 match why {
6483 auth::AuthDenied::Unknown => error_response(
6484 StatusCode::UNAUTHORIZED,
6485 "invalid api key",
6486 "authentication_error",
6487 None,
6488 ),
6489 auth::AuthDenied::Disabled => error_response(
6490 StatusCode::FORBIDDEN,
6491 "api key is disabled",
6492 "authentication_error",
6493 None,
6494 ),
6495 }
6496}
6497
6498fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
6499 auth::authenticate_with(
6500 api_auth.keyring,
6501 api_auth.single_key.as_deref(),
6502 bearer_token(headers),
6503 )
6504 .map_err(authentication_error)
6505}
6506
6507#[derive(Debug, Clone, PartialEq, Eq)]
6508enum MetricsScope {
6509 All,
6510 CompletionDomain,
6511 Tenant(String),
6512}
6513
6514impl MetricsScope {
6515 fn operator(&self) -> bool {
6516 matches!(self, MetricsScope::All)
6517 }
6518
6519 fn process_wide(&self) -> bool {
6520 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
6521 }
6522
6523 fn includes(&self, tenant_row: &str) -> bool {
6524 match self {
6525 MetricsScope::All | MetricsScope::CompletionDomain => true,
6526 MetricsScope::Tenant(tenant) => tenant == tenant_row,
6527 }
6528 }
6529}
6530
6531fn authorize_metrics(
6532 api_auth: &ApiAuth,
6533 metrics_auth: &MetricsAuth,
6534 headers: &HeaderMap,
6535) -> Result<MetricsScope, Response> {
6536 if !metrics_auth.required {
6537 return Ok(MetricsScope::All);
6538 }
6539 let Some(candidate) = bearer_token(headers) else {
6540 return Err(authentication_error(auth::AuthDenied::Unknown));
6541 };
6542 if let Some(token) = metrics_auth.token.as_deref() {
6543 if auth::constant_time_secret_eq(token, candidate) {
6544 return Ok(MetricsScope::All);
6545 }
6546 if api_auth.configured() {
6547 return match auth::authenticate_with(
6548 api_auth.keyring,
6549 api_auth.single_key.as_deref(),
6550 Some(candidate),
6551 ) {
6552 Ok(_) => Err(error_response(
6553 StatusCode::FORBIDDEN,
6554 "completion api keys do not authorize metrics while \
6555 MEMRA_METRICS_TOKEN is configured",
6556 "authentication_error",
6557 None,
6558 )),
6559 Err(why) => Err(authentication_error(why)),
6560 };
6561 }
6562 return Err(authentication_error(auth::AuthDenied::Unknown));
6563 }
6564 if api_auth.configured() {
6565 let tenant = authenticate(api_auth, headers)?;
6566 return Ok(if api_auth.keyring.is_some() {
6567 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
6568 } else {
6569 MetricsScope::CompletionDomain
6573 });
6574 }
6575 Err(authentication_error(auth::AuthDenied::Unknown))
6576}
6577
6578fn lane_for_tenant(
6584 headers: &axum::http::HeaderMap,
6585 tenant: &auth::TenantCtx,
6586) -> Result<lanes::Lane, Response> {
6587 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
6588 None => None,
6589 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
6594 error_response_coded(
6595 StatusCode::BAD_REQUEST,
6596 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
6597 "invalid_request_error",
6598 Some("x-lane"),
6599 Some("invalid_lane"),
6600 )
6601 })?),
6602 };
6603 match tenant.lane_class {
6604 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
6605 auth::LaneClass::Batch => match requested {
6606 None => Ok(lanes::Lane::Harvest),
6607 Some(lanes::Lane::Interactive) => Err(error_response(
6608 StatusCode::FORBIDDEN,
6609 "this api key is batch-class: x-lane interactive is not permitted \
6610 (use judge or harvest)",
6611 "authentication_error",
6612 Some("x-lane"),
6613 )),
6614 Some(l) => Ok(l),
6615 },
6616 }
6617}
6618
6619fn tenant_namespace(
6623 tenant: &auth::TenantCtx,
6624 cache_salt: &Option<String>,
6625) -> Result<String, &'static str> {
6626 let keyring_configured = auth::global().is_some();
6627 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
6628 if keyring_configured {
6629 Ok(auth::scope_namespace(&tenant.tenant, &raw))
6630 } else {
6631 Ok(raw)
6632 }
6633}
6634
6635fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
6640 eprintln!(
6641 "[meter] admit id={} tenant={} lane={} model={:?}",
6642 env.id,
6643 tenant.tenant,
6644 lane.as_str(),
6645 model
6646 );
6647}
6648
6649fn apply_model_request_limits(
6650 request: &mut Request,
6651 metadata: Option<&OpenRouterModelMetadata>,
6652 caps: Option<&ModelCaps>,
6653) -> Result<(), (String, &'static str)> {
6654 let Some(metadata) = metadata else {
6655 return Ok(());
6656 };
6657 let max_prompt = metadata
6658 .max_prompt_length
6659 .map(usize::try_from)
6660 .transpose()
6661 .map_err(|_| {
6662 (
6663 "configured model prompt limit does not fit this platform".into(),
6664 "model",
6665 )
6666 })?;
6667 let max_output = metadata
6668 .max_output_length
6669 .map(usize::try_from)
6670 .transpose()
6671 .map_err(|_| {
6672 (
6673 "configured model output limit does not fit this platform".into(),
6674 "model",
6675 )
6676 })?;
6677
6678 request.max_prompt_tokens = max_prompt;
6679 if let Some(max_output) = max_output {
6680 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
6681 request.params.max_new = metadata
6682 .default_output_length
6683 .map(usize::try_from)
6684 .transpose()
6685 .map_err(|_| {
6686 (
6687 "configured default output length does not fit this platform".into(),
6688 "model",
6689 )
6690 })?
6691 .unwrap_or(max_output);
6692 } else if request.params.max_new > max_output {
6693 return Err((
6694 format!(
6695 "max_tokens {} exceeds configured model maximum {max_output}",
6696 request.params.max_new
6697 ),
6698 "max_tokens",
6699 ));
6700 }
6701 }
6702
6703 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
6707 (max_prompt, max_output, request.params.max_ctx)
6708 {
6709 let operational_ctx = max_prompt
6710 .checked_add(max_output)
6711 .and_then(|value| value.checked_add(8))
6712 .ok_or_else(|| {
6713 (
6714 "configured model context envelope overflowed".into(),
6715 "model",
6716 )
6717 })?;
6718 let operational_ctx = caps
6719 .map(|caps| caps.context_length)
6720 .filter(|&context| context > 0)
6721 .map_or(operational_ctx, |context| operational_ctx.min(context));
6722 if requested_ctx > operational_ctx {
6723 return Err((
6724 format!(
6725 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
6726 ),
6727 "max_ctx",
6728 ));
6729 }
6730 }
6731 Ok(())
6732}
6733
6734#[allow(clippy::too_many_arguments)]
6735fn start_request_receipt(
6736 st: &AppState,
6737 env: &Envelope,
6738 tenant: &auth::TenantCtx,
6739 model: &str,
6740 route: &'static str,
6741 lane: lanes::Lane,
6742 stream: bool,
6743 budget_permit: Option<metering::Permit>,
6744) -> Option<Box<dyn metering::Receipt>> {
6745 st.metering.as_ref().map(|accounting| {
6746 accounting.open(
6747 &metering::RequestMeta {
6748 request_id: &env.id,
6749 tenant: &tenant.tenant,
6750 model,
6751 route,
6752 lane: lane.as_str(),
6753 stream,
6754 },
6755 budget_permit,
6756 )
6757 })
6758}
6759
6760fn arm_capture(
6766 mut receipt: Option<Box<dyn metering::Receipt>>,
6767 prompt: impl FnOnce() -> serde_json::Value,
6768) -> Option<Box<dyn metering::Receipt>> {
6769 if let Some(receipt) = receipt.as_mut()
6770 && receipt.wants_capture()
6771 {
6772 receipt.arm_capture(prompt());
6773 }
6774 receipt
6775}
6776
6777fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
6781 serde_json::Value::Array(
6782 messages
6783 .iter()
6784 .map(|message| {
6785 let mut row = json!({ "role": message.role, "content": message.content });
6786 if !message.tool_calls.is_empty() {
6787 row["tool_calls"] = serde_json::Value::Array(
6788 message
6789 .tool_calls
6790 .iter()
6791 .map(|call| {
6792 json!({
6793 "id": call.id,
6794 "function": {
6795 "name": call.function.name,
6796 "arguments": call.function.arguments,
6797 },
6798 })
6799 })
6800 .collect(),
6801 );
6802 }
6803 row
6804 })
6805 .collect(),
6806 )
6807}
6808
6809enum BudgetRejection {
6810 Invalid(String),
6811 Insufficient,
6812 Unenrolled,
6813 Unavailable(String),
6814}
6815
6816impl BudgetRejection {
6817 fn into_response(self) -> (Response, &'static str) {
6818 match self {
6819 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
6820 Self::Insufficient => (
6821 error_response_coded(
6822 StatusCode::PAYMENT_REQUIRED,
6823 "tenant prepaid balance is insufficient for this request",
6824 "insufficient_balance",
6825 None,
6826 Some("insufficient_balance"),
6827 ),
6828 "insufficient_balance",
6829 ),
6830 Self::Unenrolled => (
6831 error_response_coded(
6832 StatusCode::PAYMENT_REQUIRED,
6833 "tenant is not enrolled for prepaid billing",
6834 "tenant_not_enrolled",
6835 None,
6836 Some("tenant_not_enrolled"),
6837 ),
6838 "tenant_not_enrolled",
6839 ),
6840 Self::Unavailable(err) => {
6841 eprintln!("[budget] ERROR: admission unavailable: {err}");
6842 (
6843 error_response_coded(
6844 StatusCode::SERVICE_UNAVAILABLE,
6845 "tenant budget accounting is unavailable",
6846 "server_error",
6847 None,
6848 Some("tenant_budget_unavailable"),
6849 ),
6850 "tenant_budget_unavailable",
6851 )
6852 }
6853 }
6854 }
6855}
6856
6857fn prepare_budget_prompt(
6858 request: &mut Request,
6859 tokenizer: Option<&Tokenizer>,
6860) -> Result<usize, String> {
6861 if let Some(error) = worker::prompt_source_limit_error(request) {
6862 return Err(error);
6863 }
6864 if request.prepared_prompt.is_none() {
6865 if let Some(trace) = request.ttft.as_ref() {
6866 trace.mark_tokenize_start();
6867 }
6868 let prompt = if !request.prompt_ids.is_empty() {
6869 request.prompt_ids.clone()
6870 } else if !request.chat_turns.is_empty() {
6871 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
6872 let plain = worker::plain_chat_render_path(
6879 &request.tools_json,
6880 &request.think,
6881 request.reasoning_effort.as_deref(),
6882 &request.chat_turns,
6883 tokenizer.has_qwen_effort_ladder(),
6884 );
6885 let rendered = if plain {
6886 let messages: Vec<_> = request
6887 .chat_turns
6888 .iter()
6889 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
6890 .collect();
6891 tokenizer.apply_chat_template(&messages, true)
6892 } else {
6893 tokenizer
6894 .apply_chat_template_tools_ex(
6895 &request.chat_turns,
6896 true,
6897 &request.tools_json,
6898 &request.tools_struct,
6899 request.think,
6900 request.reasoning_effort.as_deref(),
6901 )
6902 .map_err(|err| format!("chat template: {err}"))?
6903 };
6904 tokenizer.encode(&rendered, true)
6905 } else if request.chat {
6906 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
6907 let rendered =
6908 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
6909 tokenizer.encode(&rendered, true)
6910 } else {
6911 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
6912 tokenizer.encode(&request.prompt_text, true)
6913 };
6914 if prompt.is_empty() {
6915 return Err("empty prompt after tokenization".into());
6916 }
6917 if let Some(trace) = request.ttft.as_ref() {
6918 trace.mark_tokenize_end(prompt.len());
6919 }
6920 request.prepared_prompt = Some(prompt);
6921 }
6922 let prompt_tokens = request
6923 .prepared_prompt
6924 .as_ref()
6925 .expect("budget prompt was prepared")
6926 .len();
6927 if let Some(limit) = request.max_prompt_tokens
6928 && prompt_tokens > limit
6929 {
6930 return Err(format!(
6931 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
6932 ));
6933 }
6934 Ok(prompt_tokens)
6935}
6936
6937fn budget_completion_bound(
6938 request: &Request,
6939 prompt_tokens: usize,
6940 caps: Option<&ModelCaps>,
6941) -> Result<usize, String> {
6942 let max_new = request.params.max_new;
6943 let requested_ctx = match (request.params.max_ctx, max_new) {
6944 (Some(cap), _) => cap,
6945 (None, worker::MAX_NEW_CTX_BOUNDED) => {
6946 let server_ctx = std::env::var("MEMRA_CTX")
6947 .ok()
6948 .and_then(|value| value.parse().ok())
6949 .unwrap_or(8192usize);
6950 let mut cap = server_ctx;
6951 if prompt_tokens.saturating_add(16) > cap {
6952 cap = prompt_tokens.saturating_add(server_ctx);
6953 }
6954 cap
6955 }
6956 (None, max_new) => prompt_tokens
6957 .checked_add(max_new)
6958 .and_then(|value| value.checked_add(8))
6959 .ok_or_else(|| "request context bound overflowed".to_string())?,
6960 };
6961 let ctx_cap = caps
6962 .map(|caps| caps.context_length)
6963 .filter(|&context| context > 0)
6964 .map_or(requested_ctx, |context| requested_ctx.min(context));
6965 if prompt_tokens >= ctx_cap {
6966 return Err(format!(
6967 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
6968 ));
6969 }
6970 Ok(max_new.min(ctx_cap - prompt_tokens))
6971}
6972
6973fn admit_tenant_budget(
6974 st: &AppState,
6975 tenant: &auth::TenantCtx,
6976 request: &mut Request,
6977) -> Result<Option<metering::Permit>, BudgetRejection> {
6978 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
6979 return Ok(None);
6980 };
6981 match accounting.is_limited(&tenant.tenant) {
6982 Ok(false) => return Err(BudgetRejection::Unenrolled),
6983 Ok(true) => {}
6984 Err(metering::AdmitError::Unavailable(err)) => {
6985 return Err(BudgetRejection::Unavailable(err));
6986 }
6987 Err(other) => {
6988 return Err(BudgetRejection::Unavailable(format!(
6989 "unexpected budget enrollment result: {other:?}"
6990 )));
6991 }
6992 }
6993 let tokenizer = st
6994 .budget_tokenizers
6995 .as_ref()
6996 .and_then(|tokenizers| tokenizers.get(&request.model))
6997 .map(Arc::as_ref);
6998 if request.prompt_ids.is_empty() && tokenizer.is_none() {
6999 return Err(BudgetRejection::Unavailable(format!(
7000 "no reservation tokenizer for model {:?}",
7001 request.model
7002 )));
7003 }
7004 let prompt_tokens =
7005 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
7006 let completion_tokens =
7007 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
7008 .map_err(BudgetRejection::Invalid)?;
7009 let prompt_tokens = u64::try_from(prompt_tokens)
7010 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
7011 let completion_tokens = u64::try_from(completion_tokens)
7012 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
7013 match accounting.reserve(
7014 &tenant.tenant,
7015 &request.model,
7016 prompt_tokens,
7017 completion_tokens,
7018 ) {
7019 Ok(permit) => Ok(permit),
7020 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
7021 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
7025 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
7026 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
7027 }
7028}
7029
7030fn request_ledger_error_response() -> Response {
7031 error_response_coded(
7032 StatusCode::INTERNAL_SERVER_ERROR,
7033 "request completion could not be committed to the billing ledger",
7034 "server_error",
7035 None,
7036 Some("request_ledger_unavailable"),
7037 )
7038}
7039
7040fn request_ledger_error_body() -> serde_json::Value {
7041 error_body(
7042 "request completion could not be committed to the billing ledger",
7043 "server_error",
7044 None,
7045 Some("request_ledger_unavailable"),
7046 )
7047}
7048
7049fn ledger_rejected(
7050 mut receipt: Option<Box<dyn metering::Receipt>>,
7051 response: Response,
7052 error_code: &str,
7053 request_id: &str,
7054) -> Response {
7055 let status = response.status().as_u16();
7056 if let Some(receipt) = receipt.as_mut()
7057 && let Err(err) = receipt.reject(status, error_code)
7058 {
7059 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
7060 return with_request_id(request_id, request_ledger_error_response());
7061 }
7062 with_request_id(request_id, response)
7063}
7064
7065fn ledger_unbilled(
7070 mut receipt: Option<Box<dyn metering::Receipt>>,
7071 response: Response,
7072 outcome: &'static str,
7073 error_code: &str,
7074 request_id: &str,
7075) -> Response {
7076 let status = response.status().as_u16();
7077 if let Some(receipt) = receipt.as_mut()
7078 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
7079 {
7080 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
7081 return with_request_id(request_id, request_ledger_error_response());
7082 }
7083 with_request_id(request_id, response)
7084}
7085
7086fn engine_error_code(class: worker::ErrClass) -> &'static str {
7087 use worker::ErrClass as C;
7088 match class {
7089 C::InvalidRequest => "invalid_request",
7090 C::ContextLength => "context_length_exceeded",
7091 C::ModelNotFound => "model_not_found",
7092 C::RateLimit => "rate_limit_exceeded",
7093 C::Overloaded => "overloaded",
7094 C::Engine => "engine_error",
7095 }
7096}
7097
7098fn model_not_found_response(models: &[String], requested: &str) -> Response {
7118 error_response_coded(
7119 StatusCode::BAD_REQUEST,
7120 &format!("unknown model {requested:?}; loaded: {models:?}"),
7121 "invalid_request_error",
7122 Some("model"),
7123 Some("model_not_found"),
7124 )
7125}
7126
7127fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
7136 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
7137 return Ok(());
7138 };
7139 if let Some((pos, &id)) = ids
7140 .iter()
7141 .enumerate()
7142 .find(|&(_, &id)| id as usize >= n_vocab)
7143 {
7144 return Err(format!(
7145 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
7146 ));
7147 }
7148 Ok(())
7149}
7150
7151#[cfg(test)]
7152mod prompt_ids_tests {
7153 use super::*;
7154
7155 #[test]
7156 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
7157 let caps = ModelCaps {
7158 n_vocab: 8,
7159 ..Default::default()
7160 };
7161 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
7163 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
7164 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
7166 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
7167 assert!(err.contains("vocab size 8"), "{err}");
7168 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
7169 assert!(err.contains("4294967295"), "{err}");
7170 let unknown = ModelCaps::default();
7172 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
7173 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
7174 }
7175}
7176
7177fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
7178 if models.iter().any(|m| m == requested) {
7179 return Some(requested.to_string());
7180 }
7181 if requested.is_empty() || requested.contains('/') {
7182 return None;
7183 }
7184 let mut matches = models.iter().filter(|m| {
7185 m.rsplit('/')
7186 .next()
7187 .is_some_and(|suffix| suffix == requested)
7188 });
7189 match (matches.next(), matches.next()) {
7190 (Some(only), None) => Some(only.clone()),
7191 _ => None,
7192 }
7193}
7194
7195async fn completions(
7196 State(st): State<AppState>,
7197 headers: axum::http::HeaderMap,
7198 trace: Option<Extension<TtftRequestTrace>>,
7199 Json(mut req): Json<CompletionReq>,
7200) -> Response {
7201 let env = Envelope::new(false);
7202 match canonical_model_id(&st.models, &req.model) {
7203 Some(canonical) => req.model = canonical,
7204 None => {
7205 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7206 }
7207 }
7208 let ttft = trace.and_then(|Extension(trace)| trace.0);
7211 if let Some(trace) = ttft.as_ref() {
7212 trace.mark_parsed();
7213 trace.bind_request(&env.id, &req.model);
7214 }
7215 let tenant = match authenticate(&st.api_auth, &headers) {
7216 Ok(t) => t,
7217 Err(resp) => return with_request_id(&env.id, resp),
7218 };
7219 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7220 Ok(ns) => ns,
7221 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7222 };
7223 if let Err((msg, param)) = reject_unsupported(&[
7225 (
7226 "logit_bias",
7227 req.logit_bias.is_some(),
7228 " (device-side sampling has no bias hook yet)",
7229 ),
7230 ("logprobs", req.logprobs.is_some(), ""),
7231 (
7232 "n",
7233 req.n.is_some_and(|n| n != 1),
7234 " for n != 1 (single choice only)",
7235 ),
7236 (
7237 "best_of",
7238 req.best_of.is_some_and(|n| n != 1),
7239 " (single choice only)",
7240 ),
7241 ]) {
7242 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7243 }
7244 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
7247 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
7248 }
7249 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7252 Ok(ms) => RequestDeadline::starting_now(ms),
7253 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7254 };
7255 let lane = match lane_for_tenant(&headers, &tenant) {
7256 Ok(l) => l,
7257 Err(resp) => return resp,
7258 };
7259 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7260 let model = req.model.clone();
7261 let stream = req.stream;
7262 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7263 let mut request = build_request_with_trace(
7264 &req,
7265 tx,
7266 lane,
7267 affinity,
7268 ttft.clone(),
7269 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
7273 );
7274 request.cache_ns = cache_ns;
7275 if let Err((message, param)) = apply_model_request_limits(
7276 &mut request,
7277 st.openrouter_metadata.get(&model),
7278 st.caps.get(&model),
7279 ) {
7280 return with_request_id(&env.id, bad_request(&message, Some(param)));
7281 }
7282 if let Err(msg) = nonstream_deadline_gate(
7287 &request,
7288 req.stream,
7289 deadline,
7290 req.max_tokens.is_some(),
7291 st.budget_tokenizers
7292 .as_ref()
7293 .and_then(|t| t.get(&req.model))
7294 .map(Arc::as_ref),
7295 ) {
7296 return with_request_id(
7297 &env.id,
7298 error_response_coded(
7299 StatusCode::BAD_REQUEST,
7300 &msg,
7301 "invalid_request_error",
7302 Some("max_tokens"),
7303 Some("nonstream_deadline_infeasible"),
7304 ),
7305 );
7306 }
7307 if draining() {
7310 let receipt = start_request_receipt(
7311 &st,
7312 &env,
7313 &tenant,
7314 &req.model,
7315 "/v1/completions",
7316 lane,
7317 req.stream,
7318 None,
7319 );
7320 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7321 }
7322 let budget_permit = match admit_tenant_budget(&st, &tenant, &mut request) {
7323 Ok(permit) => permit,
7324 Err(rejection) => {
7325 let (response, error_code) = rejection.into_response();
7326 let receipt = start_request_receipt(
7327 &st,
7328 &env,
7329 &tenant,
7330 &req.model,
7331 "/v1/completions",
7332 lane,
7333 req.stream,
7334 None,
7335 );
7336 return ledger_rejected(receipt, response, error_code, &env.id);
7337 }
7338 };
7339 let receipt = start_request_receipt(
7340 &st,
7341 &env,
7342 &tenant,
7343 &req.model,
7344 "/v1/completions",
7345 lane,
7346 req.stream,
7347 budget_permit,
7348 );
7349 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
7350 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7353 Ok(slot) => slot,
7354 Err(resp) => {
7355 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7356 }
7357 };
7358 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7361 Ok(guard) => guard,
7362 Err((resp, outcome)) => {
7363 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7364 }
7365 };
7366 meter_admit(&env, &tenant, &model, lane);
7367 let stop_strings = request.stop_strings.clone();
7368
7369 if let Some(trace) = ttft.as_ref() {
7374 trace.mark_submitted();
7375 }
7376 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
7377 drop(pending_admit);
7378 return ledger_rejected(
7379 receipt,
7380 rl.attach(worker_unavailable_response()),
7381 "worker_unavailable",
7382 &env.id,
7383 );
7384 }
7385 pending_admit.commit();
7386 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7390 Ok(Ok(rx)) => rx,
7391 Ok(Err((resp, error_code))) => {
7392 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7393 }
7394 Err(_) => {
7395 return ledger_unbilled(
7396 receipt,
7397 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7398 "deadline_exceeded",
7399 "deadline_exceeded",
7400 &env.id,
7401 );
7402 }
7403 };
7404
7405 let resp = if stream {
7406 let rx = match peek_first_token(rx, deadline).await {
7410 Ok(rx) => rx,
7411 Err(()) => {
7412 return ledger_unbilled(
7413 receipt,
7414 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7415 "deadline_exceeded",
7416 "deadline_exceeded",
7417 &env.id,
7418 );
7419 }
7420 };
7421 sse_response_with_receipt(
7422 rx,
7423 model,
7424 false,
7425 None,
7426 env.clone(),
7427 stop_strings,
7428 Some(guard),
7429 receipt,
7430 )
7431 .into_response()
7432 } else {
7433 let mut receipt = receipt;
7439 let resp = blocking_response_with_receipt(
7440 rx,
7441 model,
7442 false,
7443 stop_strings,
7444 None,
7445 env.clone(),
7446 &mut receipt,
7447 Some(deadline),
7448 )
7449 .await;
7450 drop(guard); resp.into_response()
7452 };
7453 rl.attach(with_request_id(&env.id, resp))
7454}
7455
7456async fn chat_completions(
7457 State(st): State<AppState>,
7458 headers: axum::http::HeaderMap,
7459 trace: Option<Extension<TtftRequestTrace>>,
7460 Json(mut req): Json<ChatCompletionReq>,
7461) -> Response {
7462 let env = Envelope::new(true);
7463 match canonical_model_id(&st.models, &req.model) {
7468 Some(canonical) => req.model = canonical,
7469 None => {
7470 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7471 }
7472 }
7473 let ttft = trace.and_then(|Extension(trace)| trace.0);
7474 if let Some(trace) = ttft.as_ref() {
7475 trace.mark_parsed();
7476 trace.bind_request(&env.id, &req.model);
7477 }
7478 let tenant = match authenticate(&st.api_auth, &headers) {
7479 Ok(t) => t,
7480 Err(resp) => return with_request_id(&env.id, resp),
7481 };
7482 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7483 Ok(ns) => ns,
7484 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7485 };
7486 if req.messages.is_empty()
7487 || req.messages.iter().any(|message| {
7488 !matches!(
7489 message.role.as_str(),
7490 "system" | "developer" | "user" | "assistant" | "tool"
7491 )
7492 })
7493 {
7494 return with_request_id(
7495 &env.id,
7496 bad_request(
7497 "messages must use system/developer/user/assistant/tool roles",
7498 Some("messages"),
7499 ),
7500 );
7501 }
7502 if let Err((msg, param)) = reject_unsupported(&[
7507 (
7508 "logit_bias",
7509 req.logit_bias.is_some(),
7510 " (device-side sampling has no bias hook yet)",
7511 ),
7512 (
7513 "logprobs",
7514 req.logprobs
7515 .as_ref()
7516 .is_some_and(|v| v.as_bool() != Some(false)),
7517 "",
7518 ),
7519 ("top_logprobs", req.top_logprobs.is_some(), ""),
7520 (
7521 "n",
7522 req.n.is_some_and(|n| n != 1),
7523 " for n != 1 (single choice only)",
7524 ),
7525 ]) {
7526 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7527 }
7528 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7531 Ok(ms) => RequestDeadline::starting_now(ms),
7532 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7533 };
7534 let lane = match lane_for_tenant(&headers, &tenant) {
7535 Ok(l) => l,
7536 Err(resp) => return resp,
7537 };
7538 let model = req.model.clone();
7539 let stream = req.stream;
7540 let capture_prompt = st
7543 .metering
7544 .as_ref()
7545 .filter(|m| m.captures(&tenant.tenant))
7546 .map(|_| capture_chat_messages(&req.messages));
7547 let declared_max_tokens = req.max_tokens.is_some();
7551 let vision_preprocess_permit = if request_has_vision(&req) {
7555 match VISION_PREPROCESS_SEMAPHORE.acquire().await {
7556 Ok(permit) => Some(permit),
7557 Err(_) => {
7558 return with_request_id(
7559 &env.id,
7560 error_response(
7561 StatusCode::SERVICE_UNAVAILABLE,
7562 "vision preprocessing is unavailable",
7563 "server_error",
7564 None,
7565 ),
7566 );
7567 }
7568 }
7569 } else {
7570 None
7571 };
7572 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7573 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7574 let mut plan = match build_chat_request_with_trace(
7575 req,
7576 st.caps.get(&model),
7577 tx,
7578 lane,
7579 affinity,
7580 ttft.clone(),
7581 st.openrouter_metadata
7582 .get(&model)
7583 .and_then(|m| m.default_reasoning_effort.as_deref()),
7584 &st.sampling_defaults(&model),
7585 ) {
7586 Ok(plan) => plan,
7587 Err(err) => {
7588 return with_request_id(&env.id, bad_request(&err, None));
7589 }
7590 };
7591 plan.request.cache_ns = cache_ns;
7592 if let Err((message, param)) = apply_model_request_limits(
7593 &mut plan.request,
7594 st.openrouter_metadata.get(&model),
7595 st.caps.get(&model),
7596 ) {
7597 return with_request_id(&env.id, bad_request(&message, Some(param)));
7598 }
7599 if let Err(msg) = nonstream_deadline_gate(
7602 &plan.request,
7603 stream,
7604 deadline,
7605 declared_max_tokens,
7606 st.budget_tokenizers
7607 .as_ref()
7608 .and_then(|t| t.get(&model))
7609 .map(Arc::as_ref),
7610 ) {
7611 return with_request_id(
7612 &env.id,
7613 error_response_coded(
7614 StatusCode::BAD_REQUEST,
7615 &msg,
7616 "invalid_request_error",
7617 Some("max_tokens"),
7618 Some("nonstream_deadline_infeasible"),
7619 ),
7620 );
7621 }
7622 plan.vision_memory = match reserve_vision_memory(&plan) {
7623 Ok(permit) => permit,
7624 Err(err) => {
7625 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
7626 }
7627 };
7628 if draining() {
7631 let receipt = start_request_receipt(
7632 &st,
7633 &env,
7634 &tenant,
7635 &model,
7636 "/v1/chat/completions",
7637 lane,
7638 stream,
7639 None,
7640 );
7641 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7642 }
7643 let budget_permit = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
7644 Ok(permit) => permit,
7645 Err(rejection) => {
7646 let (response, error_code) = rejection.into_response();
7647 let receipt = start_request_receipt(
7648 &st,
7649 &env,
7650 &tenant,
7651 &model,
7652 "/v1/chat/completions",
7653 lane,
7654 stream,
7655 None,
7656 );
7657 return ledger_rejected(receipt, response, error_code, &env.id);
7658 }
7659 };
7660 let receipt = start_request_receipt(
7661 &st,
7662 &env,
7663 &tenant,
7664 &model,
7665 "/v1/chat/completions",
7666 lane,
7667 stream,
7668 budget_permit,
7669 );
7670 let receipt = if let Some(prompt) = capture_prompt {
7671 arm_capture(receipt, move || prompt)
7672 } else {
7673 receipt
7674 };
7675 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7679 Ok(slot) => slot,
7680 Err(resp) => {
7681 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7682 }
7683 };
7684 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7687 Ok(guard) => guard,
7688 Err((resp, outcome)) => {
7689 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7690 }
7691 };
7692 if let Err(err) = decode_pending_vision(&mut plan) {
7697 return ledger_rejected(
7698 receipt,
7699 rl.attach(bad_request(&err, Some("messages"))),
7700 "invalid_request_error",
7701 &env.id,
7702 );
7703 }
7704 plan.request.vision_memory = plan.vision_memory.take();
7705 drop(vision_preprocess_permit);
7706 let constraint_ready = if plan.request.grammar.is_some() {
7707 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
7708 plan.request.constraint_ready = Some(ready_tx);
7709 Some(ready_rx)
7710 } else {
7711 None
7712 };
7713 meter_admit(&env, &tenant, &model, lane);
7714 let stop_strings = plan.request.stop_strings.clone();
7715 if let Some(trace) = ttft.as_ref() {
7717 trace.mark_submitted();
7718 }
7719 if st
7720 .cmd_tx
7721 .send(Cmd::Generate(Box::new(plan.request)))
7722 .is_err()
7723 {
7724 drop(pending_admit);
7725 return ledger_rejected(
7726 receipt,
7727 rl.attach(worker_unavailable_response()),
7728 "worker_unavailable",
7729 &env.id,
7730 );
7731 }
7732 pending_admit.commit();
7733 if let Some(ready) = constraint_ready {
7739 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
7740 match tokio::time::timeout(bound, ready).await {
7741 Ok(Ok(Ok(()))) => {}
7742 Ok(Ok(Err(err))) => {
7743 return ledger_rejected(
7744 receipt,
7745 rl.attach(engine_error_response(&err)),
7746 engine_error_code(err.class),
7747 &env.id,
7748 );
7749 }
7750 Ok(Err(_)) => {
7751 return ledger_rejected(
7752 receipt,
7753 rl.attach(worker_unavailable_response()),
7754 "worker_unavailable",
7755 &env.id,
7756 );
7757 }
7758 Err(_) if deadline.remaining().is_zero() => {
7759 return ledger_unbilled(
7760 receipt,
7761 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7762 "deadline_exceeded",
7763 "deadline_exceeded",
7764 &env.id,
7765 );
7766 }
7767 Err(_) => {
7768 return ledger_rejected(
7769 receipt,
7770 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
7771 "constraint_compile_timeout",
7772 &env.id,
7773 );
7774 }
7775 }
7776 }
7777 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7779 Ok(Ok(rx)) => rx,
7780 Ok(Err((resp, error_code))) => {
7781 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7782 }
7783 Err(_) => {
7784 return ledger_unbilled(
7785 receipt,
7786 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7787 "deadline_exceeded",
7788 "deadline_exceeded",
7789 &env.id,
7790 );
7791 }
7792 };
7793 let resp = if stream {
7794 let rx = match peek_first_token(rx, deadline).await {
7796 Ok(rx) => rx,
7797 Err(()) => {
7798 return ledger_unbilled(
7799 receipt,
7800 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7801 "deadline_exceeded",
7802 "deadline_exceeded",
7803 &env.id,
7804 );
7805 }
7806 };
7807 sse_response_with_receipt(
7808 rx,
7809 model,
7810 true,
7811 plan.parser,
7812 env.clone(),
7813 stop_strings,
7814 Some(guard),
7815 receipt,
7816 )
7817 .into_response()
7818 } else {
7819 let mut receipt = receipt;
7822 let resp = blocking_response_with_receipt(
7823 rx,
7824 model,
7825 true,
7826 stop_strings,
7827 plan.parser,
7828 env.clone(),
7829 &mut receipt,
7830 Some(deadline),
7831 )
7832 .await;
7833 drop(guard); resp.into_response()
7835 };
7836 rl.attach(with_request_id(&env.id, resp))
7837}
7838
7839#[cfg(test)]
7848fn sse_response(
7849 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
7850 model: String,
7851 chat: bool,
7852 parser: Option<ToolStreamParser>,
7853 env: Envelope,
7854 stop_strings: Vec<String>,
7855 guard: Option<InflightGuard>,
7856) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
7857 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
7858}
7859
7860fn sse_response_with_receipt(
7861 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
7862 model: String,
7863 chat: bool,
7864 mut parser: Option<ToolStreamParser>,
7865 env: Envelope,
7866 stop_strings: Vec<String>,
7867 guard: Option<InflightGuard>,
7868 mut receipt: Option<Box<dyn metering::Receipt>>,
7869) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
7870 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
7874 .then(|| StopScrubber::new(stop_strings));
7875 let stream = async_stream::stream! {
7876 let _guard = guard;
7879 let mut call_index: usize = 0;
7880 let mut role_sent = false;
7883 macro_rules! chat_chunk {
7884 ($delta:expr, $finish:expr) => {{
7885 let mut delta = $delta;
7886 if chat && !role_sent {
7887 role_sent = true;
7888 delta["role"] = json!("assistant");
7889 }
7890 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
7891 "choices": [{ "index": 0, "delta": delta,
7892 "finish_reason": $finish }] }))
7893 .to_string()
7894 }};
7895 }
7896 macro_rules! piece_chunks {
7898 ($piece:expr) => {{
7899 let mut payloads: Vec<String> = Vec::new();
7900 match $piece {
7901 Piece::Content(text) => {
7902 let text = match scrub.as_mut() {
7903 Some(sc) => sc.push(&text),
7904 None => text,
7905 };
7906 if !text.is_empty() {
7907 payloads.push(chat_chunk!(json!({ "content": text }),
7908 serde_json::Value::Null));
7909 }
7910 }
7911 Piece::Reasoning(text) => payloads.push(
7915 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
7916 Piece::Call(call) => {
7917 payloads.push(chat_chunk!(json!({ "tool_calls": [{
7918 "index": call_index, "id": call.id, "type": "function",
7919 "function": { "name": call.name, "arguments": "" } }] }),
7920 serde_json::Value::Null));
7921 payloads.push(chat_chunk!(json!({ "tool_calls": [{
7922 "index": call_index,
7923 "function": { "arguments": call.arguments } }] }),
7924 serde_json::Value::Null));
7925 call_index += 1;
7926 }
7927 }
7928 payloads
7929 }};
7930 }
7931 let mut terminal = false;
7935 while let Some(ev) = rx.recv().await {
7936 match ev {
7937 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
7939 if let Some(receipt) = receipt.as_mut()
7940 && let Err(err) = receipt.record_prompt_usage(
7941 n_prompt as u64,
7942 n_cached as u64,
7943 )
7944 {
7945 eprintln!(
7946 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
7947 env.id
7948 );
7949 let _ = receipt.reject(500, "request_ledger_unavailable");
7952 let payload = request_ledger_error_body().to_string();
7953 if chat || openai_compat() {
7954 yield Ok(SseEvent::default().data(payload));
7955 yield Ok(SseEvent::default().data("[DONE]"));
7956 } else {
7957 yield Ok(SseEvent::default().event("error").data(payload));
7958 }
7959 terminal = true;
7960 break;
7961 }
7962 }
7963 Event::Token { id, text } => {
7964 if let Some(receipt) = receipt.as_mut()
7965 && let Err(err) = receipt.record_completion_token()
7966 {
7967 eprintln!(
7968 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
7969 env.id
7970 );
7971 let _ = receipt.reject(500, "request_ledger_unavailable");
7972 let payload = request_ledger_error_body().to_string();
7973 if chat || openai_compat() {
7974 yield Ok(SseEvent::default().data(payload));
7975 yield Ok(SseEvent::default().data("[DONE]"));
7976 } else {
7977 yield Ok(SseEvent::default().event("error").data(payload));
7978 }
7979 terminal = true;
7980 break;
7981 }
7982 if let Some(receipt) = receipt.as_mut() {
7985 receipt.capture_completion_delta(&text);
7986 }
7987 if let Some(p) = parser.as_mut() {
7988 for piece in p.push(&text) {
7989 for payload in piece_chunks!(piece) {
7990 yield Ok(SseEvent::default().data(payload));
7991 }
7992 }
7993 continue;
7994 }
7995 let text = match scrub.as_mut() {
7996 Some(sc) => sc.push(&text),
7997 None => text,
7998 };
7999 if text.is_empty() && scrub.is_some() {
8000 continue; }
8002 let payload = if chat {
8003 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
8004 } else if openai_compat() {
8005 env.stamp(json!({ "object": "text_completion", "model": model,
8006 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
8007 .to_string()
8008 } else {
8009 json!({ "model": model, "id": id, "text": text }).to_string()
8010 };
8011 yield Ok(SseEvent::default().data(payload));
8012 }
8013 Event::TokenSnapshot(_) => {}
8017 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
8018 let mut finish = stop_reason_to_finish(&stop_reason);
8019 if let Some(p) = parser.as_mut() {
8020 for piece in p.finish() {
8021 for payload in piece_chunks!(piece) {
8022 yield Ok(SseEvent::default().data(payload));
8023 }
8024 }
8025 if p.n_calls() > 0 { finish = "tool_calls"; }
8026 }
8027 if let Some(sc) = scrub.as_mut() {
8029 let tail = sc.finish();
8030 if !tail.is_empty() {
8031 let payload = if chat {
8032 chat_chunk!(json!({ "content": tail }),
8033 serde_json::Value::Null)
8034 } else {
8035 env.stamp(json!({ "object": "text_completion",
8036 "model": model,
8037 "choices": [{ "index": 0, "text": tail,
8038 "finish_reason": null }] })).to_string()
8039 };
8040 yield Ok(SseEvent::default().data(payload));
8041 }
8042 }
8043 if let Some(receipt) = receipt.as_mut()
8044 && let Err(err) = receipt.complete(
8045 metering::UsageCounts {
8046 prompt_tokens: n_prompt as u64,
8047 cached_prompt_tokens: n_cached as u64,
8048 completion_tokens: n_tokens as u64,
8049 },
8050 elapsed_s,
8051 )
8052 {
8053 eprintln!(
8054 "[ledger] ERROR: request {} completion receipt failed: {err}",
8055 env.id
8056 );
8057 let _ = receipt.reject(500, "request_ledger_unavailable");
8061 let payload = request_ledger_error_body().to_string();
8062 if chat || openai_compat() {
8063 yield Ok(SseEvent::default().data(payload));
8064 yield Ok(SseEvent::default().data("[DONE]"));
8065 } else {
8066 yield Ok(SseEvent::default().event("error").data(payload));
8067 }
8068 terminal = true;
8069 break;
8070 }
8071 if chat || openai_compat() {
8072 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
8073 let fin = if chat {
8074 let mut v = env.stamp(json!({
8075 "object": "chat.completion.chunk", "model": model,
8076 "choices": [{ "index": 0, "delta": {},
8077 "finish_reason": finish }],
8078 "usage": usage }));
8079 if !role_sent {
8081 v["choices"][0]["delta"]["role"] = json!("assistant");
8082 }
8083 v
8084 } else {
8085 env.stamp(json!({ "object": "text_completion", "model": model,
8086 "choices": [{ "index": 0, "text": "",
8087 "finish_reason": finish }],
8088 "usage": usage }))
8089 }.to_string();
8090 yield Ok(SseEvent::default().data(fin));
8091 yield Ok(SseEvent::default().data("[DONE]"));
8092 } else {
8093 let payload = json!({
8094 "stop_reason": stop_reason, "n_tokens": n_tokens,
8095 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
8096 "elapsed_s": elapsed_s
8097 }).to_string();
8098 yield Ok(SseEvent::default().event("done").data(payload));
8099 }
8100 terminal = true;
8101 break;
8102 }
8103 Event::Error(err) => {
8104 let ledger_error = if let Some(receipt) = receipt.as_mut() {
8114 receipt
8115 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
8116 .err()
8117 } else {
8118 None
8119 };
8120 if let Some(ref ledger_error) = ledger_error {
8121 eprintln!(
8122 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
8123 env.id
8124 );
8125 }
8126 let payload = if ledger_error.is_some() {
8127 request_ledger_error_body().to_string()
8128 } else {
8129 engine_error_body(&err).to_string()
8130 };
8131 if chat || openai_compat() {
8132 yield Ok(SseEvent::default().data(payload));
8135 yield Ok(SseEvent::default().data("[DONE]"));
8136 } else {
8137 yield Ok(SseEvent::default().event("error").data(payload));
8140 }
8141 terminal = true;
8142 break;
8143 }
8144 }
8145 }
8146 if !terminal {
8147 let e = worker::EngineError::overloaded(
8153 "worker closed the stream without completing (worker restart in progress)",
8154 );
8155 if let Some(receipt) = receipt.as_mut()
8156 && let Err(ledger_err) = receipt.reject(
8157 class_http(e.class).0.as_u16(),
8158 engine_error_code(e.class),
8159 )
8160 {
8161 eprintln!(
8162 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8163 env.id
8164 );
8165 }
8166 let payload = engine_error_body(&e).to_string();
8167 if chat || openai_compat() {
8168 yield Ok(SseEvent::default().data(payload));
8169 yield Ok(SseEvent::default().data("[DONE]"));
8170 } else {
8171 yield Ok(SseEvent::default().event("error").data(payload));
8172 }
8173 }
8174 };
8175 Sse::new(stream).keep_alive(
8176 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
8179 )
8180}
8181
8182fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
8184 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
8185 text.truncate(offset);
8186 }
8187}
8188
8189fn partial_stop_suffix(s: &str, tag: &str) -> usize {
8192 let mut best = 0;
8193 for (k, _) in tag.char_indices().skip(1) {
8194 if k <= s.len() && s.ends_with(&tag[..k]) {
8195 best = k;
8196 }
8197 }
8198 best
8199}
8200
8201struct StopScrubber {
8207 stops: Vec<String>,
8208 buf: String,
8209 done: bool,
8210}
8211
8212impl StopScrubber {
8213 fn new(stops: Vec<String>) -> Self {
8214 Self {
8215 stops,
8216 buf: String::new(),
8217 done: false,
8218 }
8219 }
8220
8221 fn push(&mut self, text: &str) -> String {
8223 if self.done {
8224 return String::new();
8225 }
8226 self.buf.push_str(text);
8227 if let Some(i) = self
8228 .stops
8229 .iter()
8230 .filter_map(|s| self.buf.find(s.as_str()))
8231 .min()
8232 {
8233 self.done = true;
8234 let out = self.buf[..i].to_string();
8235 self.buf.clear();
8236 return out;
8237 }
8238 let keep = self
8239 .stops
8240 .iter()
8241 .map(|s| partial_stop_suffix(&self.buf, s))
8242 .max()
8243 .unwrap_or(0);
8244 let emit_to = self.buf.len() - keep;
8245 let out = self.buf[..emit_to].to_string();
8246 self.buf.drain(..emit_to);
8247 out
8248 }
8249
8250 fn finish(&mut self) -> String {
8252 if self.done {
8253 self.buf.clear();
8254 return String::new();
8255 }
8256 std::mem::take(&mut self.buf)
8257 }
8258}
8259
8260#[cfg(test)]
8261async fn blocking_response(
8262 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8263 model: String,
8264 chat: bool,
8265 stop_strings: Vec<String>,
8266 parser: Option<ToolStreamParser>,
8267 env: Envelope,
8268) -> Response {
8269 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
8270 .await
8271}
8272
8273struct BlockingPayload<'a> {
8277 env: &'a Envelope,
8278 model: String,
8279 chat: bool,
8280 finish: &'static str,
8281 text: String,
8282 reasoning: String,
8283 calls: Vec<ParsedToolCall>,
8284 tokens: Vec<u32>,
8285 stop_reason: String,
8286 n_prompt: usize,
8287 n_tokens: usize,
8288 n_cached: usize,
8289 elapsed_s: f64,
8290 spec: Option<worker::SpecUsage>,
8291 deadline_error: Option<serde_json::Value>,
8297}
8298
8299fn blocking_payload(p: BlockingPayload<'_>) -> Response {
8300 let BlockingPayload {
8301 env,
8302 model,
8303 chat,
8304 finish,
8305 text,
8306 reasoning,
8307 calls,
8308 tokens,
8309 stop_reason,
8310 n_prompt,
8311 n_tokens,
8312 n_cached,
8313 elapsed_s,
8314 spec,
8315 deadline_error,
8316 } = p;
8317 if chat {
8318 let content = if !calls.is_empty() && text.is_empty() {
8320 serde_json::Value::Null
8321 } else {
8322 serde_json::Value::String(text)
8323 };
8324 let mut message = json!({ "role": "assistant", "content": content });
8325 if !reasoning.is_empty() {
8328 message["reasoning"] = json!(reasoning);
8329 message["reasoning_details"] = json!([{
8330 "type": "reasoning.text", "text": reasoning }]);
8331 }
8332 if !calls.is_empty() {
8333 message["tool_calls"] =
8334 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
8335 }
8336 let mut body = json!({
8337 "object": "chat.completion", "model": model,
8338 "choices": [{ "index": 0,
8339 "message": message,
8340 "finish_reason": finish }],
8341 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8342 });
8343 if let Some(err) = deadline_error {
8344 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8345 body["error"] = err;
8346 }
8347 return Json(env.stamp(body)).into_response();
8348 }
8349 if openai_compat() {
8350 let mut body = json!({
8351 "object": "text_completion", "model": model,
8352 "choices": [{ "index": 0, "text": text,
8353 "finish_reason": finish }],
8354 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8355 });
8356 if let Some(err) = deadline_error {
8357 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8358 body["error"] = err;
8359 }
8360 return Json(env.stamp(body)).into_response();
8361 }
8362 Json(CompletionResp {
8363 model,
8364 text,
8365 tokens,
8366 stop_reason,
8367 error: deadline_error,
8368 n_tokens,
8369 prompt_tokens: n_prompt,
8370 cached_tokens: n_cached,
8371 elapsed_s,
8372 })
8373 .into_response()
8374}
8375
8376async fn blocking_response_with_receipt(
8401 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8402 model: String,
8403 chat: bool,
8404 stop_strings: Vec<String>,
8405 mut parser: Option<ToolStreamParser>,
8406 env: Envelope,
8407 receipt: &mut Option<Box<dyn metering::Receipt>>,
8408 deadline: Option<RequestDeadline>,
8409) -> Response {
8410 let mut text = String::new();
8411 let mut reasoning = String::new();
8412 let mut tokens: Vec<u32> = Vec::new();
8413 let mut calls: Vec<ParsedToolCall> = Vec::new();
8414 let consume = |pieces: Vec<Piece>,
8415 text: &mut String,
8416 reasoning: &mut String,
8417 calls: &mut Vec<ParsedToolCall>| {
8418 for piece in pieces {
8419 match piece {
8420 Piece::Content(t) => text.push_str(&t),
8421 Piece::Reasoning(t) => reasoning.push_str(&t),
8422 Piece::Call(c) => calls.push(c),
8423 }
8424 }
8425 };
8426 let started = std::time::Instant::now();
8428 let mut seen_prompt: usize = 0;
8429 let mut seen_cached: usize = 0;
8430 let mut seen_tokens: usize = 0;
8431 loop {
8432 let ev = match deadline {
8433 Some(d) => tokio::select! {
8434 biased;
8435 ev = rx.recv() => ev,
8436 () = tokio::time::sleep_until(d.at) => {
8437 drop(rx);
8440 if seen_tokens == 0 {
8441 if let Some(receipt) = receipt.as_mut()
8445 && let Err(err) = receipt.settle_unbilled(
8446 "deadline_exceeded",
8447 StatusCode::REQUEST_TIMEOUT.as_u16(),
8448 "deadline_exceeded",
8449 )
8450 {
8451 eprintln!(
8452 "[ledger] ERROR: request {} deadline receipt failed: {err}",
8453 env.id
8454 );
8455 return request_ledger_error_response();
8456 }
8457 return deadline_exceeded_response(d.ms, false);
8458 }
8459 if let Some(p) = parser.as_mut() {
8460 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8461 }
8462 truncate_at_stop(&mut text, &stop_strings);
8463 let elapsed_s = started.elapsed().as_secs_f64();
8464 if let Some(receipt) = receipt.as_mut()
8467 && let Err(err) = receipt.complete_deadline_partial(
8468 metering::UsageCounts {
8469 prompt_tokens: seen_prompt as u64,
8470 cached_prompt_tokens: seen_cached as u64,
8471 completion_tokens: seen_tokens as u64,
8472 },
8473 elapsed_s,
8474 )
8475 {
8476 eprintln!(
8477 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
8478 env.id
8479 );
8480 let _ = receipt.reject(500, "request_ledger_unavailable");
8481 return request_ledger_error_response();
8482 }
8483 eprintln!(
8484 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
8485 {} ms deadline (prompt {}); non-streaming caller advised to stream",
8486 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
8487 );
8488 let err_obj = json!({
8489 "message": format!(
8490 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
8491 the {} tokens produced before the cut are delivered above and are \
8492 billed. Set \"stream\": true for work this long — a stream's \
8493 deadline bounds only the time to first token — or lower max_tokens.",
8494 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
8495 ),
8496 "code": "deadline_exceeded",
8497 "metadata": { "error_type": "timeout", "provider_name": "memra" }
8498 });
8499 return blocking_payload(BlockingPayload {
8500 env: &env,
8501 model,
8502 chat,
8503 finish: "error",
8504 text,
8505 reasoning,
8506 calls,
8507 tokens,
8508 stop_reason: "Deadline".to_string(),
8509 n_prompt: seen_prompt,
8510 n_tokens: seen_tokens,
8511 n_cached: seen_cached,
8512 elapsed_s,
8513 spec: None,
8514 deadline_error: Some(err_obj),
8515 });
8516 }
8517 },
8518 None => rx.recv().await,
8519 };
8520 let Some(ev) = ev else { break };
8521 match ev {
8522 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8524 if let Some(receipt) = receipt.as_mut()
8525 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
8526 {
8527 eprintln!(
8528 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8529 env.id
8530 );
8531 let _ = receipt.reject(500, "request_ledger_unavailable");
8534 return request_ledger_error_response();
8535 }
8536 seen_prompt = n_prompt;
8537 seen_cached = n_cached;
8538 }
8539 Event::Token { id, text: delta } => {
8540 if let Some(receipt) = receipt.as_mut()
8541 && let Err(err) = receipt.record_completion_token()
8542 {
8543 eprintln!(
8544 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8545 env.id
8546 );
8547 let _ = receipt.reject(500, "request_ledger_unavailable");
8548 return request_ledger_error_response();
8549 }
8550 if let Some(receipt) = receipt.as_mut() {
8552 receipt.capture_completion_delta(&delta);
8553 }
8554 tokens.push(id);
8555 seen_tokens += 1;
8556 match parser.as_mut() {
8557 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
8558 None => text.push_str(&delta),
8559 }
8560 }
8561 Event::TokenSnapshot(ids) => tokens = ids,
8562 Event::Done {
8563 stop_reason,
8564 n_tokens,
8565 n_prompt,
8566 n_cached,
8567 elapsed_s,
8568 spec,
8569 } => {
8570 if let Some(p) = parser.as_mut() {
8571 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8572 }
8573 truncate_at_stop(&mut text, &stop_strings);
8574 let finish = if calls.is_empty() {
8575 stop_reason_to_finish(&stop_reason)
8576 } else {
8577 "tool_calls"
8578 };
8579 if let Some(receipt) = receipt.as_mut()
8580 && let Err(err) = receipt.complete(
8581 metering::UsageCounts {
8582 prompt_tokens: n_prompt as u64,
8583 cached_prompt_tokens: n_cached as u64,
8584 completion_tokens: n_tokens as u64,
8585 },
8586 elapsed_s,
8587 )
8588 {
8589 eprintln!(
8590 "[ledger] ERROR: request {} completion receipt failed: {err}",
8591 env.id
8592 );
8593 let _ = receipt.reject(500, "request_ledger_unavailable");
8596 return request_ledger_error_response();
8597 }
8598 return blocking_payload(BlockingPayload {
8599 env: &env,
8600 model,
8601 chat,
8602 finish,
8603 text,
8604 reasoning,
8605 calls,
8606 tokens,
8607 stop_reason,
8608 n_prompt,
8609 n_tokens,
8610 n_cached,
8611 elapsed_s,
8612 spec,
8613 deadline_error: None,
8614 });
8615 }
8616 Event::Error(err) => {
8617 if let Some(receipt) = receipt.as_mut()
8621 && let Err(ledger_err) = receipt.reject(
8622 class_http(err.class).0.as_u16(),
8623 engine_error_code(err.class),
8624 )
8625 {
8626 eprintln!(
8627 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
8628 env.id
8629 );
8630 return request_ledger_error_response();
8631 }
8632 return engine_error_response(&err);
8633 }
8634 }
8635 }
8636 let e = worker::EngineError::overloaded(
8641 "worker closed the stream without completing (worker restart in progress)",
8642 );
8643 if let Some(receipt) = receipt.as_mut()
8644 && let Err(ledger_err) =
8645 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
8646 {
8647 eprintln!(
8648 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8649 env.id
8650 );
8651 return request_ledger_error_response();
8652 }
8653 engine_error_response(&e)
8654}
8655
8656#[cfg(test)]
8657mod tests {
8658 use super::*;
8659
8660 #[derive(Debug, Clone, PartialEq)]
8668 enum MeterEvent {
8669 Reserve {
8670 tenant: String,
8671 model: String,
8672 },
8673 Open {
8674 request_id: String,
8675 tenant: String,
8676 model: String,
8677 route: &'static str,
8678 stream: bool,
8679 with_permit: bool,
8680 },
8681 PromptUsage {
8682 prompt: u64,
8683 cached: u64,
8684 },
8685 Token,
8686 CapturePrompt(serde_json::Value),
8687 CaptureDelta(String),
8688 Complete {
8689 prompt: u64,
8690 cached: u64,
8691 completion: u64,
8692 },
8693 DeadlinePartial {
8694 prompt: u64,
8695 cached: u64,
8696 completion: u64,
8697 },
8698 Reject {
8699 status: u16,
8700 code: String,
8701 },
8702 Unbilled {
8703 outcome: &'static str,
8704 status: u16,
8705 code: String,
8706 },
8707 Dropped {
8710 prompt: u64,
8711 cached: u64,
8712 completion: u64,
8713 },
8714 }
8715
8716 enum ReserveScript {
8719 Admit { with_permit: bool },
8720 Insufficient,
8721 Blocked,
8722 Unenrolled,
8723 }
8724
8725 struct MockMetering {
8726 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
8727 limits: bool,
8728 limited: bool,
8729 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
8730 captures: bool,
8731 }
8732
8733 impl MockMetering {
8734 fn admit_all() -> Arc<Self> {
8735 Arc::new(MockMetering {
8736 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8737 limits: false,
8738 limited: true,
8739 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
8740 captures: false,
8741 })
8742 }
8743
8744 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
8745 Arc::new(MockMetering {
8746 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8747 limits: true,
8748 limited: true,
8749 reserve_script: std::sync::Mutex::new(script.into()),
8750 captures: false,
8751 })
8752 }
8753
8754 fn capturing() -> Arc<Self> {
8755 Arc::new(MockMetering {
8756 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8757 limits: false,
8758 limited: true,
8759 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
8760 captures: true,
8761 })
8762 }
8763
8764 fn events(&self) -> Vec<MeterEvent> {
8765 self.events.lock().unwrap().clone()
8766 }
8767 }
8768
8769 impl metering::Metering for MockMetering {
8770 fn enforces_limits(&self) -> bool {
8771 self.limits
8772 }
8773
8774 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
8775 Ok(self.limited)
8776 }
8777
8778 fn reserve(
8779 &self,
8780 tenant: &str,
8781 model: &str,
8782 _prompt_tokens: u64,
8783 _completion_bound: u64,
8784 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
8785 self.events.lock().unwrap().push(MeterEvent::Reserve {
8786 tenant: tenant.into(),
8787 model: model.into(),
8788 });
8789 match self.reserve_script.lock().unwrap().pop_front() {
8790 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
8791 Some(ReserveScript::Admit { with_permit: true }) => {
8792 Ok(Some(Box::new(()) as metering::Permit))
8793 }
8794 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
8795 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
8796 Some(ReserveScript::Unenrolled) => Err(metering::AdmitError::Unenrolled),
8797 }
8798 }
8799
8800 fn open(
8801 &self,
8802 meta: &metering::RequestMeta<'_>,
8803 permit: Option<metering::Permit>,
8804 ) -> Box<dyn metering::Receipt> {
8805 self.events.lock().unwrap().push(MeterEvent::Open {
8806 request_id: meta.request_id.into(),
8807 tenant: meta.tenant.into(),
8808 model: meta.model.into(),
8809 route: meta.route,
8810 stream: meta.stream,
8811 with_permit: permit.is_some(),
8812 });
8813 Box::new(MockReceipt {
8814 events: self.events.clone(),
8815 wants_capture: self.captures,
8816 prompt: 0,
8817 cached: 0,
8818 completion: 0,
8819 finalized: false,
8820 })
8821 }
8822
8823 fn captures(&self, _tenant: &str) -> bool {
8824 self.captures
8825 }
8826
8827 fn limits_health(&self) -> Option<metering::LimitsHealth> {
8828 self.limits.then_some(metering::LimitsHealth {
8829 source_reload_failed: 0,
8830 source_reload_consecutive: 0,
8831 source_available: true,
8832 })
8833 }
8834 }
8835
8836 struct MockReceipt {
8837 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
8838 wants_capture: bool,
8839 prompt: u64,
8840 cached: u64,
8841 completion: u64,
8842 finalized: bool,
8843 }
8844
8845 impl metering::Receipt for MockReceipt {
8846 fn wants_capture(&self) -> bool {
8847 self.wants_capture
8848 }
8849
8850 fn arm_capture(&mut self, prompt: serde_json::Value) {
8851 self.events
8852 .lock()
8853 .unwrap()
8854 .push(MeterEvent::CapturePrompt(prompt));
8855 }
8856
8857 fn capture_completion_delta(&mut self, text: &str) {
8858 if self.wants_capture {
8859 self.events
8860 .lock()
8861 .unwrap()
8862 .push(MeterEvent::CaptureDelta(text.into()));
8863 }
8864 }
8865
8866 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
8867 self.prompt = prompt;
8868 self.cached = cached;
8869 self.events
8870 .lock()
8871 .unwrap()
8872 .push(MeterEvent::PromptUsage { prompt, cached });
8873 Ok(())
8874 }
8875
8876 fn record_completion_token(&mut self) -> Result<(), String> {
8877 self.completion += 1;
8878 self.events.lock().unwrap().push(MeterEvent::Token);
8879 Ok(())
8880 }
8881
8882 fn complete(
8883 &mut self,
8884 usage: metering::UsageCounts,
8885 _worker_elapsed_s: f64,
8886 ) -> Result<(), String> {
8887 self.finalized = true;
8888 self.events.lock().unwrap().push(MeterEvent::Complete {
8889 prompt: usage.prompt_tokens,
8890 cached: usage.cached_prompt_tokens,
8891 completion: usage.completion_tokens,
8892 });
8893 Ok(())
8894 }
8895
8896 fn complete_deadline_partial(
8897 &mut self,
8898 usage: metering::UsageCounts,
8899 _worker_elapsed_s: f64,
8900 ) -> Result<(), String> {
8901 self.finalized = true;
8902 self.events
8903 .lock()
8904 .unwrap()
8905 .push(MeterEvent::DeadlinePartial {
8906 prompt: usage.prompt_tokens,
8907 cached: usage.cached_prompt_tokens,
8908 completion: usage.completion_tokens,
8909 });
8910 Ok(())
8911 }
8912
8913 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
8914 self.finalized = true;
8915 self.events.lock().unwrap().push(MeterEvent::Reject {
8916 status,
8917 code: error_code.into(),
8918 });
8919 Ok(())
8920 }
8921
8922 fn settle_unbilled(
8923 &mut self,
8924 outcome: &'static str,
8925 status: u16,
8926 error_code: &str,
8927 ) -> Result<(), String> {
8928 self.finalized = true;
8929 self.events.lock().unwrap().push(MeterEvent::Unbilled {
8930 outcome,
8931 status,
8932 code: error_code.into(),
8933 });
8934 Ok(())
8935 }
8936 }
8937
8938 impl Drop for MockReceipt {
8939 fn drop(&mut self) {
8940 if !self.finalized {
8941 self.events.lock().unwrap().push(MeterEvent::Dropped {
8942 prompt: self.prompt,
8943 cached: self.cached,
8944 completion: self.completion,
8945 });
8946 }
8947 }
8948 }
8949
8950 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
8955
8956 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
8959 let req: CompletionReq = serde_json::from_value(json!({
8960 "model": "qwen/qwen3.8-27b",
8961 "prompt_ids": vec![7u32; prompt_ids],
8962 }))
8963 .unwrap();
8964 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
8965 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
8966 request.params.max_new = max_new;
8967 request
8968 }
8969
8970 #[test]
8977 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
8978 let prompt = 30_278u64;
8979 let deadline_ms = TIMEOUT_MS_DEFAULT;
8980 let margin = |max_new: u64| {
8981 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
8982 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
8983 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
8984 };
8985 for allowed in [64u64, 2048, 4096, 5120, 6144] {
8986 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
8987 }
8988 for refused in [8192u64, 16384, 262_144] {
8989 assert!(
8990 !margin(refused),
8991 "{refused} measured as a 408 and must be refused"
8992 );
8993 }
8994 }
8995
8996 #[test]
8997 fn the_gate_names_a_max_tokens_that_actually_fits() {
8998 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
9001 assert!(
9002 fits > 0 && fits < 7_800,
9003 "advice {fits} must fit the measured ceiling"
9004 );
9005 assert_eq!(
9007 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
9008 None
9009 );
9010 }
9011
9012 #[test]
9013 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
9014 let req = gate_request(262_144, 30_000);
9015 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
9016 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
9018 assert!(
9019 err.contains("stream"),
9020 "message must name the streaming alternative: {err}"
9021 );
9022 assert!(
9023 err.contains("max_tokens"),
9024 "message must name the knob: {err}"
9025 );
9026 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
9028 let _l = GATE_ENV_LOCK.lock().unwrap(); for off in ["0", "off", "false"] {
9035 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
9036 assert!(
9037 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
9038 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
9039 );
9040 }
9041 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
9042 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
9043 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
9044 assert!(
9045 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
9046 "unset means ON (the documented default)"
9047 );
9048 }
9049
9050 #[test]
9057 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
9058 let strip = |src: &str| -> String {
9065 src.lines()
9066 .map(|line| match line.find("//") {
9067 Some(i) => line[..i].to_string(),
9068 None => line.to_string(),
9069 })
9070 .collect::<Vec<_>>()
9071 .join("\n")
9072 };
9073 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
9075 let start = src
9076 .find(signature)
9077 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
9078 let rest = &src[start + signature.len()..];
9079 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
9080 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
9081 &rest[..end]
9082 }
9083 let main_src = strip(include_str!("lib.rs"));
9084 let surfaces_src = strip(include_str!("surfaces.rs"));
9085 for (surface, src, signature) in [
9086 ("/v1/completions", &main_src, "async fn completions("),
9087 (
9088 "/v1/chat/completions",
9089 &main_src,
9090 "async fn chat_completions(",
9091 ),
9092 (
9093 "/v1/messages + /v1/responses (shared admission)",
9094 &surfaces_src,
9095 "pub(crate) async fn admit_translated(",
9096 ),
9097 ] {
9098 let handler = body(src, signature);
9099 assert!(
9100 handler.contains("nonstream_deadline_gate("),
9101 "{surface} must CALL the feasibility gate inside {signature}"
9102 );
9103 let limits = handler
9106 .find("apply_model_request_limits(")
9107 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
9108 let gate = handler.find("nonstream_deadline_gate(").unwrap();
9109 assert!(
9110 limits < gate,
9111 "{surface}: the gate must run after apply_model_request_limits"
9112 );
9113 }
9114 }
9115
9116 #[test]
9120 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
9121 let err = json!({"code": "deadline_exceeded",
9122 "metadata": {"error_type": "timeout"}});
9123 let cut = CompletionResp {
9124 model: "m".into(),
9125 text: "partial".into(),
9126 tokens: vec![1, 2],
9127 stop_reason: "Deadline".into(),
9128 error: Some(err.clone()),
9129 n_tokens: 2,
9130 prompt_tokens: 9,
9131 cached_tokens: 0,
9132 elapsed_s: 1.0,
9133 };
9134 let v = serde_json::to_value(&cut).unwrap();
9135 assert_eq!(v["stop_reason"], "Deadline");
9136 assert_eq!(v["error"]["code"], "deadline_exceeded");
9137 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
9138 let whole = CompletionResp {
9140 error: None,
9141 stop_reason: "Eos".into(),
9142 ..cut
9143 };
9144 let v = serde_json::to_value(&whole).unwrap();
9145 assert!(
9146 v.get("error").is_none(),
9147 "a complete response must not grow an error key: {v}"
9148 );
9149 }
9150
9151 #[test]
9152 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
9153 let _l = GATE_ENV_LOCK.lock().unwrap();
9154 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
9158 assert!(
9159 nonstream_deadline_gate(
9160 &req,
9161 false,
9162 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9163 false,
9164 None,
9165 )
9166 .is_ok(),
9167 "an omitted max_tokens is never gated — context is its only limit"
9168 );
9169 let resolved = gate_request(32_768, 30_000);
9174 assert!(
9175 nonstream_deadline_gate(
9176 &resolved,
9177 false,
9178 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9179 false,
9180 None,
9181 )
9182 .is_ok(),
9183 "a resolved-but-undeclared cap is not the caller's number to be refused over"
9184 );
9185 assert!(
9187 nonstream_deadline_gate(
9188 &resolved,
9189 false,
9190 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9191 true,
9192 None,
9193 )
9194 .is_err()
9195 );
9196 }
9197
9198 #[test]
9199 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
9200 let req = gate_request(64, 1234);
9201 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
9202 let mut text = gate_request(64, 0);
9203 text.prompt_ids.clear();
9204 text.prompt_text = "x".repeat(6_000);
9205 assert_eq!(
9206 prompt_tokens_estimate(&text, None),
9207 1_000,
9208 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
9209 that would have succeeded"
9210 );
9211 }
9212
9213 #[test]
9214 fn vision_memory_reservation_is_bounded_and_released() {
9215 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
9216 let Err(capacity) = try_reserve_vision_memory(1) else {
9217 panic!("a full process vision budget admitted another request");
9218 };
9219 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
9220 let response = vision_memory_error_response(capacity, Some("messages"));
9221 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
9222 assert_eq!(response.headers()["retry-after"], "5");
9223 assert_eq!(response.headers()["retry-after-ms"], "5000");
9224 drop(permit);
9225 assert!(try_reserve_vision_memory(1).is_ok());
9226 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
9227 panic!("an over-limit vision request was admitted");
9228 };
9229 assert!(matches!(request, VisionMemoryError::Request(_)));
9230 let response = vision_memory_error_response(request, Some("messages"));
9231 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
9232 assert_eq!(response.headers()["x-should-retry"], "false");
9233 let _ = try_reserve_vision_memory(1);
9234 }
9235
9236 #[test]
9237 fn header_auth_gate_covers_only_inference_dialects() {
9238 for path in [
9239 "/v1/auth/check",
9240 "/v1/completions",
9241 "/v1/chat/completions",
9242 "/v1/messages",
9243 "/v1/responses",
9244 "/v1/embeddings",
9245 "/v1/rerank",
9246 ] {
9247 assert!(protected_inference_path(path), "{path}");
9248 }
9249 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
9250 assert!(!protected_inference_path(path), "{path}");
9251 }
9252 }
9253 #[tokio::test]
9260 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
9261 use crate::metering::Metering as _;
9262 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
9263
9264 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
9265 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
9266 tx.send(Event::PromptUsage {
9267 n_prompt: 7,
9268 n_cached: 0,
9269 })
9270 .unwrap();
9271 tx.send(Event::Token {
9272 id: 1,
9273 text: "Hel".into(),
9274 })
9275 .unwrap();
9276 tx.send(Event::Token {
9277 id: 2,
9278 text: "lo".into(),
9279 })
9280 .unwrap();
9281 tx.send(Event::Done {
9282 stop_reason: "eos".into(),
9283 n_tokens: 2,
9284 n_prompt: 7,
9285 n_cached: 0,
9286 elapsed_s: 0.05,
9287 spec: None,
9288 })
9289 .unwrap();
9290 drop(tx);
9291 let mut receipt = receipt;
9292 blocking_response_with_receipt(
9293 rx,
9294 "m".into(),
9295 true,
9296 Vec::new(),
9297 None,
9298 Envelope::new(true),
9299 &mut receipt,
9300 None,
9301 )
9302 .await
9303 };
9304
9305 let plain = MockMetering::admit_all();
9307 let receipt = plain.open(
9308 &metering::RequestMeta {
9309 request_id: "cap-unmarked",
9310 tenant: "unmarked",
9311 model: "m",
9312 route: "/v1/chat/completions",
9313 lane: "interactive",
9314 stream: false,
9315 },
9316 None,
9317 );
9318 let response = drive(Some(receipt)).await;
9319 assert_eq!(response.status(), StatusCode::OK);
9320 assert!(
9321 !plain.events().iter().any(|e| matches!(
9322 e,
9323 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
9324 )),
9325 "an unarmed receipt must see no capture traffic: {:?}",
9326 plain.events()
9327 );
9328
9329 let capturing = MockMetering::capturing();
9332 let mut receipt = capturing.open(
9333 &metering::RequestMeta {
9334 request_id: "cap-marked",
9335 tenant: "marked",
9336 model: "m",
9337 route: "/v1/chat/completions",
9338 lane: "interactive",
9339 stream: false,
9340 },
9341 None,
9342 );
9343 assert!(receipt.wants_capture());
9344 receipt.arm_capture(prompt.clone());
9345 let response = drive(Some(receipt)).await;
9346 assert_eq!(response.status(), StatusCode::OK);
9347 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
9348 .await
9349 .unwrap();
9350 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
9351 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
9352
9353 let events = capturing.events();
9354 assert!(
9355 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
9356 "prompt must arm byte-exact: {events:?}"
9357 );
9358 let completion: String = events
9359 .iter()
9360 .filter_map(|e| match e {
9361 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
9362 _ => None,
9363 })
9364 .collect();
9365 assert_eq!(
9366 completion, "Hello",
9367 "the deltas must reassemble the served completion byte-exact: {events:?}"
9368 );
9369 assert!(
9370 events.contains(&MeterEvent::Complete {
9371 prompt: 7,
9372 cached: 0,
9373 completion: 2,
9374 }),
9375 "worker-truth usage settles alongside the capture: {events:?}"
9376 );
9377 }
9378
9379 fn tool_caps() -> ModelCaps {
9380 ModelCaps {
9381 tools_branch: true,
9382 qwen_think: true,
9383 think_switch: true,
9384 chat_ok: true,
9385 ..Default::default()
9386 }
9387 }
9388
9389 fn ladder_caps() -> ModelCaps {
9394 ModelCaps {
9395 qwen_effort: true,
9396 ..tool_caps()
9397 }
9398 }
9399
9400 fn gemma_tool_caps() -> ModelCaps {
9401 ModelCaps {
9402 tools_branch: true,
9403 gemma_think: true,
9404 chat_ok: true,
9405 instruct_type: Some("gemma".into()),
9406 ..Default::default()
9407 }
9408 }
9409
9410 fn gemma_template(kind: &str) -> String {
9411 let file = match kind {
9412 "qat" => "qat-trunk-template.jinja",
9413 _ => "official-tooluse-template.jinja",
9414 };
9415 let path = format!(
9416 "{}/../../research/gemma4-tools-20260817/{file}",
9417 env!("CARGO_MANIFEST_DIR")
9418 );
9419 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9420 }
9421
9422 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
9427 let tools_arr = request
9428 .get("tools")
9429 .and_then(|t| t.as_array())
9430 .cloned()
9431 .unwrap_or_default();
9432 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
9433 (Vec::new(), Vec::new(), HashMap::new())
9434 } else {
9435 prepare_tools(&tools_arr).unwrap()
9436 };
9437 let effort = request
9438 .get("reasoning_effort")
9439 .and_then(|v| v.as_str())
9440 .map(String::from);
9441 let (think, _lvl, _explicit) =
9442 parse_think(&effort, &None, None, None, None, false).unwrap();
9443
9444 let mut turns: Vec<TmplTurn> = Vec::new();
9445 for msg in request["messages"].as_array().unwrap() {
9446 let role = msg["role"].as_str().unwrap();
9447 let role = if role == "developer" { "system" } else { role };
9448 let content =
9449 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9450 let tool_calls = msg
9451 .get("tool_calls")
9452 .and_then(|a| a.as_array())
9453 .map(|a| {
9454 a.iter()
9455 .map(|tc| {
9456 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9457 render_req_tool_call(&rtc).unwrap()
9458 })
9459 .collect()
9460 })
9461 .unwrap_or_default();
9462 let tool_responses = msg
9463 .get("tool_responses")
9464 .and_then(|a| a.as_array())
9465 .map(|a| {
9466 a.iter()
9467 .map(|tr| {
9468 (
9469 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
9470 json_to_val(&tr["response"]),
9471 )
9472 })
9473 .collect()
9474 })
9475 .unwrap_or_default();
9476 turns.push(TmplTurn {
9477 role: role.to_string(),
9478 content,
9479 tool_calls,
9480 reasoning: msg
9481 .get("reasoning")
9482 .and_then(|r| r.as_str())
9483 .map(String::from)
9484 .filter(|s| !s.is_empty()),
9485 tool_call_id: msg
9486 .get("tool_call_id")
9487 .and_then(|s| s.as_str())
9488 .map(String::from),
9489 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9490 tool_responses,
9491 task: None,
9492 tools: Vec::new(),
9493 });
9494 }
9495 chat::apply_chat_template_tools_ex(
9496 Some(template),
9497 &turns,
9498 true,
9499 &tools_json,
9500 &tools_struct,
9501 think,
9502 None,
9503 None,
9504 )
9505 .unwrap()
9506 }
9507
9508 #[test]
9512 fn gemma4_tools_fixtures_match_the_official_jinja() {
9513 let dir = format!(
9514 "{}/../../research/gemma4-tools-20260817/fixtures",
9515 env!("CARGO_MANIFEST_DIR")
9516 );
9517 let mut entries: Vec<_> = std::fs::read_dir(&dir)
9518 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
9519 .map(|e| e.unwrap().path())
9520 .filter(|p| p.is_dir())
9521 .collect();
9522 entries.sort();
9523 assert!(
9524 entries.len() >= 14,
9525 "expected >=14 fixtures, found {}",
9526 entries.len()
9527 );
9528 let (mut official, mut qat) = (0u32, 0u32);
9529 for d in entries {
9530 let input: serde_json::Value =
9531 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
9532 .unwrap();
9533 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
9534 let kind = input
9535 .get("template")
9536 .and_then(|t| t.as_str())
9537 .unwrap_or("official");
9538 match kind {
9539 "qat" => qat += 1,
9540 _ => official += 1,
9541 }
9542 let tmpl = gemma_template(kind);
9543 let got = render_fixture(&input["request"], &tmpl);
9544 assert_eq!(
9545 got, expected,
9546 "fixture {:?} diverged from the jinja oracle",
9547 d
9548 );
9549 }
9550 assert!(
9551 official >= 12 && qat >= 2,
9552 "coverage: {official} official, {qat} qat"
9553 );
9554 }
9555
9556 #[test]
9562 fn gemma4_tools_flow_through_build_chat_request() {
9563 let tmpl = gemma_template("official");
9564 for name in [
9565 "01-system-tools-basic",
9566 "04-single-call-cycle",
9567 "07-multi-cycle-agentic",
9568 ] {
9569 let path = format!(
9570 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
9571 env!("CARGO_MANIFEST_DIR")
9572 );
9573 let input: serde_json::Value =
9574 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
9575 let expected_path = format!(
9576 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
9577 env!("CARGO_MANIFEST_DIR")
9578 );
9579 let expected = std::fs::read_to_string(&expected_path).unwrap();
9580 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
9581 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9582 let plan = build_chat_request(
9583 req,
9584 Some(&gemma_tool_caps()),
9585 tx,
9586 lanes::Lane::Interactive,
9587 None,
9588 )
9589 .unwrap();
9590 let got = chat::apply_chat_template_tools_ex(
9591 Some(&tmpl),
9592 &plan.request.chat_turns,
9593 true,
9594 &plan.request.tools_json,
9595 &plan.request.tools_struct,
9596 plan.request.think,
9597 plan.request.reasoning_effort.as_deref(),
9598 None,
9599 )
9600 .unwrap();
9601 assert_eq!(got, expected, "pipeline render diverged for {name}");
9602 }
9603 }
9604
9605 fn dsv4_sentinel() -> String {
9613 let path = format!(
9614 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
9615 env!("CARGO_MANIFEST_DIR")
9616 );
9617 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9618 }
9619
9620 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
9625 let role = msg["role"].as_str().unwrap().to_string();
9626 let content =
9627 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9628 let reasoning = msg
9629 .get("reasoning")
9630 .or_else(|| msg.get("reasoning_content"))
9631 .and_then(|r| r.as_str())
9632 .map(String::from)
9633 .filter(|s| !s.is_empty());
9634 let tool_calls = msg
9635 .get("tool_calls")
9636 .and_then(|a| a.as_array())
9637 .map(|a| {
9638 a.iter()
9639 .map(|tc| {
9640 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9641 render_req_tool_call(&rtc).unwrap()
9642 })
9643 .collect()
9644 })
9645 .unwrap_or_default();
9646 let tools = msg
9647 .get("tools")
9648 .and_then(|a| a.as_array())
9649 .map(|a| {
9650 a.iter()
9651 .filter_map(|t| t.get("function").map(json_to_val))
9652 .collect()
9653 })
9654 .unwrap_or_default();
9655 TmplTurn {
9656 role,
9657 content,
9658 tool_calls,
9659 reasoning,
9660 tool_call_id: msg
9661 .get("tool_call_id")
9662 .and_then(|s| s.as_str())
9663 .map(String::from),
9664 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9665 tool_responses: Vec::new(),
9666 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
9667 tools,
9668 }
9669 }
9670
9671 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
9672 v.and_then(|t| t.as_array())
9673 .map(|a| {
9674 a.iter()
9675 .filter_map(|t| t.get("function").map(json_to_val))
9676 .collect()
9677 })
9678 .unwrap_or_default()
9679 }
9680
9681 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
9685 let dir = format!(
9686 "{}/../../research/dsv4-template-20260818/{subdir}",
9687 env!("CARGO_MANIFEST_DIR")
9688 );
9689 let tmpl = dsv4_sentinel();
9690 let mut entries: Vec<_> = std::fs::read_dir(&dir)
9691 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
9692 .map(|e| e.unwrap().path())
9693 .filter(|p| p.is_dir())
9694 .collect();
9695 entries.sort();
9696 assert!(
9697 entries.len() >= min_fixtures,
9698 "expected >={min_fixtures} fixtures, found {}",
9699 entries.len()
9700 );
9701 for d in &entries {
9702 let input: serde_json::Value =
9703 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
9704 .unwrap();
9705 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
9706 let turns: Vec<TmplTurn> = input["turns"]
9707 .as_array()
9708 .unwrap()
9709 .iter()
9710 .map(dsv4_turn)
9711 .collect();
9712 let think = match input["think"].as_str().unwrap() {
9713 "chat" => ThinkMode::NoThink,
9714 _ => ThinkMode::Think,
9715 };
9716 let effort = input
9717 .get("reasoning_effort")
9718 .and_then(|v| v.as_str())
9719 .map(String::from);
9720 let req_tools = dsv4_req_tools(input.get("req_tools"));
9721 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
9722 let got = chat::apply_chat_template_tools_ex(
9723 Some(&tmpl),
9724 &turns,
9725 agp,
9726 &[],
9727 &req_tools,
9728 think,
9729 effort.as_deref(),
9730 Some(encoding),
9731 )
9732 .unwrap();
9733 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
9734 }
9735 }
9736
9737 #[test]
9738 fn dsv4_template_fixtures_match_the_oracle() {
9739 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
9740 }
9741
9742 #[test]
9748 fn dsv4_0731_fixtures_match_the_oracle() {
9749 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
9750 }
9751
9752 #[test]
9753 fn dsv4_artifact_fixtures_are_byte_identical() {
9754 let base = format!(
9758 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
9759 env!("CARGO_MANIFEST_DIR")
9760 );
9761 let tmpl = dsv4_sentinel();
9762 for (n, think) in [
9763 (1u32, ThinkMode::Think),
9764 (2, ThinkMode::Think),
9765 (3, ThinkMode::Think),
9766 (4, ThinkMode::NoThink),
9767 ] {
9768 let td: serde_json::Value = serde_json::from_str(
9769 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
9770 )
9771 .unwrap();
9772 let (messages, tools) = if td.is_object() {
9773 (td["messages"].clone(), td.get("tools").cloned())
9774 } else {
9775 (td.clone(), None)
9776 };
9777 let mut turns: Vec<TmplTurn> = Vec::new();
9778 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
9779 let mut t = dsv4_turn(msg);
9780 if i == 0 {
9781 if let Some(tl) = &tools {
9782 t.tools = tl
9783 .as_array()
9784 .unwrap()
9785 .iter()
9786 .filter_map(|x| x.get("function").map(json_to_val))
9787 .collect();
9788 }
9789 }
9790 turns.push(t);
9791 }
9792 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
9793 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
9797 let got = chat::apply_chat_template_tools_ex(
9798 Some(&tmpl),
9799 &turns,
9800 true,
9801 &[],
9802 &[],
9803 think,
9804 None,
9805 Some(encoding),
9806 )
9807 .unwrap();
9808 assert_eq!(
9809 got, expected,
9810 "artifact fixture {n} diverged from the oracle under {encoding:?}"
9811 );
9812 }
9813 }
9814 }
9815
9816 #[test]
9817 fn dsv4_default_thinkmode_renders_thinking() {
9818 let tmpl = dsv4_sentinel();
9821 let turns = vec![TmplTurn {
9822 role: "user".into(),
9823 content: "Hi".into(),
9824 ..Default::default()
9825 }];
9826 let dflt = chat::apply_chat_template_tools_ex(
9827 Some(&tmpl),
9828 &turns,
9829 true,
9830 &[],
9831 &[],
9832 ThinkMode::Default,
9833 None,
9834 None,
9835 )
9836 .unwrap();
9837 let think = chat::apply_chat_template_tools_ex(
9838 Some(&tmpl),
9839 &turns,
9840 true,
9841 &[],
9842 &[],
9843 ThinkMode::Think,
9844 None,
9845 None,
9846 )
9847 .unwrap();
9848 assert_eq!(dflt, think);
9849 assert!(
9850 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
9851 "{dflt:?}"
9852 );
9853 let chat_mode = chat::apply_chat_template_tools_ex(
9854 Some(&tmpl),
9855 &turns,
9856 true,
9857 &[],
9858 &[],
9859 ThinkMode::NoThink,
9860 None,
9861 None,
9862 )
9863 .unwrap();
9864 assert!(
9865 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
9866 "{chat_mode:?}"
9867 );
9868 }
9869
9870 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
9875 let base = format!(
9876 "{}/../../research/dsv4-template-20260818",
9877 env!("CARGO_MANIFEST_DIR")
9878 );
9879 let refdir = std::path::Path::new(&base).join("ref");
9880 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
9881 .expect("load dsv4 tokenizer from ref dir");
9882 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
9883 let banked: serde_json::Value = serde_json::from_str(
9884 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
9885 .unwrap(),
9886 )
9887 .unwrap();
9888 let obj = banked.as_object().unwrap();
9889 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
9890 for (name, ids_v) in obj {
9891 let rendered =
9892 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
9893 let want: Vec<u32> = ids_v
9894 .as_array()
9895 .unwrap()
9896 .iter()
9897 .map(|v| v.as_u64().unwrap() as u32)
9898 .collect();
9899 let got = tok.encode(&rendered, true);
9900 assert_eq!(got, want, "tokenization diverged for {name}");
9901 }
9902 }
9903
9904 #[test]
9905 fn dsv4_tokenization_crosscheck_matches_official_ids() {
9906 dsv4_run_tokenization_crosscheck("fixtures");
9907 }
9908
9909 #[test]
9913 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
9914 dsv4_run_tokenization_crosscheck("fixtures-0731");
9915 }
9916
9917 #[test]
9918 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
9919 let base = format!(
9929 "{}/../../research/dsv4-template-20260818",
9930 env!("CARGO_MANIFEST_DIR")
9931 );
9932 let refdir = std::path::Path::new(&base).join("ref");
9933 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
9934 .expect("load dsv4 tokenizer from ref dir");
9935 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
9936 let tmpl = dsv4_sentinel();
9937 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
9938 {"type": "function", "function": {
9939 "name": "get_data",
9940 "description": "Fetch a blob",
9941 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
9942 "required": ["key"]}
9943 }}
9944 ])));
9945
9946 let cases: Vec<(&str, String)> = vec![
9947 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
9949 ("space-131k", " ".repeat(131_072)),
9950 ("digit-131k", "7".repeat(131_072)),
9951 (
9952 "mixed-runs",
9953 format!(
9954 "{}{}{}{}",
9955 "Z".repeat(65_536),
9956 " ".repeat(65_536),
9957 "7".repeat(65_536),
9958 "\n".repeat(65_536)
9959 ),
9960 ),
9961 ("cjk-64k", "中".repeat(65_536)),
9962 ("accented-letter-64k", "é".repeat(65_536)),
9963 ];
9964 for (name, blob) in &cases {
9965 let msgs = serde_json::json!([
9966 {"role": "system", "content": "You are a tool-using assistant."},
9967 {"role": "user", "content": "Fetch the blob."},
9968 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
9969 "tool_calls": [{"id": "call_001", "type": "function",
9970 "function": {"name": "get_data",
9971 "arguments": "{\"key\": \"blob\"}"}}]},
9972 {"role": "tool", "tool_call_id": "call_001", "content": blob}
9973 ]);
9974 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
9975 let rendered = chat::apply_chat_template_tools_ex(
9976 Some(&tmpl),
9977 &turns,
9978 true,
9979 &[],
9980 &req_tools,
9981 ThinkMode::Think,
9982 None,
9983 None,
9984 )
9985 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
9986 assert!(
9987 rendered.contains(blob.as_str()),
9988 "{name}: tool result missing from render"
9989 );
9990 let t0 = std::time::Instant::now();
9991 let ids = tok.encode(&rendered, true);
9992 let encode_dt = t0.elapsed();
9993 assert!(!ids.is_empty(), "{name}: empty encode");
9994 let back = tok.decode(&ids);
9995 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
9996 assert!(
10000 encode_dt < std::time::Duration::from_secs(60),
10001 "{name}: encode took {encode_dt:?}"
10002 );
10003 if *name == "ascii-letter-131k" {
10007 if let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR") {
10008 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
10009 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
10010 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
10011 }
10012 }
10013 }
10014 }
10015
10016 #[test]
10017 fn models_v1_entry_advertises_thinking_support() {
10018 let step_caps = ModelCaps {
10021 effort_levels: true,
10022 ..tool_caps()
10023 };
10024 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
10025 assert_eq!(entry["capabilities"]["reasoning"], true);
10026 assert_eq!(entry["capabilities"]["tools"], true);
10027
10028 let plain = ModelCaps {
10030 chat_ok: true,
10031 ..Default::default()
10032 };
10033 let entry = model_entry_v1("plain", Some(&plain), None);
10034 assert_eq!(entry["capabilities"]["reasoning"], false);
10035 assert_eq!(entry["capabilities"]["tools"], false);
10036 let entry = model_entry_v1("unknown", None, None);
10038 assert_eq!(entry["capabilities"]["reasoning"], false);
10039 assert_eq!(entry["capabilities"]["streaming"], true);
10040 }
10041
10042 #[test]
10043 fn chat_request_preserves_turns_and_openai_stop_forms() {
10044 let payload = serde_json::json!({
10045 "model": "plain_quant",
10046 "messages": [
10047 {"role": "system", "content": "rules"},
10048 {"role": "developer", "content": "dev rules"},
10049 {"role": "user", "content": "task"},
10050 {"role": "assistant", "content": "work"}
10051 ],
10052 "max_tokens": 64,
10053 "temperature": 0.0,
10054 "stop": "<stop>"
10055 });
10056 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10057 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10058 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10059 let request = plan.request;
10060 assert!(
10061 plan.parser.is_none(),
10062 "no tools -> no parser (isolation contract)"
10063 );
10064 assert!(request.tools_json.is_empty());
10065 assert_eq!(request.think, ThinkMode::Default);
10066 assert_eq!(request.model, "plain_quant");
10067 assert_eq!(request.params.max_new, 64);
10068 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10070 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
10071 }))
10072 .unwrap();
10073 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10074 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10075 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
10076 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10078 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10079 "max_completion_tokens": 7
10080 }))
10081 .unwrap();
10082 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10083 assert_eq!(
10084 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
10085 .unwrap()
10086 .request
10087 .params
10088 .max_new,
10089 7
10090 );
10091 let req: CompletionReq = serde_json::from_value(serde_json::json!({
10093 "model": "plain_quant", "prompt": "task"
10094 }))
10095 .unwrap();
10096 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10097 assert_eq!(
10098 build_request(&req, tx, lanes::Lane::Interactive, None)
10099 .params
10100 .max_new,
10101 worker::MAX_NEW_CTX_BOUNDED
10102 );
10103 let turns: Vec<(String, String)> = request
10104 .chat_turns
10105 .iter()
10106 .map(|t| (t.role.clone(), t.content.clone()))
10107 .collect();
10108 assert_eq!(
10109 turns,
10110 vec![
10111 ("system".into(), "rules".into()),
10112 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
10114 ("assistant".into(), "work".into()),
10115 ]
10116 );
10117 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
10118 assert_eq!(request.stop_strings, vec!["<stop>"]);
10119
10120 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10121 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10122 "stop": ["a", "b"]
10123 }))
10124 .unwrap();
10125 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
10126
10127 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10131 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10132 "stop": ["", "real", ""]
10133 }))
10134 .unwrap();
10135 assert_eq!(req.stop.into_vec(), vec!["real"]);
10136 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10137 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10138 "stop": ""
10139 }))
10140 .unwrap();
10141 assert!(req.stop.into_vec().is_empty());
10142
10143 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10144 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10145 "stop": null
10146 }))
10147 .unwrap();
10148 assert!(req.stop.into_vec().is_empty());
10149 }
10150
10151 #[tokio::test]
10152 async fn chat_response_has_openai_message_shape() {
10153 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10154 tx.send(Event::Token {
10155 id: 1,
10156 text: "hello".into(),
10157 })
10158 .unwrap();
10159 tx.send(Event::Done {
10160 stop_reason: "Eos".into(),
10161 n_tokens: 1,
10162 n_prompt: 42,
10163 n_cached: 30,
10164 elapsed_s: 0.5,
10165 spec: None,
10166 })
10167 .unwrap();
10168 drop(tx);
10169 let response = blocking_response(
10170 rx,
10171 "plain_quant".into(),
10172 true,
10173 Vec::new(),
10174 None,
10175 Envelope::new(true),
10176 )
10177 .await;
10178 assert_eq!(response.status(), StatusCode::OK);
10179 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10180 .await
10181 .unwrap();
10182 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10183 assert_eq!(payload["object"], "chat.completion");
10184 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
10186 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
10187 assert!(
10188 payload["system_fingerprint"]
10189 .as_str()
10190 .unwrap()
10191 .starts_with("memra-")
10192 );
10193 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
10194 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
10195 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
10196 assert_eq!(payload["usage"]["prompt_tokens"], 42);
10198 assert_eq!(payload["usage"]["completion_tokens"], 1);
10199 assert_eq!(payload["usage"]["total_tokens"], 43);
10200 assert_eq!(
10201 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
10202 30
10203 );
10204 assert!(payload["usage"].get("spec").is_none());
10207 }
10208
10209 #[tokio::test]
10210 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
10211 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10212 tx.send(Event::Token {
10214 id: 4,
10215 text: "hello".into(),
10216 })
10217 .unwrap();
10218 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
10219 tx.send(Event::Done {
10220 stop_reason: "MaxNew".into(),
10221 n_tokens: 4,
10222 n_prompt: 2,
10223 n_cached: 0,
10224 elapsed_s: 0.5,
10225 spec: None,
10226 })
10227 .unwrap();
10228 drop(tx);
10229
10230 let response = blocking_response(
10231 rx,
10232 "plain_quant".into(),
10233 false,
10234 Vec::new(),
10235 None,
10236 Envelope::new(false),
10237 )
10238 .await;
10239 assert_eq!(response.status(), StatusCode::OK);
10240 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10241 .await
10242 .unwrap();
10243 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10244 assert_eq!(payload["text"], "hello");
10245 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
10246 assert_eq!(payload["n_tokens"], 4);
10247 }
10248
10249 #[tokio::test]
10252 async fn chat_usage_carries_spec_acceptance_summary() {
10253 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10254 tx.send(Event::Token {
10255 id: 1,
10256 text: "hello".into(),
10257 })
10258 .unwrap();
10259 tx.send(Event::Done {
10260 stop_reason: "Eos".into(),
10261 n_tokens: 1,
10262 n_prompt: 42,
10263 n_cached: 0,
10264 elapsed_s: 0.5,
10265 spec: Some(worker::SpecUsage {
10266 rounds: 10,
10267 drafted: 30,
10268 accepted: 21,
10269 }),
10270 })
10271 .unwrap();
10272 drop(tx);
10273 let response = blocking_response(
10274 rx,
10275 "plain_quant".into(),
10276 true,
10277 Vec::new(),
10278 None,
10279 Envelope::new(true),
10280 )
10281 .await;
10282 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10283 .await
10284 .unwrap();
10285 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10286 let sp = &payload["usage"]["spec"];
10287 assert_eq!(sp["rounds"], 10);
10288 assert_eq!(sp["drafted"], 30);
10289 assert_eq!(sp["accepted"], 21);
10290 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
10291 assert_eq!(payload["usage"]["total_tokens"], 43);
10293 }
10294
10295 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
10296 let mut payload = serde_json::json!({
10297 "model": "m",
10298 "messages": [{"role": "user", "content": "Weather in Paris?"}],
10299 "tools": [{"type": "function", "function": {
10300 "name": "get_weather",
10301 "description": "Get current weather",
10302 "parameters": {"type": "object",
10303 "properties": {"city": {"type": "string"},
10304 "days": {"type": "integer"}},
10305 "required": ["city"]}}}],
10306 });
10307 if let Some(obj) = extra.as_object() {
10308 for (k, v) in obj {
10309 payload[k] = v.clone();
10310 }
10311 }
10312 serde_json::from_value(payload).unwrap()
10313 }
10314
10315 #[test]
10316 fn vision_decode_is_deferred_and_grid_pinned() {
10317 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10322 let req: ChatCompletionReq = serde_json::from_value(json!({
10323 "model": "m", "messages": [{"role": "user", "content": "hi"}],
10324 }))
10325 .unwrap();
10326 let mut plan = build_chat_request(
10327 req,
10328 Some(&ModelCaps {
10329 chat_ok: true,
10330 ..Default::default()
10331 }),
10332 tx,
10333 lanes::Lane::Interactive,
10334 None,
10335 )
10336 .unwrap();
10337 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
10341 let mut b = Vec::new();
10342 b.extend_from_slice(b"BM");
10343 b.extend_from_slice(&54u32.to_le_bytes());
10344 b.extend_from_slice(&0u32.to_le_bytes());
10345 b.extend_from_slice(&54u32.to_le_bytes());
10346 b.extend_from_slice(&40u32.to_le_bytes());
10347 b.extend_from_slice(&w.to_le_bytes());
10348 b.extend_from_slice(&h.to_le_bytes());
10349 b.extend_from_slice(&1u16.to_le_bytes());
10350 b.extend_from_slice(&24u16.to_le_bytes());
10351 b.extend_from_slice(&[0u8; 24]);
10352 if with_pixels {
10353 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
10354 }
10355 b
10356 };
10357 let bytes = bmp(64, 64, true);
10358 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
10359 plan.pending_images.push(PendingVisionUnit::Still {
10360 bytes: bytes.clone(),
10361 gh,
10362 gw,
10363 });
10364 decode_pending_vision(&mut plan).unwrap();
10365 assert_eq!(plan.request.images.len(), 1);
10366 assert_eq!(
10367 (
10368 plan.request.images[0].prep.gh,
10369 plan.request.images[0].prep.gw
10370 ),
10371 (gh, gw),
10372 "decoded grid must equal the header-planned grid the pad run was rendered from"
10373 );
10374 plan.request.images.clear();
10376 plan.pending_images.push(PendingVisionUnit::Still {
10377 bytes,
10378 gh: gh + 2,
10379 gw,
10380 });
10381 let err = decode_pending_vision(&mut plan).unwrap_err();
10382 assert!(err.contains("header-planned"), "got: {err}");
10383 let bomb = bmp(16_000, 16_000, false);
10386 plan.pending_images.clear();
10387 plan.pending_images.push(PendingVisionUnit::Still {
10388 bytes: bomb,
10389 gh: 2,
10390 gw: 2,
10391 });
10392 let err = decode_pending_vision(&mut plan).unwrap_err();
10393 assert!(err.contains("exceeds the decode budget"), "got: {err}");
10394 }
10395
10396 #[test]
10397 fn tools_request_renders_client_key_order_and_arms_parser() {
10398 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10399 let plan = build_chat_request(
10400 weather_request(json!({})),
10401 Some(&tool_caps()),
10402 tx,
10403 lanes::Lane::Interactive,
10404 None,
10405 )
10406 .unwrap();
10407 assert!(plan.parser.is_some());
10408 assert_eq!(plan.request.tools_json.len(), 1);
10409 assert_eq!(
10411 plan.request.tools_json[0],
10412 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
10413 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
10414 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
10415 \"integer\"}}, \"required\": [\"city\"]}}}"
10416 );
10417 }
10418
10419 #[test]
10420 fn tool_choice_none_strips_tools_and_parser() {
10421 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10422 let plan = build_chat_request(
10423 weather_request(json!({"tool_choice": "none"})),
10424 Some(&tool_caps()),
10425 tx,
10426 lanes::Lane::Interactive,
10427 None,
10428 )
10429 .unwrap();
10430 let mut p = plan
10433 .parser
10434 .expect("think-open chat arms the reasoning splitter");
10435 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
10436 assert_eq!(
10437 pieces,
10438 vec![
10439 Piece::Reasoning("x".into()),
10440 Piece::Content("<tool_call> stays prose".into()),
10441 ]
10442 );
10443 assert!(plan.request.tools_json.is_empty());
10444 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10446 assert!(
10447 build_chat_request(
10448 weather_request(json!({"tool_choice": "required"})),
10449 Some(&tool_caps()),
10450 tx,
10451 lanes::Lane::Interactive,
10452 None
10453 )
10454 .is_err()
10455 );
10456 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10457 assert!(
10458 build_chat_request(
10459 weather_request(json!({"tool_choice":
10460 {"type": "function", "function": {"name": "get_weather"}}})),
10461 Some(&tool_caps()),
10462 tx,
10463 lanes::Lane::Interactive,
10464 None
10465 )
10466 .is_err()
10467 );
10468 }
10469
10470 #[test]
10471 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
10472 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
10473 let _ = std::fs::remove_dir_all(&root);
10474
10475 let st = root.join("st_single");
10477 std::fs::create_dir_all(&st).unwrap();
10478 std::fs::write(st.join("config.json"), "{}").unwrap();
10479 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
10480 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
10481
10482 let sh = root.join("st_sharded");
10484 std::fs::create_dir_all(&sh).unwrap();
10485 std::fs::write(sh.join("config.json"), "{}").unwrap();
10486 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
10487 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
10488
10489 let rp = root.join("repack");
10491 std::fs::create_dir_all(&rp).unwrap();
10492 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
10493 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
10494
10495 let bogus = root.join("bogus");
10497 std::fs::create_dir_all(&bogus).unwrap();
10498 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
10499 assert!(
10500 err.contains("model.safetensors"),
10501 "error should say what is missing: {err}"
10502 );
10503 assert!(
10504 err.contains("manifest.json"),
10505 "error should mention the repack form: {err}"
10506 );
10507
10508 let nc = root.join("no_config");
10510 std::fs::create_dir_all(&nc).unwrap();
10511 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
10512 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
10513 assert!(
10514 err.contains("config.json"),
10515 "error should name config.json: {err}"
10516 );
10517
10518 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
10520 assert!(err.contains("does not exist"), "{err}");
10521
10522 let f = root.join("model.gguf");
10524 std::fs::write(&f, b"g").unwrap();
10525 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
10526
10527 let _ = std::fs::remove_dir_all(&root);
10528 }
10529
10530 #[test]
10531 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
10532 let caps = ModelCaps {
10535 tools_branch: false,
10536 qwen_think: false,
10537 think_switch: false,
10538 chat_ok: false,
10539 ..Default::default()
10540 };
10541 let payload = serde_json::json!({
10542 "model": "st_model",
10543 "messages": [{"role": "user", "content": "hello"}],
10544 });
10545 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10546 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10547 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
10548 Err(e) => e,
10549 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
10550 };
10551 assert!(
10552 err.contains("no chat template"),
10553 "message should name the cause: {err}"
10554 );
10555 assert!(
10556 err.contains("/v1/completions"),
10557 "message should point at the raw-prompt escape hatch: {err}"
10558 );
10559 }
10560
10561 #[test]
10562 fn tools_on_model_without_tools_branch_is_rejected() {
10563 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10564 let caps = ModelCaps {
10565 chat_ok: true,
10566 ..Default::default()
10567 };
10568 assert!(
10569 build_chat_request(
10570 weather_request(json!({})),
10571 Some(&caps),
10572 tx,
10573 lanes::Lane::Interactive,
10574 None
10575 )
10576 .is_err()
10577 );
10578 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10579 assert!(
10580 build_chat_request(
10581 weather_request(json!({})),
10582 None,
10583 tx,
10584 lanes::Lane::Interactive,
10585 None
10586 )
10587 .is_err()
10588 );
10589 }
10590
10591 #[test]
10592 fn reasoning_effort_maps_to_think_switch() {
10593 for (extra, want) in [
10599 (json!({}), ThinkMode::Default),
10600 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
10601 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
10602 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
10603 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
10604 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
10605 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
10606 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
10607 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
10608 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
10612 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
10613 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
10614 (
10618 json!({"reasoning": {"enabled": true, "effort": "none"}}),
10619 ThinkMode::Think,
10620 ),
10621 (
10622 json!({"reasoning": {"enabled": false, "effort": "high"}}),
10623 ThinkMode::NoThink,
10624 ),
10625 ] {
10626 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10627 let plan = build_chat_request(
10628 weather_request(extra.clone()),
10629 Some(&ladder_caps()),
10634 tx,
10635 lanes::Lane::Interactive,
10636 None,
10637 )
10638 .unwrap();
10639 assert_eq!(plan.request.think, want, "extra={extra}");
10640 }
10641 for extra in [
10645 json!({"reasoning_effort": "extreme"}),
10646 json!({"reasoning": {"effort": "banana"}}),
10647 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
10648 json!({"reasoning": {"enabled": true, "effort": ""}}),
10649 ] {
10650 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10651 assert!(
10652 build_chat_request(
10653 weather_request(extra.clone()),
10654 Some(&tool_caps()),
10655 tx,
10656 lanes::Lane::Interactive,
10657 None
10658 )
10659 .is_err(),
10660 "extra={extra} must be rejected by the one allowlist"
10661 );
10662 }
10663 for (raw, want) in [
10666 ("none", Some("none")),
10667 ("minimal", Some("minimal")),
10668 ("low", Some("low")),
10669 ("medium", Some("medium")),
10670 ("high", Some("high")),
10671 ("xhigh", Some("high")),
10672 ("max", Some("high")),
10673 ("ultra", Some("high")),
10674 ("banana", None),
10675 ("", None),
10676 ("HIGH", None),
10677 ] {
10678 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
10679 }
10680 for (raw, want) in [
10683 ("none", Some("none")),
10684 ("minimal", Some("minimal")),
10685 ("low", Some("low")),
10686 ("medium", Some("medium")),
10687 ("high", Some("high")),
10688 ("xhigh", Some("max")),
10689 ("max", Some("max")),
10690 ("ultra", Some("max")),
10691 ("banana", None),
10692 ("", None),
10693 ("MAX", None),
10694 ] {
10695 assert_eq!(
10696 canonical_effort_for(raw, true),
10697 want,
10698 "canonical_effort_for({raw:?}, dsv4)"
10699 );
10700 }
10701 }
10702
10703 #[test]
10704 fn dsv4_reasoning_effort_max_survives_canonicalization() {
10705 let dsv4_caps = ModelCaps {
10711 chat_ok: true,
10712 dsv4: true,
10713 ..Default::default()
10714 };
10715 let build = |caps: &ModelCaps, effort: &str| {
10716 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10717 let req: ChatCompletionReq = serde_json::from_value(json!({
10718 "model": "m",
10719 "messages": [{"role": "user", "content": "hi"}],
10720 "reasoning_effort": effort,
10721 }))
10722 .unwrap();
10723 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
10724 };
10725 for raw in ["max", "xhigh", "ultra"] {
10726 let plan = build(&dsv4_caps, raw).unwrap();
10727 assert_eq!(
10728 plan.request.reasoning_effort.as_deref(),
10729 Some("max"),
10730 "dsv4 {raw:?} must reach the renderer as the max rung"
10731 );
10732 assert_eq!(plan.request.think, chat::ThinkMode::Think);
10733 }
10734 let plan = build(&dsv4_caps, "high").unwrap();
10736 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
10737 let step_caps = ModelCaps {
10739 chat_ok: true,
10740 effort_levels: true,
10741 ..Default::default()
10742 };
10743 let plan = build(&step_caps, "max").unwrap();
10744 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
10745 }
10746
10747 #[test]
10748 fn default_reasoning_effort_flips_only_the_unset_request() {
10749 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
10756 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10757 build_chat_request_with_trace(
10758 weather_request(extra),
10759 Some(&ladder_caps()),
10760 tx,
10761 lanes::Lane::Interactive,
10762 None,
10763 None,
10764 default_effort,
10765 &ModelSamplingDefaults::default(),
10766 )
10767 .unwrap()
10768 };
10769 for (extra, want) in [
10770 (json!({}), ThinkMode::Think),
10772 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
10776 (json!({"include_reasoning": false}), ThinkMode::NoThink),
10777 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
10779 (json!({"include_reasoning": true}), ThinkMode::Think),
10780 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
10782 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
10783 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
10784 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
10786 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
10787 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
10788 ] {
10789 let plan = build(extra.clone(), Some("high"));
10790 assert_eq!(plan.request.think, want, "extra={extra}");
10791 }
10792 assert_eq!(
10794 build(json!({}), Some("none")).request.think,
10795 ThinkMode::NoThink
10796 );
10797 assert_eq!(
10798 build(json!({"reasoning_effort": "high"}), Some("none"))
10799 .request
10800 .think,
10801 ThinkMode::Think
10802 );
10803 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
10807 }
10808
10809 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
10815 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
10816 {%- else %}'<think>\\n'{%- endif %}";
10817
10818 #[test]
10819 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
10820 let build = |extra: serde_json::Value| {
10826 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10827 build_chat_request(
10828 weather_request(extra),
10829 Some(&tool_caps()),
10830 tx,
10831 lanes::Lane::Interactive,
10832 None,
10833 )
10834 };
10835 for (extra, want) in [
10836 (json!({"enable_thinking": false}), ThinkMode::NoThink),
10837 (json!({"enable_thinking": true}), ThinkMode::Think),
10838 (
10839 json!({"chat_template_kwargs": {"enable_thinking": false}}),
10840 ThinkMode::NoThink,
10841 ),
10842 (
10843 json!({"chat_template_kwargs": {"enable_thinking": true}}),
10844 ThinkMode::Think,
10845 ),
10846 (
10849 json!({"enable_thinking": false, "reasoning_effort": "high"}),
10850 ThinkMode::NoThink,
10851 ),
10852 (
10854 json!({"enable_thinking": false,
10855 "chat_template_kwargs": {"enable_thinking": false}}),
10856 ThinkMode::NoThink,
10857 ),
10858 ] {
10859 let plan = build(extra.clone()).unwrap_or_else(|e| {
10860 panic!("{extra} must be accepted and honored, got 400: {e}");
10861 });
10862 assert_eq!(
10863 plan.request.think, want,
10864 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
10865 );
10866 }
10867 let render = |extra: serde_json::Value| -> String {
10870 let plan = build(extra).unwrap();
10871 chat::apply_chat_template_tools_ex(
10872 Some(SWITCHED_QWEN_TMPL),
10873 &plan.request.chat_turns,
10874 true,
10875 &plan.request.tools_json,
10876 &plan.request.tools_struct,
10877 plan.request.think,
10878 plan.request.reasoning_effort.as_deref(),
10879 None,
10880 )
10881 .unwrap()
10882 };
10883 let off = render(json!({"enable_thinking": false}));
10884 assert!(
10885 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
10886 "enable_thinking:false must render the CLOSED think pair: {off:?}"
10887 );
10888 let on = render(json!({}));
10889 assert!(
10890 on.ends_with("<|im_start|>assistant\n<think>\n"),
10891 "an unset request must still render the template's OPEN think tail: {on:?}"
10892 );
10893 assert_eq!(
10894 off,
10895 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
10896 "both vLLM spellings must render byte-identically"
10897 );
10898 assert_eq!(
10899 off,
10900 render(json!({"reasoning_effort": "none"})),
10901 "the vLLM spelling must render byte-identically to the OpenAI spelling"
10902 );
10903 }
10904
10905 #[test]
10906 fn unknown_chat_template_kwarg_refuses_by_name() {
10907 let build = |extra: serde_json::Value| {
10910 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10911 build_chat_request(
10912 weather_request(extra),
10913 Some(&tool_caps()),
10914 tx,
10915 lanes::Lane::Interactive,
10916 None,
10917 )
10918 };
10919 let refusal = |extra: serde_json::Value, why: &str| -> String {
10920 build(extra).err().unwrap_or_else(|| panic!("{why}"))
10921 };
10922 let err = refusal(
10923 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
10924 "an unimplementable template kwarg must not be accepted",
10925 );
10926 assert!(
10927 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
10928 "the refusal must name the offending key AND the supported one: {err}"
10929 );
10930 let err = refusal(
10931 json!({"chat_template_kwargs": "enable_thinking=false"}),
10932 "a non-object chat_template_kwargs must not be accepted",
10933 );
10934 assert!(
10935 err.contains("must be an object"),
10936 "refusal must say what shape is expected: {err}"
10937 );
10938 let err = refusal(
10939 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
10940 "a stringly-typed switch must not be accepted",
10941 );
10942 assert!(
10943 err.contains("true or false"),
10944 "refusal must name the expected type: {err}"
10945 );
10946 let plan = build(json!({"chat_template_kwargs": null}))
10948 .expect("null chat_template_kwargs is the unset case");
10949 assert_eq!(plan.request.think, ThinkMode::Default);
10950 }
10951
10952 const Q38_TMPL: &str =
10967 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
10968
10969 fn render_with(
10972 tmpl: &str,
10973 caps: &ModelCaps,
10974 extra: serde_json::Value,
10975 default_effort: Option<&str>,
10976 ) -> Result<String, String> {
10977 let mut payload = serde_json::json!({
10978 "model": "m",
10979 "messages": [{"role": "user", "content": "hi"}],
10980 });
10981 if let Some(obj) = extra.as_object() {
10982 for (k, v) in obj {
10983 payload[k] = v.clone();
10984 }
10985 }
10986 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10987 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10988 let plan = build_chat_request_with_trace(
10989 req,
10990 Some(caps),
10991 tx,
10992 lanes::Lane::Interactive,
10993 None,
10994 None,
10995 default_effort,
10996 &ModelSamplingDefaults::default(),
10997 )?;
10998 Ok(chat::apply_chat_template_tools_ex(
10999 Some(tmpl),
11000 &plan.request.chat_turns,
11001 true,
11002 &plan.request.tools_json,
11003 &plan.request.tools_struct,
11004 plan.request.think,
11005 plan.request.reasoning_effort.as_deref(),
11006 None,
11007 )
11008 .unwrap())
11009 }
11010
11011 #[test]
11012 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
11013 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
11020 let xhigh = "Reasoning effort is set to xhigh.";
11021 let low = "Reasoning effort is set to low.";
11022 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
11024 assert!(
11025 r(json!({"reasoning_effort": "high"}))
11026 .unwrap()
11027 .contains(xhigh)
11028 );
11029 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
11032 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
11033 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
11035 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
11036 assert_ne!(low_p, high_p);
11037 assert_ne!(low_p, medium);
11038 assert_ne!(high_p, medium);
11039 for alias in ["xhigh", "max", "ultra"] {
11042 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
11043 }
11044 assert_eq!(r(json!({})).unwrap(), high_p);
11047 assert_eq!(
11050 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
11051 medium
11052 );
11053 let off = r(json!({"reasoning_effort": "none"})).unwrap();
11056 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
11057 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
11058 }
11059
11060 #[test]
11061 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
11062 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
11075focused, moving directly to the conclusion without unnecessary elaboration.";
11076 let expected = format!(
11077 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
11078 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
11079 );
11080 let after_fix = render_with(
11082 Q38_TMPL,
11083 &ladder_caps(),
11084 json!({"reasoning_effort": "low"}),
11085 None,
11086 )
11087 .unwrap();
11088 assert_eq!(
11089 after_fix, expected,
11090 "the shipped prompt for reasoning_effort:\"low\""
11091 );
11092 const ORNITH_TMPL: &str = include_str!(
11095 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11096 );
11097 let on_deployed_binary = render_with(
11098 ORNITH_TMPL,
11099 &tool_caps(),
11100 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
11101 {"role": "user", "content": "hi"}]}),
11102 None,
11103 )
11104 .unwrap();
11105 assert_eq!(
11106 on_deployed_binary, expected,
11107 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
11108 level, or its reasoning-volume numbers do not describe the shipped prompt"
11109 );
11110 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
11113 assert!(
11114 !ladderless_unset.contains("Reasoning effort is set to"),
11115 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
11116 );
11117 assert_eq!(
11118 ladderless_unset,
11119 render_with(
11120 Q38_TMPL,
11121 &ladder_caps(),
11122 json!({"reasoning_effort": "medium"}),
11123 None
11124 )
11125 .unwrap(),
11126 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
11127 );
11128 }
11129
11130 #[test]
11131 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
11132 let off = render_with(
11139 Q38_TMPL,
11140 &ladder_caps(),
11141 json!({"reasoning_effort": "none"}),
11142 None,
11143 )
11144 .unwrap();
11145 for extra in [
11146 json!({"include_reasoning": false}),
11147 json!({"reasoning": {"exclude": true}}),
11148 ] {
11149 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
11150 assert!(
11151 got.ends_with("<think>\n\n</think>\n\n"),
11152 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
11153 );
11154 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
11155 }
11156 for extra in [
11161 json!({"enable_thinking": true, "include_reasoning": false}),
11162 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
11163 json!({"reasoning": {"enabled": true, "exclude": true}}),
11164 ] {
11165 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
11166 .err()
11167 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
11168 assert!(e.contains("contradictory"), "{extra}: {e}");
11169 assert!(
11170 e.contains("include_reasoning") || e.contains("exclude"),
11171 "{extra}: the refusal must name the suppression field the caller sent: {e}"
11172 );
11173 }
11174 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
11177 for extra in [
11178 json!({"include_reasoning": true}),
11179 json!({"reasoning": {"exclude": false}}),
11180 ] {
11181 assert_eq!(
11182 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
11183 dflt,
11184 "{extra} must not perturb the model's default"
11185 );
11186 }
11187 let switchless = ModelCaps {
11191 think_switch: false,
11192 ..tool_caps()
11193 };
11194 let err = render_with(
11195 Q38_TMPL,
11196 &switchless,
11197 json!({"include_reasoning": false}),
11198 None,
11199 )
11200 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
11201 assert!(err.contains("cannot disable reasoning"), "{err}");
11202 }
11203
11204 #[test]
11205 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
11206 let build = |extra: serde_json::Value| {
11207 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11208 build_chat_request(
11209 weather_request(extra),
11210 Some(&ladder_caps()),
11211 tx,
11212 lanes::Lane::Interactive,
11213 None,
11214 )
11215 };
11216 let err = |extra: serde_json::Value, why: &str| -> String {
11217 build(extra).err().unwrap_or_else(|| panic!("{why}"))
11218 };
11219 let e = err(
11223 json!({"reasoning": {"max_tokens": 1024}}),
11224 "reasoning.max_tokens must not be accepted-and-ignored",
11225 );
11226 assert!(e.contains("reasoning.max_tokens"), "{e}");
11227 assert!(e.contains("ONE output budget"), "{e}");
11228 for extra in [
11233 json!({"reasoning": {"max_tokens": null}}),
11234 json!({"reasoning": {"banana": null}}),
11235 ] {
11236 let e = err(
11237 extra.clone(),
11238 "a null-valued unhonourable key must still refuse",
11239 );
11240 assert!(
11241 e.contains("max_tokens") || e.contains("banana"),
11242 "{extra}: {e}"
11243 );
11244 }
11245 let e = err(
11247 json!({"reasoning": {"budget": 5}}),
11248 "an unknown reasoning key must not be accepted",
11249 );
11250 assert!(
11251 e.contains("reasoning.budget") && e.contains("enabled"),
11252 "{e}"
11253 );
11254 for (extra, want) in [
11258 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
11259 (json!({"reasoning": {"exclude": 1}}), "true or false"),
11260 (json!({"reasoning": {"effort": 3}}), "must be a string"),
11261 ] {
11262 let e = err(
11263 extra.clone(),
11264 "a wrong-typed reasoning key must not be ignored",
11265 );
11266 assert!(e.contains(want), "{extra}: {e}");
11267 }
11268 for extra in [
11273 json!({"reasoning": {"enabled": true}}),
11274 json!({"reasoning": {"effort": "low"}}),
11275 json!({"reasoning": {"exclude": false}}),
11276 json!({"reasoning": null}),
11277 json!({"reasoning": {"effort": null}}),
11278 json!({"reasoning": {"enabled": null, "exclude": null}}),
11279 ] {
11280 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
11281 }
11282 }
11283
11284 #[test]
11285 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
11286 const ORNITH_TMPL: &str = include_str!(
11294 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11295 );
11296 let explicit_on = render_with(
11300 ORNITH_TMPL,
11301 &tool_caps(),
11302 json!({"reasoning": {"enabled": true}}),
11303 None,
11304 )
11305 .unwrap();
11306 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
11307 for extra in [
11308 json!({"reasoning_effort": "low"}),
11309 json!({"reasoning_effort": "medium"}),
11310 json!({"reasoning_effort": "high"}),
11311 json!({"reasoning_effort": "xhigh"}),
11313 json!({"reasoning": {"effort": "xhigh"}}),
11314 ] {
11315 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11316 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
11317 assert_eq!(
11318 got, explicit_on,
11319 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
11320 documented translation, not a decorative accept"
11321 );
11322 }
11323 for extra in [
11325 json!({}),
11326 json!({"reasoning_effort": "none"}),
11327 json!({"reasoning_effort": "minimal"}),
11328 json!({"enable_thinking": false}),
11329 ] {
11330 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11331 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
11332 }
11333 let minimal = render_with(
11336 ORNITH_TMPL,
11337 &tool_caps(),
11338 json!({"reasoning_effort": "minimal"}),
11339 None,
11340 )
11341 .unwrap();
11342 assert!(
11343 minimal.ends_with("<think>\n\n</think>\n\n"),
11344 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
11345 );
11346 let ladder_low = render_with(
11349 Q38_TMPL,
11350 &ladder_caps(),
11351 json!({"reasoning_effort": "low"}),
11352 None,
11353 )
11354 .unwrap();
11355 assert!(
11356 ladder_low.contains("Reasoning effort is set to low."),
11357 "{ladder_low:?}"
11358 );
11359 assert_ne!(
11360 ladder_low,
11361 render_with(
11362 Q38_TMPL,
11363 &ladder_caps(),
11364 json!({"reasoning_effort": "high"}),
11365 None
11366 )
11367 .unwrap(),
11368 "the ladder model's rungs stay distinct prompts"
11369 );
11370 }
11371
11372 #[test]
11373 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
11374 let render_chat = |body: serde_json::Value| -> Result<String, String> {
11385 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11386 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11387 let plan = build_chat_request(
11388 req,
11389 Some(&ladder_caps()),
11390 tx,
11391 lanes::Lane::Interactive,
11392 None,
11393 )?;
11394 Ok(chat::apply_chat_template_tools_ex(
11395 Some(Q38_TMPL),
11396 &plan.request.chat_turns,
11397 true,
11398 &plan.request.tools_json,
11399 &plan.request.tools_struct,
11400 plan.request.think,
11401 plan.request.reasoning_effort.as_deref(),
11402 None,
11403 )
11404 .unwrap())
11405 };
11406 for (intent, chat_body, responses_body, messages_body) in [
11411 (
11412 "reasoning OFF",
11413 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11414 "reasoning_effort": "none"}),
11415 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
11416 json!({"model": "m", "max_tokens": 16,
11417 "messages": [{"role": "user", "content": "hi"}],
11418 "thinking": {"type": "disabled"}}),
11419 ),
11420 (
11421 "reasoning ON at the top rung",
11422 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11423 "reasoning_effort": "xhigh"}),
11424 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
11425 json!({"model": "m", "max_tokens": 16,
11426 "messages": [{"role": "user", "content": "hi"}],
11427 "output_config": {"effort": "xhigh"}}),
11428 ),
11429 (
11430 "reasoning ON at the bottom rung",
11431 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11432 "reasoning_effort": "low"}),
11433 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
11434 json!({"model": "m", "max_tokens": 16,
11435 "messages": [{"role": "user", "content": "hi"}],
11436 "output_config": {"effort": "low"}}),
11437 ),
11438 (
11439 "the model's own default",
11440 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11441 json!({"model": "m", "input": "hi"}),
11442 json!({"model": "m", "max_tokens": 16,
11443 "messages": [{"role": "user", "content": "hi"}]}),
11444 ),
11445 ] {
11446 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
11447 let via_responses = responses_api::translate(&responses_body)
11448 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
11449 let via_messages = anthropic::translate(&messages_body)
11450 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
11451 for (surface, translated) in [
11452 ("/v1/responses", via_responses),
11453 ("/v1/messages", via_messages),
11454 ] {
11455 let got = render_chat(translated)
11456 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
11457 assert_eq!(
11458 got, chat,
11459 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
11460 /v1/chat/completions — the parameter is honoured on one format and not \
11461 the other"
11462 );
11463 }
11464 }
11465 let switchless = ModelCaps {
11468 think_switch: false,
11469 ..ladder_caps()
11470 };
11471 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
11472 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11473 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11474 let plan =
11475 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
11476 Ok(format!("{:?}", plan.request.think))
11477 };
11478 for (surface, body) in [
11479 (
11480 "/v1/responses",
11481 responses_api::translate(&json!({
11482 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
11483 .unwrap(),
11484 ),
11485 (
11486 "/v1/messages",
11487 anthropic::translate(&json!({
11488 "model": "m", "max_tokens": 16,
11489 "messages": [{"role": "user", "content": "hi"}],
11490 "thinking": {"type": "disabled"}}))
11491 .unwrap(),
11492 ),
11493 ] {
11494 let err = render_switchless(body)
11495 .err()
11496 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
11497 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
11498 }
11499 }
11500
11501 #[test]
11502 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
11503 let build = |extra: serde_json::Value| {
11511 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11512 build_chat_request(
11513 weather_request(extra),
11514 Some(&ladder_caps()),
11515 tx,
11516 lanes::Lane::Interactive,
11517 None,
11518 )
11519 };
11520 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
11521 .expect("preserve_thinking:true is the vendor default the renderer implements");
11522 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
11523 .err()
11524 .expect("preserve_thinking:false (the strip arm) must refuse");
11525 assert!(e.contains("preserve_thinking"), "{e}");
11526 assert!(e.contains("strip"), "{e}");
11527 assert_eq!(
11530 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
11531 .unwrap()
11532 .request
11533 .think,
11534 ThinkMode::NoThink
11535 );
11536 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
11538 .err()
11539 .expect("a stringly-typed preserve_thinking must not be accepted");
11540 assert!(e.contains("true or false"), "{e}");
11541 }
11542
11543 #[test]
11544 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
11545 let dsv4_caps = ModelCaps {
11551 qwen_think: true,
11552 think_switch: false,
11553 dsv4: true,
11554 ..tool_caps()
11555 };
11556 for extra in [
11557 json!({"reasoning_effort": "none"}),
11558 json!({"reasoning": {"enabled": false}}),
11559 json!({"enable_thinking": false}),
11560 json!({"include_reasoning": false}),
11561 ] {
11562 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11563 let plan = build_chat_request(
11564 weather_request(extra.clone()),
11565 Some(&dsv4_caps),
11566 tx,
11567 lanes::Lane::Interactive,
11568 None,
11569 )
11570 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
11571 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
11572 }
11573 }
11574
11575 #[test]
11576 fn contradictory_think_switches_refuse_instead_of_picking_one() {
11577 let build = |extra: serde_json::Value| {
11580 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11581 build_chat_request(
11582 weather_request(extra),
11583 Some(&tool_caps()),
11584 tx,
11585 lanes::Lane::Interactive,
11586 None,
11587 )
11588 };
11589 for extra in [
11590 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
11591 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
11592 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
11593 ] {
11594 match build(extra.clone()) {
11595 Err(err) => assert!(
11596 err.contains("contradictory"),
11597 "the refusal must say the switches contradict: {err}"
11598 ),
11599 Ok(plan) => panic!(
11600 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
11601 plan.request.think
11602 ),
11603 }
11604 }
11605 for extra in [
11607 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
11608 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
11609 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
11610 ] {
11611 build(extra.clone())
11612 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
11613 }
11614 }
11615
11616 #[test]
11617 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
11618 let switchless = ModelCaps {
11623 tools_branch: true,
11624 qwen_think: true,
11625 think_switch: false,
11626 chat_ok: true,
11627 ..Default::default()
11628 };
11629 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
11630 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11631 build_chat_request_with_trace(
11632 weather_request(extra),
11633 Some(caps),
11634 tx,
11635 lanes::Lane::Interactive,
11636 None,
11637 None,
11638 default_effort,
11639 &ModelSamplingDefaults::default(),
11640 )
11641 };
11642 for extra in [
11643 json!({"reasoning_effort": "none"}),
11644 json!({"reasoning_effort": "minimal"}),
11645 json!({"reasoning": {"enabled": false}}),
11646 json!({"enable_thinking": false}),
11647 json!({"chat_template_kwargs": {"enable_thinking": false}}),
11648 ] {
11649 let err = build(extra.clone(), &switchless, None)
11650 .err()
11651 .unwrap_or_else(|| {
11652 panic!(
11653 "{extra} on a switchless think template must not be accepted-and-ignored"
11654 )
11655 });
11656 assert!(
11657 err.contains("cannot disable reasoning"),
11658 "the refusal must say the model cannot disable reasoning: {err}"
11659 );
11660 }
11661 for (extra, default_effort) in [
11665 (json!({}), None),
11666 (json!({"reasoning_effort": "high"}), None),
11669 (json!({"reasoning": {"enabled": true}}), None),
11670 (json!({"enable_thinking": true}), None),
11671 (json!({}), Some("none")),
11672 (json!({}), Some("minimal")),
11673 (json!({}), Some("high")),
11674 ] {
11675 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
11676 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
11677 });
11678 }
11679 assert_eq!(
11682 build(json!({"enable_thinking": false}), &tool_caps(), None)
11683 .unwrap()
11684 .request
11685 .think,
11686 ThinkMode::NoThink
11687 );
11688 }
11689
11690 #[test]
11691 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
11692 let gemma_caps = ModelCaps {
11699 tools_branch: true,
11700 chat_ok: true,
11701 gemma_think: true,
11702 instruct_type: Some("gemma".into()),
11703 ..Default::default()
11704 };
11705 let render =
11706 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
11707 let mut payload = serde_json::json!({
11708 "model": "google/gemma-4-31b-it",
11709 "messages": [{"role": "user", "content": "Weather in Paris?"}],
11710 });
11711 if let Some(obj) = extra.as_object() {
11712 for (k, v) in obj {
11713 payload[k] = v.clone();
11714 }
11715 }
11716 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11717 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11718 let plan = build_chat_request_with_trace(
11719 req,
11720 Some(&gemma_caps),
11721 tx,
11722 lanes::Lane::Interactive,
11723 None,
11724 None,
11725 default_effort,
11726 &ModelSamplingDefaults::default(),
11727 )
11728 .unwrap();
11729 chat::apply_chat_template_tools_ex(
11730 Some(tmpl),
11731 &plan.request.chat_turns,
11732 true,
11733 &plan.request.tools_json,
11734 &plan.request.tools_struct,
11735 plan.request.think,
11736 plan.request.reasoning_effort.as_deref(),
11737 None, )
11739 .unwrap()
11740 };
11741 let official = gemma_template("official");
11742 let unset_with_knob = render(&official, json!({}), Some("high"));
11743 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
11744 assert_eq!(
11745 unset_with_knob, explicit_on,
11746 "knob render must be byte-identical to the explicit think-on render"
11747 );
11748 assert!(
11749 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
11750 "think-on injects the <|think|> system token: {unset_with_knob:?}"
11751 );
11752 assert!(
11753 unset_with_knob.ends_with("<|turn>model\n"),
11754 "think-on generation turn is OPEN: {unset_with_knob:?}"
11755 );
11756 let explicit_off_with_knob =
11760 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
11761 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
11762 assert_eq!(explicit_off_with_knob, explicit_off);
11763 assert!(
11764 !explicit_off_with_knob.contains("<|think|>")
11765 && explicit_off_with_knob.ends_with("<|turn>model\n"),
11766 "explicit off keeps the official template's thinking-off bytes: \
11767 {explicit_off_with_knob:?}"
11768 );
11769 let unset_no_knob = render(&official, json!({}), None);
11771 assert_eq!(
11772 unset_no_knob, explicit_off,
11773 "knobless unset stays the template's own thinking-off default"
11774 );
11775 assert_ne!(unset_no_knob, unset_with_knob);
11776 let qat = gemma_template("qat");
11779 assert!(
11780 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
11781 "QAT knobless unset keeps the closed-channel default"
11782 );
11783 assert_eq!(
11784 render(&qat, json!({}), Some("high")),
11785 render(&qat, json!({"reasoning_effort": "high"}), None),
11786 "QAT knob render must equal the explicit think-on render"
11787 );
11788 }
11789
11790 #[test]
11791 fn default_reasoning_effort_is_validated_at_metadata_load() {
11792 let parsed = OpenRouterMetadataFile::from_toml(
11794 r#"
11795[models.g]
11796default_reasoning_effort = "high"
11797"#,
11798 )
11799 .unwrap();
11800 assert_eq!(
11801 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
11802 Some("high")
11803 );
11804 let err = OpenRouterMetadataFile::from_toml(
11805 r#"
11806[models.g]
11807default_reasoning_effort = "always"
11808"#,
11809 )
11810 .unwrap_err();
11811 assert!(err.contains("default_reasoning_effort"), "{err}");
11812 }
11813
11814 #[test]
11815 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
11816 let effort_caps = ModelCaps {
11828 effort_levels: true,
11829 think_switch: false,
11830 ..tool_caps()
11831 };
11832 for (extra, want) in [
11833 (json!({}), None),
11834 (json!({"reasoning_effort": "low"}), Some("low")),
11835 (json!({"reasoning_effort": "medium"}), Some("medium")),
11836 (json!({"reasoning_effort": "high"}), Some("high")),
11837 (json!({"reasoning": {"effort": "high"}}), Some("high")),
11838 (json!({"reasoning_effort": "xhigh"}), Some("high")),
11840 (json!({"reasoning": {"effort": "max"}}), Some("high")),
11841 ] {
11842 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11843 let plan = build_chat_request(
11844 weather_request(extra.clone()),
11845 Some(&effort_caps),
11846 tx,
11847 lanes::Lane::Interactive,
11848 None,
11849 )
11850 .unwrap();
11851 assert_eq!(
11852 plan.request.reasoning_effort.as_deref(),
11853 want,
11854 "extra={extra}"
11855 );
11856 }
11857 for extra in [
11864 json!({"reasoning_effort": "none"}),
11865 json!({"reasoning_effort": "minimal"}),
11866 json!({"reasoning": {"enabled": false}}),
11867 json!({"enable_thinking": false}),
11868 json!({"include_reasoning": false}),
11869 ] {
11870 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11871 let err = build_chat_request(
11872 weather_request(extra.clone()),
11873 Some(&effort_caps),
11874 tx,
11875 lanes::Lane::Interactive,
11876 None,
11877 )
11878 .err()
11879 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
11880 assert!(
11881 err.contains("cannot disable reasoning"),
11882 "extra={extra}: {err}"
11883 );
11884 }
11885 for extra in [
11892 json!({"reasoning_effort": "high"}),
11893 json!({"reasoning": {"effort": "low"}}),
11894 ] {
11895 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11896 let plan = build_chat_request(
11897 weather_request(extra.clone()),
11898 Some(&tool_caps()),
11899 tx,
11900 lanes::Lane::Interactive,
11901 None,
11902 )
11903 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
11904 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
11905 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
11906 }
11907 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11909 let plan = build_chat_request(
11910 weather_request(json!({})),
11911 Some(&tool_caps()),
11912 tx,
11913 lanes::Lane::Interactive,
11914 None,
11915 )
11916 .unwrap();
11917 assert_eq!(plan.request.reasoning_effort, None);
11918 }
11919
11920 #[test]
11921 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
11922 let payload = serde_json::json!({
11923 "model": "m",
11924 "messages": [
11925 {"role": "user", "content": "Weather in Paris?"},
11926 {"role": "assistant", "content": null, "tool_calls": [
11927 {"id": "call_x", "type": "function", "function": {
11928 "name": "get_weather",
11929 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
11930 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
11931 ],
11932 });
11933 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11934 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11935 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
11936 .unwrap();
11937 let turns = &plan.request.chat_turns;
11938 assert_eq!(turns[1].tool_calls.len(), 1);
11939 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
11940 assert_eq!(
11941 turns[1].tool_calls[0].params,
11942 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
11943 );
11944 assert_eq!(turns[2].role, "tool");
11945 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
11946 let mut p = plan
11949 .parser
11950 .expect("think-open chat arms the reasoning splitter");
11951 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
11952 assert_eq!(
11953 pieces,
11954 vec![
11955 Piece::Reasoning("thought".into()),
11956 Piece::Content("answer <tool_call> is prose here".into()),
11957 ]
11958 );
11959 }
11960
11961 #[tokio::test]
11962 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
11963 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
11964 tx.send(Event::Token {
11965 id: 1,
11966 text: "plan</think>\n\n".into(),
11967 })
11968 .unwrap();
11969 tx.send(Event::Token {
11970 id: 2,
11971 text: "<tool_call>\n<function=get_weather>\n\
11972<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
11973 .into(),
11974 })
11975 .unwrap();
11976 tx.send(Event::Done {
11977 stop_reason: "Eos".into(),
11978 n_tokens: 2,
11979 n_prompt: 40,
11980 n_cached: 0,
11981 elapsed_s: 0.5,
11982 spec: None,
11983 })
11984 .unwrap();
11985 drop(tx);
11986 let parser = ToolStreamParser::new(HashMap::new(), true);
11987 let response = blocking_response(
11988 rx,
11989 "m".into(),
11990 true,
11991 Vec::new(),
11992 Some(parser),
11993 Envelope::new(true),
11994 )
11995 .await;
11996 assert_eq!(response.status(), StatusCode::OK);
11997 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
11998 .await
11999 .unwrap();
12000 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12001 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
12002 assert_eq!(
12005 payload["choices"][0]["message"]["content"],
12006 serde_json::Value::Null
12007 );
12008 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
12009 assert_eq!(
12010 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
12011 "plan"
12012 );
12013 let call = &payload["choices"][0]["message"]["tool_calls"][0];
12014 assert_eq!(call["type"], "function");
12015 assert_eq!(call["function"]["name"], "get_weather");
12016 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
12017 assert_eq!(payload["usage"]["prompt_tokens"], 40);
12020 assert_eq!(payload["usage"]["completion_tokens"], 2);
12021 assert_eq!(payload["usage"]["total_tokens"], 42);
12022 assert_eq!(
12023 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12024 0
12025 );
12026 }
12027
12028 #[test]
12029 fn cache_salt_plumbs_to_the_worker_namespace() {
12030 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12032 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
12033 }))
12034 .unwrap();
12035 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12036 assert_eq!(
12037 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12038 "tenant-a"
12039 );
12040
12041 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12042 "model": "m", "messages": [{"role": "user", "content": "task"}],
12043 "cache_salt": "tenant-b"
12044 }))
12045 .unwrap();
12046 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12047 assert_eq!(
12048 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12049 .unwrap()
12050 .request
12051 .cache_ns,
12052 "tenant-b"
12053 );
12054
12055 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12057 "model": "m", "prompt": "task"
12058 }))
12059 .unwrap();
12060 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12061 assert_eq!(
12062 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12063 ""
12064 );
12065 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12066 "model": "m", "messages": [{"role": "user", "content": "task"}]
12067 }))
12068 .unwrap();
12069 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12070 assert_eq!(
12071 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12072 .unwrap()
12073 .request
12074 .cache_ns,
12075 ""
12076 );
12077 }
12078
12079 #[test]
12080 fn cache_salt_validation_rejects_oversized_value() {
12081 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
12082 assert_eq!(
12083 validate_cache_namespace(&salt, false),
12084 Err("cache_salt must be at most 64 bytes")
12085 );
12086 }
12087
12088 #[test]
12089 fn cache_salt_validation_rejects_reserved_open_namespace() {
12090 let salt = Some("t:acme\u{1f}private".to_string());
12091 assert_eq!(
12092 validate_cache_namespace(&salt, false),
12093 Err("cache_salt must not use the reserved t: prefix without a keyring")
12094 );
12095 }
12096
12097 #[test]
12098 fn cache_salt_validation_accepts_normal_value() {
12099 let salt = Some("tenant-A_7.c2VjcmV0LXNjb3Bl+/=".to_string());
12100 assert_eq!(
12101 validate_cache_namespace(&salt, false).unwrap(),
12102 salt.unwrap()
12103 );
12104 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
12105 let max = Some("a".repeat(CACHE_SALT_MAX_BYTES));
12106 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max.unwrap());
12107 }
12108
12109 #[test]
12110 fn cache_salt_validation_rejects_unsupported_characters() {
12111 let salt = Some("tenant salt".to_string());
12112 assert_eq!(
12113 validate_cache_namespace(&salt, false),
12114 Err("cache_salt contains unsupported characters")
12115 );
12116 }
12117
12118 #[test]
12119 fn affinity_key_honors_both_client_conventions_in_priority_order() {
12120 use axum::http::HeaderMap;
12121 let hdr = |v: &str| {
12122 let mut h = HeaderMap::new();
12123 h.insert("x-session-id", v.parse().unwrap());
12124 h
12125 };
12126 let empty = HeaderMap::new();
12127 let s = |v: &str| Some(v.to_string());
12128 assert_eq!(affinity_key(&s("explicit"), &None, &empty), s("explicit"));
12130 assert_eq!(
12131 affinity_key(&None, &s("openai-user"), &empty),
12132 s("openai-user")
12133 );
12134 assert_eq!(affinity_key(&None, &None, &hdr("hdr-id")), s("hdr-id"));
12135 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")), s("a"));
12138 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")), s("b"));
12139 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")), None);
12142 assert_eq!(affinity_key(&s(""), &s("real"), &empty), s("real"));
12143 assert_eq!(affinity_key(&s(" padded "), &None, &empty), s("padded"));
12145 assert_eq!(affinity_key(&None, &None, &empty), None);
12147 }
12148
12149 #[test]
12150 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
12151 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12152 "model": "m", "prompt": "task", "session_id": "conv-1"
12153 }))
12154 .unwrap();
12155 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12156 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12157 assert_eq!(
12158 build_request(&req, tx, lanes::Lane::Interactive, key)
12159 .affinity
12160 .as_deref(),
12161 Some("conv-1")
12162 );
12163 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12165 "model": "m", "messages": [{"role": "user", "content": "task"}],
12166 "user": "conv-2"
12167 }))
12168 .unwrap();
12169 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12170 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12171 assert_eq!(
12172 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
12173 .unwrap()
12174 .request
12175 .affinity
12176 .as_deref(),
12177 Some("conv-2")
12178 );
12179 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12181 "model": "m", "prompt": "task"
12182 }))
12183 .unwrap();
12184 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12185 assert!(
12186 build_request(&req, tx, lanes::Lane::Interactive, None)
12187 .affinity
12188 .is_none()
12189 );
12190 }
12191
12192 async fn sse_data_lines(resp: Response) -> Vec<String> {
12194 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12195 .await
12196 .unwrap();
12197 String::from_utf8(bytes.to_vec())
12198 .unwrap()
12199 .lines()
12200 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
12201 .collect()
12202 }
12203
12204 #[tokio::test]
12205 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
12206 let feed = |think: bool| {
12212 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12213 let body = if think {
12214 "a plan</think>\n\nanswer"
12215 } else {
12216 "answer"
12217 };
12218 tx.send(Event::Token {
12219 id: 1,
12220 text: body.into(),
12221 })
12222 .unwrap();
12223 tx.send(Event::Done {
12224 stop_reason: "Eos".into(),
12225 n_tokens: 3,
12226 n_prompt: 10,
12227 n_cached: 0,
12228 elapsed_s: 0.1,
12229 spec: None,
12230 })
12231 .unwrap();
12232 drop(tx);
12233 rx
12234 };
12235 let resp = blocking_response(
12237 feed(true),
12238 "m".into(),
12239 true,
12240 Vec::new(),
12241 Some(ToolStreamParser::reasoning_only()),
12242 Envelope::new(true),
12243 )
12244 .await;
12245 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12246 .await
12247 .unwrap();
12248 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12249 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
12250 assert_eq!(
12251 v["choices"][0]["message"]["reasoning_details"][0]["text"],
12252 "a plan"
12253 );
12254 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12255 let resp = blocking_response(
12258 feed(false),
12259 "m".into(),
12260 true,
12261 Vec::new(),
12262 None,
12263 Envelope::new(true),
12264 )
12265 .await;
12266 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12267 .await
12268 .unwrap();
12269 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12270 assert!(
12271 v["choices"][0]["message"].get("reasoning").is_none(),
12272 "a reasoning-off response must carry no reasoning field: {v}"
12273 );
12274 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12275 let resp = sse_response(
12277 feed(true),
12278 "m".into(),
12279 true,
12280 Some(ToolStreamParser::reasoning_only()),
12281 Envelope::new(true),
12282 Vec::new(),
12283 None,
12284 )
12285 .into_response();
12286 let lines = sse_data_lines(resp).await;
12287 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12288 .iter()
12289 .map(|l| serde_json::from_str(l).unwrap())
12290 .collect();
12291 let reasoning: String = chunks
12292 .iter()
12293 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
12294 .collect();
12295 assert_eq!(
12296 reasoning, "a plan",
12297 "think text must stream as delta.reasoning"
12298 );
12299 let content: String = chunks
12300 .iter()
12301 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
12302 .collect();
12303 assert_eq!(content, "answer", "content must exclude the think segment");
12304 let resp = sse_response(
12306 feed(false),
12307 "m".into(),
12308 true,
12309 None,
12310 Envelope::new(true),
12311 Vec::new(),
12312 None,
12313 )
12314 .into_response();
12315 let lines = sse_data_lines(resp).await;
12316 for l in &lines[..lines.len() - 1] {
12317 let c: serde_json::Value = serde_json::from_str(l).unwrap();
12318 assert!(
12319 c["choices"][0]["delta"].get("reasoning").is_none(),
12320 "a reasoning-off stream must carry no reasoning deltas: {c}"
12321 );
12322 }
12323 }
12324
12325 #[tokio::test]
12326 async fn stream_chunks_carry_envelope_and_first_delta_role() {
12327 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12328 tx.send(Event::Token {
12329 id: 1,
12330 text: "he".into(),
12331 })
12332 .unwrap();
12333 tx.send(Event::Token {
12334 id: 2,
12335 text: "llo".into(),
12336 })
12337 .unwrap();
12338 tx.send(Event::Done {
12339 stop_reason: "Eos".into(),
12340 n_tokens: 2,
12341 n_prompt: 10,
12342 n_cached: 0,
12343 elapsed_s: 0.1,
12344 spec: None,
12345 })
12346 .unwrap();
12347 drop(tx);
12348 let resp = sse_response(
12349 rx,
12350 "m".into(),
12351 true,
12352 None,
12353 Envelope::new(true),
12354 Vec::new(),
12355 None,
12356 )
12357 .into_response();
12358 let lines = sse_data_lines(resp).await;
12359 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12360 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12361 .iter()
12362 .map(|l| serde_json::from_str(l).unwrap())
12363 .collect();
12364 let id = chunks[0]["id"].as_str().unwrap().to_string();
12366 assert!(id.starts_with("chatcmpl-"));
12367 for c in &chunks {
12368 assert_eq!(c["id"], id.as_str());
12369 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
12370 assert!(
12371 c["system_fingerprint"]
12372 .as_str()
12373 .unwrap()
12374 .starts_with("memra-")
12375 );
12376 assert_eq!(c["object"], "chat.completion.chunk");
12377 }
12378 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
12380 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
12381 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
12382 let fin = chunks.last().unwrap();
12384 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
12385 assert_eq!(fin["usage"]["prompt_tokens"], 10);
12386 }
12387
12388 #[tokio::test]
12389 async fn stream_token_events_equal_usage_on_every_finish_path() {
12390 for (stop_reason, expected_finish) in [
12391 ("Eos", "stop"),
12392 ("Callback", "stop"),
12393 ("MaxNew", "length"),
12394 ("ContextFull", "length"),
12395 ] {
12396 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12397 tx.send(Event::Token {
12400 id: 248_046,
12401 text: String::new(),
12402 })
12403 .unwrap();
12404 tx.send(Event::Done {
12405 stop_reason: stop_reason.into(),
12406 n_tokens: 1,
12407 n_prompt: 8,
12408 n_cached: 8,
12409 elapsed_s: 0.1,
12410 spec: None,
12411 })
12412 .unwrap();
12413 drop(tx);
12414
12415 let resp = sse_response(
12416 rx,
12417 "m".into(),
12418 true,
12419 None,
12420 Envelope::new(true),
12421 Vec::new(),
12422 None,
12423 )
12424 .into_response();
12425 let lines = sse_data_lines(resp).await;
12426 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12427 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12428 .iter()
12429 .map(|line| serde_json::from_str(line).unwrap())
12430 .collect();
12431 let token_events = chunks
12432 .iter()
12433 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
12434 .count();
12435 let terminal = chunks.last().unwrap();
12436 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
12437 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
12438 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
12439 }
12440 }
12441
12442 #[tokio::test]
12443 async fn stream_excludes_stop_text_like_non_stream_does() {
12444 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12448 tx.send(Event::Token {
12449 id: 1,
12450 text: "answer\nPro".into(),
12451 })
12452 .unwrap();
12453 tx.send(Event::Token {
12454 id: 2,
12455 text: "blem: leaked prompt".into(),
12456 })
12457 .unwrap();
12458 tx.send(Event::Done {
12459 stop_reason: "Callback".into(),
12460 n_tokens: 2,
12461 n_prompt: 8,
12462 n_cached: 0,
12463 elapsed_s: 0.1,
12464 spec: None,
12465 })
12466 .unwrap();
12467 drop(tx);
12468 let resp = sse_response(
12469 rx,
12470 "m".into(),
12471 true,
12472 None,
12473 Envelope::new(true),
12474 vec!["Problem:".into()],
12475 None,
12476 )
12477 .into_response();
12478 let lines = sse_data_lines(resp).await;
12479 let content: String = lines
12480 .iter()
12481 .filter(|l| *l != "[DONE]")
12482 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12483 .filter_map(|c| {
12484 c["choices"][0]["delta"]["content"]
12485 .as_str()
12486 .map(str::to_string)
12487 })
12488 .collect();
12489 assert_eq!(content, "answer\n");
12490
12491 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12493 tx.send(Event::Token {
12494 id: 1,
12495 text: "ends in Pro".into(),
12496 })
12497 .unwrap();
12498 tx.send(Event::Done {
12499 stop_reason: "Eos".into(),
12500 n_tokens: 1,
12501 n_prompt: 8,
12502 n_cached: 0,
12503 elapsed_s: 0.1,
12504 spec: None,
12505 })
12506 .unwrap();
12507 drop(tx);
12508 let resp = sse_response(
12509 rx,
12510 "m".into(),
12511 true,
12512 None,
12513 Envelope::new(true),
12514 vec!["Problem:".into()],
12515 None,
12516 )
12517 .into_response();
12518 let lines = sse_data_lines(resp).await;
12519 let content: String = lines
12520 .iter()
12521 .filter(|l| *l != "[DONE]")
12522 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12523 .filter_map(|c| {
12524 c["choices"][0]["delta"]["content"]
12525 .as_str()
12526 .map(str::to_string)
12527 })
12528 .collect();
12529 assert_eq!(content, "ends in Pro");
12530 }
12531
12532 #[tokio::test]
12533 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
12534 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12535 tx.send(Event::Error(worker::EngineError::engine("boom")))
12536 .unwrap();
12537 drop(tx);
12538 let resp = sse_response(
12539 rx,
12540 "m".into(),
12541 true,
12542 None,
12543 Envelope::new(true),
12544 Vec::new(),
12545 None,
12546 )
12547 .into_response();
12548 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12549 .await
12550 .unwrap();
12551 let body = String::from_utf8(bytes.to_vec()).unwrap();
12552 assert!(
12554 !body.contains("event: error"),
12555 "named SSE event leaked: {body}"
12556 );
12557 let lines: Vec<&str> = body
12558 .lines()
12559 .filter_map(|l| l.strip_prefix("data: "))
12560 .collect();
12561 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
12562 assert_eq!(err["error"]["message"], "boom");
12563 assert_eq!(err["error"]["type"], "server_error");
12564 assert_eq!(err["error"]["code"], "engine_error");
12565 assert_eq!(lines.last(), Some(&"[DONE]"));
12566 }
12567
12568 #[test]
12569 fn ttft_sse_marker_ignores_keepalive_comments() {
12570 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
12571 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
12572 assert!(is_sse_data_frame(
12573 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
12574 ));
12575 }
12576
12577 #[tokio::test]
12578 async fn error_bodies_use_the_openai_object_shape() {
12579 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12580 tx.send(Event::Error(worker::EngineError::model_not_found(
12581 "unknown model \"x\"",
12582 )))
12583 .unwrap();
12584 drop(tx);
12585 let response =
12586 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
12587 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
12588 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12589 .await
12590 .unwrap();
12591 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12592 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
12594 assert_eq!(payload["error"]["type"], "invalid_request_error");
12595 assert_eq!(payload["error"]["param"], "model");
12596 assert_eq!(payload["error"]["code"], "model_not_found");
12597 }
12598
12599 fn retry_after(resp: &Response) -> Option<String> {
12606 resp.headers()
12607 .get(axum::http::header::RETRY_AFTER)
12608 .and_then(|v| v.to_str().ok())
12609 .map(str::to_string)
12610 }
12611
12612 async fn body_value(resp: Response) -> serde_json::Value {
12615 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12616 .await
12617 .expect("body");
12618 serde_json::from_slice(&bytes).expect("json body")
12619 }
12620
12621 #[test]
12622 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
12623 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
12625 assert_eq!(
12626 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
12627 TIMEOUT_MS_DEFAULT
12628 );
12629 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
12633 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
12634 }
12635 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
12637 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
12638 assert!(err.contains("timeout_ms"), "{err}");
12639 assert!(
12640 err.contains(&TIMEOUT_MS_MIN.to_string())
12641 && err.contains(&TIMEOUT_MS_MAX.to_string()),
12642 "the message must state the range: {err}"
12643 );
12644 assert!(
12645 err.contains("stream"),
12646 "the message must point at streaming for longer work: {err}"
12647 );
12648 }
12649 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
12651 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
12652 assert!(
12653 err.contains("timeout_ms") && err.contains("stream"),
12654 "{err}"
12655 );
12656 }
12657 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
12659 }
12660
12661 #[tokio::test]
12664 async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
12665 let _l = DRAIN_LOCK.lock().unwrap();
12666 let st = fake_worker_state();
12667
12668 let comp = completions(
12669 State(st.clone()),
12670 HeaderMap::new(),
12671 None,
12672 Json(
12673 serde_json::from_value(json!({
12674 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
12675 .unwrap(),
12676 ),
12677 )
12678 .await;
12679 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
12680 let chat = chat_completions(
12681 State(st.clone()),
12682 HeaderMap::new(),
12683 None,
12684 Json(
12685 serde_json::from_value(json!({
12686 "model": "m", "messages": [{"role": "user", "content": "t"}],
12687 "timeout_ms": 90_001}))
12688 .unwrap(),
12689 ),
12690 )
12691 .await;
12692 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
12693 let resp_api = responses_api::responses(
12694 State(st.clone()),
12695 HeaderMap::new(),
12696 None,
12697 axum::body::Bytes::from(
12698 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
12699 ),
12700 )
12701 .await;
12702 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
12703 let msgs = anthropic::messages(
12704 State(st.clone()),
12705 HeaderMap::new(),
12706 None,
12707 axum::body::Bytes::from(
12708 json!({"model": "m", "max_tokens": 16,
12709 "messages": [{"role": "user", "content": "t"}],
12710 "timeout_ms": 90_001})
12711 .to_string(),
12712 ),
12713 )
12714 .await;
12715 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
12716
12717 for (surface, resp) in [
12719 ("/v1/completions", comp),
12720 ("/v1/chat/completions", chat),
12721 ("/v1/responses", resp_api),
12722 ] {
12723 let body = body_value(resp).await;
12724 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
12725 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
12726 let m = body["error"]["message"].as_str().unwrap();
12727 assert!(
12728 m.contains("90000") && m.contains("stream"),
12729 "{surface}: {m}"
12730 );
12731 }
12732 let body = body_value(msgs).await;
12734 assert_eq!(body["error"]["type"], "invalid_request_error");
12735 let m = body["error"]["message"].as_str().unwrap();
12736 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
12737 }
12738
12739 #[tokio::test]
12742 async fn a_non_integer_timeout_ms_is_a_named_400() {
12743 let _l = DRAIN_LOCK.lock().unwrap();
12744 let st = fake_worker_state();
12745 let resp = chat_completions(
12746 State(st),
12747 HeaderMap::new(),
12748 None,
12749 Json(
12750 serde_json::from_value(json!({
12751 "model": "m", "messages": [{"role": "user", "content": "t"}],
12752 "timeout_ms": "30s"}))
12753 .unwrap(),
12754 ),
12755 )
12756 .await;
12757 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
12758 let body = body_value(resp).await;
12759 assert_eq!(body["error"]["param"], "timeout_ms");
12760 }
12761
12762 #[tokio::test]
12768 async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
12769 let _l = DRAIN_LOCK.lock().unwrap();
12770 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
12775 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
12776 let worker_cancel = cancel_seen.clone();
12777 let health = health::WorkerHealth::new();
12778 let h = health.clone();
12779 std::thread::spawn(move || {
12780 h.mark_ready();
12781 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
12782 worker::release_pending_admit();
12783 worker::release_admission_reservation(req.lane);
12784 let _ = req.tx.send(Event::PromptUsage {
12785 n_prompt: 1,
12786 n_cached: 0,
12787 });
12788 let _ = req.tx.send(Event::Token {
12789 id: 1,
12790 text: "partial".into(),
12791 });
12792 for _ in 0..5_000 {
12795 if req.tx.is_closed() {
12796 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
12797 break;
12798 }
12799 std::thread::sleep(std::time::Duration::from_millis(1));
12800 }
12801 }
12802 });
12803 for _ in 0..2_000 {
12804 if health.live().is_ok() {
12805 break;
12806 }
12807 std::thread::sleep(std::time::Duration::from_millis(1));
12808 }
12809 let mut st = fake_worker_state();
12810 st.cmd_tx = cmd_tx;
12811 st.health = health;
12812 let mock = MockMetering::admit_all();
12813 st.metering = Some(mock.clone());
12814
12815 let resp = chat_completions(
12816 State(st),
12817 HeaderMap::new(),
12818 None,
12819 Json(
12820 serde_json::from_value(json!({
12821 "model": "m", "messages": [{"role": "user", "content": "t"}],
12822 "timeout_ms": 1_000}))
12823 .unwrap(),
12824 ),
12825 )
12826 .await;
12827
12828 assert_eq!(resp.status(), StatusCode::OK);
12833 let body = body_value(resp).await;
12834 assert!(
12835 body["choices"][0]["message"]["content"]
12836 .as_str()
12837 .unwrap()
12838 .contains("partial"),
12839 "the tokens generated before the cut must be delivered: {body}"
12840 );
12841 assert_eq!(body["choices"][0]["finish_reason"], "error");
12845 assert_eq!(
12846 body["choices"][0]["native_finish_reason"],
12847 "deadline_exceeded"
12848 );
12849 assert_eq!(body["error"]["code"], "deadline_exceeded");
12850 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
12851 let message = body["error"]["message"].as_str().unwrap();
12852 assert!(
12853 message.contains("1000") && message.contains("stream"),
12854 "the partial must name the deadline and the streaming alternative: {message}"
12855 );
12856 assert_eq!(body["usage"]["completion_tokens"], 1);
12857
12858 let mut cancelled = false;
12863 for _ in 0..500 {
12864 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
12865 cancelled = true;
12866 break;
12867 }
12868 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
12869 }
12870 assert!(
12871 cancelled,
12872 "the deadline must CANCEL generation (worker's event channel closed)"
12873 );
12874
12875 let events = mock.events();
12880 assert!(
12881 events.contains(&MeterEvent::DeadlinePartial {
12882 prompt: 1,
12883 cached: 0,
12884 completion: 1,
12885 }),
12886 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
12887 );
12888 assert!(
12889 !events
12890 .iter()
12891 .any(|e| matches!(e, MeterEvent::Complete { .. })),
12892 "a deadline cut must stay distinguishable from a full answer: {events:?}"
12893 );
12894 }
12895
12896 #[tokio::test]
12900 async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
12901 let _l = DRAIN_LOCK.lock().unwrap();
12902 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
12903 let health = health::WorkerHealth::new();
12904 let h = health.clone();
12905 std::thread::spawn(move || {
12906 h.mark_ready();
12907 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
12910 worker::release_pending_admit();
12911 worker::release_admission_reservation(req.lane);
12912 let _ = req.tx.send(Event::PromptUsage {
12913 n_prompt: 1,
12914 n_cached: 0,
12915 });
12916 for _ in 0..5_000 {
12917 if req.tx.is_closed() {
12918 break;
12919 }
12920 std::thread::sleep(std::time::Duration::from_millis(1));
12921 }
12922 }
12923 });
12924 for _ in 0..2_000 {
12925 if health.live().is_ok() {
12926 break;
12927 }
12928 std::thread::sleep(std::time::Duration::from_millis(1));
12929 }
12930 let mut st = fake_worker_state();
12931 st.cmd_tx = cmd_tx;
12932 st.health = health;
12933 let mock = MockMetering::admit_all();
12934 st.metering = Some(mock.clone());
12935 let resp = chat_completions(
12936 State(st),
12937 HeaderMap::new(),
12938 None,
12939 Json(
12940 serde_json::from_value(json!({
12941 "model": "m", "messages": [{"role": "user", "content": "t"}],
12942 "timeout_ms": 1_000}))
12943 .unwrap(),
12944 ),
12945 )
12946 .await;
12947 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
12948 assert!(resp.headers().get("x-should-retry").is_none());
12950 assert_eq!(retry_after(&resp), None);
12951 let body = body_value(resp).await;
12952 assert_eq!(body["error"]["code"], "deadline_exceeded");
12953 assert!(
12954 body["error"]["message"]
12955 .as_str()
12956 .unwrap()
12957 .contains("not billed"),
12958 "the zero-token 408 keeps the billing promise: {body}"
12959 );
12960 let events = mock.events();
12961 assert!(
12962 events.contains(&MeterEvent::Unbilled {
12963 outcome: "deadline_exceeded",
12964 status: 408,
12965 code: "deadline_exceeded".into(),
12966 }),
12967 "the named zero-debit census outcome, not the generic reject — every sibling \
12968 deadline path settles this one: {events:?}"
12969 );
12970 }
12971
12972 #[tokio::test]
12975 async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
12976 let _l = DRAIN_LOCK.lock().unwrap();
12977 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
12979 let health = health::WorkerHealth::new();
12980 let h = health.clone();
12981 std::thread::spawn(move || {
12982 h.mark_ready();
12983 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
12984 worker::release_pending_admit();
12985 worker::release_admission_reservation(req.lane);
12986 let _ = req.tx.send(Event::PromptUsage {
12987 n_prompt: 1,
12988 n_cached: 0,
12989 });
12990 while !req.tx.is_closed() {
12991 std::thread::sleep(std::time::Duration::from_millis(1));
12992 }
12993 }
12994 });
12995 for _ in 0..2_000 {
12996 if health.live().is_ok() {
12997 break;
12998 }
12999 std::thread::sleep(std::time::Duration::from_millis(1));
13000 }
13001 let mut st = fake_worker_state();
13002 st.cmd_tx = cmd_tx;
13003 st.health = health;
13004 let mock = MockMetering::admit_all();
13005 st.metering = Some(mock.clone());
13006
13007 let resp = chat_completions(
13008 State(st),
13009 HeaderMap::new(),
13010 None,
13011 Json(
13012 serde_json::from_value(json!({
13013 "model": "m", "messages": [{"role": "user", "content": "t"}],
13014 "stream": true, "timeout_ms": 1_000}))
13015 .unwrap(),
13016 ),
13017 )
13018 .await;
13019 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
13022 let body = body_value(resp).await;
13023 assert_eq!(body["error"]["code"], "deadline_exceeded");
13024 assert!(
13025 body["error"]["message"]
13026 .as_str()
13027 .unwrap()
13028 .contains("first token"),
13029 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
13030 );
13031 let events = mock.events();
13032 assert!(
13033 events.contains(&MeterEvent::Unbilled {
13034 outcome: "deadline_exceeded",
13035 status: 408,
13036 code: "deadline_exceeded".into(),
13037 }),
13038 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
13039 );
13040 }
13041
13042 #[tokio::test]
13046 async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
13047 let _l = DRAIN_LOCK.lock().unwrap();
13048 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
13051 let mock = MockMetering::admit_all();
13052 st.metering = Some(mock.clone());
13053 let resp = chat_completions(
13054 State(st),
13055 HeaderMap::new(),
13056 None,
13057 Json(
13058 serde_json::from_value(json!({
13059 "model": "m", "messages": [{"role": "user", "content": "t"}],
13060 "stream": true, "timeout_ms": 1_000}))
13061 .unwrap(),
13062 ),
13063 )
13064 .await;
13065 assert_eq!(
13066 resp.status(),
13067 StatusCode::OK,
13068 "TTFT was met — 200 is correct"
13069 );
13070 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13071 .await
13072 .expect("the stream must run to completion past the deadline");
13073 let text = String::from_utf8(bytes.to_vec()).unwrap();
13074 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
13075 let events = mock.events();
13076 assert!(
13077 events
13078 .iter()
13079 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
13080 "a stream past its deadline after first token still settles as COMPLETE with \
13081 all four tokens: {events:?}"
13082 );
13083 }
13084
13085 #[test]
13088 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
13089 let st = fake_worker_state();
13090 let cap = lane_cap(lanes::Lane::Interactive);
13092 {
13093 let mut m = st.metrics.lock().unwrap();
13094 m.completed = 10;
13095 m.tokens_out = 1_000; m.step_p50_ms = 10.0; m.queued_requests = (cap * 4 + 1) as u64;
13098 }
13099 let rl = RateLimit {
13100 limit: cap,
13101 remaining: 0,
13102 reset_s: 1,
13103 };
13104 let (resp, outcome) = admission_backpressure(
13105 &st,
13106 lanes::Lane::Interactive,
13107 &rl,
13108 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13109 )
13110 .expect_err("a backlog past the bound must shed");
13111 assert_eq!(outcome, "shed_queue");
13112 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13113 assert!(
13114 retry_after(&resp).is_some(),
13115 "a shed must carry Retry-After so the router's spill can act on it"
13116 );
13117 let stamped = rl.attach(resp);
13119 for h in [
13120 "x-ratelimit-limit",
13121 "x-ratelimit-remaining",
13122 "x-ratelimit-reset",
13123 ] {
13124 assert!(stamped.headers().get(h).is_some(), "missing {h}");
13125 }
13126 }
13127
13128 #[test]
13132 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
13133 let st = fake_worker_state();
13134 let cap = lane_cap(lanes::Lane::Interactive);
13135 {
13136 let mut m = st.metrics.lock().unwrap();
13137 m.completed = 10;
13138 m.tokens_out = 1_000;
13139 m.step_p50_ms = 10.0; m.queued_requests = cap as u64; }
13142 let rl = RateLimit {
13143 limit: cap,
13144 remaining: 0,
13145 reset_s: 1,
13146 };
13147 assert!(
13149 admission_backpressure(
13150 &st,
13151 lanes::Lane::Interactive,
13152 &rl,
13153 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13154 )
13155 .is_ok(),
13156 "a request whose deadline covers the estimate must be admitted"
13157 );
13158 let (resp, outcome) = admission_backpressure(
13160 &st,
13161 lanes::Lane::Interactive,
13162 &rl,
13163 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13164 )
13165 .expect_err("a deadline shorter than the estimated wait must shed");
13166 assert_eq!(outcome, "shed_deadline");
13167 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13168 assert!(retry_after(&resp).is_some());
13169 }
13170
13171 #[test]
13174 fn admission_backpressure_is_interactive_only_and_silent_with_free_slots() {
13175 let st = fake_worker_state();
13176 let cap = lane_cap(lanes::Lane::Interactive);
13177 {
13178 let mut m = st.metrics.lock().unwrap();
13179 m.completed = 10;
13180 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
13182 m.queued_requests = (cap * 100) as u64;
13183 }
13184 let free = RateLimit {
13186 limit: cap,
13187 remaining: 1,
13188 reset_s: 0,
13189 };
13190 assert!(
13191 admission_backpressure(
13192 &st,
13193 lanes::Lane::Interactive,
13194 &free,
13195 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13196 )
13197 .is_ok()
13198 );
13199 let full = RateLimit {
13201 limit: cap,
13202 remaining: 0,
13203 reset_s: 5,
13204 };
13205 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
13206 assert!(
13207 admission_backpressure(
13208 &st,
13209 lane,
13210 &full,
13211 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13212 )
13213 .is_ok(),
13214 "{lane:?} must not be shed by the interactive gate"
13215 );
13216 }
13217 }
13218
13219 #[test]
13220 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
13221 let st = fake_worker_state();
13222 let cap = lane_cap(lanes::Lane::Interactive);
13223 let bound = max_queue_depth(cap);
13224 assert!(bound > 0, "the queue bound must admit at least one request");
13225 let rl = RateLimit {
13226 limit: cap,
13227 remaining: 0,
13228 reset_s: 1,
13229 };
13230 let _ = worker::PENDING_ADMITS.fetch_update(
13231 std::sync::atomic::Ordering::AcqRel,
13232 std::sync::atomic::Ordering::Acquire,
13233 |_| Some(0),
13234 );
13235 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
13236 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
13237 let guard = reserve_pending_admit(
13238 &st,
13239 lanes::Lane::Interactive,
13240 &rl,
13241 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13242 )
13243 .expect("the final queue slot should be reservable");
13244 assert_eq!(
13245 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13246 1
13247 );
13248 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
13249 drop(guard);
13250 assert_eq!(
13251 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13252 0
13253 );
13254 assert_eq!(
13255 counter.load(std::sync::atomic::Ordering::Acquire),
13256 bound - 1
13257 );
13258
13259 counter.store(bound, std::sync::atomic::Ordering::Release);
13260 let rejected = reserve_pending_admit(
13261 &st,
13262 lanes::Lane::Interactive,
13263 &rl,
13264 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13265 );
13266 assert!(matches!(rejected, Err((_, "shed_queue"))));
13267 counter.store(0, std::sync::atomic::Ordering::Release);
13268 }
13269
13270 #[test]
13271 fn admission_reservations_are_lane_scoped() {
13272 let st = fake_worker_state();
13273 let harvest = lanes::Lane::Harvest;
13274 let interactive = lanes::Lane::Interactive;
13275 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
13276 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
13277 harvest_counter.store(
13278 max_queue_depth(lane_cap(harvest)),
13279 std::sync::atomic::Ordering::Release,
13280 );
13281 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
13282 let free = RateLimit {
13283 limit: lane_cap(interactive),
13284 remaining: 1,
13285 reset_s: 0,
13286 };
13287 let guard = reserve_pending_admit(
13288 &st,
13289 interactive,
13290 &free,
13291 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13292 )
13293 .expect("a full harvest queue must not consume interactive capacity");
13294 drop(guard);
13295 let harvest_rl = RateLimit {
13296 limit: lane_cap(harvest),
13297 remaining: 0,
13298 reset_s: 1,
13299 };
13300 assert!(matches!(
13301 reserve_pending_admit(
13302 &st,
13303 harvest,
13304 &harvest_rl,
13305 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13306 ),
13307 Err((_, "shed_queue"))
13308 ));
13309 harvest_counter.store(0, std::sync::atomic::Ordering::Release);
13310 }
13311
13312 #[test]
13313 fn taxonomy_maps_every_class_to_its_status_and_code() {
13314 use worker::{EngineError as E, ErrClass as C};
13315 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
13316 (
13317 E::invalid_param("bad json", "response_format"),
13318 StatusCode::BAD_REQUEST,
13319 "invalid_request_error",
13320 "",
13321 ),
13322 (
13323 E::context_length("prompt (9000 tok) >= context cap (8192)"),
13324 StatusCode::BAD_REQUEST,
13325 "invalid_request_error",
13326 "context_length_exceeded",
13327 ),
13328 (
13329 E::model_not_found("unknown model \"nope\""),
13330 StatusCode::BAD_REQUEST,
13331 "invalid_request_error",
13332 "model_not_found",
13333 ),
13334 (
13335 E::rate_limit("lane judge is at capacity, retry"),
13336 StatusCode::TOO_MANY_REQUESTS,
13337 "rate_limit_error",
13338 "rate_limit_exceeded",
13339 ),
13340 (
13341 E::overloaded("no VRAM for a new session"),
13342 StatusCode::SERVICE_UNAVAILABLE,
13343 "server_error",
13344 "overloaded",
13345 ),
13346 (
13347 E::engine("graph step failed: launch error"),
13348 StatusCode::INTERNAL_SERVER_ERROR,
13349 "server_error",
13350 "engine_error",
13351 ),
13352 ];
13353 for (err, want_status, want_type, want_code) in cases {
13354 let (status, etype, code) = class_http(err.class);
13355 assert_eq!(status, want_status, "{:?}", err);
13356 assert_eq!(etype, want_type, "{:?}", err);
13357 if !want_code.is_empty() {
13358 assert_eq!(code, Some(want_code), "{:?}", err);
13359 }
13360 let body = engine_error_body(&err);
13362 assert_eq!(body["error"]["message"], err.message);
13363 assert_eq!(body["error"]["type"], want_type);
13364 }
13365 for c in [
13367 C::InvalidRequest,
13368 C::ContextLength,
13369 C::ModelNotFound,
13370 C::RateLimit,
13371 C::Overloaded,
13372 C::Engine,
13373 ] {
13374 let (s, t, _) = class_http(c);
13375 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
13376 assert!(!t.is_empty());
13377 }
13378 }
13379
13380 #[test]
13381 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
13382 let e = worker::EngineError::engine(
13387 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
13388 );
13389 let resp = engine_error_response(&e);
13390 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13391 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13392 }
13393
13394 #[test]
13395 fn retry_headers_follow_the_sdk_contract() {
13396 for e in [
13400 worker::EngineError::rate_limit("shed"),
13401 worker::EngineError::overloaded("no VRAM"),
13402 ] {
13403 let resp = engine_error_response(&e);
13404 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
13405 let secs: u64 = ra
13406 .parse()
13407 .expect("Retry-After must be integer delay-seconds");
13408 assert!(
13409 secs > 0 && secs <= 60,
13410 "Retry-After {secs}s outside the honored window"
13411 );
13412 let ms = resp
13413 .headers()
13414 .get("retry-after-ms")
13415 .unwrap()
13416 .to_str()
13417 .unwrap();
13418 assert_eq!(
13419 ms.parse::<u64>().unwrap(),
13420 secs * 1000,
13421 "the two headers disagree"
13422 );
13423 assert!(
13424 resp.headers().get("x-should-retry").is_none(),
13425 "a retryable class must not say x-should-retry: false"
13426 );
13427 }
13428 }
13429
13430 #[tokio::test]
13431 async fn command_send_failure_obeys_the_retry_contract() {
13432 let _l = DRAIN_LOCK.lock().unwrap();
13433 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
13434 let mut st = fake_worker_state();
13435 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13436 drop(cmd_rx);
13437 st.cmd_tx = cmd_tx;
13438
13439 let completion = completions(
13440 State(st.clone()),
13441 axum::http::HeaderMap::new(),
13442 None,
13443 Json(
13444 serde_json::from_value(serde_json::json!({
13445 "model": "m", "prompt": "test"
13446 }))
13447 .unwrap(),
13448 ),
13449 )
13450 .await;
13451 let chat = chat_completions(
13452 State(st),
13453 axum::http::HeaderMap::new(),
13454 None,
13455 Json(
13456 serde_json::from_value(serde_json::json!({
13457 "model": "m", "messages": [{"role": "user", "content": "test"}]
13458 }))
13459 .unwrap(),
13460 ),
13461 )
13462 .await;
13463
13464 for resp in [completion, chat] {
13465 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13466 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13467 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
13468 assert_ne!(
13469 resp.headers()
13470 .get("x-should-retry")
13471 .and_then(|v| v.to_str().ok()),
13472 Some("false")
13473 );
13474 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13475 .await
13476 .unwrap();
13477 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13478 assert_eq!(payload["error"]["type"], "server_error");
13479 assert_eq!(payload["error"]["code"], "overloaded");
13480 }
13481 }
13482
13483 #[test]
13484 fn unfixable_client_errors_say_x_should_retry_false() {
13485 for e in [
13488 worker::EngineError::model_not_found("unknown model \"x\""),
13489 worker::EngineError::context_length("prompt too long"),
13490 worker::EngineError::invalid_param("bad", "messages"),
13491 ] {
13492 let resp = engine_error_response(&e);
13493 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13494 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
13495 assert!(
13496 retry_after(&resp).is_none(),
13497 "a 400 must not promise a retry window"
13498 );
13499 }
13500 }
13501
13502 #[tokio::test]
13503 async fn a_closed_worker_channel_is_503_not_500() {
13504 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
13508 drop(tx);
13509 let resp =
13510 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
13511 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13512 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13513 }
13514
13515 #[tokio::test]
13516 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
13517 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13520 tx.send(Event::Error(worker::EngineError::rate_limit(
13521 "lane judge shed: interactive p99 over budget, retry",
13522 )))
13523 .unwrap();
13524 let (resp, error_code) = peek_admission(rx)
13525 .await
13526 .expect_err("a shed must not be forwarded into the stream");
13527 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13528 assert_eq!(error_code, "rate_limit_exceeded");
13529 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13530 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13531 .await
13532 .unwrap();
13533 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13534 assert!(
13535 payload["error"].is_object(),
13536 "bare-string error body: {payload}"
13537 );
13538 assert_eq!(payload["error"]["type"], "rate_limit_error");
13539 assert!(
13540 payload["error"]["message"]
13541 .as_str()
13542 .unwrap()
13543 .contains("shed")
13544 );
13545 }
13546
13547 #[tokio::test]
13548 async fn interactive_admission_error_is_a_preheader_429() {
13549 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13552 tx.send(Event::Error(worker::EngineError::rate_limit(
13553 "KV capacity unavailable",
13554 )))
13555 .unwrap();
13556 let (resp, error_code) = peek_admission(rx)
13557 .await
13558 .expect_err("admission error must stay pre-header");
13559 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13560 assert_eq!(error_code, "rate_limit_exceeded");
13561 }
13562
13563 #[tokio::test]
13564 async fn admission_peek_preserves_context_error_for_the_ledger() {
13565 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13566 tx.send(Event::Error(worker::EngineError::context_length(
13567 "prompt exceeds configured model maximum",
13568 )))
13569 .unwrap();
13570 let (resp, error_code) = peek_admission(rx)
13571 .await
13572 .expect_err("context rejection must stay pre-header");
13573 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13574 assert_eq!(error_code, "context_length_exceeded");
13575 }
13576
13577 #[tokio::test]
13578 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
13579 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13580 tx.send(Event::PromptUsage {
13581 n_prompt: 262_143,
13582 n_cached: 0,
13583 })
13584 .unwrap();
13585 let mut replay = peek_admission(rx).await.expect("successful admission");
13586 assert!(matches!(
13587 replay.recv().await,
13588 Some(Event::PromptUsage {
13589 n_prompt: 262_143,
13590 n_cached: 0
13591 }),
13592 ));
13593 }
13594
13595 #[test]
13596 fn penalties_plumb_from_http_to_sampler_config() {
13597 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13600 "model": "m", "messages": [{"role": "user", "content": "task"}],
13601 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
13602 }))
13603 .unwrap();
13604 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13605 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13606 .unwrap()
13607 .request
13608 .sampler_cfg;
13609 assert_eq!(cfg.penalty_freq, 0.5);
13610 assert_eq!(cfg.penalty_present, 0.25);
13611 assert_eq!(cfg.penalty_repeat, 1.1);
13612 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
13613
13614 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13615 "model": "m", "prompt": "task", "frequency_penalty": 1.5
13616 }))
13617 .unwrap();
13618 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13619 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13620 assert_eq!(cfg.penalty_freq, 1.5);
13621 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
13622
13623 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13625 "model": "m", "prompt": "task"
13626 }))
13627 .unwrap();
13628 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13629 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13630 assert_eq!(cfg.penalty_last_n, 0);
13631 assert_eq!(cfg.penalty_repeat, 1.0);
13632 }
13633
13634 #[test]
13635 fn omitted_temperature_is_openai_default_not_greedy() {
13636 let chat_temp = |body: serde_json::Value| {
13652 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13653 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13654 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13655 .unwrap()
13656 .request
13657 .sampler_cfg
13658 .temperature
13659 };
13660 let comp_temp = |body: serde_json::Value| {
13661 let req: CompletionReq = serde_json::from_value(body).unwrap();
13662 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13663 build_request(&req, tx, lanes::Lane::Interactive, None)
13664 .sampler_cfg
13665 .temperature
13666 };
13667
13668 assert_eq!(
13670 chat_temp(serde_json::json!({
13671 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
13672 1.0,
13673 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
13674 );
13675 assert_eq!(
13676 comp_temp(serde_json::json!({
13677 "model": "m", "prompt": "t"})),
13678 1.0,
13679 "omitted completions temperature must be the OpenAI 1.0 default"
13680 );
13681
13682 assert_eq!(
13684 chat_temp(serde_json::json!({
13685 "model": "m", "messages": [{"role": "user", "content": "t"}],
13686 "temperature": 0.0})),
13687 0.0,
13688 "explicit temperature 0 must stay greedy"
13689 );
13690 assert_eq!(
13691 comp_temp(serde_json::json!({
13692 "model": "m", "prompt": "t", "temperature": 0})),
13693 0.0,
13694 "explicit temperature 0 must stay greedy"
13695 );
13696 assert!(
13698 memra_engine::sampler::Sampler::new(sampler_config(
13699 0.0,
13700 0,
13701 1.0,
13702 0.0,
13703 0.0,
13704 0.0,
13705 1.0,
13706 Some(0)
13707 ))
13708 .is_greedy()
13709 );
13710 assert!(
13711 !memra_engine::sampler::Sampler::new(sampler_config(
13712 1.0,
13713 0,
13714 1.0,
13715 0.0,
13716 0.0,
13717 0.0,
13718 1.0,
13719 Some(0)
13720 ))
13721 .is_greedy()
13722 );
13723
13724 assert_eq!(
13726 chat_temp(serde_json::json!({
13727 "model": "m", "messages": [{"role": "user", "content": "t"}],
13728 "temperature": 0.7})),
13729 0.7
13730 );
13731
13732 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13736 "model": "m", "prompt": "t"}))
13737 .unwrap();
13738 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13739 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13740 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
13741 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
13742 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
13743 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
13744 assert!(
13749 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
13750 "the omitted-temperature default must ride sampled spec's pure-temp regime"
13751 );
13752 }
13753
13754 #[test]
13755 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
13756 let caps = ModelCaps {
13757 chat_temperature_default: Some(0.5),
13758 chat_top_p_default: Some(0.9),
13759 chat_ok: true,
13760 ..Default::default()
13761 };
13762 let cfg = |extra: serde_json::Value| {
13763 let mut body = serde_json::json!({
13764 "model": "step35",
13765 "messages": [{"role": "user", "content": "task"}]
13766 });
13767 body.as_object_mut()
13768 .unwrap()
13769 .extend(extra.as_object().unwrap().clone());
13770 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13771 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13772 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
13773 .unwrap()
13774 .request
13775 .sampler_cfg
13776 };
13777
13778 let omitted = cfg(serde_json::json!({}));
13779 assert_eq!(omitted.temperature, 0.5);
13780 assert_eq!(omitted.top_p, 0.9);
13781
13782 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
13783 assert_eq!(explicit_temp.temperature, 0.7);
13784 assert_eq!(
13785 explicit_temp.top_p, 0.9,
13786 "omitting top_p must retain StepFun's nucleus default"
13787 );
13788
13789 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
13790 assert_eq!(
13791 explicit.temperature, 0.0,
13792 "explicit greedy must remain authoritative"
13793 );
13794 assert_eq!(
13795 explicit.top_p, 1.0,
13796 "explicit untruncated sampling must remain authoritative"
13797 );
13798 }
13799
13800 fn qwen38_vendor_defaults() -> SamplingDefaults {
13804 SamplingDefaults {
13805 temperature: Some(1.0),
13806 top_p: Some(0.95),
13807 top_k: Some(20),
13808 min_p: Some(0.0),
13809 presence_penalty: Some(0.0),
13810 repetition_penalty: Some(1.0),
13811 frequency_penalty: None,
13812 }
13813 }
13814
13815 fn gemma4_vendor_defaults() -> SamplingDefaults {
13820 SamplingDefaults {
13821 temperature: Some(1.0),
13822 top_p: Some(0.95),
13823 top_k: Some(64),
13824 ..Default::default()
13825 }
13826 }
13827
13828 #[test]
13829 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
13830 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
13835 let chat = |extra: serde_json::Value| {
13836 let mut body = serde_json::json!({
13837 "model": "google/gemma-4-31b-it",
13838 "messages": [{"role": "user", "content": "task"}],
13839 "seed": 7
13841 });
13842 body.as_object_mut()
13843 .unwrap()
13844 .extend(extra.as_object().unwrap().clone());
13845 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13846 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13847 build_chat_request_with_trace(
13848 req,
13849 Some(&ModelCaps {
13850 chat_ok: true,
13851 ..Default::default()
13852 }),
13853 tx,
13854 lanes::Lane::Interactive,
13855 None,
13856 None,
13857 None,
13858 &d,
13859 )
13860 .unwrap()
13861 .request
13862 .sampler_cfg
13863 };
13864
13865 let omitted = chat(serde_json::json!({}));
13867 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
13868 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
13869 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
13870 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
13872 assert_eq!(omitted.penalty_repeat, 1.0);
13873 assert_eq!(omitted.penalty_freq, 0.0);
13874 assert_eq!(omitted.penalty_present, 0.0);
13875 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
13876 assert!(
13877 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
13878 "the vendor default must NOT be greedy — that is the whole point of the lane"
13879 );
13880
13881 let greedy = chat(serde_json::json!({"temperature": 0}));
13884 assert_eq!(
13885 greedy.temperature, 0.0,
13886 "explicit temperature 0 stays greedy"
13887 );
13888 assert!(
13889 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
13890 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
13891 spec/graph exactness arms"
13892 );
13893
13894 let one_field = chat(serde_json::json!({"top_k": 3}));
13896 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
13897 assert_eq!(
13898 one_field.temperature, 1.0,
13899 "omitting temperature still takes the vendor value"
13900 );
13901 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
13902
13903 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
13906 assert_eq!(
13907 disabled.top_k, 0,
13908 "an explicit top_k 0 means KEEP ALL, not 'unset'"
13909 );
13910 assert_eq!(
13911 disabled.top_p, 1.0,
13912 "an explicit top_p 1.0 means untruncated"
13913 );
13914
13915 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
13917 assert_eq!(penal.penalty_present, 1.5);
13918 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
13919 }
13920
13921 #[test]
13922 fn vendor_sampling_defaults_are_identical_on_every_surface() {
13923 let d = qwen38_vendor_defaults();
13935 let md = ModelSamplingDefaults::single(d);
13936 let comp = |extra: serde_json::Value| {
13937 let mut body = serde_json::json!({
13938 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
13939 body.as_object_mut()
13940 .unwrap()
13941 .extend(extra.as_object().unwrap().clone());
13942 let req: CompletionReq = serde_json::from_value(body).unwrap();
13943 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13944 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
13945 };
13946 let chat = |extra: serde_json::Value| {
13947 let mut body = serde_json::json!({
13948 "model": "qwen/qwen3.8-27b",
13949 "messages": [{"role": "user", "content": "task"}],
13950 "seed": 11 });
13951 body.as_object_mut()
13952 .unwrap()
13953 .extend(extra.as_object().unwrap().clone());
13954 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13955 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13956 build_chat_request_with_trace(
13957 req,
13958 Some(&ModelCaps {
13959 chat_ok: true,
13960 ..Default::default()
13961 }),
13962 tx,
13963 lanes::Lane::Interactive,
13964 None,
13965 None,
13966 None,
13967 &md,
13968 )
13969 .unwrap()
13970 .request
13971 .sampler_cfg
13972 };
13973
13974 for extra in [
13975 serde_json::json!({}),
13976 serde_json::json!({"temperature": 0}),
13977 serde_json::json!({"temperature": 0.0}),
13978 serde_json::json!({"temperature": 0.7}),
13979 serde_json::json!({"top_p": 1.0}),
13980 serde_json::json!({"top_k": 0}),
13981 serde_json::json!({"min_p": 0.05}),
13982 serde_json::json!({"repetition_penalty": 1.1}),
13983 serde_json::json!({"frequency_penalty": 0.5}),
13984 serde_json::json!({"presence_penalty": 1.5}),
13985 serde_json::json!({
13986 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
13987 "frequency_penalty": 0.1, "presence_penalty": 0.2,
13988 "repetition_penalty": 1.05 }),
13989 ] {
13990 let c = comp(extra.clone());
13991 let h = chat(extra.clone());
13992 assert_eq!(
13993 (
13994 c.temperature,
13995 c.top_p,
13996 c.top_k,
13997 c.min_p,
13998 c.penalty_repeat,
13999 c.penalty_freq,
14000 c.penalty_present,
14001 c.penalty_last_n,
14002 c.seed
14003 ),
14004 (
14005 h.temperature,
14006 h.top_p,
14007 h.top_k,
14008 h.min_p,
14009 h.penalty_repeat,
14010 h.penalty_freq,
14011 h.penalty_present,
14012 h.penalty_last_n,
14013 h.seed
14014 ),
14015 "/v1/completions and /v1/chat/completions disagree on {extra} — \
14016 standard-surface-law violation"
14017 );
14018 }
14019
14020 let omitted = comp(serde_json::json!({}));
14022 assert_eq!(
14023 omitted.temperature, 1.0,
14024 "qwen3.8 card thinking temperature"
14025 );
14026 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
14027 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
14028 assert!(
14030 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
14031 .is_greedy()
14032 );
14033 }
14034
14035 #[tokio::test]
14048 async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
14049 let _l = DRAIN_LOCK.lock().unwrap();
14050 let step_caps = ModelCaps {
14051 chat_ok: true,
14052 chat_temperature_default: Some(0.5),
14053 chat_top_p_default: Some(0.9),
14054 ..Default::default()
14055 };
14056 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14057 let st = fake_worker_state_full(
14058 1,
14059 std::time::Duration::ZERO,
14060 HashMap::from([("m".to_string(), step_caps)]),
14061 Some(cfg_tx),
14062 );
14063 let fields = |saw: &WorkerSaw| {
14068 let c = &saw.sampler_cfg;
14069 (
14070 c.temperature,
14071 c.top_p,
14072 c.top_k,
14073 c.min_p,
14074 c.penalty_repeat,
14075 c.penalty_freq,
14076 c.penalty_present,
14077 c.penalty_last_n,
14078 )
14079 };
14080 let worker_saw = |surface: &str| {
14081 cfg_rx
14082 .recv_timeout(std::time::Duration::from_secs(10))
14083 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
14084 };
14085
14086 let resp = completions(
14087 State(st.clone()),
14088 axum::http::HeaderMap::new(),
14089 None,
14090 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
14091 )
14092 .await;
14093 assert_eq!(
14094 resp.status(),
14095 StatusCode::OK,
14096 "/v1/completions rejected the omitted-sampling request"
14097 );
14098 let comp = worker_saw("/v1/completions");
14099
14100 let resp = chat_completions(
14101 State(st.clone()),
14102 axum::http::HeaderMap::new(),
14103 None,
14104 Json(
14105 serde_json::from_value(serde_json::json!({
14106 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
14107 .unwrap(),
14108 ),
14109 )
14110 .await;
14111 assert_eq!(
14112 resp.status(),
14113 StatusCode::OK,
14114 "/v1/chat/completions rejected the omitted-sampling request"
14115 );
14116 let chat = worker_saw("/v1/chat/completions");
14117
14118 let resp = anthropic::messages(
14119 State(st.clone()),
14120 axum::http::HeaderMap::new(),
14121 None,
14122 axum::body::Bytes::from(
14123 serde_json::json!({
14124 "model": "m", "max_tokens": 16,
14125 "messages": [{"role": "user", "content": "t"}]})
14126 .to_string(),
14127 ),
14128 )
14129 .await;
14130 assert_eq!(
14131 resp.status(),
14132 StatusCode::OK,
14133 "/v1/messages rejected the omitted-sampling request"
14134 );
14135 let msg = worker_saw("/v1/messages");
14136
14137 let resp = responses_api::responses(
14138 State(st.clone()),
14139 axum::http::HeaderMap::new(),
14140 None,
14141 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
14142 )
14143 .await;
14144 assert_eq!(
14145 resp.status(),
14146 StatusCode::OK,
14147 "/v1/responses rejected the omitted-sampling request"
14148 );
14149 let rsp = worker_saw("/v1/responses");
14150
14151 for (surface, cfg) in [
14152 ("/v1/completions", &comp),
14153 ("/v1/messages", &msg),
14154 ("/v1/responses", &rsp),
14155 ] {
14156 assert_eq!(
14157 fields(cfg),
14158 fields(&chat),
14159 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
14160 for the same omitted-sampling request — standard-surface-law violation \
14161 (hermes d991b51699218285)"
14162 );
14163 }
14164 assert_eq!(
14167 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
14168 (0.5, 0.9),
14169 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
14170 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
14171 );
14172 }
14173
14174 #[tokio::test]
14184 async fn same_effort_value_resolves_identically_on_every_surface() {
14185 let _l = DRAIN_LOCK.lock().unwrap();
14186 let caps = ModelCaps {
14189 chat_ok: true,
14190 effort_levels: true,
14191 ..Default::default()
14192 };
14193 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14194 let st = fake_worker_state_full(
14195 1,
14196 std::time::Duration::ZERO,
14197 HashMap::from([("m".to_string(), caps)]),
14198 Some(saw_tx),
14199 );
14200 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
14201 match surface {
14202 "/v1/chat/completions" => {
14203 chat_completions(
14204 State(st),
14205 axum::http::HeaderMap::new(),
14206 None,
14207 Json(
14208 serde_json::from_value(serde_json::json!({
14209 "model": "m", "max_tokens": 8,
14210 "reasoning_effort": effort,
14211 "messages": [{"role": "user", "content": "t"}]}))
14212 .unwrap(),
14213 ),
14214 )
14215 .await
14216 }
14217 "/v1/responses" => {
14218 responses_api::responses(
14219 State(st),
14220 axum::http::HeaderMap::new(),
14221 None,
14222 axum::body::Bytes::from(
14223 serde_json::json!({
14224 "model": "m", "max_output_tokens": 8, "input": "t",
14225 "reasoning": {"effort": effort}})
14226 .to_string(),
14227 ),
14228 )
14229 .await
14230 }
14231 "/v1/messages" => {
14232 anthropic::messages(
14233 State(st),
14234 axum::http::HeaderMap::new(),
14235 None,
14236 axum::body::Bytes::from(
14237 serde_json::json!({
14238 "model": "m", "max_tokens": 8,
14239 "messages": [{"role": "user", "content": "t"}],
14240 "output_config": {"effort": effort}})
14241 .to_string(),
14242 ),
14243 )
14244 .await
14245 }
14246 other => panic!("unknown surface {other}"),
14247 }
14248 };
14249 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
14250
14251 for (effort, want_think, want_level) in [
14254 ("none", ThinkMode::NoThink, Some("low")),
14255 ("minimal", ThinkMode::NoThink, Some("low")),
14256 ("low", ThinkMode::Think, Some("low")),
14257 ("medium", ThinkMode::Think, Some("medium")),
14258 ("high", ThinkMode::Think, Some("high")),
14259 ("xhigh", ThinkMode::Think, Some("high")),
14261 ] {
14262 for surface in SURFACES {
14263 let resp = send(st.clone(), surface, effort).await;
14264 assert_eq!(
14265 resp.status(),
14266 StatusCode::OK,
14267 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
14268 diverged again (issue #31)"
14269 );
14270 let saw = saw_rx
14271 .recv_timeout(std::time::Duration::from_secs(10))
14272 .unwrap_or_else(|_| {
14273 panic!("{surface}: effort {effort:?} request never reached the worker")
14274 });
14275 assert_eq!(
14276 (saw.think, saw.reasoning_effort.as_deref()),
14277 (want_think, want_level),
14278 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
14279 reasoning surface — the parameter was dropped or remapped before \
14280 parse_think (issue #31 regression)"
14281 );
14282 }
14283 }
14284
14285 for effort in ["bogus", "banana", ""] {
14288 for surface in SURFACES {
14289 let resp = send(st.clone(), surface, effort).await;
14290 assert_eq!(
14291 resp.status(),
14292 StatusCode::BAD_REQUEST,
14293 "{surface} accepted effort {effort:?} — silent-accept regression \
14294 (issue #31: the value never reached parse_think's allowlist)"
14295 );
14296 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
14298 .await
14299 .unwrap();
14300 let v: serde_json::Value = serde_json::from_slice(&body)
14301 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
14302 match surface {
14303 "/v1/messages" => {
14304 assert_eq!(v["type"], "error", "{surface} error envelope");
14305 assert_eq!(
14306 v["error"]["type"], "invalid_request_error",
14307 "{surface} error type"
14308 );
14309 }
14310 _ => {
14311 assert!(
14312 v["error"]["message"].is_string(),
14313 "{surface} OpenAI-shaped error body: {v}"
14314 );
14315 }
14316 }
14317 }
14318 }
14319
14320 let resp = anthropic::messages(
14324 State(st.clone()),
14325 axum::http::HeaderMap::new(),
14326 None,
14327 axum::body::Bytes::from(
14328 serde_json::json!({
14329 "model": "m", "max_tokens": 8,
14330 "messages": [{"role": "user", "content": "t"}],
14331 "thinking": {"type": "enabled"},
14332 "output_config": {"effort": "none"}})
14333 .to_string(),
14334 ),
14335 )
14336 .await;
14337 assert_eq!(resp.status(), StatusCode::OK);
14338 let saw = saw_rx
14339 .recv_timeout(std::time::Duration::from_secs(10))
14340 .expect("thinking+effort request never reached the worker");
14341 assert_eq!(
14342 saw.think,
14343 ThinkMode::Think,
14344 "thinking.type (the documented Anthropic lever) must win the switch over \
14345 output_config.effort"
14346 );
14347 let resp = anthropic::messages(
14348 State(st.clone()),
14349 axum::http::HeaderMap::new(),
14350 None,
14351 axum::body::Bytes::from(
14352 serde_json::json!({
14353 "model": "m", "max_tokens": 8,
14354 "messages": [{"role": "user", "content": "t"}],
14355 "thinking": {"type": "enabled"},
14356 "output_config": {"effort": "banana"}})
14357 .to_string(),
14358 ),
14359 )
14360 .await;
14361 assert_eq!(
14362 resp.status(),
14363 StatusCode::BAD_REQUEST,
14364 "an invalid effort must 400 even next to an explicit thinking.type — \
14365 precedence must not re-open the silent-accept hole"
14366 );
14367 }
14368
14369 #[test]
14370 fn vendor_sampling_defaults_are_boot_validated() {
14371 let parsed = OpenRouterMetadataFile::from_toml(
14374 r#"
14375[models.g]
14376default_temperature = 1.0
14377default_top_p = 0.95
14378default_top_k = 64
14379default_min_p = 0.0
14380default_presence_penalty = 0.0
14381default_frequency_penalty = 0.0
14382default_repetition_penalty = 1.0
14383"#,
14384 )
14385 .unwrap();
14386 let g = parsed.get("g").unwrap();
14387 assert_eq!(g.default_temperature, Some(1.0));
14388 assert_eq!(g.default_top_p, Some(0.95));
14389 assert_eq!(g.default_top_k, Some(64));
14390
14391 let err = OpenRouterMetadataFile::from_toml(
14395 r#"
14396[models.g]
14397default_temperature = 0.0
14398"#,
14399 )
14400 .unwrap_err();
14401 assert!(err.contains("default_temperature"), "{err}");
14402 assert!(
14403 err.contains("greedy"),
14404 "the refusal must say WHY a zero default is refused: {err}"
14405 );
14406
14407 for bad in [
14408 "default_temperature = 2.5",
14409 "default_temperature = -1.0",
14410 "default_top_p = 0.0",
14411 "default_top_p = 1.5",
14412 "default_min_p = 1.0",
14413 "default_min_p = -0.1",
14414 "default_presence_penalty = 3.0",
14415 "default_frequency_penalty = -2.5",
14416 "default_repetition_penalty = 0.0",
14417 ] {
14418 let err =
14419 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
14420 let key = bad.split(' ').next().unwrap();
14421 assert!(err.contains(key), "{bad} must be refused by name: {err}");
14422 }
14423
14424 let err = OpenRouterMetadataFile::from_toml(
14428 r#"
14429[models.g]
14430default_temperture = 1.0
14431"#,
14432 )
14433 .unwrap_err();
14434 assert!(
14435 err.contains("unknown field"),
14436 "an unknown key must be fatal, which is what makes binary-first ordering \
14437 mandatory: {err}"
14438 );
14439 }
14440
14441 #[test]
14442 fn non_thinking_sampling_arm_is_boot_validated() {
14443 let parsed = OpenRouterMetadataFile::from_toml(
14447 r#"
14448[models.q]
14449default_temperature = 1.0
14450default_top_p = 0.95
14451default_top_k = 20
14452
14453[models.q.non_thinking_sampling]
14454temperature = 0.7
14455top_p = 0.8
14456top_k = 20
14457presence_penalty = 1.5
14458"#,
14459 )
14460 .unwrap();
14461 let arm = parsed
14462 .get("q")
14463 .unwrap()
14464 .non_thinking_sampling
14465 .as_ref()
14466 .unwrap();
14467 assert_eq!(arm.temperature, Some(0.7));
14468 assert_eq!(arm.top_p, Some(0.8));
14469 assert_eq!(arm.top_k, Some(20));
14470 assert_eq!(arm.presence_penalty, Some(1.5));
14471 assert_eq!(
14472 arm.min_p, None,
14473 "undeclared arm fields stay undeclared, never invented"
14474 );
14475
14476 let err = OpenRouterMetadataFile::from_toml(
14480 r#"
14481[models.q]
14482[models.q.non_thinking_sampling]
14483temperature = 0.0
14484"#,
14485 )
14486 .unwrap_err();
14487 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
14488 assert!(err.contains("greedy"), "{err}");
14489
14490 let err = OpenRouterMetadataFile::from_toml(
14493 r#"
14494[models.q]
14495[models.q.non_thinking_sampling]
14496"#,
14497 )
14498 .unwrap_err();
14499 assert!(err.contains("non_thinking_sampling"), "{err}");
14500 assert!(err.contains("declare"), "{err}");
14501
14502 for bad in [
14504 "temperature = 2.5",
14505 "top_p = 0.0",
14506 "top_p = 1.5",
14507 "min_p = 1.0",
14508 "presence_penalty = 3.0",
14509 "frequency_penalty = -2.5",
14510 "repetition_penalty = 0.0",
14511 ] {
14512 let err = OpenRouterMetadataFile::from_toml(&format!(
14513 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
14514 ))
14515 .unwrap_err();
14516 let key = bad.split(' ').next().unwrap();
14517 assert!(
14518 err.contains(&format!("non_thinking_sampling.{key}")),
14519 "the refusal for {bad:?} must name the nested key: {err}"
14520 );
14521 }
14522
14523 let err = OpenRouterMetadataFile::from_toml(
14527 r#"
14528[models.q]
14529[models.q.non_thinking_sampling]
14530temperture = 0.7
14531"#,
14532 )
14533 .unwrap_err();
14534 assert!(err.contains("unknown field"), "{err}");
14535 }
14536
14537 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
14542 SamplingDefaults {
14543 temperature: Some(0.7),
14544 top_p: Some(0.8),
14545 top_k: Some(20),
14546 presence_penalty: Some(1.5),
14547 ..Default::default()
14548 }
14549 }
14550
14551 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
14552 ModelSamplingDefaults {
14553 thinking: qwen38_vendor_defaults(),
14554 non_thinking: Some(qwen38_non_thinking_defaults()),
14555 }
14556 }
14557
14558 fn qwen38_caps() -> ModelCaps {
14562 ModelCaps {
14563 chat_ok: true,
14564 qwen_think: true,
14565 think_switch: true,
14566 ..Default::default()
14567 }
14568 }
14569
14570 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
14573 (
14574 c.temperature,
14575 c.top_p,
14576 c.top_k,
14577 c.min_p,
14578 c.penalty_present,
14579 c.penalty_freq,
14580 c.penalty_repeat,
14581 c.penalty_last_n,
14582 c.seed,
14583 )
14584 }
14585
14586 fn build_with_arms(
14587 defaults: &ModelSamplingDefaults,
14588 caps: &ModelCaps,
14589 default_effort: Option<&str>,
14590 extra: serde_json::Value,
14591 ) -> Request {
14592 let mut body = serde_json::json!({
14593 "model": "m",
14594 "messages": [{"role": "user", "content": "task"}],
14595 "seed": 3
14597 });
14598 body.as_object_mut()
14599 .unwrap()
14600 .extend(extra.as_object().unwrap().clone());
14601 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14602 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14603 build_chat_request_with_trace(
14604 req,
14605 Some(caps),
14606 tx,
14607 lanes::Lane::Interactive,
14608 None,
14609 None,
14610 default_effort,
14611 defaults,
14612 )
14613 .unwrap()
14614 .request
14615 }
14616
14617 #[test]
14618 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
14619 let two_arm = qwen38_two_arm_defaults();
14624 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
14625 let caps = qwen38_caps();
14626
14627 let off_spellings = [
14629 serde_json::json!({"reasoning_effort": "none"}),
14630 serde_json::json!({"enable_thinking": false}),
14631 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
14632 serde_json::json!({"reasoning": {"enabled": false}}),
14633 ];
14634 for extra in &off_spellings {
14635 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
14636 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
14637 let c = &r.sampler_cfg;
14638 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
14639 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
14640 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
14641 assert_eq!(
14642 c.penalty_present, 1.5,
14643 "{extra}: non-thinking presence_penalty"
14644 );
14645 assert_eq!(
14646 c.penalty_last_n,
14647 memra_engine::spec::PEN_WINDOW_MAX,
14648 "{extra}: the arm's presence penalty uses the cross-path history window"
14649 );
14650 assert_eq!(
14651 c.min_p, 0.0,
14652 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
14653 );
14654
14655 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
14658 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
14659 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
14660 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
14661 assert_eq!(
14662 s.sampler_cfg.penalty_present, 0.0,
14663 "{extra}: single-arm model"
14664 );
14665 }
14666
14667 for extra in [
14670 serde_json::json!({}),
14671 serde_json::json!({"enable_thinking": true}),
14672 serde_json::json!({"reasoning_effort": "high"}),
14673 serde_json::json!({"reasoning": {"enabled": true}}),
14674 ] {
14675 for defaults in [&two_arm, &single_arm] {
14676 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
14677 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
14678 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
14679 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
14680 assert_eq!(
14681 c.penalty_present, 0.0,
14682 "{extra}: thinking arm has no presence"
14683 );
14684 }
14685 }
14686
14687 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
14690 assert_eq!(
14691 c.temperature, 0.7,
14692 "deployment-default off = non-thinking arm"
14693 );
14694 let c = build_with_arms(
14696 &two_arm,
14697 &caps,
14698 Some("none"),
14699 serde_json::json!({"enable_thinking": true}),
14700 )
14701 .sampler_cfg;
14702 assert_eq!(
14703 c.temperature, 1.0,
14704 "explicit ON beats the deployment default"
14705 );
14706
14707 let c = build_with_arms(
14709 &two_arm,
14710 &caps,
14711 None,
14712 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
14713 )
14714 .sampler_cfg;
14715 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
14716 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
14717 let c = build_with_arms(
14718 &two_arm,
14719 &caps,
14720 None,
14721 serde_json::json!({
14722 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
14723 )
14724 .sampler_cfg;
14725 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
14726 assert_eq!(
14727 c.penalty_present, 0.0,
14728 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
14729 is a value, not an absence"
14730 );
14731 assert_eq!(
14732 c.penalty_last_n, 0,
14733 "all penalties off => no history window"
14734 );
14735 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
14736
14737 let c = build_with_arms(
14740 &two_arm,
14741 &caps,
14742 None,
14743 serde_json::json!({"enable_thinking": false, "temperature": 0}),
14744 )
14745 .sampler_cfg;
14746 assert!(
14747 memra_engine::sampler::Sampler::new(c).is_greedy(),
14748 "explicit temperature 0 must stay greedy on the non-thinking arm"
14749 );
14750
14751 let c = build_with_arms(
14754 &single_arm,
14755 &caps,
14756 None,
14757 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
14758 )
14759 .sampler_cfg;
14760 assert_eq!(c.temperature, 0.55);
14761 assert_eq!(
14762 c.top_p, 0.95,
14763 "single-arm model: unset top_p takes its one arm"
14764 );
14765 }
14766
14767 #[test]
14768 fn sampling_arms_never_blend_field_by_field() {
14769 let parsed = OpenRouterMetadataFile::from_toml(
14774 r#"
14775[models.m]
14776default_temperature = 1.0
14777default_min_p = 0.05
14778
14779[models.m.non_thinking_sampling]
14780temperature = 0.6
14781"#,
14782 )
14783 .unwrap();
14784 let caps = ModelCaps {
14785 chat_temperature_default: Some(0.5),
14786 chat_top_p_default: Some(0.9),
14787 ..Default::default()
14788 };
14789 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
14790 let client = ClientSampling {
14791 seed: Some(1),
14792 ..Default::default()
14793 };
14794
14795 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
14796 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
14797 assert_eq!(
14798 off.min_p, 0.0,
14799 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
14800 );
14801 assert_eq!(
14802 off.top_p, 1.0,
14803 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
14804 );
14805
14806 for mode in [ThinkMode::Default, ThinkMode::Think] {
14808 let on = resolve_sampler_config(client, d.for_mode(mode));
14809 assert_eq!(on.temperature, 1.0);
14810 assert_eq!(on.min_p, 0.05);
14811 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
14812 }
14813 }
14814
14815 #[test]
14816 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
14817 let caps = qwen38_caps();
14826 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
14827 let two_arm = qwen38_two_arm_defaults();
14828
14829 let bodies = [
14830 serde_json::json!({}),
14831 serde_json::json!({"enable_thinking": true}),
14832 serde_json::json!({"reasoning_effort": "high"}),
14833 serde_json::json!({"reasoning_effort": "none"}),
14834 serde_json::json!({"enable_thinking": false}),
14835 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
14836 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
14837 serde_json::json!({"enable_thinking": false, "temperature": 0}),
14838 ];
14839 for extra in &bodies {
14840 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
14842 let mut client = ClientSampling {
14843 seed: Some(3),
14844 ..Default::default()
14845 };
14846 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
14847 client.temperature = Some(t as f32);
14848 }
14849 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
14850 client.top_p = Some(p as f32);
14851 }
14852 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
14853 assert_eq!(
14854 sampler_key(&r.sampler_cfg),
14855 sampler_key(&pre_arm),
14856 "{extra}: single-arm model diverged from the pre-arm resolution law"
14857 );
14858
14859 if r.think != ThinkMode::NoThink {
14862 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
14863 assert_eq!(t.think, r.think, "{extra}");
14864 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
14865 assert_eq!(
14866 sampler_key(&t.sampler_cfg),
14867 sampler_key(&r.sampler_cfg),
14868 "{extra}: a thinking-on request must not feel the non-thinking arm"
14869 );
14870 }
14871 }
14872 }
14873
14874 #[test]
14875 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
14876 let r = build_with_arms(
14882 &qwen38_two_arm_defaults(),
14883 &qwen38_caps(),
14884 None,
14885 serde_json::json!({"response_format": {"type": "json_object"}}),
14886 );
14887 assert_eq!(
14888 r.think,
14889 ThinkMode::NoThink,
14890 "constraint forces the switch off"
14891 );
14892 assert_eq!(
14893 r.sampler_cfg.temperature, 0.7,
14894 "and the arm follows the real mode"
14895 );
14896 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
14897 }
14898
14899 #[test]
14900 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
14901 let caps = ModelCaps {
14906 chat_temperature_default: Some(0.5),
14907 chat_top_p_default: Some(0.9),
14908 chat_ok: true,
14909 ..Default::default()
14910 };
14911 let metadata = OpenRouterModelMetadata {
14912 default_temperature: Some(1.0),
14913 default_top_p: Some(0.95),
14914 default_top_k: Some(64),
14915 ..Default::default()
14916 };
14917
14918 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
14919 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
14920 assert_eq!(caps_only.top_p, Some(0.9));
14921 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
14922
14923 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
14924 assert_eq!(
14925 both.temperature,
14926 Some(1.0),
14927 "metadata outranks the arch cap"
14928 );
14929 assert_eq!(both.top_p, Some(0.95));
14930 assert_eq!(both.top_k, Some(64));
14931
14932 let partial = SamplingDefaults::resolve(
14934 Some(&OpenRouterModelMetadata {
14935 default_temperature: Some(0.7),
14936 ..Default::default()
14937 }),
14938 Some(&caps),
14939 );
14940 assert_eq!(partial.temperature, Some(0.7));
14941 assert_eq!(
14942 partial.top_p,
14943 Some(0.9),
14944 "an undeclared metadata field must fall through to the cap, not to 1.0"
14945 );
14946
14947 assert_eq!(
14949 SamplingDefaults::resolve(None, None),
14950 SamplingDefaults::default()
14951 );
14952 }
14953
14954 #[test]
14955 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
14956 let resolved = |d: &SamplingDefaults| {
14969 resolve_sampler_config(
14970 ClientSampling {
14971 seed: Some(1),
14972 ..Default::default()
14973 },
14974 d,
14975 )
14976 };
14977
14978 assert!(
14980 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
14981 .is_spec_sampling(),
14982 "the API-standard default must stay in the fast pure-temp regime"
14983 );
14984
14985 for (name, d) in [
14986 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
14987 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
14988 ] {
14989 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
14990 assert!(
14991 !sampler.is_greedy(),
14992 "{name}: vendor default must not be greedy"
14993 );
14994 assert!(
14995 !sampler.is_spec_sampling(),
14996 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
14997 starts passing, either the vendor numbers changed or the in-graph draft \
14998 learned filters, and the perf note in docs/SERVING.md needs revisiting"
14999 );
15000 }
15001
15002 let opted_out = resolve_sampler_config(
15004 ClientSampling {
15005 top_p: Some(1.0),
15006 top_k: Some(0),
15007 seed: Some(1),
15008 ..Default::default()
15009 },
15010 &qwen38_vendor_defaults(),
15011 );
15012 assert!(
15013 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
15014 "explicitly disabling the filters must restore the pure-temp regime"
15015 );
15016 }
15017
15018 #[test]
15019 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
15020 let comp_seed = |body: serde_json::Value| {
15027 let req: CompletionReq = serde_json::from_value(body).unwrap();
15028 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15029 build_request(&req, tx, lanes::Lane::Interactive, None)
15030 .sampler_cfg
15031 .seed
15032 };
15033 let chat_seed = |body: serde_json::Value| {
15034 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15035 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15036 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15037 .unwrap()
15038 .request
15039 .sampler_cfg
15040 .seed
15041 };
15042
15043 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15046 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15047 let c = chat_seed(serde_json::json!({
15048 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
15049 assert_ne!(
15050 a, 0,
15051 "omitted seed must not be the pinned 0 that caused the loop"
15052 );
15053 assert_ne!(b, 0);
15054 assert_ne!(c, 0);
15055 assert_ne!(
15056 a, b,
15057 "two seed-omitting requests must get DIFFERENT streams"
15058 );
15059 assert_ne!(a, c);
15060
15061 assert_eq!(
15064 comp_seed(serde_json::json!({
15065 "model": "m", "prompt": "t", "seed": 0})),
15066 0,
15067 "explicit seed 0 must stay 0 — the determinism gates depend on it"
15068 );
15069 assert_eq!(
15070 comp_seed(serde_json::json!({
15071 "model": "m", "prompt": "t", "seed": 12345})),
15072 12345
15073 );
15074 assert_eq!(
15075 chat_seed(serde_json::json!({
15076 "model": "m", "messages": [{"role": "user", "content": "t"}],
15077 "seed": 777})),
15078 777
15079 );
15080 assert_eq!(
15082 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
15083 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
15084 );
15085
15086 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
15089 assert_eq!(
15090 seeds.len(),
15091 256,
15092 "fresh_seed must not collide across rapid calls"
15093 );
15094 assert!(!seeds.contains(&0));
15095 }
15096
15097 #[test]
15098 fn response_format_builds_grammar_only_when_present() {
15099 let mk = |rf: Option<serde_json::Value>| {
15103 let mut body = serde_json::json!({
15104 "model": "m", "messages": [{"role": "user", "content": "t"}]});
15105 if let Some(rf) = rf {
15106 body["response_format"] = rf;
15107 }
15108 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15109 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15110 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15111 };
15112 assert!(mk(None).unwrap().request.grammar.is_none());
15113 assert!(
15114 mk(Some(serde_json::json!({"type": "text"})))
15115 .unwrap()
15116 .request
15117 .grammar
15118 .is_none()
15119 );
15120 assert!(matches!(
15121 mk(Some(serde_json::json!({"type": "json_object"})))
15122 .unwrap()
15123 .request
15124 .grammar,
15125 Some(constrained::GrammarSpec::JsonObject)
15126 ));
15127 assert!(matches!(
15128 mk(Some(serde_json::json!({"type": "json_schema",
15129 "json_schema": {"schema": {"type": "object"}}})))
15130 .unwrap()
15131 .request
15132 .grammar,
15133 Some(constrained::GrammarSpec::JsonSchema(_))
15134 ));
15135 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
15137 }
15138
15139 #[test]
15140 fn unsupported_semantic_params_are_named_rejections() {
15141 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15143 "model": "m", "messages": [{"role": "user", "content": "t"}],
15144 "response_format": {"type": "json_object"}
15145 }))
15146 .unwrap();
15147 assert!(req.response_format.is_some());
15148 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15149 "model": "m", "messages": [{"role": "user", "content": "t"}],
15150 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
15151 "user": "u-1", "stream_options": {"include_usage": true}
15152 }))
15153 .unwrap();
15154 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
15156 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
15157 assert_eq!(req.n, Some(1));
15158 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
15160 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
15161 assert_eq!(param, "logit_bias");
15162 assert_eq!(msg, "logit_bias is not supported (why)");
15163 }
15164
15165 #[test]
15166 fn completions_accept_openai_stop_forms() {
15167 for (value, expected) in [
15168 (serde_json::json!("Problem:"), vec!["Problem:"]),
15169 (
15170 serde_json::json!(["Question:", "Problem:"]),
15171 vec!["Question:", "Problem:"],
15172 ),
15173 (serde_json::Value::Null, Vec::<&str>::new()),
15174 ] {
15175 let req: CompletionReq = serde_json::from_value(serde_json::json!({
15176 "model": "plain_quant", "prompt": "task", "stop": value
15177 }))
15178 .unwrap();
15179 assert_eq!(req.stop.into_vec(), expected);
15180 }
15181 }
15182
15183 fn fake_worker_state() -> AppState {
15190 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
15191 }
15192
15193 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
15194 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
15195 }
15196
15197 struct WorkerSaw {
15202 sampler_cfg: SamplerConfig,
15203 think: ThinkMode,
15204 reasoning_effort: Option<String>,
15205 }
15206
15207 fn fake_worker_state_full(
15214 steps: usize,
15215 step_delay: std::time::Duration,
15216 caps: HashMap<String, ModelCaps>,
15217 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
15218 ) -> AppState {
15219 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15220 let health = health::WorkerHealth::new();
15221 let h = health.clone();
15222 std::thread::spawn(move || {
15223 h.mark_ready();
15224 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
15225 if let Some(tx) = &saw_tx {
15226 let _ = tx.send(WorkerSaw {
15227 sampler_cfg: req.sampler_cfg.clone(),
15228 think: req.think,
15229 reasoning_effort: req.reasoning_effort.clone(),
15230 });
15231 }
15232 worker::release_pending_admit();
15236 worker::release_admission_reservation(req.lane);
15237 h.beat_busy();
15238 if let Some(ready) = req.constraint_ready.take() {
15239 let _ = ready.send(Ok(()));
15240 }
15241 let _ = req.tx.send(Event::PromptUsage {
15242 n_prompt: 1,
15243 n_cached: 0,
15244 });
15245 for step in 0..steps {
15246 h.beat_busy();
15247 let text = if steps == 1 { "ok" } else { "x" };
15248 let _ = req.tx.send(Event::Token {
15249 id: step as u32 + 1,
15250 text: text.into(),
15251 });
15252 if !step_delay.is_zero() {
15253 std::thread::sleep(step_delay);
15254 }
15255 }
15256 let _ = req.tx.send(Event::Done {
15257 stop_reason: "Eos".into(),
15258 n_tokens: steps,
15259 n_prompt: 1,
15260 n_cached: 0,
15261 elapsed_s: 0.01,
15262 spec: None,
15263 });
15264 h.set_phase(health::PHASE_IDLE);
15265 }
15266 });
15267 for _ in 0..2000 {
15270 if health.live().is_ok() {
15271 break;
15272 }
15273 std::thread::sleep(std::time::Duration::from_millis(1));
15274 }
15275 AppState {
15276 cmd_tx,
15277 models: Arc::new(vec!["m".into()]),
15278 caps: Arc::new(caps),
15279 openrouter_metadata: Arc::new(HashMap::new()),
15280 provider_metadata: Arc::new(None),
15281 metering: None,
15282
15283 budget_tokenizers: None,
15284 api_auth: ApiAuth::default(),
15285 metrics_auth: MetricsAuth::default(),
15286 metrics: SharedMetrics::default(),
15287 started: 1,
15288 inflight: Arc::new(Default::default()),
15289 tenant_inflight: Arc::new(Default::default()),
15290 health,
15291 bg: None,
15292 }
15293 }
15294
15295 #[tokio::test]
15296 async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
15297 let _l = DRAIN_LOCK.lock().unwrap();
15298 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
15299 let normal_state = st.clone();
15300 let normal = tokio::spawn(async move {
15301 chat_completions(
15302 State(normal_state),
15303 axum::http::HeaderMap::new(),
15304 None,
15305 Json(
15306 serde_json::from_value(serde_json::json!({
15307 "model": "m",
15308 "messages": [{"role": "user", "content": "keep decoding"}],
15309 }))
15310 .unwrap(),
15311 ),
15312 )
15313 .await
15314 });
15315 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
15316
15317 let mut deep = serde_json::json!({"type": "string"});
15318 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
15319 deep = serde_json::json!({"allOf": [deep]});
15320 }
15321 let bad = chat_completions(
15322 State(st.clone()),
15323 axum::http::HeaderMap::new(),
15324 None,
15325 Json(
15326 serde_json::from_value(serde_json::json!({
15327 "model": "m",
15328 "messages": [{"role": "user", "content": "bad schema"}],
15329 "response_format": {
15330 "type": "json_schema",
15331 "json_schema": {"schema": deep},
15332 },
15333 }))
15334 .unwrap(),
15335 ),
15336 )
15337 .await;
15338 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
15339 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
15340 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
15341 .await
15342 .unwrap();
15343 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15344 assert!(
15345 payload["error"]["message"]
15346 .as_str()
15347 .unwrap()
15348 .contains("maximum nesting depth")
15349 );
15350 assert!(
15351 !normal.is_finished(),
15352 "bad schema stalled or replaced the normal decode"
15353 );
15354
15355 let normal_response = normal.await.unwrap();
15356 assert_eq!(normal_response.status(), StatusCode::OK);
15357 let snapshot = st.health.snapshot();
15358 assert!(
15359 st.health.live().is_ok(),
15360 "normal decode left health stalled"
15361 );
15362 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
15363 }
15364
15365 #[tokio::test]
15366 async fn valid_response_format_preflight_preserves_generation() {
15367 let _l = DRAIN_LOCK.lock().unwrap();
15368 let response = chat_completions(
15369 State(fake_worker_state()),
15370 axum::http::HeaderMap::new(),
15371 None,
15372 Json(
15373 serde_json::from_value(serde_json::json!({
15374 "model": "m",
15375 "messages": [{"role": "user", "content": "valid schema"}],
15376 "response_format": {"type": "json_object"},
15377 }))
15378 .unwrap(),
15379 ),
15380 )
15381 .await;
15382 assert_eq!(response.status(), StatusCode::OK);
15383 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15384 .await
15385 .unwrap();
15386 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15387 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
15388 }
15389
15390 #[tokio::test]
15391 async fn unknown_model_refuses_model_not_found_before_admission() {
15392 let _l = DRAIN_LOCK.lock().unwrap();
15393 let response = chat_completions(
15398 State(fake_worker_state()),
15399 axum::http::HeaderMap::new(),
15400 None,
15401 Json(
15402 serde_json::from_value(serde_json::json!({
15403 "model": "qwen/qwen3.8-27b-typo",
15404 "messages": [{"role": "user", "content": "hi"}],
15405 }))
15406 .unwrap(),
15407 ),
15408 )
15409 .await;
15410 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15411 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15412 .await
15413 .unwrap();
15414 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15415 assert_eq!(payload["error"]["code"], "model_not_found");
15416 assert_eq!(payload["error"]["type"], "invalid_request_error");
15417
15418 let response = completions(
15420 State(fake_worker_state()),
15421 axum::http::HeaderMap::new(),
15422 None,
15423 Json(
15424 serde_json::from_value(serde_json::json!({
15425 "model": "nope",
15426 "prompt": "hi",
15427 }))
15428 .unwrap(),
15429 ),
15430 )
15431 .await;
15432 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15433 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15434 .await
15435 .unwrap();
15436 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15437 assert_eq!(payload["error"]["code"], "model_not_found");
15438 }
15439
15440 const METRICS_KEY_ACME: &str = "completion-acme-secret";
15441 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
15442
15443 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
15444 let spec = format!(
15445 "acme:{},blue:{}",
15446 auth::sha256_hex(METRICS_KEY_ACME),
15447 auth::sha256_hex(METRICS_KEY_BLUE),
15448 );
15449 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
15450 let mut st = fake_worker_state();
15451 st.api_auth.keyring = Some(keyring);
15452 st.metrics_auth = MetricsAuth::new(
15453 true,
15454 st.api_auth.configured(),
15455 metrics_token.map(str::to_string),
15456 );
15457 {
15458 let mut metrics = st.metrics.lock().unwrap();
15459 metrics.admitted = 17;
15460 metrics.prompt_tokens_in = 400;
15461 metrics.cached_tokens_in = 60;
15462 metrics.prefix_hits = 2;
15463 metrics.prefix_misses = 3;
15464 metrics.prefix_inserts = 5;
15465 metrics.prefix_evictions = 7;
15466 metrics.prefix_skips_budget = 9;
15467 metrics.prefix_skips_pinned = 10;
15468 metrics.prefix_hit_tokens = 11;
15469 metrics.lcp_hist[4] = 13;
15470 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
15471 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
15472 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
15473 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
15474 metrics.prefix_entries = 29;
15475 metrics.prefix_bytes = 31;
15476 metrics.active_sessions = 3;
15477 metrics.queued_requests = 5;
15478 metrics.continuation_pool_entries = 7;
15479 metrics.spec_pool_entries = 11;
15480 metrics.cuda_driver_free_bytes = 13;
15481 metrics.cuda_pool_reserved_bytes = 17;
15482 metrics.cuda_pool_used_bytes = 19;
15483 metrics.cuda_pool_cached_bytes = 23;
15484 metrics.batch_size_last = 37;
15485 metrics.spec.insert(
15486 "m".into(),
15487 memra_engine::spec::SpecTelemetry {
15488 rounds: 2,
15489 drafted: 6,
15490 accepted: 4,
15491 ..Default::default()
15492 },
15493 );
15494 let mut spec_window = memra_engine::spec::SpecTelemetry {
15495 rounds: 4,
15496 drafted: 12,
15497 accepted: 6,
15498 ..Default::default()
15499 };
15500 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
15501 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
15502 metrics.spec_window.insert("m".into(), spec_window);
15503 metrics.constraint_compiler_fail_closed.insert(
15504 "m".into(),
15505 Arc::new(std::sync::atomic::AtomicBool::new(true)),
15506 );
15507 }
15508 st
15509 }
15510
15511 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15512 let mut headers = HeaderMap::new();
15513 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15514 let response = get_metrics(State(st), headers).await;
15515 assert_eq!(response.status(), StatusCode::OK);
15516 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15517 .await
15518 .unwrap();
15519 serde_json::from_slice(&bytes).unwrap()
15520 }
15521
15522 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15523 let mut headers = HeaderMap::new();
15524 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15525 let response = yield_metrics(State(st), headers).await;
15526 assert_eq!(response.status(), StatusCode::OK);
15527 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15528 .await
15529 .unwrap();
15530 serde_json::from_slice(&bytes).unwrap()
15531 }
15532
15533 #[test]
15534 fn exposed_open_bind_is_refused_before_server_start() {
15535 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
15536 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
15537
15538 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
15539 assert!(err.contains("refusing unauthenticated non-loopback bind"));
15540 assert!(err.contains("MEMRA_API_KEY"));
15541 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
15542 assert!(validate_bind_security("[::]:8000", false, false).is_err());
15543
15544 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
15545 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
15546 }
15547
15548 #[tokio::test]
15549 async fn keyed_metrics_require_and_accept_api_bearer() {
15550 let mut st = fake_worker_state();
15551 st.api_auth.single_key = Some(Arc::from("completion-secret"));
15552 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
15553
15554 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
15555 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
15556 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
15557 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
15558
15559 let mut headers = HeaderMap::new();
15560 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
15561 assert_eq!(
15562 get_metrics(State(st.clone()), headers.clone())
15563 .await
15564 .status(),
15565 StatusCode::OK,
15566 );
15567 let body = metrics_json(st.clone(), "completion-secret").await;
15568 assert!(
15569 body.get("admitted").is_some(),
15570 "the legacy single-key domain keeps cumulative counters",
15571 );
15572 assert!(
15573 body.get("active_sessions").is_none(),
15574 "a static completion key is not an operator metrics principal",
15575 );
15576 assert_eq!(
15577 yield_metrics(State(st), headers).await.status(),
15578 StatusCode::OK
15579 );
15580 }
15581
15582 #[tokio::test]
15583 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
15584 let st = multi_key_metrics_state(None);
15585 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
15586 assert_eq!(
15587 body.as_object().unwrap().len(),
15588 2,
15589 "completion metrics must contain only tenant-scoped rows",
15590 );
15591 let tenants = body["tenants"].as_object().unwrap();
15592 assert_eq!(tenants.len(), 1);
15593 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
15594 assert!(!tenants.contains_key("t:blue"));
15595 let adsd = body["adsd_suspect_total"].as_object().unwrap();
15596 assert_eq!(adsd.len(), 1);
15597 assert_eq!(adsd["t:acme"], 1);
15598 assert!(!adsd.contains_key("t:blue"));
15599
15600 let mut headers = HeaderMap::new();
15601 headers.insert(
15602 "authorization",
15603 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
15604 );
15605 assert_eq!(
15606 yield_metrics(State(st), headers).await.status(),
15607 StatusCode::FORBIDDEN,
15608 "the process-wide yield view requires an operator metrics token",
15609 );
15610 }
15611
15612 #[tokio::test]
15613 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
15614 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
15615 for operator_only in [
15616 "prefix_cache_entries",
15617 "prefix_cache_bytes",
15618 "prefix_cache_skips_budget",
15619 "prefix_cache_skips_pinned",
15620 "active_sessions",
15621 "queued_requests",
15622 "continuation_pool_entries",
15623 "spec_pool_entries",
15624 "cuda_driver_free_bytes",
15625 "cuda_pool_reserved_bytes",
15626 "cuda_pool_used_bytes",
15627 "cuda_pool_cached_bytes",
15628 "constraint_compiler_fail_closed",
15629 "serve_idle_seconds",
15630 "spec",
15631 "spec_tau",
15632 "spec_accept_by_position",
15633 "dual_pp",
15634 "peer_probe_bypassed",
15635 "peer_probe_boundary_copies",
15636 "peer_probe_runtime_reprobes",
15637 "peer_probe_runtime_failures",
15638 "peer_probe_deferred_total",
15639 "peer_probe_integrity_degraded",
15640 "peer_probe_degraded_to_host_bounce",
15641 ] {
15642 assert!(
15643 body.get(operator_only).is_none(),
15644 "tenant metrics must not expose operator field {operator_only}",
15645 );
15646 }
15647 }
15648
15649 #[test]
15650 fn populated_spec_acceptance_metrics_are_operator_only() {
15651 for scope in [
15652 MetricsScope::CompletionDomain,
15653 MetricsScope::Tenant("t:acme".into()),
15654 ] {
15655 let mut body = json!({});
15656 insert_spec_acceptance_metrics(&mut body, &scope, || {
15657 panic!("tenant scope evaluated the process-wide spec snapshot")
15658 });
15659 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
15660 assert!(
15661 body.get("spec_accept_by_position").is_none(),
15662 "{scope:?} leaked the accept histogram"
15663 );
15664 }
15665
15666 let mut telemetry = memra_engine::spec::SpecTelemetry {
15667 rounds: 4,
15668 drafted: 12,
15669 accepted: 6,
15670 ..Default::default()
15671 };
15672 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
15673 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
15674 let mut body = json!({});
15675 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
15676 HashMap::from([("model-a".to_string(), telemetry)])
15677 });
15678 assert_eq!(body["spec_tau"]["model-a"], 1.5);
15679 let histogram = &body["spec_accept_by_position"]["model-a"];
15680 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
15681 assert_eq!(histogram["rounds"], 4);
15682 assert_eq!(histogram["offered"], json!([4, 4, 4]));
15683 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
15684 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
15685 }
15686
15687 #[test]
15688 fn populated_dual_pp_metrics_are_operator_only() {
15689 let populated = DualPpMetricsSnapshot {
15690 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
15691 stage_samples: [1, 1, 1, 1],
15692 dropped_timing_samples: 0,
15693 overlaps: 17,
15694 slot_pairs: 19,
15695 slot_uses: [19, 19],
15696 slot_collisions: 0,
15697 };
15698 for scope in [
15699 MetricsScope::CompletionDomain,
15700 MetricsScope::Tenant("t:acme".into()),
15701 ] {
15702 let mut body = json!({});
15703 insert_dual_pp_metrics(&mut body, &scope, || populated);
15704 assert!(
15705 body.get("dual_pp").is_none(),
15706 "{scope:?} leaked dual PP topology"
15707 );
15708 }
15709
15710 let mut body = json!({});
15711 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
15712 assert_eq!(body["dual_pp"]["overlaps"], 17);
15713 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
15714 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
15715 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
15716 assert_eq!(
15717 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
15718 1.0
15719 );
15720 }
15721
15722 #[test]
15723 fn peer_probe_metrics_are_operator_only() {
15724 let populated = memra_engine::pp::PeerProbeMetrics {
15725 bypassed: 1,
15726 boundary_copies: 8_192,
15727 runtime_probes: 1,
15728 runtime_failures: 0,
15729 deferred_total: 4,
15730 integrity_degraded: true,
15731 degraded_to_host_bounce: true,
15732 };
15733 for scope in [
15734 MetricsScope::CompletionDomain,
15735 MetricsScope::Tenant("t:acme".into()),
15736 ] {
15737 let mut body = json!({});
15738 insert_peer_probe_metrics(&mut body, &scope, || populated);
15739 assert!(body.get("peer_probe_bypassed").is_none());
15740 }
15741
15742 let mut body = json!({});
15743 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
15744 assert_eq!(body["peer_probe_bypassed"], 1);
15745 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
15746 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
15747 assert_eq!(body["peer_probe_runtime_failures"], 0);
15748 assert_eq!(body["peer_probe_deferred_total"], 4);
15749 assert_eq!(body["peer_probe_integrity_degraded"], true);
15750 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
15751 }
15752
15753 #[tokio::test]
15754 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
15755 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
15756 for operator_only in [
15757 "lcp_histogram",
15758 "cache_hit_token_ratio",
15759 "prefix_cache_hits",
15760 "prefix_cache_misses",
15761 "prefix_cache_inserts",
15762 "prefix_cache_evictions",
15763 "prefix_cache_skips_budget",
15764 "prefix_cache_skips_pinned",
15765 "prefix_cache_hit_tokens",
15766 ] {
15767 assert!(
15768 tenant_body.get(operator_only).is_none(),
15769 "tenant metrics must not expose global prefix field {operator_only}",
15770 );
15771 }
15772 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
15773 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
15774 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
15775 assert_eq!(
15776 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
15777 0.4
15778 );
15779
15780 let operator_body = metrics_json(
15781 multi_key_metrics_state(Some("scrape-secret")),
15782 "scrape-secret",
15783 )
15784 .await;
15785 assert_eq!(operator_body["prefix_cache_hits"], 2);
15786 assert_eq!(operator_body["prefix_cache_misses"], 3);
15787 assert_eq!(operator_body["prefix_cache_inserts"], 5);
15788 assert_eq!(operator_body["prefix_cache_evictions"], 7);
15789 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
15790 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
15791 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
15792 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
15793 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
15794 }
15795
15796 #[tokio::test]
15797 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
15798 let st = multi_key_metrics_state(Some("scrape-secret"));
15799 let mut completion_headers = HeaderMap::new();
15800 completion_headers.insert(
15801 "authorization",
15802 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
15803 );
15804 assert_eq!(
15805 get_metrics(State(st.clone()), completion_headers.clone())
15806 .await
15807 .status(),
15808 StatusCode::FORBIDDEN,
15809 );
15810 assert_eq!(
15811 yield_metrics(State(st.clone()), completion_headers)
15812 .await
15813 .status(),
15814 StatusCode::FORBIDDEN,
15815 );
15816
15817 let body = metrics_json(st.clone(), "scrape-secret").await;
15818 let tenants = body["tenants"].as_object().unwrap();
15819 assert_eq!(tenants.len(), 2);
15820 assert!(tenants.contains_key("t:acme"));
15821 assert!(tenants.contains_key("t:blue"));
15822 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
15823 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
15824 assert_eq!(body["active_sessions"], 3);
15825 assert_eq!(body["queued_requests"], 5);
15826 assert_eq!(body["prefix_cache_bytes"], 31);
15827 assert_eq!(body["cuda_driver_free_bytes"], 13);
15828 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
15829 assert_eq!(body["spec"]["m"]["drafted"], 6);
15830 assert_eq!(body["spec_tau"]["m"], 1.5);
15831 assert_eq!(
15832 body["spec_accept_by_position"]["m"]["accepted"],
15833 json!([3, 2, 1])
15834 );
15835 let yield_body = yield_metrics_json(st, "scrape-secret").await;
15836 assert_eq!(yield_body["batch_size_last"], 37);
15837 }
15838
15839 #[tokio::test]
15840 async fn metrics_token_protects_public_override_without_api_keys() {
15841 let mut st = fake_worker_state();
15842 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
15843
15844 assert_eq!(
15845 get_metrics(State(st.clone()), HeaderMap::new())
15846 .await
15847 .status(),
15848 StatusCode::UNAUTHORIZED,
15849 );
15850 let mut headers = HeaderMap::new();
15851 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
15852 assert_eq!(
15853 get_metrics(State(st.clone()), headers.clone())
15854 .await
15855 .status(),
15856 StatusCode::OK,
15857 );
15858 assert_eq!(
15859 yield_metrics(State(st), headers).await.status(),
15860 StatusCode::OK
15861 );
15862 }
15863
15864 #[tokio::test]
15865 async fn no_key_loopback_metrics_remain_open_for_development() {
15866 let mut st = fake_worker_state();
15867 st.metrics_auth = MetricsAuth::new(true, false, None);
15868 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
15869 assert_eq!(response.status(), StatusCode::OK);
15870 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15871 .await
15872 .unwrap();
15873 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15874 assert!(
15875 body.get("active_sessions").is_some(),
15876 "no-key loopback development keeps full operator visibility",
15877 );
15878 assert_eq!(
15879 yield_metrics(State(st), HeaderMap::new()).await.status(),
15880 StatusCode::OK,
15881 );
15882 }
15883
15884 #[test]
15885 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
15886 let metrics = SharedMetrics::default();
15887 let rl = RateLimit::compute(4, 1, &metrics);
15889 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
15890 let rl = RateLimit::compute(4, 3, &metrics);
15891 assert_eq!(rl.remaining, 1);
15892 let rl = RateLimit::compute(4, 4, &metrics);
15894 assert_eq!(rl.remaining, 0);
15895 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
15896 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
15898 let m = worker::Metrics {
15900 completed: 2,
15901 tokens_out: 200,
15902 step_p50_ms: 20.0,
15903 ..Default::default()
15904 };
15905 assert_eq!(reset_estimate_s(&m), 2); }
15907
15908 #[test]
15909 fn inflight_guard_counts_up_and_frees_on_drop() {
15910 let counts: InflightCounts = Arc::new(Default::default());
15911 let tenants: TenantGauge = Arc::new(Default::default());
15912 let (g1, n1, t1) = InflightGuard::try_acquire(
15913 counts.clone(),
15914 lanes::Lane::Interactive,
15915 tenants.clone(),
15916 "acme",
15917 None,
15918 )
15919 .unwrap();
15920 let (g2, n2, t2) = InflightGuard::try_acquire(
15921 counts.clone(),
15922 lanes::Lane::Interactive,
15923 tenants.clone(),
15924 "acme",
15925 None,
15926 )
15927 .unwrap();
15928 assert_eq!((n1, n2), (1, 2));
15929 assert_eq!((t1, t2), (1, 2));
15931 let (gj, nj, tj) = InflightGuard::try_acquire(
15933 counts.clone(),
15934 lanes::Lane::Judge,
15935 tenants.clone(),
15936 "blue",
15937 None,
15938 )
15939 .unwrap();
15940 assert_eq!((nj, tj), (1, 1));
15941 drop(g1);
15942 drop(gj);
15943 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
15944 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
15945 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
15946 assert!(tenants.lock().unwrap().get("blue").is_none());
15948 drop(g2);
15949 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
15950 assert!(tenants.lock().unwrap().is_empty());
15951 }
15952
15953 #[test]
15954 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
15955 let counts: InflightCounts = Arc::new(Default::default());
15956 let tenants: TenantGauge = Arc::new(Default::default());
15957 let start = Arc::new(std::sync::Barrier::new(3));
15958 let attempted = Arc::new(std::sync::Barrier::new(3));
15959 let mut joins = Vec::new();
15960 for _ in 0..2 {
15961 let counts = counts.clone();
15962 let tenants = tenants.clone();
15963 let start = start.clone();
15964 let attempted = attempted.clone();
15965 joins.push(std::thread::spawn(move || {
15966 start.wait();
15967 let result = InflightGuard::try_acquire(
15968 counts,
15969 lanes::Lane::Interactive,
15970 tenants,
15971 "preview_001",
15972 Some(1),
15973 );
15974 let won = result.is_ok();
15975 attempted.wait(); drop(result);
15977 won
15978 }));
15979 }
15980 start.wait();
15981 attempted.wait();
15982 let wins = joins
15983 .into_iter()
15984 .map(|join| join.join().unwrap())
15985 .filter(|won| *won)
15986 .count();
15987 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
15988 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
15989 assert!(tenants.lock().unwrap().is_empty());
15990 }
15991
15992 #[tokio::test]
15993 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
15994 let st = fake_worker_state();
15995 let tenant = auth::TenantCtx {
15996 tenant: "preview_001".into(),
15997 lane_class: auth::LaneClass::Interactive,
15998 rate_limit: Some(1),
15999 };
16000 let first_env = Envelope::new(true);
16001 let (guard, first_rl) =
16002 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
16003 Ok(slot) => slot,
16004 Err(_) => panic!("the first request must acquire the tenant slot"),
16005 };
16006 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
16007
16008 let second_env = Envelope::new(true);
16009 let response =
16010 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
16011 Err(response) => response,
16012 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
16013 };
16014 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
16015 assert_eq!(response.headers()["retry-after"], "2");
16016 assert_eq!(response.headers()["retry-after-ms"], "2000");
16017 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
16018 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
16019 assert_eq!(response.headers()["x-request-id"], second_env.id);
16020 assert_eq!(
16021 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16022 1,
16023 "rejected request must not consume a lane slot"
16024 );
16025 assert_eq!(
16026 st.tenant_inflight
16027 .lock()
16028 .unwrap()
16029 .get("preview_001")
16030 .copied(),
16031 Some(1),
16032 "rejected request must not increment the tenant gauge"
16033 );
16034 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16035 .await
16036 .unwrap();
16037 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16038 assert_eq!(payload["error"]["type"], "rate_limit_error");
16039 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
16040 assert!(
16041 payload["error"]["message"]
16042 .as_str()
16043 .unwrap()
16044 .contains("concurrent request limit")
16045 );
16046
16047 drop(guard);
16048 let _ = InflightGuard::try_acquire(
16049 st.inflight.clone(),
16050 lanes::Lane::Interactive,
16051 st.tenant_inflight.clone(),
16052 "preview_001",
16053 Some(1),
16054 )
16055 .expect("slot must reopen after the in-flight request completes");
16056 }
16057
16058 #[test]
16059 fn tenant_rate_limit_override_is_min_with_global_cap() {
16060 let metrics = SharedMetrics::default();
16061 let unlimited = auth::TenantCtx::default_tenant();
16062 let capped = auth::TenantCtx {
16063 tenant: "acme".into(),
16064 lane_class: auth::LaneClass::Interactive,
16065 rate_limit: Some(2),
16066 };
16067 let global = lane_cap(lanes::Lane::Interactive);
16068 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
16070 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16071 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
16073 assert_eq!((rl.limit, rl.remaining), (2, 1));
16074 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
16075 assert_eq!(rl.remaining, 0);
16076 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
16077 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
16081 assert_eq!(rl.remaining, 0);
16082 let wide = auth::TenantCtx {
16083 rate_limit: Some(global + 100),
16084 ..capped.clone()
16085 };
16086 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
16087 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16088 }
16089
16090 #[test]
16091 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
16092 let batch = auth::TenantCtx {
16093 tenant: "bulk".into(),
16094 lane_class: auth::LaneClass::Batch,
16095 rate_limit: None,
16096 };
16097 let interactive = auth::TenantCtx::default_tenant();
16098 let hdr = |v: Option<&str>| {
16099 let mut h = axum::http::HeaderMap::new();
16100 if let Some(v) = v {
16101 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16102 }
16103 h
16104 };
16105 assert_eq!(
16107 lane_for_tenant(&hdr(None), &interactive).unwrap(),
16108 lanes::Lane::Interactive
16109 );
16110 assert_eq!(
16111 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
16112 lanes::Lane::Judge
16113 );
16114 assert_eq!(
16116 lane_for_tenant(&hdr(None), &batch).unwrap(),
16117 lanes::Lane::Harvest
16118 );
16119 assert_eq!(
16120 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
16121 lanes::Lane::Judge
16122 );
16123 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
16124 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16125 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
16127 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16128 }
16129
16130 #[tokio::test]
16131 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
16132 let hdr = |v: &str| {
16137 let mut h = axum::http::HeaderMap::new();
16138 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16139 h
16140 };
16141 let body = |resp: Response| async move {
16142 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16143 .await
16144 .unwrap();
16145 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
16146 };
16147
16148 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
16149 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16150 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16151 let payload = body(resp).await;
16152 assert!(
16153 payload["error"].is_object(),
16154 "bare-string error body: {payload}"
16155 );
16156 assert_eq!(payload["error"]["type"], "invalid_request_error");
16157 assert_eq!(payload["error"]["param"], "x-lane");
16158 assert_eq!(payload["error"]["code"], "invalid_lane");
16159
16160 let batch = auth::TenantCtx {
16161 tenant: "bulk".into(),
16162 lane_class: auth::LaneClass::Batch,
16163 rate_limit: None,
16164 };
16165 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
16166 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16167 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16168 let payload = body(resp).await;
16169 assert_eq!(payload["error"]["type"], "authentication_error");
16170 assert_eq!(payload["error"]["param"], "x-lane");
16171 }
16172
16173 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
16176
16177 #[tokio::test]
16178 async fn responses_carry_rate_limit_headers_and_slot_frees() {
16179 let _l = DRAIN_LOCK.lock().unwrap();
16180 let st = fake_worker_state();
16181 let resp = chat_completions(
16184 State(st.clone()),
16185 axum::http::HeaderMap::new(),
16186 None,
16187 Json(
16188 serde_json::from_value(serde_json::json!({
16189 "model": "m", "messages": [{"role": "user", "content": "t"}]
16190 }))
16191 .unwrap(),
16192 ),
16193 )
16194 .await;
16195 assert_eq!(resp.status(), StatusCode::OK);
16196 let h = resp.headers();
16197 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
16198 let remaining: usize = h["x-ratelimit-remaining"]
16199 .to_str()
16200 .unwrap()
16201 .parse()
16202 .unwrap();
16203 assert_eq!(remaining, limit - 1);
16204 assert_eq!(h["x-ratelimit-reset"], "0");
16205 assert_eq!(
16206 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16207 0,
16208 "slot must free at completion"
16209 );
16210 let resp = completions(
16213 State(st.clone()),
16214 axum::http::HeaderMap::new(),
16215 None,
16216 Json(
16217 serde_json::from_value(serde_json::json!({
16218 "model": "m", "prompt": "t", "stream": true
16219 }))
16220 .unwrap(),
16221 ),
16222 )
16223 .await;
16224 assert_eq!(resp.status(), StatusCode::OK);
16225 assert!(resp.headers().contains_key("x-ratelimit-limit"));
16226 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
16227 assert!(resp.headers().contains_key("x-ratelimit-reset"));
16228 assert_eq!(
16229 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16230 1,
16231 "stream in flight holds the slot"
16232 );
16233 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
16234 .await
16235 .unwrap();
16236 assert_eq!(
16237 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16238 0,
16239 "slot must free when the stream completes"
16240 );
16241 }
16242
16243 #[tokio::test]
16244 async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
16245 let _l = DRAIN_LOCK.lock().unwrap();
16246 let mut st = fake_worker_state();
16247 let mock = MockMetering::admit_all();
16248 st.metering = Some(mock.clone());
16249
16250 let nonstream = chat_completions(
16251 State(st.clone()),
16252 HeaderMap::new(),
16253 None,
16254 Json(
16255 serde_json::from_value(json!({
16256 "model": "m",
16257 "messages": [{"role": "user", "content": "t"}],
16258 }))
16259 .unwrap(),
16260 ),
16261 )
16262 .await;
16263 assert_eq!(nonstream.status(), StatusCode::OK);
16264 let nonstream_id = nonstream.headers()["x-request-id"]
16265 .to_str()
16266 .unwrap()
16267 .to_string();
16268
16269 let stream = completions(
16270 State(st),
16271 HeaderMap::new(),
16272 None,
16273 Json(
16274 serde_json::from_value(json!({
16275 "model": "m",
16276 "prompt": "t",
16277 "stream": true,
16278 }))
16279 .unwrap(),
16280 ),
16281 )
16282 .await;
16283 assert_eq!(stream.status(), StatusCode::OK);
16284 let stream_id = stream.headers()["x-request-id"]
16285 .to_str()
16286 .unwrap()
16287 .to_string();
16288 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
16289 .await
16290 .unwrap();
16291
16292 let events = mock.events();
16296 let opened: Vec<&str> = events
16297 .iter()
16298 .filter_map(|e| match e {
16299 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
16300 _ => None,
16301 })
16302 .collect();
16303 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
16304 let completes = events
16305 .iter()
16306 .filter(|e| {
16307 matches!(
16308 e,
16309 MeterEvent::Complete {
16310 prompt: 1,
16311 cached: 0,
16312 completion: 1,
16313 }
16314 )
16315 })
16316 .count();
16317 assert_eq!(
16318 completes, 2,
16319 "both surfaces settle complete with worker-truth usage: {events:?}"
16320 );
16321 }
16322
16323 #[tokio::test]
16324 async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
16325 let _l = DRAIN_LOCK.lock().unwrap();
16326 let mock = MockMetering::with_limits(vec![
16332 ReserveScript::Insufficient,
16333 ReserveScript::Admit { with_permit: false },
16334 ReserveScript::Blocked,
16335 ReserveScript::Admit { with_permit: true },
16336 ]);
16337 let mut st = fake_worker_state();
16338 st.metering = Some(mock.clone());
16339
16340 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
16342 assert_eq!(metrics.status(), StatusCode::OK);
16343 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
16344 .await
16345 .unwrap();
16346 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
16347 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
16348 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
16349 assert_eq!(metrics_body["budget_source_available"], true);
16350
16351 let request = || {
16352 Json(
16353 serde_json::from_value::<CompletionReq>(json!({
16354 "model": "m",
16355 "prompt_ids": [1],
16356 "max_tokens": 1,
16357 }))
16358 .unwrap(),
16359 )
16360 };
16361
16362 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16363 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
16364 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
16365 .await
16366 .unwrap();
16367 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
16368 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
16369 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
16370
16371 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16372 assert_eq!(included.status(), StatusCode::OK);
16373
16374 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16377 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
16378
16379 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16380 assert_eq!(admitted.status(), StatusCode::OK);
16381
16382 let events = mock.events();
16383 let terminal: Vec<&MeterEvent> = events
16384 .iter()
16385 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
16386 .collect();
16387 assert_eq!(
16388 terminal.len(),
16389 4,
16390 "four requests, four terminal settles: {events:?}"
16391 );
16392 assert!(matches!(
16393 terminal[0],
16394 MeterEvent::Reject { status: 402, .. }
16395 ));
16396 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
16397 assert!(matches!(
16398 terminal[2],
16399 MeterEvent::Reject { status: 402, .. }
16400 ));
16401 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
16402 let permits: Vec<bool> = events
16404 .iter()
16405 .filter_map(|e| match e {
16406 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
16407 _ => None,
16408 })
16409 .collect();
16410 assert_eq!(
16411 permits,
16412 vec![false, false, false, true],
16413 "the permit rides the receipt exactly when reserve minted one: {events:?}"
16414 );
16415 }
16416
16417 #[tokio::test]
16418 async fn streaming_client_disconnect_records_partial_usage_and_cost() {
16419 let _l = DRAIN_LOCK.lock().unwrap();
16420 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
16421 let mock = MockMetering::admit_all();
16422 st.metering = Some(mock.clone());
16423
16424 let response = completions(
16425 State(st),
16426 HeaderMap::new(),
16427 None,
16428 Json(
16429 serde_json::from_value(json!({
16430 "model": "m",
16431 "prompt": "disconnect after one delta",
16432 "stream": true,
16433 }))
16434 .unwrap(),
16435 ),
16436 )
16437 .await;
16438 assert_eq!(response.status(), StatusCode::OK);
16439 let request_id = response.headers()["x-request-id"]
16440 .to_str()
16441 .unwrap()
16442 .to_string();
16443 let mut body = Box::pin(response.into_body().into_data_stream());
16444 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
16445 .await
16446 .expect("stream ended before first delta")
16447 .expect("stream body failed");
16448 assert!(
16449 is_sse_data_frame(&first),
16450 "first frame was not SSE data: {first:?}"
16451 );
16452 drop(body);
16453
16454 let mut dropped = None;
16457 for _ in 0..500 {
16458 if let Some(event) = mock
16459 .events()
16460 .into_iter()
16461 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
16462 {
16463 dropped = Some(event);
16464 break;
16465 }
16466 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
16467 }
16468 let events = mock.events();
16469 assert!(
16470 events
16471 .iter()
16472 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
16473 "the receipt was opened under the caller-visible request id: {events:?}"
16474 );
16475 assert_eq!(
16476 dropped,
16477 Some(MeterEvent::Dropped {
16478 prompt: 1,
16479 cached: 0,
16480 completion: 1,
16481 }),
16482 "a client disconnect must leave the partial counts on the dropped receipt \
16483 (the implementation prices that drop): {events:?}"
16484 );
16485 }
16486
16487 #[tokio::test]
16488 async fn draining_rejects_new_requests_with_503_and_retry_after() {
16489 let _l = DRAIN_LOCK.lock().unwrap();
16490 let st = fake_worker_state();
16491 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
16492 let resp = chat_completions(
16494 State(st.clone()),
16495 axum::http::HeaderMap::new(),
16496 None,
16497 Json(
16498 serde_json::from_value(serde_json::json!({
16499 "model": "m", "messages": [{"role": "user", "content": "t"}]
16500 }))
16501 .unwrap(),
16502 ),
16503 )
16504 .await;
16505 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16506 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
16511 let ra_s: u64 = ra
16512 .parse()
16513 .expect("Retry-After must be integer delay-seconds");
16514 assert!(
16515 ra_s > 0 && ra_s <= 60,
16516 "Retry-After {ra_s}s is outside the honored window"
16517 );
16518 let ra_ms: u64 = resp.headers()["retry-after-ms"]
16519 .to_str()
16520 .unwrap()
16521 .parse()
16522 .unwrap();
16523 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
16524 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16525 .await
16526 .unwrap();
16527 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16528 assert!(
16529 payload["error"]["message"]
16530 .as_str()
16531 .unwrap()
16532 .contains("draining")
16533 );
16534 assert_eq!(payload["error"]["type"], "server_error");
16535 assert_eq!(payload["error"]["code"], "draining");
16536 let resp = completions(
16537 State(st.clone()),
16538 axum::http::HeaderMap::new(),
16539 None,
16540 Json(
16541 serde_json::from_value(serde_json::json!({
16542 "model": "m", "prompt": "t"
16543 }))
16544 .unwrap(),
16545 ),
16546 )
16547 .await;
16548 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16549 assert!(resp.headers().contains_key("retry-after"));
16550 assert_eq!(
16551 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16552 0,
16553 "rejected requests must not hold slots"
16554 );
16555 let resp = health_live(State(st.clone())).await.into_response();
16558 assert_eq!(
16559 resp.status(),
16560 StatusCode::OK,
16561 "a drain must not look like a liveness fault"
16562 );
16563 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16564 .await
16565 .unwrap();
16566 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16567 assert_eq!(payload["status"], "draining");
16568 let resp = health_ready(State(st.clone())).await.into_response();
16570 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16571 let retry_s = drain_deadline_s().clamp(1, 60);
16572 let retry_s_text = retry_s.to_string();
16573 let retry_ms_text = (retry_s * 1000).to_string();
16574 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
16575 assert_eq!(
16576 resp.headers().get("retry-after-ms").unwrap(),
16577 retry_ms_text.as_str()
16578 );
16579 assert_ne!(
16580 resp.headers()
16581 .get("x-should-retry")
16582 .and_then(|v| v.to_str().ok()),
16583 Some("false")
16584 );
16585 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16586 .await
16587 .unwrap();
16588 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16589 assert_eq!(payload["status"], "not_ready");
16590 assert!(payload["detail"].as_str().unwrap().contains("draining"));
16591 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16592 let resp = chat_completions(
16594 State(st.clone()),
16595 axum::http::HeaderMap::new(),
16596 None,
16597 Json(
16598 serde_json::from_value(serde_json::json!({
16599 "model": "m", "messages": [{"role": "user", "content": "t"}]
16600 }))
16601 .unwrap(),
16602 ),
16603 )
16604 .await;
16605 assert_eq!(resp.status(), StatusCode::OK);
16606 }
16607
16608 #[tokio::test]
16611 async fn health_is_green_only_while_the_worker_is_alive() {
16612 let _l = DRAIN_LOCK.lock().unwrap();
16615 let st = fake_worker_state();
16616 let resp = health_live(State(st.clone())).await.into_response();
16619 assert_eq!(resp.status(), StatusCode::OK);
16620 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16621 .await
16622 .unwrap();
16623 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16624 assert_eq!(payload["status"], "ok");
16625 assert_eq!(payload["worker"]["phase"], "idle");
16626 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
16627 let ready = health_ready(State(st.clone())).await.into_response();
16628 assert_eq!(ready.status(), StatusCode::OK);
16629
16630 st.health.mark_dead("worker thread panicked: test-injected");
16634 let resp = health_live(State(st.clone())).await.into_response();
16635 assert_eq!(
16636 resp.status(),
16637 StatusCode::SERVICE_UNAVAILABLE,
16638 "a dead worker MUST NOT report a healthy liveness"
16639 );
16640 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16641 .await
16642 .unwrap();
16643 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16644 assert_eq!(payload["status"], "unhealthy");
16645 assert!(
16647 payload["detail"]
16648 .as_str()
16649 .unwrap()
16650 .contains("test-injected"),
16651 "cause not surfaced: {payload}"
16652 );
16653 let ready = health_ready(State(st.clone())).await.into_response();
16654 assert_eq!(
16655 ready.status(),
16656 StatusCode::SERVICE_UNAVAILABLE,
16657 "dead is also not ready"
16658 );
16659
16660 st.health.mark_ready();
16663 assert_eq!(
16664 health_live(State(st.clone()))
16665 .await
16666 .into_response()
16667 .status(),
16668 StatusCode::OK,
16669 "mark_ready must clear the latch (a successful respawn)"
16670 );
16671 }
16672
16673 #[tokio::test]
16674 async fn readyz_peer_probe_integrity_is_present_and_advisory() {
16675 let _l = DRAIN_LOCK.lock().unwrap();
16676 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16677 let st = fake_worker_state();
16678
16679 let ready = health_ready(State(st.clone())).await.into_response();
16680 assert_eq!(ready.status(), StatusCode::OK);
16681 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
16682 .await
16683 .unwrap();
16684 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16685 assert_eq!(payload["peer_probe_integrity"], "ok");
16686
16687 st.health.note_peer_probe_deferral(2, false);
16688 let deferred = health_ready(State(st.clone())).await.into_response();
16689 assert_eq!(deferred.status(), StatusCode::OK);
16690 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
16691 .await
16692 .unwrap();
16693 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16694 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
16695
16696 st.health.note_peer_probe_deferral(4, true);
16697 let degraded = health_ready(State(st.clone())).await.into_response();
16698 assert_eq!(
16699 degraded.status(),
16700 StatusCode::OK,
16701 "peer degradation is advisory while plain serving remains healthy"
16702 );
16703 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
16704 .await
16705 .unwrap();
16706 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16707 assert_eq!(payload["peer_probe_integrity"], "degraded");
16708
16709 st.health.mark_dead("test-injected worker failure");
16710 let unready = health_ready(State(st)).await.into_response();
16711 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
16712 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
16713 .await
16714 .unwrap();
16715 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16716 assert_eq!(
16717 payload["peer_probe_integrity"], "degraded",
16718 "the advisory field must also survive an unrelated readiness failure"
16719 );
16720 }
16721
16722 #[tokio::test]
16723 async fn liveness_failure_obeys_the_retry_contract() {
16724 let _l = DRAIN_LOCK.lock().unwrap();
16728 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16729 let st = fake_worker_state();
16730 st.health
16731 .mark_dead("worker thread panicked: retry-contract-test");
16732
16733 let resp = health_live(State(st)).await.into_response();
16734 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16735 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16736 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
16737 assert_ne!(
16738 resp.headers()
16739 .get("x-should-retry")
16740 .and_then(|v| v.to_str().ok()),
16741 Some("false")
16742 );
16743 }
16744
16745 #[tokio::test]
16746 async fn readiness_failure_obeys_the_retry_contract() {
16747 let _l = DRAIN_LOCK.lock().unwrap();
16748 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16749 let st = fake_worker_state();
16750 st.health
16751 .mark_dead("worker thread panicked: retry-contract-test");
16752
16753 let resp = health_ready(State(st)).await.into_response();
16754 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16755 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16756 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
16757 assert_ne!(
16758 resp.headers()
16759 .get("x-should-retry")
16760 .and_then(|v| v.to_str().ok()),
16761 Some("false")
16762 );
16763 }
16764
16765 #[tokio::test]
16766 async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
16767 let _l = DRAIN_LOCK.lock().unwrap();
16777 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16778 let st = fake_worker_state();
16779 assert_eq!(
16780 health_live(State(st.clone()))
16781 .await
16782 .into_response()
16783 .status(),
16784 StatusCode::OK
16785 );
16786 st.health
16787 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
16788 let resp = health_live(State(st.clone())).await.into_response();
16789 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16790 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16791 .await
16792 .unwrap();
16793 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16794 assert!(
16795 payload["detail"]
16796 .as_str()
16797 .unwrap()
16798 .contains("probe exceeded")
16799 );
16800 st.health.mark_ready();
16803 assert_eq!(
16804 health_live(State(st.clone()))
16805 .await
16806 .into_response()
16807 .status(),
16808 StatusCode::SERVICE_UNAVAILABLE,
16809 "a GPU fault must not be cleared by an in-process respawn"
16810 );
16811 }
16812
16813 #[test]
16814 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
16815 let caps = ModelCaps {
16817 tools_branch: true,
16818 qwen_think: true,
16819 think_switch: true,
16820 chat_ok: true,
16821 context_length: 262144,
16822 tokenizer: "qwen2".into(),
16823 instruct_type: Some("chatml".into()),
16824 effort_levels: false,
16825 qwen_effort: false,
16826 gemma_think: false,
16827 dsv4: false,
16828 chat_temperature_default: None,
16829 chat_top_p_default: None,
16830 n_vocab: 151_936,
16831 };
16832 let e = model_entry_v1("main", Some(&caps), None);
16833 assert_eq!(e["id"], "main");
16834 assert_eq!(e["name"], "main");
16835 assert_eq!(e["object"], "model");
16836 assert_eq!(e["context_length"], 262144);
16837 assert!(e["pricing"]["input"].is_null());
16839 assert!(e["pricing"]["output"].is_null());
16840
16841 let meta = OpenRouterModelMetadata {
16844 pricing: OpenRouterPricing {
16845 prompt: Some("0.00000038".into()),
16846 cached_prompt: Some("0.0000002".into()),
16847 completion: Some("0.0000026".into()),
16848 ..Default::default()
16849 },
16850 input_modalities: vec!["image".into(), "video".into()],
16851 max_output_length: Some(32768),
16852 ..Default::default()
16853 };
16854 let e = model_entry_v1("main", Some(&caps), Some(&meta));
16855 assert_eq!(e["pricing"]["currency"], "USD");
16858 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
16859 assert_eq!(e["pricing"]["input"], "0.38");
16860 assert_eq!(e["pricing"]["output"], "2.60");
16861 assert_eq!(e["pricing"]["cached_input"], "0.20");
16862 assert!(e["pricing"]["cache_write"].is_null());
16863 assert_eq!(e["pricing"]["minimum_request"], "0");
16864 assert_eq!(e["owned_by"], "main");
16865 assert_eq!(e["type"], "chat");
16866 assert_eq!(e["max_output_tokens"], 32768);
16867 assert_eq!(e["endpoints"], json!(["chat/completions"]));
16868 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
16869 assert_eq!(e["output_modalities"], json!(["text"]));
16870 assert_eq!(e["capabilities"]["streaming"], true);
16871 assert_eq!(e["capabilities"]["tools"], true);
16872 assert_eq!(e["lifecycle"]["status"], "active");
16873 assert!(e["lifecycle"]["deprecation_at"].is_null());
16874 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
16875 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
16876 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
16880 keys.sort_unstable();
16881 assert_eq!(
16882 keys,
16883 [
16884 "capabilities",
16885 "context_length",
16886 "endpoints",
16887 "id",
16888 "input_modalities",
16889 "lifecycle",
16890 "max_output_tokens",
16891 "name",
16892 "object",
16893 "output_modalities",
16894 "owned_by",
16895 "pricing",
16896 "reliability",
16897 "type",
16898 ],
16899 "unexpected /v1/models entry keys"
16900 );
16901 let mut price_keys: Vec<&str> = e["pricing"]
16902 .as_object()
16903 .unwrap()
16904 .keys()
16905 .map(String::as_str)
16906 .collect();
16907 price_keys.sort_unstable();
16908 assert_eq!(
16909 price_keys,
16910 [
16911 "cache_write",
16912 "cached_input",
16913 "currency",
16914 "input",
16915 "minimum_request",
16916 "output",
16917 "unit",
16918 ],
16919 "unexpected /v1/models pricing keys"
16920 );
16921
16922 let e = model_entry_v1("m", None, None);
16924 assert!(e["context_length"].is_null());
16925 assert!(e["max_output_tokens"].is_null());
16926 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
16928 assert!(e["context_length"].is_null());
16929 }
16930
16931 #[test]
16937 fn catalog_row_follows_the_declared_surface() {
16938 let caps = ModelCaps {
16939 tools_branch: true,
16940 ..Default::default()
16941 };
16942
16943 let embed = OpenRouterModelMetadata {
16944 surface: Some("embedding".into()),
16945 max_output_length: Some(1),
16946 ..Default::default()
16947 };
16948 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
16949 assert_eq!(e["type"], "embedding");
16950 assert_eq!(e["endpoints"], json!(["embeddings"]));
16951 assert_eq!(e["output_modalities"], json!(["embeddings"]));
16952 assert_eq!(e["capabilities"]["streaming"], false);
16953 assert_eq!(
16954 e["capabilities"]["tools"], false,
16955 "an embedder has no tools"
16956 );
16957 assert_eq!(e["capabilities"]["reasoning"], false);
16958 assert_eq!(e["capabilities"]["structured_output"], false);
16959 assert_eq!(e["capabilities"]["prompt_caching"], false);
16960 assert!(
16961 e["max_output_tokens"].is_null(),
16962 "a surface that emits no completion tokens must not advertise a ceiling"
16963 );
16964
16965 let rerank = OpenRouterModelMetadata {
16966 surface: Some("rerank".into()),
16967 ..Default::default()
16968 };
16969 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
16970 assert_eq!(r["type"], "rerank");
16971 assert_eq!(r["endpoints"], json!(["rerank"]));
16972 assert_eq!(r["output_modalities"], json!(["rerank"]));
16973 assert_eq!(r["capabilities"]["tools"], false);
16974 assert_eq!(r["capabilities"]["reasoning"], false);
16975
16976 let chat = OpenRouterModelMetadata {
16979 max_output_length: Some(32768),
16980 ..Default::default()
16981 };
16982 let c = model_entry_v1("main", Some(&caps), Some(&chat));
16983 assert_eq!(c["type"], "chat");
16984 assert_eq!(c["endpoints"], json!(["chat/completions"]));
16985 assert_eq!(c["output_modalities"], json!(["text"]));
16986 assert_eq!(c["capabilities"]["tools"], true);
16987 assert_eq!(c["max_output_tokens"], 32768);
16988 }
16989
16990 #[test]
16993 fn unknown_surface_is_rejected_at_config_load() {
16994 let bad = OpenRouterModelMetadata {
16995 surface: Some("embeddings".into()), ..Default::default()
16997 };
16998 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
16999 .expect_err("an unknown surface must not load");
17000 assert!(err.contains("surface"), "{err}");
17001
17002 for good in ["chat", "embedding", "rerank"] {
17003 let ok = OpenRouterModelMetadata {
17004 surface: Some(good.into()),
17005 ..Default::default()
17006 };
17007 assert!(
17008 validate_openrouter_metadata("m", &ok).is_ok(),
17009 "{good} must load"
17010 );
17011 }
17012 }
17013
17014 #[test]
17015 fn per_million_price_is_exact_decimal_shift() {
17016 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
17018 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
17019 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
17020 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
17021 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
17022 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
17023 assert_eq!(per_million_price("not-a-price"), None);
17024 assert_eq!(per_million_price(""), None);
17025 }
17026
17027 #[test]
17028 fn metadata_provider_block_parses_and_validates() {
17029 let (_, provider) = OpenRouterMetadataFile::parse(
17030 r#"
17031 [provider]
17032 id = "tiyuvta"
17033 status_url = "https://status.tiyuvta.ai"
17034 support_contact = "mailto:support@tiyuvta.ai"
17035 incident_contact = "mailto:incidents@tiyuvta.ai"
17036 regions = ["eu-central"]
17037 "#,
17038 )
17039 .unwrap();
17040 let provider = provider.unwrap();
17041 assert_eq!(provider.id, "tiyuvta");
17042 assert_eq!(provider.regions, vec!["eu-central"]);
17043 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
17045 assert!(err.contains("provider.id"), "{err}");
17046 let err = OpenRouterMetadataFile::parse(
17048 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
17049 )
17050 .unwrap_err();
17051 assert!(err.contains("must be a URI"), "{err}");
17052 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
17054 assert!(provider.is_none());
17055 }
17056
17057 #[test]
17058 fn models_openai_default_body_stays_byte_identical() {
17059 let body = models_openai_body(&["main".into(), "judge".into()]);
17060 let bytes = serde_json::to_vec(&body).unwrap();
17061 assert_eq!(
17062 bytes,
17063 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
17064 );
17065 }
17066
17067 #[test]
17068 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
17069 let loaded = vec![
17071 "qwen/qwen3.6-27b".to_string(),
17072 "qwen/qwen3.6-35b-a3b".to_string(),
17073 ];
17074 assert_eq!(
17075 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
17076 Some("qwen/qwen3.6-35b-a3b"),
17077 );
17078 assert_eq!(
17079 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
17080 Some("qwen/qwen3.6-27b"),
17081 );
17082 assert_eq!(
17084 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
17085 Some("qwen/qwen3.6-35b-a3b"),
17086 );
17087 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
17089 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
17090 assert_eq!(canonical_model_id(&loaded, ""), None);
17091 }
17092
17093 #[test]
17094 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
17095 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
17098 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
17099 assert_eq!(
17101 canonical_model_id(&loaded, "a/shared-name").as_deref(),
17102 Some("a/shared-name")
17103 );
17104 assert_eq!(
17105 canonical_model_id(&loaded, "b/shared-name").as_deref(),
17106 Some("b/shared-name")
17107 );
17108 let bare = vec!["solo".to_string()];
17110 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
17111 }
17112
17113 #[test]
17114 fn openrouter_models_entry_serializes_complete_metadata() {
17115 let metadata = OpenRouterMetadataFile::from_toml(
17116 r#"
17117[models.main]
17118hugging_face_id = "Qwen/Qwen3.6-27B"
17119created = 1786032000
17120quantization = "nvfp4"
17121description = "Qwen3.6 27B served by memra."
17122max_prompt_length = 245760
17123max_output_length = 16384
17124default_output_length = 4096
17125is_ready = true
17126is_free = false
17127discount_to_user = 0.1
17128openrouter_slug = "qwen/qwen3.6-27b"
17129datacenters = [{ country_code = "US", region = "us-east-1" }]
17130zdr = true
17131hipaa = false
17132
17133[models.main.pricing]
17134prompt = "0.000000234"
17135cached_prompt = "0.0000000585"
17136cache_write = "0.000000234"
17137completion = "0.000001872"
17138internal_reasoning = "0.000001872"
17139request = "0.01"
17140
17141[models.main.capacity]
17142prompt_tpm = 1000000
17143cached_prompt_tpm = 2000000
17144completion_tpm = 500000
17145request_rpm = 1000
17146concurrency = 64
17147"#,
17148 )
17149 .unwrap();
17150 let caps = ModelCaps {
17151 tools_branch: true,
17152 qwen_think: true,
17153 think_switch: true,
17154 chat_ok: true,
17155 context_length: 262144,
17156 tokenizer: "qwen2".into(),
17157 instruct_type: Some("chatml".into()),
17158 ..Default::default()
17159 };
17160 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
17161
17162 assert_eq!(entry["schema_version"], "2.4");
17163 assert_eq!(entry["id"], "main");
17164 assert_eq!(entry["name"], "main");
17165 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
17166 assert_eq!(entry["created"], 1786032000u64);
17167 assert_eq!(entry["quantization"], "nvfp4");
17168 assert_eq!(entry["tokenizer"], "qwen2");
17169 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
17170 assert!(
17171 entry.get("object").is_none(),
17172 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
17173 );
17174
17175 let input = &entry["input_modalities"][0];
17176 assert_eq!(input["type"], "text");
17177 assert_eq!(
17178 input["supported_inputs"]["max_context_length"]["value"],
17179 262144
17180 );
17181 assert_eq!(
17182 input["supported_inputs"]["max_prompt_length"]["value"],
17183 245760
17184 );
17185 let input_prices = input["pricing"].as_array().unwrap();
17186 let input_price = |kind: &str| {
17187 input_prices
17188 .iter()
17189 .find(|price| price["type"] == kind)
17190 .unwrap()
17191 };
17192 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
17193 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
17194 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
17195 assert_eq!(input["capacity"][0]["value"], 1000000);
17196 assert_eq!(input["capacity"][1]["value"], 2000000);
17197
17198 let output = &entry["output_modalities"][0];
17199 assert_eq!(output["type"], "text");
17200 assert_eq!(output["max_length"]["value"], 16384);
17201 assert_eq!(output["streaming"], true);
17202 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
17203 assert_eq!(
17204 output["supported_parameters"]["structured_outputs"]["type"],
17205 "boolean"
17206 );
17207 assert_eq!(
17208 output["supported_parameters"]["reasoning"]["type"],
17209 "boolean"
17210 );
17211 assert_eq!(output["pricing"][0]["type"], "completion");
17212 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
17213 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
17214 assert_eq!(output["capacity"][0]["value"], 500000);
17215 assert_eq!(output["capacity"][1]["type"], "concurrency");
17216 assert_eq!(output["capacity"][1]["value"], 64);
17217
17218 assert_eq!(entry["pricing"][0]["type"], "request");
17219 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
17220 assert_eq!(entry["capacity"][0]["value"], 1000);
17221 assert_eq!(entry["is_ready"], true);
17222 assert_eq!(entry["is_free"], false);
17223 assert_eq!(entry["discount_to_user"], 0.1);
17224 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
17225 assert_eq!(entry["datacenters"][0]["country_code"], "US");
17226 assert_eq!(entry["compliance"]["zdr"], true);
17227 assert_eq!(entry["compliance"]["hipaa"], false);
17228 }
17229
17230 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
17235[models."qwen/qwen3.6-35b-a3b"]
17236hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
17237created = 1777260255
17238quantization = "int4"
17239description = "Qwen3.6 35B-A3B fixture entry."
17240max_prompt_length = 262144
17241max_output_length = 262144
17242default_output_length = 8192
17243is_ready = true
17244is_free = false
17245discount_to_user = 0.0
17246openrouter_slug = "qwen/qwen3.6-35b-a3b"
17247zdr = false
17248hipaa = false
17249
17250[[models."qwen/qwen3.6-35b-a3b".datacenters]]
17251country_code = "CA"
17252region = "Ontario"
17253
17254[models."qwen/qwen3.6-35b-a3b".pricing]
17255prompt = "0.0000000931"
17256cached_prompt = "0.0000000652"
17257completion = "0.0000009025"
17258
17259[models."qwen/qwen3.6-35b-a3b".capacity]
17260prompt_tpm = 780000
17261cached_prompt_tpm = 310000
17262completion_tpm = 9600
17263request_rpm = 160
17264concurrency = 16
17265
17266[planned_models."qwen/qwen3.8-27b"]
17267description = "Planned fixture entry; must never be emitted."
17268max_prompt_length = 262144
17269max_output_length = 262144
17270default_output_length = 8192
17271is_ready = false
17272is_free = false
17273discount_to_user = 0.0
17274openrouter_slug = "qwen/qwen3.8-27b"
17275zdr = false
17276hipaa = false
17277
17278[planned_models."qwen/qwen3.8-27b".pricing]
17279prompt = "0.0000002745"
17280cached_prompt = "0.0000001922"
17281completion = "0.0000022800"
17282
17283[planned_models."google/gemma-4-26b-a4b-it"]
17284hugging_face_id = "google/gemma-4-26B-A4B-it"
17285created = 1775227989
17286quantization = "int4"
17287description = "Planned fixture entry; must never be emitted."
17288max_prompt_length = 262144
17289max_output_length = 262144
17290default_output_length = 8192
17291is_ready = false
17292is_free = false
17293discount_to_user = 0.0
17294openrouter_slug = "google/gemma-4-26b-a4b-it"
17295zdr = false
17296hipaa = false
17297
17298[planned_models."google/gemma-4-26b-a4b-it".pricing]
17299prompt = "0.0000000665"
17300cached_prompt = "0.0000000466"
17301completion = "0.0000003230"
17302"#;
17303
17304 #[test]
17305 fn gateway_registry_generates_the_staged_active_shape() {
17306 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17307 let caps = ModelCaps {
17308 tools_branch: true,
17309 qwen_think: true,
17310 think_switch: true,
17311 chat_ok: true,
17312 context_length: 262144,
17313 tokenizer: "qwen2".into(),
17314 instruct_type: Some("chatml".into()),
17315 ..Default::default()
17316 };
17317 let q35_entry = model_entry_openrouter(
17318 "qwen/qwen3.6-35b-a3b",
17319 Some(&caps),
17320 metadata.get("qwen/qwen3.6-35b-a3b"),
17321 );
17322 assert_eq!(q35_entry["created"], 1777260255u64);
17323 assert_eq!(q35_entry["quantization"], "int4");
17324 assert_eq!(q35_entry["is_ready"], true);
17325 assert_eq!(
17326 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
17327 262144
17328 );
17329 assert_eq!(
17330 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
17331 262144
17332 );
17333 assert_eq!(
17334 q35_entry["output_modalities"][0]["max_length"]["value"],
17335 262144
17336 );
17337 let prices = q35_entry["input_modalities"][0]["pricing"]
17338 .as_array()
17339 .unwrap();
17340 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
17341 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
17342 assert_eq!(
17346 q35_entry["input_modalities"][0]["capacity"][0]["value"],
17347 780000
17348 );
17349 assert_eq!(
17350 q35_entry["input_modalities"][0]["capacity"][1]["value"],
17351 310000
17352 );
17353 assert_eq!(
17354 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
17355 262144
17356 );
17357 assert_eq!(
17358 q35_entry["output_modalities"][0]["capacity"][0]["value"],
17359 9600
17360 );
17361 assert_eq!(
17362 q35_entry["output_modalities"][0]["capacity"][1]["value"],
17363 16
17364 );
17365 assert_eq!(
17366 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
17367 "0.0000009025"
17368 );
17369 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
17371
17372 assert_eq!(
17373 metadata.len(),
17374 1,
17375 "planned models must never enter the active map"
17376 );
17377 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
17378 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
17379 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
17380
17381 let openmodels = model_entry_openmodels(
17382 "qwen/qwen3.6-35b-a3b",
17383 Some(&caps),
17384 metadata.get("qwen/qwen3.6-35b-a3b"),
17385 )
17386 .unwrap();
17387 assert_eq!(openmodels["currency"], "USD");
17388 assert_eq!(openmodels["max_output_length"], 262144);
17389 assert_eq!(openmodels["is_ready"], true);
17390 assert_eq!(openmodels["is_free"], false);
17391 assert_eq!(openmodels["discount_to_user"], 0.0);
17392 }
17393
17394 #[test]
17395 fn gateway_registry_limits_are_live_request_limits() {
17396 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17397 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
17398 let caps = ModelCaps {
17399 context_length: 262_144,
17400 ..Default::default()
17401 };
17402 let build = |value: serde_json::Value| {
17403 let req: CompletionReq = serde_json::from_value(value).unwrap();
17404 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
17405 build_request(&req, tx, lanes::Lane::Interactive, None)
17406 };
17407
17408 let mut omitted = build(json!({
17409 "model": "qwen/qwen3.6-35b-a3b",
17410 "prompt_ids": [1, 2, 3]
17411 }));
17412 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
17413 assert_eq!(omitted.params.max_new, 8_192);
17414 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
17415
17416 let mut field_top = build(json!({
17417 "model": "qwen/qwen3.6-35b-a3b",
17418 "prompt_ids": [1],
17419 "max_tokens": 262144
17420 }));
17421 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
17422 assert_eq!(field_top.params.max_new, 262_144);
17423 assert_eq!(
17424 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
17425 262_044,
17426 "the field-top output request is accepted but bounded by remaining trained context",
17427 );
17428
17429 let mut too_much_output = build(json!({
17430 "model": "qwen/qwen3.6-35b-a3b",
17431 "prompt_ids": [1],
17432 "max_tokens": 262145
17433 }));
17434 let (message, param) =
17435 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
17436 .unwrap_err();
17437 assert_eq!(param, "max_tokens");
17438 assert!(message.contains("262145"));
17439
17440 let mut oversized_allocation = build(json!({
17441 "model": "qwen/qwen3.6-35b-a3b",
17442 "prompt_ids": [1],
17443 "max_tokens": 1,
17444 "max_ctx": 262145
17445 }));
17446 let (_, param) =
17447 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
17448 .unwrap_err();
17449 assert_eq!(param, "max_ctx");
17450 }
17451
17452 #[test]
17453 fn planned_registry_entries_are_validated_but_never_activated() {
17454 let parsed = OpenRouterMetadataFile::from_toml(
17455 r#"
17456[planned_models.future]
17457max_output_length = 262144
17458default_output_length = 8192
17459
17460[planned_models.future.pricing]
17461prompt = "0.0000001"
17462"#,
17463 )
17464 .unwrap();
17465 assert!(parsed.is_empty());
17466
17467 let error = OpenRouterMetadataFile::from_toml(
17468 r#"
17469[planned_models.future]
17470default_output_length = 8192
17471"#,
17472 )
17473 .unwrap_err();
17474 assert!(error.contains("requires max_output_length"));
17475 }
17476
17477 #[test]
17482 fn every_catalog_feed_honours_the_declared_surface() {
17483 let metadata = OpenRouterMetadataFile::from_toml(
17484 r#"
17485[models."qwen/qwen3-embedding-8b"]
17486surface = "embedding"
17487created = 1787961600
17488max_output_length = 1
17489is_ready = true
17490is_free = false
17491discount_to_user = 0.0
17492
17493[models."qwen/qwen3-embedding-8b".pricing]
17494prompt = "0.00000001"
17495cached_prompt = "0.0"
17496completion = "0.0"
17497
17498[models."main"]
17499created = 1787443200
17500max_output_length = 32768
17501is_ready = true
17502is_free = false
17503discount_to_user = 0.0
17504
17505[models."main".pricing]
17506prompt = "0.00000025"
17507cached_prompt = "0.00000009"
17508completion = "0.0000012"
17509"#,
17510 )
17511 .unwrap();
17512 let caps = ModelCaps {
17513 tools_branch: true,
17514 qwen_think: true,
17515 chat_ok: true,
17516 context_length: 32768,
17517 ..Default::default()
17518 };
17519 let embed = metadata.get("qwen/qwen3-embedding-8b");
17520 let chat = metadata.get("main");
17521
17522 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
17524 let out = &or["output_modalities"][0];
17525 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
17526 assert!(
17527 out.get("streaming").is_none(),
17528 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
17529 );
17530 let params = &out["supported_parameters"];
17535 assert_eq!(
17536 params.as_object().map(|o| o.len()),
17537 Some(0),
17538 "no completion parameter belongs on an embedder row: {params}"
17539 );
17540 for field in [
17541 "tools",
17542 "tool_choice",
17543 "reasoning",
17544 "max_tokens",
17545 "json_mode",
17546 "structured_outputs",
17547 "stop",
17548 "temperature",
17549 "seed",
17550 ] {
17551 assert!(params[field].is_null(), "{field} leaked onto an embedder");
17552 }
17553 assert!(
17554 out["max_length"].is_null(),
17555 "a surface emitting no completion tokens advertises no ceiling: {out}"
17556 );
17557
17558 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
17560 .expect("openmodels entry builds");
17561 assert_eq!(om["output_modalities"], json!(["embeddings"]));
17562 let features = om["supported_features"].as_array().unwrap();
17563 assert!(
17564 !features
17565 .iter()
17566 .any(|f| f == "tool_calling" || f == "reasoning"),
17567 "chat-only features leaked onto an embedder: {features:?}"
17568 );
17569
17570 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
17572 assert_eq!(v1["type"], "embedding");
17573 assert_eq!(v1["capabilities"]["tools"], false);
17574
17575 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
17577 let out_chat = &or_chat["output_modalities"][0];
17578 assert_eq!(out_chat["type"], "text");
17579 assert_eq!(out_chat["streaming"], true);
17580 assert!(!out_chat["supported_parameters"]["tools"].is_null());
17581 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
17582 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
17583 assert_eq!(out_chat["max_length"]["value"], 32768u64);
17584 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
17585 assert_eq!(om_chat["output_modalities"], json!(["text"]));
17586 assert!(
17587 om_chat["supported_features"]
17588 .as_array()
17589 .unwrap()
17590 .iter()
17591 .any(|f| f == "tool_calling")
17592 );
17593 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
17594 }
17595
17596 #[test]
17603 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
17604 let raw = std::fs::read_to_string(concat!(
17605 env!("CARGO_MANIFEST_DIR"),
17606 "/../../research/gateway-20260812/raw/sources/",
17607 "openrouter-provider-schema-v2.4-20260812.json"
17608 ))
17609 .expect("vendored Provider Monitor 2.4 schema is in-tree");
17610 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
17611 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
17612 .as_array()
17613 .expect("OutputModality is a oneOf");
17614
17615 let metadata = OpenRouterMetadataFile::from_toml(
17616 r#"
17617[models."embed"]
17618surface = "embedding"
17619created = 1787961600
17620max_output_length = 1
17621is_ready = true
17622is_free = false
17623discount_to_user = 0.0
17624
17625[models."embed".pricing]
17626prompt = "0.00000001"
17627cached_prompt = "0.0"
17628completion = "0.0"
17629
17630[models."rr"]
17631surface = "rerank"
17632created = 1787961600
17633max_output_length = 1
17634is_ready = true
17635is_free = false
17636discount_to_user = 0.0
17637
17638[models."rr".pricing]
17639prompt = "0.00000003"
17640cached_prompt = "0.0"
17641completion = "0.0"
17642
17643[models."chatty"]
17644created = 1787443200
17645max_output_length = 32768
17646is_ready = true
17647is_free = false
17648discount_to_user = 0.0
17649
17650[models."chatty".pricing]
17651prompt = "0.00000025"
17652cached_prompt = "0.00000009"
17653completion = "0.0000012"
17654"#,
17655 )
17656 .unwrap();
17657 let caps = ModelCaps {
17658 tools_branch: true,
17659 qwen_think: true,
17660 chat_ok: true,
17661 context_length: 32768,
17662 ..Default::default()
17663 };
17664
17665 for (alias, want_type) in [
17666 ("embed", "embeddings"),
17667 ("rr", "rerank"),
17668 ("chatty", "text"),
17669 ] {
17670 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
17671 let modality = &row["output_modalities"][0];
17672 assert_eq!(modality["type"], want_type, "{alias}: {row}");
17673
17674 let branch = branches
17676 .iter()
17677 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
17678 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
17679 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
17680 .as_object()
17681 .expect("branch properties")
17682 .keys()
17683 .map(String::as_str)
17684 .collect();
17685 for key in modality.as_object().expect("modality object").keys() {
17686 assert!(
17687 allowed.contains(key.as_str()),
17688 "{alias}: {key:?} is not a property of the {want_type:?} branch \
17689 (additionalProperties:false); allowed = {allowed:?}"
17690 );
17691 }
17692 for req in branch["required"].as_array().into_iter().flatten() {
17693 let req = req.as_str().expect("required entry is a string");
17694 assert!(
17695 modality.get(req).is_some(),
17696 "{alias}: required property {req:?} missing from the {want_type:?} branch"
17697 );
17698 }
17699 }
17700 }
17701
17702 #[test]
17703 fn openrouter_models_entry_omits_undeclared_optional_fields() {
17704 let entry = model_entry_openrouter("minimal", None, None);
17705 let object = entry.as_object().unwrap();
17706 for field in [
17707 "hugging_face_id",
17708 "created",
17709 "quantization",
17710 "tokenizer",
17711 "description",
17712 "pricing",
17713 "capacity",
17714 "is_ready",
17715 "is_free",
17716 "discount_to_user",
17717 "openrouter",
17718 "datacenters",
17719 "compliance",
17720 ] {
17721 assert!(
17722 !object.contains_key(field),
17723 "optional field {field} must be absent, not null"
17724 );
17725 }
17726 assert_eq!(entry["schema_version"], "2.4");
17727 assert_eq!(entry["input_modalities"][0]["type"], "text");
17728 assert!(
17729 entry["input_modalities"][0]
17730 .get("supported_inputs")
17731 .is_none()
17732 );
17733 assert!(entry["input_modalities"][0].get("pricing").is_none());
17734 assert!(entry["input_modalities"][0].get("capacity").is_none());
17735 assert_eq!(entry["output_modalities"][0]["type"], "text");
17736 assert_eq!(entry["output_modalities"][0]["streaming"], true);
17737 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
17738 assert!(entry["output_modalities"][0].get("max_length").is_none());
17739 assert!(entry["output_modalities"][0].get("pricing").is_none());
17740 assert!(entry["output_modalities"][0].get("capacity").is_none());
17741 }
17742
17743 #[test]
17744 fn openmodels_entry_serializes_standard_provider_shape() {
17745 let metadata = OpenRouterMetadataFile::from_toml(
17746 r#"
17747[models."qwen/qwen3.6-27b"]
17748created = 1786032000
17749max_output_length = 16384
17750is_ready = true
17751is_free = false
17752discount_to_user = 0.05
17753
17754[models."qwen/qwen3.6-27b".pricing]
17755prompt = "0.000000291"
17756cached_prompt = "0.000000291"
17757completion = "0.000002763"
17758request = "0"
17759"#,
17760 )
17761 .unwrap();
17762 let caps = ModelCaps {
17763 tools_branch: true,
17764 qwen_think: true,
17765 chat_ok: true,
17766 context_length: 262144,
17767 ..Default::default()
17768 };
17769 let entry = model_entry_openmodels(
17770 "qwen/qwen3.6-27b",
17771 Some(&caps),
17772 metadata.get("qwen/qwen3.6-27b"),
17773 )
17774 .unwrap();
17775
17776 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
17777 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
17778 assert_eq!(entry["created"], 1786032000u64);
17779 assert_eq!(entry["input_modalities"], json!(["text"]));
17780 assert_eq!(entry["output_modalities"], json!(["text"]));
17781 assert_eq!(entry["context_length"], 262144u64);
17782 assert_eq!(entry["max_output_length"], 16384u64);
17783 assert_eq!(entry["currency"], "USD");
17784 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
17785 assert_eq!(entry["pricing"]["completion"], "0.000002763");
17786 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
17787 assert_eq!(entry["pricing"]["request"], "0");
17788 assert_eq!(
17789 entry["supported_features"],
17790 json!(["tool_calling", "reasoning"])
17791 );
17792 assert_eq!(entry["is_ready"], true);
17793 assert_eq!(entry["is_free"], false);
17794 assert_eq!(entry["discount_to_user"], 0.05);
17795 assert!(entry.get("schema_version").is_none());
17796 assert!(entry.get("quantization").is_none());
17797 }
17798
17799 #[test]
17800 fn openmodels_entry_rejects_missing_operator_metadata() {
17801 let caps = ModelCaps {
17802 context_length: 262144,
17803 ..Default::default()
17804 };
17805 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
17806 assert_eq!(
17807 error,
17808 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
17809 );
17810 }
17811
17812 #[tokio::test]
17813 async fn blocking_response_excludes_stop_text_across_token_events() {
17814 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
17815 tx.send(Event::Token {
17816 id: 1,
17817 text: "answer\nPro".into(),
17818 })
17819 .unwrap();
17820 tx.send(Event::Token {
17821 id: 2,
17822 text: "blem: leaked prompt".into(),
17823 })
17824 .unwrap();
17825 tx.send(Event::Done {
17826 stop_reason: "Callback".into(),
17827 n_tokens: 2,
17828 n_prompt: 8,
17829 n_cached: 0,
17830 elapsed_s: 0.5,
17831 spec: None,
17832 })
17833 .unwrap();
17834 drop(tx);
17835 let response = blocking_response(
17836 rx,
17837 "plain_quant".into(),
17838 false,
17839 vec!["Problem:".into()],
17840 None,
17841 Envelope::new(false),
17842 )
17843 .await;
17844 assert_eq!(response.status(), StatusCode::OK);
17845 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
17846 .await
17847 .unwrap();
17848 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17849 assert_eq!(payload["text"], "answer\n");
17850 assert_eq!(payload["stop_reason"], "Callback");
17851 }
17852}