1pub mod auth;
59pub(crate) mod constrained;
60pub(crate) mod darklane;
64pub(crate) mod health;
68pub(crate) mod lanes {
69 pub use memra_lanes::*;
70}
71mod admit_predict;
76mod affinity;
83mod anthropic;
89#[allow(dead_code)] mod build_id;
94mod dsv4_serve;
95mod embed_api;
96pub mod metering;
103mod prefill_receipt;
104mod responses_api;
105mod surfaces;
106mod toolcall;
107mod ttft;
108mod worker;
109
110use std::collections::HashMap;
111use std::net::{SocketAddr, ToSocketAddrs};
112use std::sync::mpsc::Sender;
113use std::sync::{Arc, Mutex, RwLock};
114
115use axum::{
116 Extension, Json, Router,
117 body::Body,
118 extract::{DefaultBodyLimit, FromRequest, Path, Query, Request as AxumRequest, State},
119 http::{
120 HeaderMap, StatusCode,
121 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
122 },
123 middleware::{self, Next},
124 response::{
125 IntoResponse, Response,
126 sse::{Event as SseEvent, Sse},
127 },
128 routing::{get, post},
129};
130use futures_core::Stream as _;
131use serde::de::DeserializeOwned;
132use serde::{Deserialize, Serialize};
133use serde_json::json;
134use sha2::{Digest as _, Sha256};
135use tower::ServiceExt as _;
136
137use memra_engine::decode::GenParams;
138use memra_engine::sampler::SamplerConfig;
139use memra_tokenizer::{
140 Tokenizer,
141 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
142};
143use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
144use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
145
146const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
170const MAX_BODY_ADMISSIONS: usize = 4;
171const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
172#[allow(clippy::identity_op)] const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
177const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
178const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
179const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
180const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
181const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
182const MAX_STOP_SEQUENCES: usize = 16;
183const MAX_STOP_SEQUENCE_BYTES: usize = 1_024;
184const MAX_STOP_SEQUENCES_BYTES: usize = 4 * 1_024;
185const MAX_CLIENT_IDENTIFIER_BYTES: usize = 256;
186const MAX_HTTP_CONNECTIONS: usize = 1_024;
187const MAX_HTTP2_STREAMS_PER_CONNECTION: u32 = 128;
188const HTTP1_HEADER_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
189const HTTP_CONNECTION_MAX_LIFETIME: std::time::Duration = std::time::Duration::from_secs(300);
190
191fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
192 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
193 SEMAPHORE
194 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
195 .clone()
196}
197
198fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
199 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
200 SEMAPHORE
201 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
202 .clone()
203}
204
205#[derive(Clone)]
206pub(crate) struct BodyAdmissionGuard {
207 permit: Arc<Mutex<Option<tokio::sync::OwnedSemaphorePermit>>>,
208}
209
210impl BodyAdmissionGuard {
211 fn new(permit: tokio::sync::OwnedSemaphorePermit) -> Self {
212 Self {
213 permit: Arc::new(Mutex::new(Some(permit))),
214 }
215 }
216
217 pub(crate) fn release(&self) {
218 if let Ok(mut permit) = self.permit.lock() {
219 permit.take();
220 }
221 }
222}
223
224pub(crate) struct BodyAdmissionLease(Option<BodyAdmissionGuard>);
225
226impl BodyAdmissionLease {
227 fn release(&mut self) {
228 if let Some(admission) = self.0.take() {
229 admission.release();
230 }
231 }
232
233 pub(crate) fn guard(&self) -> Option<&BodyAdmissionGuard> {
234 self.0.as_ref()
235 }
236}
237
238impl Drop for BodyAdmissionLease {
239 fn drop(&mut self) {
240 self.release();
241 }
242}
243
244pub(crate) struct AdmittedJson<T>(pub(crate) T, pub(crate) BodyAdmissionLease);
245
246#[axum::async_trait]
247impl<S, T> FromRequest<S> for AdmittedJson<T>
248where
249 S: Send + Sync,
250 T: DeserializeOwned,
251{
252 type Rejection = axum::extract::rejection::JsonRejection;
253
254 async fn from_request(req: AxumRequest, state: &S) -> Result<Self, Self::Rejection> {
255 let admission = req.extensions().get::<BodyAdmissionGuard>().cloned();
256 let parsed = Json::<T>::from_request(req, state).await;
257 parsed.map(|Json(value)| Self(value, BodyAdmissionLease(admission)))
258 }
259}
260
261fn declared_body_length(req: &AxumRequest) -> Option<usize> {
262 req.headers()
263 .get(CONTENT_LENGTH)
264 .and_then(|value| value.to_str().ok())
265 .and_then(|value| value.parse().ok())
266}
267
268fn body_requires_admission(req: &AxumRequest) -> bool {
269 if req.headers().contains_key(TRANSFER_ENCODING) {
273 return true;
274 }
275 declared_body_length(req).is_none_or(|length| length > BODY_ADMISSION_BYPASS_BYTES)
276}
277
278fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
282 let Some(length) = declared_body_length(req) else {
283 return BODY_READ_TIMEOUT;
284 };
285 let bytes = length as u64;
286 let extra_seconds =
287 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
288 let seconds = BODY_READ_TIMEOUT
289 .as_secs()
290 .saturating_add(extra_seconds)
291 .min(BODY_READ_TIMEOUT_MAX.as_secs());
292 std::time::Duration::from_secs(seconds)
293}
294
295async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
300 let resp = next.run(req).await;
301 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
302 return resp;
303 }
304 error_response_coded(
305 StatusCode::PAYLOAD_TOO_LARGE,
306 &format!(
307 "request body exceeds the {} MiB limit",
308 MAX_BODY_BYTES / (1024 * 1024)
309 ),
310 "invalid_request_error",
311 None,
312 Some("request_too_large"),
313 )
314}
315
316fn apply_body_limit(app: Router) -> Router {
319 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
320 .layer(middleware::from_fn(shape_payload_too_large))
321}
322
323fn protected_inference_path(path: &str) -> bool {
324 matches!(
325 path,
326 "/v1/auth/check"
327 | "/v1/completions"
328 | "/v1/chat/completions"
329 | "/v1/messages"
330 | "/v1/responses"
331 | "/v1/embeddings"
332 | "/v1/rerank"
333 )
334}
335
336async fn shape_inference_early_response(path: &str, response: Response) -> Response {
340 let request_id = Envelope::new(path != "/v1/completions");
341 if path == "/v1/messages" {
342 anthropic::with_anthropic_request_id(
343 &request_id.id,
344 anthropic::reshape_error(response, &request_id.id).await,
345 )
346 } else {
347 with_request_id(&request_id.id, response)
348 }
349}
350
351async fn authenticate_inference_before_body(
356 State(st): State<AppState>,
357 mut req: AxumRequest,
358 next: Next,
359) -> Response {
360 if !protected_inference_path(req.uri().path()) {
361 return next.run(req).await;
362 }
363 let path = req.uri().path().to_string();
364 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
368 return shape_inference_early_response(
369 &path,
370 error_response_coded(
371 StatusCode::PAYLOAD_TOO_LARGE,
372 &format!(
373 "request body exceeds the {} MiB limit",
374 MAX_BODY_BYTES / (1024 * 1024)
375 ),
376 "invalid_request_error",
377 None,
378 Some("request_too_large"),
379 ),
380 )
381 .await;
382 }
383 let headers = req.headers();
384 let bearer = bearer_token(headers);
385 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
386 let api_key = headers
387 .get("x-api-key")
388 .and_then(|value| value.to_str().ok());
389 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
390 } else {
391 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
392 };
393 if let Err(why) = auth {
394 return shape_inference_early_response(&path, authentication_error(why)).await;
395 }
396 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
404 let body_admission = if body_requires_admission(&req) {
405 body_admission_semaphore()
406 } else {
407 small_body_admission_semaphore()
408 };
409 let body_permit = match body_admission.try_acquire_owned() {
410 Ok(permit) => permit,
411 Err(tokio::sync::TryAcquireError::Closed) => {
412 let response = retry_contract_response(
413 error_response_coded(
414 StatusCode::SERVICE_UNAVAILABLE,
415 "request body admission is unavailable",
416 "server_error",
417 None,
418 Some("body_admission_unavailable"),
419 ),
420 Some(BODY_ADMISSION_RETRY_AFTER_S),
421 );
422 return shape_inference_early_response(&path, response).await;
423 }
424 Err(tokio::sync::TryAcquireError::NoPermits) => {
425 let response = retry_contract_response(
426 error_response_coded(
427 StatusCode::TOO_MANY_REQUESTS,
428 "request body admission is busy",
429 "rate_limit_error",
430 None,
431 Some("body_admission_busy"),
432 ),
433 Some(BODY_ADMISSION_RETRY_AFTER_S),
434 );
435 return shape_inference_early_response(&path, response).await;
436 }
437 };
438 let body_admission_guard = BodyAdmissionGuard::new(body_permit);
443 req.extensions_mut().insert(body_admission_guard.clone());
444 let body = std::mem::replace(req.body_mut(), Body::empty());
445 let mut body = Box::pin(body.into_data_stream());
446 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
447 let body_timed_out_flag = body_timed_out.clone();
448 let guarded_body = async_stream::stream! {
449 loop {
450 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
451 if remaining.is_zero() {
452 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
453 yield Err(std::io::Error::new(
454 std::io::ErrorKind::TimedOut,
455 "request body read deadline exceeded",
456 ));
457 break;
458 }
459 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
460 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
461 Ok(frame) => frame,
462 Err(_) => {
463 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
464 yield Err(std::io::Error::new(
465 std::io::ErrorKind::TimedOut,
466 "request body idle timeout exceeded",
467 ));
468 break;
469 }
470 };
471 match frame {
472 Some(Ok(bytes)) => yield Ok(bytes),
473 Some(Err(error)) => {
474 yield Err(std::io::Error::other(error.to_string()));
475 break;
476 }
477 None => break,
478 }
479 }
480 };
481 *req.body_mut() = Body::from_stream(guarded_body);
482 let response = next.run(req).await;
483 body_admission_guard.release();
484 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
485 let request_id = Envelope::new(path != "/v1/completions");
486 let timeout = error_response_coded(
487 StatusCode::REQUEST_TIMEOUT,
488 "request body read timed out",
489 "invalid_request_error",
490 None,
491 Some("request_body_timeout"),
492 );
493 return if path == "/v1/messages" {
494 anthropic::with_anthropic_request_id(
495 &request_id.id,
496 anthropic::reshape_error(timeout, &request_id.id).await,
497 )
498 } else {
499 with_request_id(&request_id.id, timeout)
500 };
501 }
502 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
503 let request_id = Envelope::new(true);
504 return anthropic::with_anthropic_request_id(
505 &request_id.id,
506 anthropic::reshape_error(response, &request_id.id).await,
507 );
508 }
509 response
510}
511
512#[cfg(test)]
513mod body_limit_tests {
514 use super::*;
515
516 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
517
518 fn test_app() -> Router {
522 let app = Router::new()
523 .route(
524 "/bytes",
525 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
526 )
527 .route(
528 "/json",
529 post(
530 |AdmittedJson(v, _admission): AdmittedJson<serde_json::Value>| async move {
531 v["pad"].as_str().unwrap_or("").len().to_string()
532 },
533 ),
534 );
535 apply_body_limit(app)
536 }
537
538 fn streamed_body(chunks: usize) -> Body {
539 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
542 Body::from_stream(async_stream::stream! {
543 for _ in 0..chunks {
544 yield Ok::<_, std::io::Error>(chunk.clone());
545 }
546 })
547 }
548
549 #[tokio::test]
550 async fn bodies_past_the_old_2mib_default_are_accepted() {
551 for (path, body) in [
554 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
555 (
556 "/json",
557 Body::from(
558 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
559 ),
560 ),
561 ] {
562 let resp = test_app()
563 .oneshot(
564 axum::http::Request::post(path)
565 .header(CONTENT_TYPE, "application/json")
566 .body(body)
567 .unwrap(),
568 )
569 .await
570 .unwrap();
571 assert_eq!(resp.status(), StatusCode::OK, "{path}");
572 }
573 }
574
575 #[tokio::test]
576 async fn body_at_exactly_the_limit_is_accepted() {
577 let resp = test_app()
578 .oneshot(
579 axum::http::Request::post("/bytes")
580 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
581 .unwrap(),
582 )
583 .await
584 .unwrap();
585 assert_eq!(resp.status(), StatusCode::OK);
586 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
587 .await
588 .unwrap();
589 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
590 }
591
592 #[tokio::test]
593 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
594 for path in ["/bytes", "/json"] {
598 let resp = test_app()
599 .oneshot(
600 axum::http::Request::post(path)
601 .header(CONTENT_TYPE, "application/json")
602 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
603 .unwrap(),
604 )
605 .await
606 .unwrap();
607 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
608 assert_eq!(
609 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
610 Some(b"false".as_ref()),
611 "{path}: retrying identical bytes cannot fix a 413"
612 );
613 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
614 .await
615 .unwrap();
616 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
617 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
618 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
619 assert!(
620 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
621 "{path}: message names the limit"
622 );
623 }
624 }
625
626 #[tokio::test]
627 async fn authenticated_body_admission_is_finite() {
628 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
629 let semaphore = body_admission_semaphore();
630 let mut permits = Vec::new();
631 for _ in 0..MAX_BODY_ADMISSIONS {
632 permits.push(semaphore.clone().acquire_owned().await.unwrap());
633 }
634 assert!(
635 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
636 .await
637 .is_err(),
638 "body parser admission must not be unbounded"
639 );
640 drop(permits);
641 assert!(semaphore.acquire().await.is_ok());
642 }
643
644 #[tokio::test]
645 async fn small_body_admission_is_finite_and_separate() {
646 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
647 let large = body_admission_semaphore();
648 let small = small_body_admission_semaphore();
649 let mut small_permits = Vec::new();
650 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
651 small_permits.push(small.clone().acquire_owned().await.unwrap());
652 }
653 assert!(
654 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
655 .await
656 .is_err(),
657 "small body parser admission must be bounded"
658 );
659 assert!(
660 large.clone().try_acquire().is_ok(),
661 "small uploads must not consume large-upload permits"
662 );
663 drop(small_permits);
664 assert!(small.acquire().await.is_ok());
665 }
666
667 #[test]
668 fn small_declared_bodies_bypass_large_upload_admission() {
669 let request = axum::http::Request::post("/v1/chat/completions")
670 .header(CONTENT_LENGTH, "2048")
671 .body(Body::empty())
672 .unwrap();
673 assert!(!body_requires_admission(&request));
674
675 let request = axum::http::Request::post("/v1/chat/completions")
676 .header(
677 CONTENT_LENGTH,
678 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
679 )
680 .body(Body::empty())
681 .unwrap();
682 assert!(body_requires_admission(&request));
683
684 let request = axum::http::Request::post("/v1/chat/completions")
685 .header(CONTENT_LENGTH, "2048")
686 .header(TRANSFER_ENCODING, "chunked")
687 .body(Body::empty())
688 .unwrap();
689 assert!(body_requires_admission(&request));
690 }
691
692 #[test]
693 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
694 let unknown = axum::http::Request::post("/v1/chat/completions")
695 .body(Body::empty())
696 .unwrap();
697 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
698
699 let large = axum::http::Request::post("/v1/chat/completions")
700 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
701 .body(Body::empty())
702 .unwrap();
703 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
704 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
705
706 let absurd = axum::http::Request::post("/v1/chat/completions")
707 .header(CONTENT_LENGTH, u64::MAX.to_string())
708 .body(Body::empty())
709 .unwrap();
710 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
711 }
712
713 #[tokio::test]
714 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
715 let too_large = shape_inference_early_response(
716 "/v1/messages",
717 error_response_coded(
718 StatusCode::PAYLOAD_TOO_LARGE,
719 "request body exceeds the 192 MiB limit",
720 "invalid_request_error",
721 None,
722 Some("request_too_large"),
723 ),
724 )
725 .await;
726 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
727 let house_id = too_large.headers()["x-request-id"].clone();
728 assert_eq!(too_large.headers()["request-id"], house_id);
729 assert_eq!(too_large.headers()["x-should-retry"], "false");
730 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
731 .await
732 .unwrap();
733 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
734 assert_eq!(payload["type"], "error");
735 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
736
737 let busy = shape_inference_early_response(
738 "/v1/chat/completions",
739 retry_contract_response(
740 error_response_coded(
741 StatusCode::TOO_MANY_REQUESTS,
742 "request body admission is busy",
743 "rate_limit_error",
744 None,
745 Some("body_admission_busy"),
746 ),
747 Some(BODY_ADMISSION_RETRY_AFTER_S),
748 ),
749 )
750 .await;
751 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
752 assert!(!busy.headers()["x-request-id"].is_empty());
753 assert_eq!(busy.headers()["retry-after"], "1");
754 assert_eq!(busy.headers()["retry-after-ms"], "1000");
755 assert!(busy.headers().get("x-should-retry").is_none());
756 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
757 .await
758 .unwrap();
759 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
760 assert_eq!(payload["error"]["code"], "body_admission_busy");
761 }
762
763 #[tokio::test]
764 async fn vision_preprocess_admission_is_fail_fast_and_retryable() {
765 let semaphore = Box::leak(Box::new(tokio::sync::Semaphore::new(1)));
766 let held = semaphore.try_acquire().unwrap();
767 let busy = try_vision_preprocess_with(true, semaphore).unwrap_err();
768 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
769 assert_eq!(busy.headers()["retry-after"], "1");
770 drop(held);
771 assert!(
772 try_vision_preprocess_with(true, semaphore)
773 .unwrap()
774 .is_some()
775 );
776 assert!(
777 try_vision_preprocess_with(false, semaphore)
778 .unwrap()
779 .is_none()
780 );
781 }
782
783 #[tokio::test]
784 async fn typed_json_retains_body_admission_until_handler_validation_releases_it() {
785 #[derive(Clone)]
786 struct Signals {
787 parsed: Arc<tokio::sync::Notify>,
788 finish: Arc<tokio::sync::Notify>,
789 semaphore: Arc<tokio::sync::Semaphore>,
790 }
791
792 let semaphore = Arc::new(tokio::sync::Semaphore::new(1));
793 let guard = BodyAdmissionGuard::new(semaphore.clone().try_acquire_owned().unwrap());
794 let signals = Signals {
795 parsed: Arc::new(tokio::sync::Notify::new()),
796 finish: Arc::new(tokio::sync::Notify::new()),
797 semaphore: semaphore.clone(),
798 };
799 let app = Router::new()
800 .route(
801 "/",
802 post(
803 |Extension(signals): Extension<Signals>,
804 AdmittedJson(_, mut admission): AdmittedJson<serde_json::Value>| async move {
805 assert_eq!(
806 signals.semaphore.available_permits(),
807 0,
808 "typed deserialization alone must not release post-parse admission"
809 );
810 admission.release();
811 assert_eq!(signals.semaphore.available_permits(), 1);
812 signals.parsed.notify_one();
813 signals.finish.notified().await;
814 "ok"
815 },
816 ),
817 )
818 .layer(Extension(signals.clone()))
819 .layer(Extension(guard));
820 let response = tokio::spawn(
821 app.oneshot(
822 axum::http::Request::post("/")
823 .header(CONTENT_TYPE, "application/json")
824 .body(Body::from(r#"{"value":1}"#))
825 .unwrap(),
826 ),
827 );
828 signals.parsed.notified().await;
829 assert_eq!(
830 semaphore.available_permits(),
831 1,
832 "validated work must release admission before generation waits"
833 );
834 signals.finish.notify_one();
835 assert_eq!(response.await.unwrap().unwrap().status(), StatusCode::OK);
836 }
837
838 #[tokio::test]
839 async fn transport_closes_stalled_headers_and_caps_connections() {
840 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
841
842 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
843 let address = listener.local_addr().unwrap();
844 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
845 let server = tokio::spawn(serve_bounded_http_with_limits(
846 listener,
847 Router::new().route("/", get(|| async { "ok" })).route(
848 "/slow",
849 get(|| async {
850 tokio::time::sleep(std::time::Duration::from_millis(140)).await;
851 "slow-ok"
852 }),
853 ),
854 async move {
855 let _ = shutdown_rx.await;
856 },
857 std::time::Duration::from_millis(30),
858 1,
859 std::time::Duration::from_millis(80),
860 ));
861
862 let mut stalled = tokio::net::TcpStream::connect(address).await.unwrap();
863 stalled.write_all(b"GET / HT").await.unwrap();
864 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
865 let mut excess = tokio::net::TcpStream::connect(address).await.unwrap();
866 let mut bytes = Vec::new();
867 tokio::time::timeout(
868 std::time::Duration::from_millis(250),
869 excess.read_to_end(&mut bytes),
870 )
871 .await
872 .expect("connection beyond the cap must be closed promptly")
873 .unwrap();
874
875 bytes.clear();
876 tokio::time::timeout(
877 std::time::Duration::from_millis(500),
878 stalled.read_to_end(&mut bytes),
879 )
880 .await
881 .expect("stalled request headers must hit the configured deadline")
882 .unwrap();
883
884 let mut idle = tokio::net::TcpStream::connect(address).await.unwrap();
885 idle.write_all(b"GET / HTTP/1.1\r\nHost: local\r\n\r\n")
886 .await
887 .unwrap();
888 bytes.clear();
889 tokio::time::timeout(
890 std::time::Duration::from_millis(500),
891 idle.read_to_end(&mut bytes),
892 )
893 .await
894 .expect("an idle keep-alive connection must hit the maximum lifetime")
895 .unwrap();
896 assert!(String::from_utf8_lossy(&bytes).contains("200 OK"));
897
898 let mut active = tokio::net::TcpStream::connect(address).await.unwrap();
899 active
900 .write_all(b"GET /slow HTTP/1.1\r\nHost: local\r\n\r\n")
901 .await
902 .unwrap();
903 bytes.clear();
904 tokio::time::timeout(
905 std::time::Duration::from_millis(500),
906 active.read_to_end(&mut bytes),
907 )
908 .await
909 .expect("an active response must finish across the connection age boundary")
910 .unwrap();
911 let active_response = String::from_utf8_lossy(&bytes);
912 assert!(active_response.contains("200 OK"), "{active_response}");
913 assert!(active_response.contains("slow-ok"), "{active_response}");
914
915 let h2_stream = tokio::net::TcpStream::connect(address).await.unwrap();
919 let (mut h2_client, h2_connection) = h2::client::handshake(h2_stream).await.unwrap();
920 let h2_driver = tokio::spawn(h2_connection);
921 let request = axum::http::Request::builder()
922 .uri(format!("http://{address}/"))
923 .body(())
924 .unwrap();
925 let (response, _) = h2_client.send_request(request, true).unwrap();
926 let response = tokio::time::timeout(std::time::Duration::from_millis(500), response)
927 .await
928 .expect("HTTP/2 handshake and response must complete")
929 .expect("HTTP/2 connection must stay alive through the response");
930 assert_eq!(response.status(), StatusCode::OK);
931 drop(h2_client);
932 h2_driver.abort();
933 let _ = h2_driver.await;
934
935 let _ = shutdown_tx.send(());
936 server.await.unwrap().unwrap();
937 }
938}
939
940#[derive(Clone, Default)]
941struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
942
943fn is_sse_data_frame(bytes: &[u8]) -> bool {
944 bytes
945 .windows(b"data:".len())
946 .any(|window| window == b"data:")
947}
948
949async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
950 let trace = ttft::start(req.uri().path());
951 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
952 let response = next.run(req).await;
953 let Some(trace) = trace else {
954 return response;
955 };
956 let is_sse = response
957 .headers()
958 .get(CONTENT_TYPE)
959 .and_then(|value| value.to_str().ok())
960 .is_some_and(|value| value.starts_with("text/event-stream"));
961 if !is_sse {
962 return response;
963 }
964
965 let (parts, body) = response.into_parts();
968 let mut body = Box::pin(body.into_data_stream());
969 let stream = async_stream::stream! {
970 while let Some(frame) =
971 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
972 {
973 if frame
974 .as_ref()
975 .is_ok_and(|bytes| is_sse_data_frame(bytes))
976 {
977 trace.mark_first_sse_byte();
978 }
979 yield frame;
980 }
981 };
982 Response::from_parts(parts, Body::from_stream(stream))
983}
984
985const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
986const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
987
988#[derive(Debug, Clone, Default, Deserialize)]
989#[serde(deny_unknown_fields)]
990struct OpenRouterMetadataFile {
991 #[serde(default)]
992 models: HashMap<String, OpenRouterModelMetadata>,
993 #[serde(default)]
996 planned_models: HashMap<String, OpenRouterModelMetadata>,
997 #[serde(default)]
1000 provider: Option<ProviderMetadata>,
1001}
1002
1003#[derive(Debug, Clone, Deserialize)]
1007#[serde(deny_unknown_fields)]
1008struct ProviderMetadata {
1009 id: String,
1010 #[serde(default)]
1011 status_url: Option<String>,
1012 #[serde(default)]
1013 support_contact: Option<String>,
1014 #[serde(default)]
1015 incident_contact: Option<String>,
1016 #[serde(default)]
1017 regions: Vec<String>,
1018}
1019
1020#[derive(Debug, Clone, Default, Deserialize)]
1022#[serde(deny_unknown_fields)]
1023struct LifecycleMetadata {
1024 #[serde(default)]
1025 status: Option<String>,
1026 #[serde(default)]
1027 deprecation_at: Option<String>,
1028 #[serde(default)]
1029 retirement_at: Option<String>,
1030 #[serde(default)]
1031 replacement_model_id: Option<String>,
1032}
1033
1034#[derive(Debug, Clone, Default, Deserialize)]
1036#[serde(deny_unknown_fields)]
1037struct ReliabilityMetadata {
1038 #[serde(default)]
1039 first_token_timeout_seconds: Option<u64>,
1040 #[serde(default)]
1041 completion_timeout_seconds: Option<u64>,
1042 #[serde(default)]
1043 stream_idle_timeout_seconds: Option<u64>,
1044 #[serde(default)]
1045 capacity_scope: Option<String>,
1046}
1047
1048#[derive(Debug, Clone, Default, Deserialize)]
1049#[serde(deny_unknown_fields)]
1050struct OpenRouterModelMetadata {
1051 #[serde(default)]
1053 owned_by: Option<String>,
1054 #[serde(default)]
1055 lifecycle: Option<LifecycleMetadata>,
1056 #[serde(default)]
1057 reliability: Option<ReliabilityMetadata>,
1058 #[serde(default)]
1059 hugging_face_id: Option<String>,
1060 #[serde(default)]
1061 created: Option<u64>,
1062 #[serde(default)]
1063 quantization: Option<String>,
1064 #[serde(default)]
1065 description: Option<String>,
1066 #[serde(default)]
1067 max_prompt_length: Option<u64>,
1068 #[serde(default)]
1069 max_output_length: Option<u64>,
1070 #[serde(default)]
1073 default_output_length: Option<u64>,
1074 #[serde(default)]
1075 pricing: OpenRouterPricing,
1076 #[serde(default)]
1077 capacity: OpenRouterCapacity,
1078 #[serde(default)]
1079 is_ready: Option<bool>,
1080 #[serde(default)]
1081 is_free: Option<bool>,
1082 #[serde(default)]
1083 discount_to_user: Option<f64>,
1084 #[serde(default)]
1085 openrouter_slug: Option<String>,
1086 #[serde(default)]
1087 datacenters: Vec<OpenRouterDatacenter>,
1088 #[serde(default)]
1092 input_modalities: Vec<String>,
1093 #[serde(default)]
1108 surface: Option<String>,
1109 #[serde(default)]
1110 zdr: Option<bool>,
1111 #[serde(default)]
1112 hipaa: Option<bool>,
1113 #[serde(default)]
1123 default_reasoning_effort: Option<String>,
1124 #[serde(default)]
1142 default_temperature: Option<f32>,
1143 #[serde(default)]
1144 default_top_p: Option<f32>,
1145 #[serde(default)]
1147 default_top_k: Option<usize>,
1148 #[serde(default)]
1149 default_min_p: Option<f32>,
1150 #[serde(default)]
1151 default_presence_penalty: Option<f32>,
1152 #[serde(default)]
1153 default_frequency_penalty: Option<f32>,
1154 #[serde(default)]
1156 default_repetition_penalty: Option<f32>,
1157 #[serde(default)]
1178 non_thinking_sampling: Option<SamplingArmMetadata>,
1179}
1180
1181#[derive(Debug, Clone, Default, Deserialize)]
1187#[serde(deny_unknown_fields)]
1188struct SamplingArmMetadata {
1189 #[serde(default)]
1190 temperature: Option<f32>,
1191 #[serde(default)]
1192 top_p: Option<f32>,
1193 #[serde(default)]
1194 top_k: Option<usize>,
1195 #[serde(default)]
1196 min_p: Option<f32>,
1197 #[serde(default)]
1198 presence_penalty: Option<f32>,
1199 #[serde(default)]
1200 frequency_penalty: Option<f32>,
1201 #[serde(default)]
1202 repetition_penalty: Option<f32>,
1203}
1204
1205impl SamplingArmMetadata {
1206 fn is_empty(&self) -> bool {
1207 self.temperature.is_none()
1208 && self.top_p.is_none()
1209 && self.top_k.is_none()
1210 && self.min_p.is_none()
1211 && self.presence_penalty.is_none()
1212 && self.frequency_penalty.is_none()
1213 && self.repetition_penalty.is_none()
1214 }
1215}
1216
1217#[derive(Debug, Clone, Default, Deserialize)]
1218#[serde(deny_unknown_fields)]
1219struct OpenRouterPricing {
1220 #[serde(default)]
1221 prompt: Option<String>,
1222 #[serde(default)]
1223 cached_prompt: Option<String>,
1224 #[serde(default)]
1225 cache_write: Option<String>,
1226 #[serde(default)]
1227 completion: Option<String>,
1228 #[serde(default)]
1229 internal_reasoning: Option<String>,
1230 #[serde(default)]
1231 request: Option<String>,
1232}
1233
1234#[derive(Debug, Clone, Default, Deserialize)]
1235#[serde(deny_unknown_fields)]
1236struct OpenRouterCapacity {
1237 #[serde(default)]
1238 prompt_tpm: Option<u64>,
1239 #[serde(default)]
1240 cached_prompt_tpm: Option<u64>,
1241 #[serde(default)]
1242 completion_tpm: Option<u64>,
1243 #[serde(default)]
1244 request_rpm: Option<u64>,
1245 #[serde(default)]
1246 concurrency: Option<u64>,
1247}
1248
1249#[derive(Debug, Clone, Deserialize, Serialize)]
1250#[serde(deny_unknown_fields)]
1251struct OpenRouterDatacenter {
1252 country_code: String,
1253 #[serde(default, skip_serializing_if = "Option::is_none")]
1254 region: Option<String>,
1255}
1256
1257impl OpenRouterMetadataFile {
1258 fn parse(
1259 text: &str,
1260 ) -> Result<
1261 (
1262 HashMap<String, OpenRouterModelMetadata>,
1263 Option<ProviderMetadata>,
1264 ),
1265 String,
1266 > {
1267 let file: Self =
1268 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
1269 for (alias, metadata) in &file.models {
1270 validate_openrouter_metadata(alias, metadata)?;
1271 }
1272 for (alias, metadata) in &file.planned_models {
1273 validate_openrouter_metadata(alias, metadata)?;
1274 if file.models.contains_key(alias) {
1275 return Err(format!(
1276 "model alias {alias:?} appears in both models and planned_models"
1277 ));
1278 }
1279 }
1280 if let Some(provider) = &file.provider {
1281 if provider.id.is_empty() {
1282 return Err("provider.id must be a non-empty slug".into());
1283 }
1284 for (field, value) in [
1286 ("provider.support_contact", &provider.support_contact),
1287 ("provider.incident_contact", &provider.incident_contact),
1288 ] {
1289 if let Some(value) = value
1290 && !value.contains(':')
1291 {
1292 return Err(format!(
1293 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1294 ));
1295 }
1296 }
1297 }
1298 Ok((file.models, file.provider))
1299 }
1300
1301 #[cfg(test)]
1302 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1303 Self::parse(text).map(|(models, _)| models)
1304 }
1305}
1306
1307fn per_million_price(per_token: &str) -> Option<String> {
1311 if !valid_price_string(per_token) {
1312 return None;
1313 }
1314 let (whole, frac) = match per_token.split_once('.') {
1315 Some((whole, frac)) => (whole, frac),
1316 None => (per_token, ""),
1317 };
1318 let mut digits = format!("{whole}{frac}");
1319 let point = whole.len() + 6;
1320 while digits.len() < point {
1321 digits.push('0');
1322 }
1323 let (int_part, frac_part) = digits.split_at(point);
1324 let int_part = int_part.trim_start_matches('0');
1325 let int_part = if int_part.is_empty() { "0" } else { int_part };
1326 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1327 while frac_out.len() < 2 {
1328 frac_out.push('0');
1329 }
1330 Some(format!("{int_part}.{frac_out}"))
1331}
1332
1333fn valid_price_string(value: &str) -> bool {
1334 let mut parts = value.split('.');
1335 let whole = parts.next().unwrap_or_default();
1336 let fraction = parts.next();
1337 !whole.is_empty()
1338 && whole.bytes().all(|b| b.is_ascii_digit())
1339 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1340 && parts.next().is_none()
1341}
1342
1343fn validate_openrouter_metadata(
1344 alias: &str,
1345 metadata: &OpenRouterModelMetadata,
1346) -> Result<(), String> {
1347 if alias.is_empty() {
1348 return Err("models metadata contains an empty model alias".into());
1349 }
1350 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1353 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1354 {
1355 return Err(format!(
1356 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1357 reasoning_effort level (none|minimal|low|medium|high)"
1358 ));
1359 }
1360 validate_sampling_defaults(alias, metadata)?;
1361 for m in &metadata.input_modalities {
1362 if m != "image" && m != "video" {
1363 return Err(format!(
1364 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1365 ));
1366 }
1367 }
1368 if let Some(sfc) = metadata.surface.as_deref()
1369 && !matches!(sfc, "chat" | "embedding" | "rerank")
1370 {
1371 return Err(format!(
1372 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1373 ));
1374 }
1375 if let Some(q) = metadata.quantization.as_deref()
1376 && !matches!(
1377 q,
1378 "int4"
1379 | "int8"
1380 | "fp4"
1381 | "mxfp4"
1382 | "nvfp4"
1383 | "fp6"
1384 | "fp8"
1385 | "mxfp8"
1386 | "fp16"
1387 | "bf16"
1388 | "fp32"
1389 )
1390 {
1391 return Err(format!(
1392 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1393 ));
1394 }
1395 for (field, value) in [
1396 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1397 (
1398 "pricing.cached_prompt",
1399 metadata.pricing.cached_prompt.as_deref(),
1400 ),
1401 (
1402 "pricing.cache_write",
1403 metadata.pricing.cache_write.as_deref(),
1404 ),
1405 ("pricing.completion", metadata.pricing.completion.as_deref()),
1406 (
1407 "pricing.internal_reasoning",
1408 metadata.pricing.internal_reasoning.as_deref(),
1409 ),
1410 ("pricing.request", metadata.pricing.request.as_deref()),
1411 ] {
1412 if let Some(value) = value
1413 && !valid_price_string(value)
1414 {
1415 return Err(format!(
1416 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1417 ));
1418 }
1419 }
1420 for (field, value) in [
1421 ("created", metadata.created),
1422 ("max_prompt_length", metadata.max_prompt_length),
1423 ("max_output_length", metadata.max_output_length),
1424 ("default_output_length", metadata.default_output_length),
1425 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1426 (
1427 "capacity.cached_prompt_tpm",
1428 metadata.capacity.cached_prompt_tpm,
1429 ),
1430 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1431 ("capacity.request_rpm", metadata.capacity.request_rpm),
1432 ("capacity.concurrency", metadata.capacity.concurrency),
1433 ] {
1434 if let Some(value) = value
1435 && value > JSON_SAFE_INTEGER_MAX
1436 {
1437 return Err(format!(
1438 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1439 ));
1440 }
1441 }
1442 for (field, value) in [
1443 ("max_prompt_length", metadata.max_prompt_length),
1444 ("max_output_length", metadata.max_output_length),
1445 ("default_output_length", metadata.default_output_length),
1446 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1447 (
1448 "capacity.cached_prompt_tpm",
1449 metadata.capacity.cached_prompt_tpm,
1450 ),
1451 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1452 ("capacity.request_rpm", metadata.capacity.request_rpm),
1453 ("capacity.concurrency", metadata.capacity.concurrency),
1454 ] {
1455 if value == Some(0) {
1456 return Err(format!(
1457 "model {alias:?}: {field} must be greater than zero when declared"
1458 ));
1459 }
1460 }
1461 if let (Some(default), Some(maximum)) =
1462 (metadata.default_output_length, metadata.max_output_length)
1463 && default > maximum
1464 {
1465 return Err(format!(
1466 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1467 ));
1468 }
1469 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1470 return Err(format!(
1471 "model {alias:?}: default_output_length requires max_output_length"
1472 ));
1473 }
1474 if let Some(discount) = metadata.discount_to_user
1475 && (!discount.is_finite() || discount >= 1.0)
1476 {
1477 return Err(format!(
1478 "model {alias:?}: discount_to_user must be finite and less than 1"
1479 ));
1480 }
1481 if metadata
1482 .openrouter_slug
1483 .as_deref()
1484 .is_some_and(str::is_empty)
1485 {
1486 return Err(format!(
1487 "model {alias:?}: openrouter_slug must not be empty when declared"
1488 ));
1489 }
1490 for dc in &metadata.datacenters {
1491 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1492 return Err(format!(
1493 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1494 dc.country_code
1495 ));
1496 }
1497 }
1498 Ok(())
1499}
1500
1501fn validate_sampling_defaults(
1516 alias: &str,
1517 metadata: &OpenRouterModelMetadata,
1518) -> Result<(), String> {
1519 validate_sampling_arm(
1520 alias,
1521 &[
1522 "default_temperature",
1523 "default_top_p",
1524 "default_min_p",
1525 "default_presence_penalty",
1526 "default_frequency_penalty",
1527 "default_repetition_penalty",
1528 ],
1529 metadata.default_temperature,
1530 metadata.default_top_p,
1531 metadata.default_min_p,
1532 metadata.default_presence_penalty,
1533 metadata.default_frequency_penalty,
1534 metadata.default_repetition_penalty,
1535 )?;
1536 if let Some(arm) = &metadata.non_thinking_sampling {
1537 if arm.is_empty() {
1541 return Err(format!(
1542 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1543 least one vendor recommendation or delete the table"
1544 ));
1545 }
1546 validate_sampling_arm(
1547 alias,
1548 &[
1549 "non_thinking_sampling.temperature",
1550 "non_thinking_sampling.top_p",
1551 "non_thinking_sampling.min_p",
1552 "non_thinking_sampling.presence_penalty",
1553 "non_thinking_sampling.frequency_penalty",
1554 "non_thinking_sampling.repetition_penalty",
1555 ],
1556 arm.temperature,
1557 arm.top_p,
1558 arm.min_p,
1559 arm.presence_penalty,
1560 arm.frequency_penalty,
1561 arm.repetition_penalty,
1562 )?;
1563 }
1564 Ok(())
1565}
1566
1567#[allow(clippy::too_many_arguments)]
1573fn validate_sampling_arm(
1574 alias: &str,
1575 keys: &[&str; 6],
1576 temperature: Option<f32>,
1577 top_p: Option<f32>,
1578 min_p: Option<f32>,
1579 presence_penalty: Option<f32>,
1580 frequency_penalty: Option<f32>,
1581 repetition_penalty: Option<f32>,
1582) -> Result<(), String> {
1583 if let Some(t) = temperature
1584 && (!t.is_finite() || t <= 0.0 || t > 2.0)
1585 {
1586 return Err(format!(
1587 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1588 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1589 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1590 vendor recommendation, not greedy); clients reach greedy by sending an \
1591 explicit temperature 0.",
1592 keys[0]
1593 ));
1594 }
1595 if let Some(p) = top_p
1596 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1597 {
1598 return Err(format!(
1599 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1600 keys[1]
1601 ));
1602 }
1603 if let Some(m) = min_p
1604 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1605 {
1606 return Err(format!(
1607 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1608 keys[2]
1609 ));
1610 }
1611 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1612 if let Some(v) = value
1613 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1614 {
1615 return Err(format!(
1616 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1617 ));
1618 }
1619 }
1620 if let Some(r) = repetition_penalty
1621 && (!r.is_finite() || r <= 0.0)
1622 {
1623 return Err(format!(
1624 "model {alias:?}: {} {r} must be finite and \
1625 greater than zero (1.0 = off)",
1626 keys[5]
1627 ));
1628 }
1629 Ok(())
1630}
1631
1632#[derive(Debug, Clone, Default)]
1637struct ModelMetadataSet {
1638 models: HashMap<String, OpenRouterModelMetadata>,
1639 provider: Option<ProviderMetadata>,
1641}
1642
1643#[derive(Debug, Clone)]
1647pub struct MetadataReloadReceipt {
1648 pub path: String,
1649 pub sha256: String,
1650 pub models: usize,
1651 pub has_provider: bool,
1652}
1653
1654fn sha256_hex_bytes(bytes: &[u8]) -> String {
1655 let mut hasher = Sha256::new();
1656 hasher.update(bytes);
1657 format!("{:x}", hasher.finalize())
1658}
1659
1660fn load_model_metadata_file(
1665 path: &std::path::Path,
1666 models: &[(String, String, Option<String>)],
1667) -> Result<(ModelMetadataSet, MetadataReloadReceipt), String> {
1668 let bytes = std::fs::read(path).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1669 let sha256 = sha256_hex_bytes(&bytes);
1670 let text =
1671 String::from_utf8(bytes).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1672 let (models_map, provider) = OpenRouterMetadataFile::parse(&text)
1673 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1674 for alias in models_map.keys() {
1675 if !models.iter().any(|(name, _, _)| name == alias) {
1676 return Err(format!(
1677 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1678 ));
1679 }
1680 }
1681 let receipt = MetadataReloadReceipt {
1682 path: path.to_string_lossy().into_owned(),
1683 sha256,
1684 models: models_map.len(),
1685 has_provider: provider.is_some(),
1686 };
1687 Ok((
1688 ModelMetadataSet {
1689 models: models_map,
1690 provider,
1691 },
1692 receipt,
1693 ))
1694}
1695
1696fn load_openrouter_metadata(
1697 models: &[(String, String, Option<String>)],
1698) -> Result<ModelMetadataSet, String> {
1699 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1700 Ok(path) => path,
1701 Err(_) => return Ok(ModelMetadataSet::default()),
1702 };
1703 let p = std::path::Path::new(&path);
1704 if !p.is_file() {
1705 return Err(format!(
1706 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1707 ));
1708 }
1709 let (set, receipt) = load_model_metadata_file(p, models)?;
1710 eprintln!(
1711 "[server] OpenRouter metadata loaded: {} model(s), sha256 {} from {path}",
1712 receipt.models, receipt.sha256,
1713 );
1714 Ok(set)
1715}
1716
1717#[derive(Clone, Default)]
1735pub struct MetadataReloadHandle {
1736 cell: Arc<RwLock<Arc<ModelMetadataSet>>>,
1737 models: Arc<Vec<(String, String, Option<String>)>>,
1742 path: Option<std::path::PathBuf>,
1745}
1746
1747impl MetadataReloadHandle {
1748 pub fn reload(&self) -> Result<MetadataReloadReceipt, String> {
1751 let path = self.path.as_ref().ok_or_else(|| {
1752 "model metadata reload: MEMRA_MODEL_METADATA was not configured at boot".to_string()
1753 })?;
1754 let (set, receipt) = load_model_metadata_file(path, &self.models)?;
1755 {
1760 let mut guard = self
1761 .cell
1762 .write()
1763 .unwrap_or_else(|poisoned| poisoned.into_inner());
1764 *guard = Arc::new(set);
1765 }
1766 eprintln!(
1767 "[server] model metadata reloaded: {} model(s), sha256 {} from {}",
1768 receipt.models, receipt.sha256, receipt.path,
1769 );
1770 Ok(receipt)
1771 }
1772}
1773
1774#[derive(Clone)]
1775struct AppState {
1776 cmd_tx: Sender<Cmd>,
1777 models: Arc<Vec<String>>,
1778 caps: Arc<HashMap<String, ModelCaps>>,
1779 openrouter_metadata: Arc<RwLock<Arc<ModelMetadataSet>>>,
1784 metering: Option<Arc<dyn metering::Metering>>,
1790 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1793 api_auth: ApiAuth,
1796 metrics_auth: MetricsAuth,
1798 metrics: SharedMetrics,
1799 inflight: InflightCounts,
1803 tenant_inflight: TenantGauge,
1806 health: health::SharedHealth,
1810 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1814}
1815
1816impl AppState {
1817 fn metadata(&self) -> Arc<ModelMetadataSet> {
1822 self.openrouter_metadata
1823 .read()
1824 .map(|guard| guard.clone())
1825 .unwrap_or_else(|poisoned| poisoned.into_inner().clone())
1826 }
1827 fn sampling_defaults_in(
1845 set: &ModelMetadataSet,
1846 caps: &HashMap<String, ModelCaps>,
1847 model: &str,
1848 ) -> ModelSamplingDefaults {
1849 ModelSamplingDefaults::resolve(set.models.get(model), caps.get(model))
1850 }
1851}
1852
1853#[derive(Clone, Default)]
1854struct ApiAuth {
1855 keyring: Option<&'static auth::KeyStore>,
1856 single_key: Option<Arc<str>>,
1857}
1858
1859impl ApiAuth {
1860 fn from_env() -> Result<ApiAuth, String> {
1861 let single_key = match std::env::var("MEMRA_API_KEY") {
1862 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1863 Ok(key) => Some(Arc::from(key)),
1864 Err(std::env::VarError::NotPresent) => None,
1865 Err(std::env::VarError::NotUnicode(_)) => {
1866 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1867 }
1868 };
1869 Ok(ApiAuth {
1870 keyring: auth::global(),
1871 single_key,
1872 })
1873 }
1874
1875 fn configured(&self) -> bool {
1876 self.keyring.is_some() || self.single_key.is_some()
1877 }
1878}
1879
1880#[derive(Clone, Default)]
1881struct MetricsAuth {
1882 required: bool,
1883 token: Option<Arc<str>>,
1884}
1885
1886impl MetricsAuth {
1887 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1888 let token = token.map(Arc::from);
1889 MetricsAuth {
1890 required: !bind_loopback || api_auth_configured || token.is_some(),
1891 token,
1892 }
1893 }
1894}
1895
1896fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1897 let mut resolved = addr
1898 .to_socket_addrs()
1899 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1900 let first = resolved
1901 .next()
1902 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1903 let mut loopback = first.ip().to_canonical().is_loopback();
1904 for socket in resolved {
1905 loopback &= socket.ip().to_canonical().is_loopback();
1906 }
1907 Ok((first, loopback))
1908}
1909
1910fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1911 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1912}
1913
1914fn validate_bind_security(
1915 addr: &str,
1916 api_auth_configured: bool,
1917 allow_open_bind: bool,
1918) -> Result<bool, String> {
1919 let loopback = bind_is_loopback(addr)?;
1920 if !loopback && !api_auth_configured && !allow_open_bind {
1921 return Err(format!(
1922 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1923 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1924 ));
1925 }
1926 Ok(loopback)
1927}
1928
1929type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1947
1948type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1951
1952struct InflightGuard {
1956 counts: InflightCounts,
1957 idx: usize,
1958 tenants: TenantGauge,
1959 tenant: String,
1960}
1961
1962impl InflightGuard {
1963 fn try_acquire(
1967 counts: InflightCounts,
1968 lane: lanes::Lane,
1969 tenants: TenantGauge,
1970 tenant: &str,
1971 tenant_cap: Option<usize>,
1972 ) -> Result<(Self, usize, usize), usize> {
1973 let idx = lane.idx();
1974 let nt = {
1975 let mut m = tenants.lock().unwrap();
1976 let e = m.entry(tenant.to_string()).or_insert(0);
1977 if tenant_cap.is_some_and(|cap| *e >= cap) {
1978 return Err(*e);
1979 }
1980 *e += 1;
1981 *e
1982 };
1983 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1984 Ok((
1985 InflightGuard {
1986 counts,
1987 idx,
1988 tenants,
1989 tenant: tenant.to_string(),
1990 },
1991 n,
1992 nt,
1993 ))
1994 }
1995}
1996
1997impl Drop for InflightGuard {
1998 fn drop(&mut self) {
1999 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
2000 let mut m = self.tenants.lock().unwrap();
2001 if let Some(e) = m.get_mut(&self.tenant) {
2002 *e -= 1;
2003 if *e == 0 {
2004 m.remove(&self.tenant);
2005 }
2006 }
2007 }
2008}
2009
2010fn lane_cap(lane: lanes::Lane) -> usize {
2014 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
2015 CAPS.get_or_init(|| {
2016 let batching = std::env::var("MEMRA_SERVE_BATCH")
2017 .map(|v| v != "0")
2018 .unwrap_or(true);
2019 let interactive = if batching {
2020 std::env::var("MEMRA_MAX_SESSIONS")
2021 .ok()
2022 .and_then(|v| v.parse().ok())
2023 .unwrap_or(64)
2024 } else {
2025 worker::MAX_ACTIVE
2026 };
2027 let p = lanes::LanePolicy::from_env();
2028 [interactive, p.max_sessions[1], p.max_sessions[2]]
2029 })[lane.idx()]
2030}
2031
2032fn reset_estimate_s(m: &worker::Metrics) -> u64 {
2035 if m.completed > 0 && m.step_p50_ms > 0.0 {
2036 let mean_toks = m.tokens_out as f64 / m.completed as f64;
2037 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
2038 }
2039 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2040 *D.get_or_init(|| {
2041 std::env::var("MEMRA_RL_RESET_S")
2042 .ok()
2043 .and_then(|v| v.parse().ok())
2044 .unwrap_or(2)
2045 })
2046}
2047
2048pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
2064pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
2065pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
2066
2067pub(crate) fn timeout_ms_max() -> u64 {
2073 static V: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2074 *V.get_or_init(|| {
2075 std::env::var("MEMRA_TIMEOUT_MS_MAX")
2076 .ok()
2077 .and_then(|s| s.parse::<u64>().ok())
2078 .filter(|&ms| ms >= TIMEOUT_MS_MIN)
2079 .unwrap_or(TIMEOUT_MS_MAX)
2080 })
2081}
2082
2083pub(crate) fn stream_ttft_ms_max() -> u64 {
2088 static V: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2089 *V.get_or_init(|| {
2090 std::env::var("MEMRA_STREAM_TTFT_MS_MAX")
2091 .ok()
2092 .and_then(|s| s.parse::<u64>().ok())
2093 .filter(|&ms| ms >= TIMEOUT_MS_MIN)
2094 .unwrap_or_else(timeout_ms_max)
2095 })
2096}
2097
2098fn timeout_ms_default(stream: bool) -> u64 {
2099 if stream {
2100 stream_ttft_ms_max()
2101 } else {
2102 timeout_ms_max()
2103 }
2104}
2105
2106pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>, stream: bool) -> Result<u64, String> {
2110 let max = timeout_ms_default(stream);
2111 let Some(v) = v.filter(|v| !v.is_null()) else {
2112 return Ok(max);
2114 };
2115 let Some(ms) = v.as_u64() else {
2116 return Err(format!(
2117 "timeout_ms must be an integer number of milliseconds in \
2118 {TIMEOUT_MS_MIN}..={max}, got {v}; for long work use \"stream\": true — \
2119 the deadline then bounds only time to first token and the stream may run \
2120 as long as it needs"
2121 ));
2122 };
2123 if !(TIMEOUT_MS_MIN..=max).contains(&ms) {
2124 return Err(format!(
2125 "timeout_ms {ms} is outside the accepted range \
2126 {TIMEOUT_MS_MIN}..={max} (milliseconds) for this response mode. For long \
2127 work use \"stream\": true — the deadline then bounds only time to first \
2128 token and the stream may run as long as it needs"
2129 ));
2130 }
2131 Ok(ms)
2132}
2133
2134#[derive(Clone, Copy)]
2137pub(crate) struct RequestDeadline {
2138 pub(crate) at: tokio::time::Instant,
2139 pub(crate) ms: u64,
2140}
2141
2142impl RequestDeadline {
2143 pub(crate) fn preheader(self, stream: bool) -> Self {
2144 if !stream || self.ms <= TIMEOUT_MS_MAX {
2145 return self;
2146 }
2147 let Some(commit_ms) = sse_prefill_commit_ms() else {
2148 return self;
2149 };
2150 let ms = self.ms.min(commit_ms);
2151 Self {
2152 at: self.at - std::time::Duration::from_millis(self.ms - ms),
2153 ms,
2154 }
2155 }
2156 pub(crate) fn starting_now(ms: u64) -> Self {
2157 Self {
2158 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
2159 ms,
2160 }
2161 }
2162
2163 pub(crate) fn remaining(&self) -> std::time::Duration {
2164 self.at
2165 .saturating_duration_since(tokio::time::Instant::now())
2166 }
2167}
2168
2169pub(crate) fn deadline_exceeded_message(ms: u64, stream: bool) -> String {
2175 let what = if stream {
2176 "the first token was produced"
2177 } else {
2178 "the response completed"
2179 };
2180 format!(
2181 "deadline of {ms} ms (timeout_ms; default {}) elapsed before \
2182 {what}; generation was cancelled and this request is not billed",
2183 timeout_ms_default(stream)
2184 )
2185}
2186
2187pub(crate) fn deadline_exceeded_error(ms: u64, stream: bool) -> serde_json::Value {
2188 error_body(
2189 &deadline_exceeded_message(ms, stream),
2190 "timeout",
2191 Some("timeout_ms"),
2192 Some("deadline_exceeded"),
2193 )
2194}
2195
2196pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
2197 error_response_coded(
2198 StatusCode::REQUEST_TIMEOUT,
2199 &deadline_exceeded_message(ms, stream),
2200 "timeout",
2201 Some("timeout_ms"),
2202 Some("deadline_exceeded"),
2203 )
2204}
2205
2206pub(crate) fn admission_deadline_response(deadline: RequestDeadline, stream: bool) -> Response {
2207 let admission = deadline.preheader(stream);
2208 if admission.ms == deadline.ms {
2209 return deadline_exceeded_response(deadline.ms, stream);
2210 }
2211 error_response_coded(
2212 StatusCode::REQUEST_TIMEOUT,
2213 &format!(
2214 "admission did not complete within the {} ms pre-header budget \
2215 (streaming first-token timeout_ms={}); generation was cancelled and this request is not billed",
2216 admission.ms, deadline.ms
2217 ),
2218 "timeout",
2219 Some("timeout_ms"),
2220 Some("deadline_exceeded"),
2221 )
2222}
2223
2224pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
2264
2265pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
2270
2271pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
2275
2276fn env_flag_on(name: &'static str, default_on: bool) -> bool {
2282 match std::env::var(name) {
2283 Ok(v) => !matches!(
2284 v.trim().to_ascii_lowercase().as_str(),
2285 "0" | "off" | "false"
2286 ),
2287 Err(_) => default_on,
2288 }
2289}
2290
2291pub(crate) fn env_u64(name: &'static str, default: u64) -> u64 {
2294 std::env::var(name)
2295 .ok()
2296 .and_then(|v| v.parse::<u64>().ok())
2297 .filter(|v| *v > 0)
2298 .unwrap_or(default)
2299}
2300
2301const CHARS_PER_TOKEN_FLOOR: usize = 6;
2312
2313pub(crate) fn prompt_tokens_estimate(
2314 request: &worker::Request,
2315 tokenizer: Option<&Tokenizer>,
2316) -> u64 {
2317 if !request.prompt_ids.is_empty() {
2318 return request.prompt_ids.len() as u64;
2319 }
2320 let mut text = String::new();
2321 text.push_str(&request.prompt_text);
2322 for turn in &request.chat_turns {
2323 text.push_str(&turn.content);
2324 }
2325 for tool in &request.tools_json {
2326 text.push_str(tool);
2327 }
2328 if let Some(tokenizer) = tokenizer {
2329 return tokenizer.encode(text.as_str(), false).len() as u64;
2330 }
2331 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
2332}
2333
2334pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
2338 let prefill_ms = prompt_tokens
2339 .saturating_mul(1_000)
2340 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
2341 .unwrap_or(u64::MAX);
2342 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
2343 if decode_ms == 0 {
2344 return None;
2345 }
2346 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
2347}
2348
2349pub(crate) fn nonstream_deadline_gate(
2359 request: &worker::Request,
2360 stream: bool,
2361 deadline: RequestDeadline,
2362 caller_declared_max_tokens: bool,
2363 tokenizer: Option<&Tokenizer>,
2364) -> Result<(), String> {
2365 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
2366 return Ok(());
2367 }
2368 let max_new = request.params.max_new as u64;
2369 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
2376 {
2377 return Ok(());
2378 }
2379 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
2380 let remaining_ms = deadline.remaining().as_millis() as u64;
2381 let prefill_ms = prompt_tokens.saturating_mul(1_000)
2382 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
2383 let decode_ms = max_new.saturating_mul(1_000)
2384 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
2385 let est_ms = prefill_ms.saturating_add(decode_ms);
2386 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
2387 if est_ms <= bound_ms {
2388 return Ok(());
2389 }
2390 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
2391 let advice = match fits {
2392 Some(fits) if fits > 0 => format!(
2393 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
2394 stream's deadline bounds only the time to first token, so it may run as long \
2395 as it needs"
2396 ),
2397 _ => format!(
2398 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
2399 token, so no max_tokens fits: set \"stream\": true"
2400 ),
2401 };
2402 Err(format!(
2403 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
2404 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
2405 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
2406 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
2407 rather than after the deadline: {advice}",
2408 est_ms / 1_000,
2409 ))
2410}
2411
2412fn max_queue_depth(cap: usize) -> usize {
2416 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
2417 D.get_or_init(|| {
2418 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
2419 .ok()
2420 .and_then(|v| v.parse().ok())
2421 })
2422 .unwrap_or(cap.saturating_mul(4))
2423}
2424
2425fn queue_wait_ceiling_s() -> u64 {
2432 static S: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2433 *S.get_or_init(|| {
2434 std::env::var("MEMRA_QUEUE_WAIT_CEILING_S")
2435 .ok()
2436 .and_then(|v| v.parse().ok())
2437 .unwrap_or(0)
2438 })
2439}
2440
2441pub(crate) struct PendingAdmissionGuard {
2462 reserved: bool,
2463 lane: lanes::Lane,
2464}
2465
2466impl PendingAdmissionGuard {
2467 pub(crate) fn commit(mut self) {
2471 self.reserved = false;
2472 std::mem::forget(self);
2473 }
2474}
2475
2476impl Drop for PendingAdmissionGuard {
2477 fn drop(&mut self) {
2478 if self.reserved {
2479 worker::release_pending_admit();
2480 worker::release_admission_reservation(self.lane);
2481 }
2482 }
2483}
2484
2485#[allow(clippy::result_large_err)] pub(crate) fn reserve_pending_admit(
2487 st: &AppState,
2488 lane: lanes::Lane,
2489 rl: &RateLimit,
2490 deadline: RequestDeadline,
2491) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2492 reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())
2493}
2494
2495#[allow(clippy::result_large_err)] fn reserve_pending_admit_with_ceiling(
2500 st: &AppState,
2501 lane: lanes::Lane,
2502 rl: &RateLimit,
2503 deadline: RequestDeadline,
2504 ceiling_s: u64,
2505) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2506 let cap = lane_cap(lane).max(1);
2511 let bound = max_queue_depth(cap);
2512 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2513 loop {
2514 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2515 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2516 let backlog = reservations;
2520 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2521 if backlog >= bound {
2522 let msg = format!(
2523 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2524 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2525 coarse estimate, not a promise)",
2526 lane.as_str()
2527 );
2528 let resp = retry_contract_response(
2529 (
2530 StatusCode::TOO_MANY_REQUESTS,
2531 Json(error_body(
2532 &msg,
2533 "rate_limit_error",
2534 None,
2535 Some("shed_queue"),
2536 )),
2537 )
2538 .into_response(),
2539 Some(est_wait_s),
2540 );
2541 return Err((resp, "shed_queue"));
2542 }
2543 let remaining_ms = deadline.remaining().as_millis() as u64;
2544 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2548 if lane == lanes::Lane::Interactive
2549 && waits_for_capacity
2550 && est_wait_s.saturating_mul(1_000) > remaining_ms
2551 {
2552 let msg = format!(
2553 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2554 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2555 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2556 estimate, not a promise)"
2557 );
2558 let resp = retry_contract_response(
2559 (
2560 StatusCode::TOO_MANY_REQUESTS,
2561 Json(error_body(
2562 &msg,
2563 "rate_limit_error",
2564 None,
2565 Some("shed_deadline"),
2566 )),
2567 )
2568 .into_response(),
2569 Some(est_wait_s),
2570 );
2571 return Err((resp, "shed_deadline"));
2572 }
2573 if lane == lanes::Lane::Interactive
2581 && waits_for_capacity
2582 && ceiling_s > 0
2583 && est_wait_s > ceiling_s
2584 {
2585 let msg = format!(
2586 "estimated queue wait ~{est_wait_s}s exceeds this deployment's queue-wait \
2587 ceiling ({ceiling_s}s); this request was not admitted and is not billed; \
2588 retry after ~{est_wait_s}s (a coarse estimate, not a promise)"
2589 );
2590 let resp = retry_contract_response(
2591 (
2592 StatusCode::TOO_MANY_REQUESTS,
2593 Json(error_body(
2594 &msg,
2595 "rate_limit_error",
2596 None,
2597 Some("shed_queue_wait"),
2598 )),
2599 )
2600 .into_response(),
2601 Some(est_wait_s),
2602 );
2603 return Err((resp, "shed_queue_wait"));
2604 }
2605 if reservations_for_lane
2606 .compare_exchange(
2607 reservations,
2608 reservations.saturating_add(1),
2609 std::sync::atomic::Ordering::AcqRel,
2610 std::sync::atomic::Ordering::Acquire,
2611 )
2612 .is_ok()
2613 {
2614 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2618 return Ok(PendingAdmissionGuard {
2619 reserved: true,
2620 lane,
2621 });
2622 }
2623 }
2624}
2625
2626static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2637
2638fn draining() -> bool {
2639 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2640}
2641
2642fn drain_deadline_s() -> u64 {
2644 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2645 *D.get_or_init(|| {
2646 std::env::var("MEMRA_DRAIN_S")
2647 .ok()
2648 .and_then(|v| v.parse().ok())
2649 .unwrap_or(30)
2650 })
2651}
2652
2653fn drain_response() -> Response {
2663 let resp = (
2664 StatusCode::SERVICE_UNAVAILABLE,
2665 Json(error_body(
2666 "server is draining (shutdown in progress); retry",
2667 "server_error",
2668 None,
2669 Some("draining"),
2670 )),
2671 )
2672 .into_response();
2673 retry_contract_response(resp, Some(drain_deadline_s()))
2674}
2675
2676struct RateLimit {
2678 limit: usize,
2679 remaining: usize,
2680 reset_s: u64,
2681}
2682
2683impl RateLimit {
2684 fn at_admit(
2689 lane: lanes::Lane,
2690 n_inflight: usize,
2691 metrics: &SharedMetrics,
2692 tenant: &auth::TenantCtx,
2693 n_tenant: usize,
2694 ) -> Self {
2695 let global = lane_cap(lane);
2696 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2697 return Self::compute(global, n_inflight, metrics);
2698 };
2699 let headroom = t
2700 .saturating_sub(n_tenant)
2701 .min(global.saturating_sub(n_inflight));
2702 Self::compute(t, t - headroom, metrics)
2704 }
2705
2706 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2707 let remaining = limit.saturating_sub(n_inflight);
2708 let reset_s = if remaining > 0 {
2709 0
2710 } else {
2711 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2712 reset_estimate_s(&m)
2713 };
2714 RateLimit {
2715 limit,
2716 remaining,
2717 reset_s,
2718 }
2719 }
2720
2721 fn attach(&self, mut resp: Response) -> Response {
2723 let h = resp.headers_mut();
2724 for (k, v) in [
2725 ("x-ratelimit-limit", self.limit as u64),
2726 ("x-ratelimit-remaining", self.remaining as u64),
2727 ("x-ratelimit-reset", self.reset_s),
2728 ] {
2729 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2730 h.insert(axum::http::HeaderName::from_static(k), v);
2731 }
2732 }
2733 resp
2734 }
2735}
2736
2737#[allow(clippy::result_large_err)] fn acquire_request_slot(
2742 st: &AppState,
2743 lane: lanes::Lane,
2744 tenant: &auth::TenantCtx,
2745 env: &Envelope,
2746) -> Result<(InflightGuard, RateLimit), Response> {
2747 let global = lane_cap(lane);
2748 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2749 match InflightGuard::try_acquire(
2750 st.inflight.clone(),
2751 lane,
2752 st.tenant_inflight.clone(),
2753 &tenant.tenant,
2754 tenant_cap,
2755 ) {
2756 Ok((guard, n_inflight, n_tenant)) => {
2757 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2758 Ok((guard, rl))
2759 }
2760 Err(n_tenant) => {
2761 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2762 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2763 let error =
2764 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2765 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2766 }
2767 }
2768}
2769
2770#[derive(Deserialize)]
2772struct CompletionReq {
2773 model: String,
2774 #[serde(default)]
2775 prompt: String,
2776 #[serde(default)]
2778 prompt_ids: Vec<u32>,
2779 #[serde(default)]
2782 max_tokens: Option<usize>,
2783 #[serde(default)]
2796 temperature: Option<f32>,
2797 #[serde(default)]
2798 top_p: Option<f32>,
2799 #[serde(default)]
2801 top_k: Option<usize>,
2802 #[serde(default)]
2804 min_p: Option<f32>,
2805 #[serde(default)]
2807 frequency_penalty: Option<f32>,
2808 #[serde(default)]
2809 presence_penalty: Option<f32>,
2810 #[serde(default)]
2812 repetition_penalty: Option<f32>,
2813 #[serde(default)]
2819 seed: Option<u64>,
2820 #[serde(default)]
2821 stop: StopSequences,
2822 #[serde(default)]
2825 logit_bias: Option<serde_json::Value>,
2826 #[serde(default)]
2827 logprobs: Option<serde_json::Value>,
2828 #[serde(default)]
2829 n: Option<usize>,
2830 #[serde(default)]
2831 best_of: Option<usize>,
2832 #[serde(default)]
2834 chat: bool,
2835 #[serde(default)]
2837 stream: bool,
2838 #[serde(default)]
2840 max_ctx: Option<usize>,
2841 #[serde(default)]
2843 trace_id: Option<String>,
2844 #[serde(default)]
2848 cache_salt: Option<String>,
2849 #[serde(default)]
2853 session_id: Option<String>,
2854 #[serde(default)]
2855 user: Option<String>,
2856 #[serde(default)]
2860 timeout_ms: Option<serde_json::Value>,
2861}
2862
2863#[derive(Deserialize)]
2864struct ChatMessage {
2865 role: String,
2866 #[serde(default)]
2868 content: serde_json::Value,
2869 #[serde(default)]
2871 tool_calls: Vec<ReqToolCall>,
2872 #[serde(default)]
2875 tool_call_id: Option<String>,
2876 #[serde(default)]
2879 name: Option<String>,
2880 #[serde(default, alias = "reasoning_content")]
2890 reasoning: Option<String>,
2891}
2892
2893#[derive(Deserialize)]
2894struct ReqToolCall {
2895 #[serde(default)]
2896 #[allow(dead_code)]
2897 id: Option<String>,
2898 function: ReqToolFunction,
2899}
2900
2901#[derive(Deserialize)]
2902struct ReqToolFunction {
2903 name: String,
2904 #[serde(default)]
2906 arguments: serde_json::Value,
2907}
2908
2909#[derive(Clone, Default, Deserialize)]
2910#[serde(untagged)]
2911enum StopSequences {
2912 One(String),
2913 Many(Vec<String>),
2914 #[default]
2915 None,
2916}
2917
2918impl StopSequences {
2919 fn into_vec(self) -> Vec<String> {
2925 let stops = match self {
2926 Self::One(stop) => vec![stop],
2927 Self::Many(stops) => stops,
2928 Self::None => Vec::new(),
2929 };
2930 stops.into_iter().filter(|s| !s.is_empty()).collect()
2931 }
2932
2933 fn validate(&self) -> Result<(), String> {
2934 let stops: &[String] = match self {
2935 Self::One(stop) => std::slice::from_ref(stop),
2936 Self::Many(stops) => stops,
2937 Self::None => &[],
2938 };
2939 if stops.len() > MAX_STOP_SEQUENCES {
2940 return Err(format!(
2941 "stop accepts at most {MAX_STOP_SEQUENCES} sequences"
2942 ));
2943 }
2944 let mut total = 0usize;
2945 for stop in stops {
2946 let bytes = stop.len();
2947 if bytes > MAX_STOP_SEQUENCE_BYTES {
2948 return Err(format!(
2949 "each stop sequence must be at most {MAX_STOP_SEQUENCE_BYTES} UTF-8 bytes"
2950 ));
2951 }
2952 total = total
2953 .checked_add(bytes)
2954 .ok_or_else(|| "stop sequence byte count overflowed".to_string())?;
2955 }
2956 if total > MAX_STOP_SEQUENCES_BYTES {
2957 return Err(format!(
2958 "stop sequences must total at most {MAX_STOP_SEQUENCES_BYTES} UTF-8 bytes"
2959 ));
2960 }
2961 Ok(())
2962 }
2963}
2964
2965#[derive(Deserialize)]
2971struct ChatCompletionReq {
2972 model: String,
2973 messages: Vec<ChatMessage>,
2974 #[serde(default, alias = "max_completion_tokens")]
2977 max_tokens: Option<usize>,
2978 #[serde(default)]
2981 temperature: Option<f32>,
2982 #[serde(default)]
2983 top_p: Option<f32>,
2984 #[serde(default)]
2988 top_k: Option<usize>,
2989 #[serde(default)]
2991 min_p: Option<f32>,
2992 #[serde(default)]
2994 frequency_penalty: Option<f32>,
2995 #[serde(default)]
2996 presence_penalty: Option<f32>,
2997 #[serde(default)]
2999 repetition_penalty: Option<f32>,
3000 #[serde(default)]
3002 seed: Option<u64>,
3003 #[serde(default)]
3004 stop: StopSequences,
3005 #[serde(default)]
3006 stream: bool,
3007 #[serde(default)]
3008 max_ctx: Option<usize>,
3009 #[serde(default)]
3014 response_format: Option<serde_json::Value>,
3015 #[serde(default)]
3016 logit_bias: Option<serde_json::Value>,
3017 #[serde(default)]
3018 logprobs: Option<serde_json::Value>,
3019 #[serde(default)]
3020 top_logprobs: Option<usize>,
3021 #[serde(default)]
3022 n: Option<usize>,
3023 #[serde(default)]
3025 tools: Vec<serde_json::Value>,
3026 #[serde(default)]
3028 tool_choice: Option<serde_json::Value>,
3029 #[serde(default)]
3035 reasoning_effort: Option<String>,
3036 #[serde(default)]
3045 reasoning: Option<serde_json::Value>,
3046 #[serde(default)]
3061 include_reasoning: Option<bool>,
3062 #[serde(default)]
3071 enable_thinking: Option<bool>,
3072 #[serde(default)]
3078 chat_template_kwargs: Option<serde_json::Value>,
3079 #[serde(default)]
3083 cache_salt: Option<String>,
3084 #[serde(default)]
3086 session_id: Option<String>,
3087 #[serde(default)]
3088 user: Option<String>,
3089 #[serde(default)]
3094 timeout_ms: Option<serde_json::Value>,
3095}
3096fn one() -> f32 {
3097 1.0
3098}
3099fn default_temperature() -> f32 {
3105 1.0
3106}
3107
3108#[derive(Debug, Clone, Copy, Default, PartialEq)]
3128struct SamplingDefaults {
3129 temperature: Option<f32>,
3130 top_p: Option<f32>,
3131 top_k: Option<usize>,
3132 min_p: Option<f32>,
3133 frequency_penalty: Option<f32>,
3134 presence_penalty: Option<f32>,
3135 repetition_penalty: Option<f32>,
3136}
3137
3138impl SamplingDefaults {
3139 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
3142 SamplingDefaults {
3143 temperature: metadata
3144 .and_then(|m| m.default_temperature)
3145 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
3146 top_p: metadata
3147 .and_then(|m| m.default_top_p)
3148 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
3149 top_k: metadata.and_then(|m| m.default_top_k),
3150 min_p: metadata.and_then(|m| m.default_min_p),
3151 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
3152 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
3153 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
3154 }
3155 }
3156}
3157
3158#[derive(Debug, Clone, Copy, Default, PartialEq)]
3171struct ModelSamplingDefaults {
3172 thinking: SamplingDefaults,
3173 non_thinking: Option<SamplingDefaults>,
3174}
3175
3176impl ModelSamplingDefaults {
3177 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
3178 ModelSamplingDefaults {
3179 thinking: SamplingDefaults::resolve(metadata, caps),
3180 non_thinking: metadata
3185 .and_then(|m| m.non_thinking_sampling.as_ref())
3186 .map(|arm| SamplingDefaults {
3187 temperature: arm.temperature,
3188 top_p: arm.top_p,
3189 top_k: arm.top_k,
3190 min_p: arm.min_p,
3191 frequency_penalty: arm.frequency_penalty,
3192 presence_penalty: arm.presence_penalty,
3193 repetition_penalty: arm.repetition_penalty,
3194 }),
3195 }
3196 }
3197
3198 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
3211 match (think, &self.non_thinking) {
3212 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
3213 _ => &self.thinking,
3214 }
3215 }
3216
3217 #[cfg(test)] fn single(thinking: SamplingDefaults) -> Self {
3221 ModelSamplingDefaults {
3222 thinking,
3223 non_thinking: None,
3224 }
3225 }
3226}
3227
3228#[derive(Debug, Clone, Copy, Default)]
3234struct ClientSampling {
3235 temperature: Option<f32>,
3236 top_p: Option<f32>,
3237 top_k: Option<usize>,
3238 min_p: Option<f32>,
3239 frequency_penalty: Option<f32>,
3240 presence_penalty: Option<f32>,
3241 repetition_penalty: Option<f32>,
3242 seed: Option<u64>,
3243}
3244
3245impl From<&CompletionReq> for ClientSampling {
3246 fn from(r: &CompletionReq) -> Self {
3247 ClientSampling {
3248 temperature: r.temperature,
3249 top_p: r.top_p,
3250 top_k: r.top_k,
3251 min_p: r.min_p,
3252 frequency_penalty: r.frequency_penalty,
3253 presence_penalty: r.presence_penalty,
3254 repetition_penalty: r.repetition_penalty,
3255 seed: r.seed,
3256 }
3257 }
3258}
3259
3260impl From<&ChatCompletionReq> for ClientSampling {
3261 fn from(r: &ChatCompletionReq) -> Self {
3262 ClientSampling {
3263 temperature: r.temperature,
3264 top_p: r.top_p,
3265 top_k: r.top_k,
3266 min_p: r.min_p,
3267 frequency_penalty: r.frequency_penalty,
3268 presence_penalty: r.presence_penalty,
3269 repetition_penalty: r.repetition_penalty,
3270 seed: r.seed,
3271 }
3272 }
3273}
3274
3275fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
3281 sampler_config(
3282 client
3283 .temperature
3284 .or(defaults.temperature)
3285 .unwrap_or_else(default_temperature),
3286 client.top_k.or(defaults.top_k).unwrap_or(0),
3287 client.top_p.or(defaults.top_p).unwrap_or_else(one),
3288 client.min_p.or(defaults.min_p).unwrap_or(0.0),
3289 client
3290 .frequency_penalty
3291 .or(defaults.frequency_penalty)
3292 .unwrap_or(0.0),
3293 client
3294 .presence_penalty
3295 .or(defaults.presence_penalty)
3296 .unwrap_or(0.0),
3297 client
3298 .repetition_penalty
3299 .or(defaults.repetition_penalty)
3300 .unwrap_or_else(one),
3301 client.seed,
3302 )
3303}
3304
3305#[derive(Serialize)]
3306struct CompletionResp {
3307 model: String,
3308 text: String,
3309 tokens: Vec<u32>,
3310 stop_reason: String,
3314 #[serde(default, skip_serializing_if = "Option::is_none")]
3319 error: Option<serde_json::Value>,
3320 n_tokens: usize,
3321 prompt_tokens: usize,
3324 cached_tokens: usize,
3325 elapsed_s: f64,
3326}
3327
3328fn usage_json(
3336 n_prompt: usize,
3337 n_tokens: usize,
3338 n_cached: usize,
3339 elapsed_s: f64,
3340 spec: Option<worker::SpecUsage>,
3341) -> serde_json::Value {
3342 let mut u = json!({
3343 "prompt_tokens": n_prompt,
3344 "completion_tokens": n_tokens,
3345 "total_tokens": n_prompt + n_tokens,
3346 "prompt_tokens_details": { "cached_tokens": n_cached },
3347 "elapsed_s": elapsed_s,
3348 });
3349 if let Some(sp) = spec {
3350 u["spec"] = json!({
3351 "rounds": sp.rounds,
3352 "drafted": sp.drafted,
3353 "accepted": sp.accepted,
3354 "acceptance_rate": if sp.drafted > 0 {
3355 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
3356 });
3357 }
3358 u
3359}
3360
3361pub const SYSTEM_FINGERPRINT: &str = concat!(
3393 "memra-",
3394 env!("CARGO_PKG_VERSION"),
3395 "-",
3396 env!("MEMRA_BUILD_ID")
3397);
3398
3399pub const BUILD_ID_SRC: &str = env!("MEMRA_BUILD_ID_SRC");
3401
3402pub const BUILD_ID_NOTE: &str = env!("MEMRA_BUILD_ID_NOTE");
3404
3405pub const BUILD_GIT_SHA: &str = env!("MEMRA_BUILD_SHA");
3409
3410pub fn build_identity_line() -> String {
3413 format!("[server] build: {SYSTEM_FINGERPRINT} (id: {BUILD_ID_SRC}, git: {BUILD_GIT_SHA})")
3414}
3415
3416fn gen_hex128() -> String {
3419 use std::hash::{BuildHasher, Hasher};
3420 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3421 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3422 let t = std::time::SystemTime::now()
3423 .duration_since(std::time::UNIX_EPOCH)
3424 .map(|d| d.as_nanos() as u64)
3425 .unwrap_or(0);
3426 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
3427 h1.write_u64(n);
3428 h1.write_u64(t);
3429 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
3430 h2.write_u64(t.rotate_left(17));
3431 h2.write_u64(n);
3432 format!("{:016x}{:016x}", h1.finish(), h2.finish())
3433}
3434
3435#[derive(Clone)]
3438struct Envelope {
3439 id: String,
3440 created: u64,
3441}
3442
3443impl Envelope {
3444 fn new(chat: bool) -> Self {
3445 Envelope {
3446 id: format!(
3447 "{}-{}",
3448 if chat { "chatcmpl" } else { "cmpl" },
3449 gen_hex128()
3450 ),
3451 created: std::time::SystemTime::now()
3452 .duration_since(std::time::UNIX_EPOCH)
3453 .map(|d| d.as_secs())
3454 .unwrap_or(0),
3455 }
3456 }
3457
3458 fn capture_child(&self, index: usize) -> Self {
3475 Envelope {
3476 id: format!("{}.{index}", self.id),
3477 created: self.created,
3478 }
3479 }
3480
3481 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
3483 v["id"] = json!(self.id);
3484 v["created"] = json!(self.created);
3485 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
3486 v
3487 }
3488}
3489
3490fn with_request_id(id: &str, mut resp: Response) -> Response {
3492 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
3493 resp.headers_mut()
3494 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
3495 }
3496 resp
3497}
3498
3499fn openai_compat() -> bool {
3507 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3508 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
3509 Ok("openai") => true,
3510 Ok(_) => false,
3511 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
3512 })
3513}
3514
3515fn cache_namespace(cache_salt: &Option<String>) -> String {
3527 cache_salt.clone().unwrap_or_default()
3528}
3529
3530const CACHE_SALT_MAX_BYTES: usize = 64;
3531
3532fn validate_cache_namespace(
3533 cache_salt: &Option<String>,
3534 keyring_configured: bool,
3535) -> Result<String, &'static str> {
3536 let raw = cache_namespace(cache_salt);
3537 if raw.len() > CACHE_SALT_MAX_BYTES {
3538 return Err("cache_salt must be at most 64 bytes");
3539 }
3540 if !keyring_configured && raw.starts_with("t:") {
3541 return Err("cache_salt must not use the reserved t: prefix without a keyring");
3542 }
3543 if !raw
3544 .bytes()
3545 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
3546 {
3547 return Err("cache_salt contains unsupported characters");
3548 }
3549 Ok(raw)
3550}
3551
3552fn affinity_key(
3568 session_id: &Option<String>,
3569 user: &Option<String>,
3570 headers: &axum::http::HeaderMap,
3571) -> Result<Option<String>, String> {
3572 let clean = |s: &str| -> Result<Option<String>, String> {
3573 let t = s.trim();
3574 if t.is_empty() {
3575 Ok(None)
3576 } else if t.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3577 Err(format!(
3578 "session identity must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3579 ))
3580 } else if t.chars().any(char::is_control) {
3581 Err("session identity must not contain control characters".into())
3582 } else {
3583 Ok(Some(t.to_string()))
3584 }
3585 };
3586 if let Some(value) = session_id.as_deref()
3587 && let Some(value) = clean(value)?
3588 {
3589 return Ok(Some(value));
3590 }
3591 if let Some(value) = user.as_deref()
3592 && let Some(value) = clean(value)?
3593 {
3594 return Ok(Some(value));
3595 }
3596 match headers.get("x-session-id") {
3597 Some(value) => clean(
3598 value
3599 .to_str()
3600 .map_err(|_| "x-session-id must contain visible ASCII or UTF-8 text")?,
3601 ),
3602 None => Ok(None),
3603 }
3604}
3605
3606fn validate_client_identifier(value: Option<&str>, name: &str) -> Result<(), String> {
3607 let Some(value) = value else {
3608 return Ok(());
3609 };
3610 if value.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3611 return Err(format!(
3612 "{name} must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3613 ));
3614 }
3615 if value.chars().any(char::is_control) {
3616 return Err(format!("{name} must not contain control characters"));
3617 }
3618 Ok(())
3619}
3620
3621fn error_body(
3626 message: &str,
3627 etype: &str,
3628 param: Option<&str>,
3629 code: Option<&str>,
3630) -> serde_json::Value {
3631 json!({ "error": {
3632 "message": message,
3633 "type": etype,
3634 "param": param,
3635 "code": code,
3636 } })
3637}
3638
3639fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3640 error_response_coded(status, message, etype, param, None)
3641}
3642
3643fn error_response_coded(
3648 status: StatusCode,
3649 message: &str,
3650 etype: &str,
3651 param: Option<&str>,
3652 code: Option<&str>,
3653) -> Response {
3654 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3655 if status.is_client_error()
3656 && status != StatusCode::TOO_MANY_REQUESTS
3657 && status != StatusCode::REQUEST_TIMEOUT
3658 && status != StatusCode::CONFLICT
3659 {
3660 resp.headers_mut().insert(
3661 "x-should-retry",
3662 axum::http::HeaderValue::from_static("false"),
3663 );
3664 }
3665 resp
3666}
3667
3668fn bad_request(message: &str, param: Option<&str>) -> Response {
3669 error_response(
3670 StatusCode::BAD_REQUEST,
3671 message,
3672 "invalid_request_error",
3673 param,
3674 )
3675}
3676
3677const 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>) {
3706 use worker::ErrClass as C;
3707 match class {
3708 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3709 C::ContextLength => (
3710 StatusCode::BAD_REQUEST,
3711 "invalid_request_error",
3712 Some("context_length_exceeded"),
3713 ),
3714 C::ModelNotFound => (
3715 StatusCode::BAD_REQUEST,
3716 "invalid_request_error",
3717 Some("model_not_found"),
3718 ),
3719 C::RateLimit => (
3720 StatusCode::TOO_MANY_REQUESTS,
3721 "rate_limit_error",
3722 Some("rate_limit_exceeded"),
3723 ),
3724 C::Overloaded => (
3725 StatusCode::SERVICE_UNAVAILABLE,
3726 "server_error",
3727 Some("overloaded"),
3728 ),
3729 C::Engine => (
3730 StatusCode::INTERNAL_SERVER_ERROR,
3731 "server_error",
3732 Some("engine_error"),
3733 ),
3734 }
3735}
3736
3737fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3739 use worker::ErrClass as C;
3740 match class {
3741 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3742 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3743 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3747 }
3748}
3749
3750fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3753 let (_, etype, code) = class_http(e.class);
3754 error_body(&e.message, etype, e.param, code)
3755}
3756
3757fn engine_error_response(e: &worker::EngineError) -> Response {
3763 engine_error_response_with_retry_after(
3764 e,
3765 e.retry_after_s.or_else(|| class_retry_after_s(e.class)),
3766 )
3767}
3768
3769fn engine_error_response_with_retry_after(
3770 e: &worker::EngineError,
3771 retry_after_s: Option<u64>,
3772) -> Response {
3773 let (status, _, _) = class_http(e.class);
3774 let resp = (status, Json(engine_error_body(e))).into_response();
3775 retry_contract_response(resp, retry_after_s)
3776}
3777
3778fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3780 let status = resp.status();
3781 let h = resp.headers_mut();
3782 match retry_after_s {
3783 Some(secs) => {
3784 let secs = secs.clamp(1, 60);
3786 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3787 h.insert(axum::http::header::RETRY_AFTER, v);
3788 }
3789 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3790 h.insert("retry-after-ms", v);
3791 }
3792 }
3793 None if status.is_client_error() => {
3794 h.insert(
3797 "x-should-retry",
3798 axum::http::HeaderValue::from_static("false"),
3799 );
3800 }
3801 None => {}
3802 }
3803 resp
3804}
3805
3806fn worker_unavailable_response() -> Response {
3807 engine_error_response_with_retry_after(
3808 &worker::EngineError::overloaded("worker unavailable"),
3809 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3810 )
3811}
3812
3813fn stop_reason_to_finish(r: &str) -> &'static str {
3814 match r {
3815 "Eos" | "Callback" => "stop",
3816 "MaxNew" | "ContextFull" => "length",
3817 _ => "stop",
3818 }
3819}
3820
3821fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3825 match v {
3826 serde_json::Value::Null => Ok(String::new()),
3827 serde_json::Value::String(s) => Ok(s.clone()),
3828 serde_json::Value::Array(parts) => {
3829 let mut out = String::new();
3830 for p in parts {
3831 match p.get("type").and_then(|t| t.as_str()) {
3832 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3833 Some(t) => out.push_str(t),
3834 None => return Err("content part has no text field".into()),
3835 },
3836 Some(other) => {
3837 return Err(format!(
3838 "unsupported content part type {other:?} (text only)"
3839 ));
3840 }
3841 }
3842 }
3843 Ok(out)
3844 }
3845 _ => Err("content must be a string, null, or an array of text parts".into()),
3846 }
3847}
3848
3849pub(crate) static VISION_PLACEMENT_SERVING: std::sync::atomic::AtomicBool =
3867 std::sync::atomic::AtomicBool::new(true);
3868
3869fn vision_placement_serving() -> bool {
3870 VISION_PLACEMENT_SERVING.load(std::sync::atomic::Ordering::Acquire)
3871}
3872
3873fn vision_media_admissible(placement: bool, kind: &str) -> Result<(), String> {
3879 if placement {
3880 Ok(())
3881 } else {
3882 Err(format!(
3883 "{kind} input is not enabled on this deployment (vision overlay placement \
3884 inadmissible at boot: see the worker's IMAGE INPUT DISABLED line)"
3885 ))
3886 }
3887}
3888
3889fn vision_placement_admits(kind: &str) -> Result<(), String> {
3890 vision_media_admissible(vision_placement_serving(), kind)
3891}
3892
3893fn vision_enabled() -> bool {
3897 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3898 *ON.get_or_init(|| {
3899 std::env::var("MEMRA_VISION_DIR").is_ok()
3900 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3901 })
3902}
3903
3904fn gemma_vision_enabled() -> bool {
3909 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3910 *ON.get_or_init(|| {
3911 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3912 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3913 })
3914}
3915
3916pub(crate) static GLM5_VISION_SERVING: std::sync::atomic::AtomicBool =
3926 std::sync::atomic::AtomicBool::new(false);
3927
3928fn glm5_vision_enabled() -> bool {
3931 GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)
3932}
3933
3934fn step_vision_enabled() -> bool {
3942 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3943 *ON.get_or_init(|| {
3944 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3945 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3946 })
3947}
3948
3949const VISION_MAX_IMAGES: usize = 8;
3951
3952pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3958 std::sync::atomic::AtomicUsize::new(0);
3959pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3963 tokio::sync::Semaphore::const_new(1);
3964
3965#[allow(clippy::result_large_err)]
3968pub(crate) fn try_vision_preprocess(
3969 required: bool,
3970) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3971 try_vision_preprocess_with(required, &VISION_PREPROCESS_SEMAPHORE)
3972}
3973
3974#[allow(clippy::result_large_err)]
3975fn try_vision_preprocess_with(
3976 required: bool,
3977 semaphore: &'static tokio::sync::Semaphore,
3978) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3979 if !required {
3980 return Ok(None);
3981 }
3982 match semaphore.try_acquire() {
3983 Ok(permit) => Ok(Some(permit)),
3984 Err(tokio::sync::TryAcquireError::NoPermits) => Err(retry_contract_response(
3985 error_response_coded(
3986 StatusCode::TOO_MANY_REQUESTS,
3987 "vision preprocessing is busy",
3988 "rate_limit_error",
3989 Some("messages"),
3990 Some("vision_preprocess_busy"),
3991 ),
3992 Some(BODY_ADMISSION_RETRY_AFTER_S),
3993 )),
3994 Err(tokio::sync::TryAcquireError::Closed) => Err(error_response_coded(
3995 StatusCode::SERVICE_UNAVAILABLE,
3996 "vision preprocessing is unavailable",
3997 "server_error",
3998 Some("messages"),
3999 Some("vision_preprocess_unavailable"),
4000 )),
4001 }
4002}
4003
4004pub(crate) struct VisionMemoryPermit {
4005 bytes: usize,
4006}
4007
4008#[derive(Debug)]
4009pub(crate) enum VisionMemoryError {
4010 Request(String),
4011 Capacity(String),
4012}
4013
4014impl Drop for VisionMemoryPermit {
4015 fn drop(&mut self) {
4016 if self.bytes != 0 {
4017 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
4018 }
4019 }
4020}
4021
4022fn try_reserve_vision_memory(
4023 bytes: usize,
4024) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
4025 if bytes == 0 {
4026 return Ok(None);
4027 }
4028 if bytes > MAX_VISION_PATCH_BYTES {
4029 return Err(VisionMemoryError::Request(format!(
4030 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
4031 MAX_VISION_PATCH_BYTES / (1024 * 1024)
4032 )));
4033 }
4034 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
4035 loop {
4036 let Some(next) = in_use.checked_add(bytes) else {
4037 return Err(VisionMemoryError::Capacity(
4038 "vision patch memory reservation overflowed".into(),
4039 ));
4040 };
4041 if next > MAX_VISION_PATCH_BYTES {
4042 return Err(VisionMemoryError::Capacity(format!(
4043 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
4044 in_use / (1024 * 1024),
4045 bytes / (1024 * 1024)
4046 )));
4047 }
4048 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
4049 in_use,
4050 next,
4051 std::sync::atomic::Ordering::AcqRel,
4052 std::sync::atomic::Ordering::Acquire,
4053 ) {
4054 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
4055 Err(actual) => in_use = actual,
4056 }
4057 }
4058}
4059
4060pub(crate) fn vision_memory_error_response(
4061 error: VisionMemoryError,
4062 param: Option<&str>,
4063) -> Response {
4064 match error {
4065 VisionMemoryError::Request(message) => bad_request(&message, param),
4066 VisionMemoryError::Capacity(message) => retry_contract_response(
4067 error_response_coded(
4068 StatusCode::SERVICE_UNAVAILABLE,
4069 &message,
4070 "server_error",
4071 None,
4072 Some("vision_memory_busy"),
4073 ),
4074 Some(RETRY_AFTER_S_OVERLOADED),
4075 ),
4076 }
4077}
4078
4079enum PendingVisionUnit {
4085 Still {
4086 bytes: Vec<u8>,
4087 gh: usize,
4088 gw: usize,
4089 },
4090 Video {
4091 bytes: Vec<u8>,
4092 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
4093 video: usize,
4094 },
4095}
4096
4097struct PendingGemmaImage {
4099 bytes: Vec<u8>,
4100 gw: usize,
4101 gh: usize,
4102}
4103
4104struct PendingGlm5Image {
4107 bytes: Vec<u8>,
4108 gh: usize,
4109 gw: usize,
4110}
4111
4112struct PendingStepImage {
4115 bytes: Vec<u8>,
4116 plan: memra_engine::vision_step::StepImagePlan,
4117}
4118
4119fn content_to_text_vision_step(
4129 v: &serde_json::Value,
4130 step_images: &mut Vec<PendingStepImage>,
4131) -> Result<String, String> {
4132 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
4133 let parts = match v {
4134 serde_json::Value::Array(parts) => parts,
4135 _ => return content_to_text(v),
4136 };
4137 let mut out = String::new();
4138 let mut needs_sep = false;
4139 for p in parts {
4140 match p.get("type").and_then(|t| t.as_str()) {
4141 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
4142 Some(t) => {
4143 if needs_sep {
4144 out.push(' ');
4145 }
4146 out.push_str(t);
4147 needs_sep = true;
4148 }
4149 None => return Err("content part has no text field".into()),
4150 },
4151 Some("image_url") => {
4152 vision_placement_admits("image")?;
4153 let url = p
4154 .get("image_url")
4155 .and_then(|u| {
4156 if u.is_string() {
4157 u.as_str()
4158 } else {
4159 u.get("url").and_then(|x| x.as_str())
4160 }
4161 })
4162 .ok_or("image_url part has no url")?;
4163 if !url.starts_with("data:") {
4164 return Err(
4165 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4166 );
4167 }
4168 if step_images.len() >= VISION_MAX_IMAGES {
4169 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4170 }
4171 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4175 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
4176 let plan = memra_engine::vision_step::step_plan_image(&bytes)
4177 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
4178 for i in 0..plan.n_tiles {
4179 out.push_str("<patch_start>");
4180 for _ in 0..SV_TILE_ROWS {
4181 out.push_str("<im_patch>");
4182 }
4183 out.push_str("<patch_end>");
4184 if plan.newline_mask[i] {
4185 out.push_str("<patch_newline>");
4186 }
4187 }
4188 out.push_str("<im_start>");
4189 for _ in 0..SV_MAIN_ROWS {
4190 out.push_str("<im_patch>");
4191 }
4192 out.push_str("<im_end>");
4193 step_images.push(PendingStepImage { bytes, plan });
4194 needs_sep = false;
4195 }
4196 Some("video_url") => {
4197 return Err("step37 has no video input (image-only processor)".into());
4198 }
4199 Some(other) => {
4200 return Err(format!("unsupported content part type {other:?}"));
4201 }
4202 }
4203 }
4204 Ok(out)
4205}
4206
4207fn content_to_text_vision(
4216 v: &serde_json::Value,
4217 images: &mut Vec<PendingVisionUnit>,
4218 gemma_images: &mut Vec<PendingGemmaImage>,
4219 glm5_images: &mut Vec<PendingGlm5Image>,
4220 step_images: &mut Vec<PendingStepImage>,
4221 next_video: &mut usize,
4222) -> Result<String, String> {
4223 if step_vision_enabled() {
4227 return content_to_text_vision_step(v, step_images);
4228 }
4229 let parts = match v {
4230 serde_json::Value::Array(parts) => parts,
4231 _ => return content_to_text(v),
4232 };
4233 let mut out = String::new();
4234 for p in parts {
4235 match p.get("type").and_then(|t| t.as_str()) {
4236 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
4237 Some(t) => out.push_str(t),
4238 None => return Err("content part has no text field".into()),
4239 },
4240 Some("image_url") if glm5_vision_enabled() => {
4241 let url = p
4242 .get("image_url")
4243 .and_then(|u| {
4244 if u.is_string() {
4245 u.as_str()
4246 } else {
4247 u.get("url").and_then(|x| x.as_str())
4248 }
4249 })
4250 .ok_or("image_url part has no url")?;
4251 if !url.starts_with("data:") {
4252 return Err(
4253 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4254 );
4255 }
4256 if glm5_images.len() >= VISION_MAX_IMAGES {
4257 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4258 }
4259 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4263 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4264 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes)
4265 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4266 out.push_str("<|begin_of_image|>");
4270 for _ in 0..memra_engine::vision_glm5::n_merged_for_grid(gh, gw) {
4271 out.push_str("<|image|>");
4272 }
4273 out.push_str("<|end_of_image|>");
4274 glm5_images.push(PendingGlm5Image { bytes, gh, gw });
4275 }
4276 Some("video_url") if glm5_vision_enabled() => {
4277 return Err(
4278 "glm5 video input is not served (tensor census only; image input is the \
4279 supported surface)"
4280 .into(),
4281 );
4282 }
4283 Some("image_url") if gemma_vision_enabled() => {
4284 vision_placement_admits("image")?;
4285 let url = p
4286 .get("image_url")
4287 .and_then(|u| {
4288 if u.is_string() {
4289 u.as_str()
4290 } else {
4291 u.get("url").and_then(|x| x.as_str())
4292 }
4293 })
4294 .ok_or("image_url part has no url")?;
4295 if !url.starts_with("data:") {
4296 return Err(
4297 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4298 );
4299 }
4300 if gemma_images.len() >= VISION_MAX_IMAGES {
4301 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4302 }
4303 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
4307 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4308 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
4309 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4310 out.push_str("<|image>");
4312 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
4313 out.push_str("<|image|>");
4314 }
4315 out.push_str("<image|>");
4316 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
4317 }
4318 Some("image_url") => {
4319 if !vision_enabled() {
4320 return Err("image input is not enabled on this deployment".into());
4321 }
4322 vision_placement_admits("image")?;
4323 let url = p
4324 .get("image_url")
4325 .and_then(|u| {
4326 if u.is_string() {
4327 u.as_str()
4328 } else {
4329 u.get("url").and_then(|x| x.as_str())
4330 }
4331 })
4332 .ok_or("image_url part has no url")?;
4333 if !url.starts_with("data:") {
4334 return Err(
4335 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4336 );
4337 }
4338 if images
4339 .iter()
4340 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
4341 .count()
4342 >= VISION_MAX_IMAGES
4343 {
4344 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4345 }
4346 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4351 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4352 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
4353 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4354 out.push_str("<|vision_start|>");
4355 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
4356 out.push_str("<|image_pad|>");
4357 }
4358 out.push_str("<|vision_end|>");
4359 images.push(PendingVisionUnit::Still { bytes, gh, gw });
4360 }
4361 Some("video_url") if gemma_vision_enabled() => {
4362 return Err("gemma-4 has no video input (image-only projector)".into());
4363 }
4364 Some("video_url") => {
4365 if !vision_enabled() {
4366 return Err("video input is not enabled on this deployment".into());
4367 }
4368 vision_placement_admits("video")?;
4369 let url = p
4370 .get("video_url")
4371 .and_then(|u| {
4372 if u.is_string() {
4373 u.as_str()
4374 } else {
4375 u.get("url").and_then(|x| x.as_str())
4376 }
4377 })
4378 .ok_or("video_url part has no url")?;
4379 if !url.starts_with("data:") {
4380 return Err(
4381 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4382 );
4383 }
4384 if *next_video >= 2 {
4385 return Err("too many videos (max 2)".into());
4386 }
4387 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
4390 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
4391 .map_err(|e| format!("video: {e}"))?;
4392 let vidx = *next_video;
4393 *next_video += 1;
4394 for group in &vid.groups {
4396 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
4397 out.push_str("<|vision_start|>");
4398 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
4399 out.push_str("<|video_pad|>");
4400 }
4401 out.push_str("<|vision_end|>");
4402 }
4403 images.push(PendingVisionUnit::Video {
4406 bytes,
4407 groups: vid.groups,
4408 video: vidx,
4409 });
4410 }
4411 Some(other) => {
4412 return Err(format!("unsupported content part type {other:?}"));
4413 }
4414 }
4415 }
4416 Ok(out)
4417}
4418
4419fn pyjson(v: &serde_json::Value, out: &mut String) {
4423 match v {
4424 serde_json::Value::Object(m) => {
4425 out.push('{');
4426 for (i, (k, val)) in m.iter().enumerate() {
4427 if i > 0 {
4428 out.push_str(", ");
4429 }
4430 out.push_str(&serde_json::Value::String(k.clone()).to_string());
4431 out.push_str(": ");
4432 pyjson(val, out);
4433 }
4434 out.push('}');
4435 }
4436 serde_json::Value::Array(a) => {
4437 out.push('[');
4438 for (i, val) in a.iter().enumerate() {
4439 if i > 0 {
4440 out.push_str(", ");
4441 }
4442 pyjson(val, out);
4443 }
4444 out.push(']');
4445 }
4446 scalar => out.push_str(&scalar.to_string()),
4447 }
4448}
4449
4450fn pyjson_str(v: &serde_json::Value) -> String {
4451 let mut s = String::new();
4452 pyjson(v, &mut s);
4453 s
4454}
4455
4456#[allow(clippy::too_many_arguments)] fn sampler_config(
4464 temperature: f32,
4465 top_k: usize,
4466 top_p: f32,
4467 min_p: f32,
4468 frequency_penalty: f32,
4469 presence_penalty: f32,
4470 repetition_penalty: f32,
4471 seed: Option<u64>,
4472) -> SamplerConfig {
4473 let penalties_on =
4474 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
4475 SamplerConfig {
4476 temperature,
4477 top_k,
4478 top_p,
4479 min_p,
4480 penalty_last_n: if penalties_on {
4481 memra_engine::spec::PEN_WINDOW_MAX
4482 } else {
4483 0
4484 },
4485 penalty_repeat: repetition_penalty,
4486 penalty_freq: frequency_penalty,
4487 penalty_present: presence_penalty,
4488 seed: seed.unwrap_or_else(fresh_seed),
4491 }
4492}
4493
4494fn fresh_seed() -> u64 {
4499 use std::sync::atomic::{AtomicU64, Ordering};
4500 static COUNTER: AtomicU64 = AtomicU64::new(0);
4501 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
4502 let nanos = std::time::SystemTime::now()
4503 .duration_since(std::time::UNIX_EPOCH)
4504 .map(|d| d.as_nanos() as u64)
4505 .unwrap_or(0);
4506 let mut z = nanos
4507 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
4508 .wrapping_add(0x9E3779B97F4A7C15);
4509 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
4510 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
4511 z ^= z >> 31;
4512 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
4515}
4516
4517fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
4522 for (param, present, why) in fields {
4523 if *present {
4524 return Err((format!("{param} is not supported{why}"), param.to_string()));
4525 }
4526 }
4527 Ok(())
4528}
4529
4530#[derive(PartialEq)]
4531enum ToolChoice {
4532 Auto,
4533 None,
4534}
4535
4536fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
4537 match v {
4538 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
4539 Some(serde_json::Value::String(s)) => match s.as_str() {
4540 "auto" => Ok(ToolChoice::Auto),
4541 "none" => Ok(ToolChoice::None),
4542 "required" => Err("tool_choice \"required\" is not supported (no constrained \
4543 decoding); use \"auto\""
4544 .into()),
4545 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
4546 },
4547 Some(serde_json::Value::Object(_)) => {
4548 Err("named-function tool_choice is not supported; use \"auto\"".into())
4549 }
4550 Some(other) => Err(format!("bad tool_choice: {other}")),
4551 }
4552}
4553
4554fn parse_think(
4604 reasoning_effort: &Option<String>,
4605 reasoning: &Option<serde_json::Value>,
4606 vllm_switch: Option<bool>,
4607 suppress_switch: Option<bool>,
4608 default_effort: Option<&str>,
4609 max_tier: bool,
4610) -> Result<(ThinkMode, Option<String>, bool), String> {
4611 let mut effort = reasoning_effort.clone();
4612 let ReasoningObject {
4613 mut enabled,
4614 effort: object_effort,
4615 exclude,
4616 } = parse_reasoning_object(reasoning)?;
4617 if let Some(e) = object_effort {
4618 effort = Some(e);
4619 }
4620 match (enabled, vllm_switch) {
4625 (Some(a), Some(b)) if a != b => {
4626 return Err(format!(
4627 "contradictory reasoning switches: reasoning.enabled={a} and \
4628 enable_thinking={b} — send one"
4629 ));
4630 }
4631 (None, Some(b)) => enabled = Some(b),
4632 _ => {}
4633 }
4634 let suppress = match (exclude, suppress_switch) {
4646 (Some(true), _) | (_, Some(false)) => Some(false),
4647 _ => None,
4648 };
4649 match (enabled, suppress) {
4650 (Some(true), Some(false)) => {
4651 return Err(
4652 "contradictory reasoning switches: reasoning is enabled but \
4653 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
4654 on this server not delivering reasoning means not generating it, so send one"
4655 .into(),
4656 );
4657 }
4658 (None, Some(b)) => enabled = Some(b),
4659 _ => {}
4660 }
4661 let client_explicit = effort.is_some() || enabled.is_some();
4665 if effort.is_none() && enabled.is_none() {
4670 effort = default_effort.map(str::to_string);
4671 }
4672 let effort_arm = match effort.as_deref() {
4677 None => None,
4678 Some(raw) => {
4679 let level = canonical_effort_for(raw, max_tier).ok_or_else(|| {
4680 format!(
4681 "bad reasoning_effort {raw:?} \
4682 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
4683 highest level this model's template distinguishes)"
4684 )
4685 })?;
4686 Some(match level {
4687 "none" | "minimal" => (ThinkMode::NoThink, "low"),
4688 "low" => (ThinkMode::Think, "low"),
4689 "medium" => (ThinkMode::Think, "medium"),
4690 "max" => (ThinkMode::Think, "max"),
4691 _ => (ThinkMode::Think, "high"),
4692 })
4693 }
4694 };
4695 let (think, level) = match (enabled, effort_arm) {
4696 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
4699 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
4700 (None, Some((think, level))) => (think, Some(level.to_string())),
4701 (None, None) => (ThinkMode::Default, None),
4702 };
4703 Ok((think, level, client_explicit))
4704}
4705
4706struct ReasoningObject {
4708 enabled: Option<bool>,
4709 effort: Option<String>,
4710 exclude: Option<bool>,
4711}
4712
4713fn parse_reasoning_object(
4732 reasoning: &Option<serde_json::Value>,
4733) -> Result<ReasoningObject, String> {
4734 let mut out = ReasoningObject {
4735 enabled: None,
4736 effort: None,
4737 exclude: None,
4738 };
4739 let Some(v) = reasoning else { return Ok(out) };
4740 let obj = match v {
4741 serde_json::Value::Null => return Ok(out),
4742 serde_json::Value::Object(obj) => obj,
4743 _ => return Err("reasoning must be an object".into()),
4744 };
4745 for (key, value) in obj {
4746 match key.as_str() {
4754 "enabled" => {
4755 if !value.is_null() {
4756 out.enabled = Some(
4757 value
4758 .as_bool()
4759 .ok_or("reasoning.enabled must be true or false")?,
4760 );
4761 }
4762 }
4763 "exclude" => {
4764 if !value.is_null() {
4765 out.exclude = Some(
4766 value
4767 .as_bool()
4768 .ok_or("reasoning.exclude must be true or false")?,
4769 );
4770 }
4771 }
4772 "effort" => {
4773 if !value.is_null() {
4774 out.effort = Some(
4775 value
4776 .as_str()
4777 .ok_or("reasoning.effort must be a string")?
4778 .to_string(),
4779 );
4780 }
4781 }
4782 "max_tokens" => {
4783 return Err(
4784 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4785 are output tokens here, and max_tokens is the ONE output budget covering \
4786 reasoning and content together — there is no separate reasoning budget to \
4787 spend against, so honouring this field is impossible rather than merely \
4788 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4789 reasoning.enabled:false) to spend less of it on reasoning"
4790 .into(),
4791 );
4792 }
4793 other => {
4794 return Err(format!(
4795 "reasoning.{other} is not a field this server implements (it would change \
4796 nothing about the request); the supported keys are enabled, effort and \
4797 exclude"
4798 ));
4799 }
4800 }
4801 }
4802 Ok(out)
4803}
4804
4805fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4826 let Some(v) = kwargs else { return Ok(None) };
4827 let obj = match v {
4828 serde_json::Value::Null => return Ok(None),
4829 serde_json::Value::Object(obj) => obj,
4830 _ => return Err("chat_template_kwargs must be an object".into()),
4831 };
4832 let mut switch = None;
4833 for (key, value) in obj {
4834 match key.as_str() {
4835 "enable_thinking" => {
4836 switch = Some(
4837 value
4838 .as_bool()
4839 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4840 );
4841 }
4842 "preserve_thinking" => {
4843 let preserve = value
4844 .as_bool()
4845 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4846 if !preserve {
4847 return Err(
4848 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4849 server: the renderer implements the vendor DEFAULT (replay every prior \
4850 assistant turn's <think> block, empty when no reasoning was sent) but \
4851 not the strip arm — serving replay bytes under a strip request would \
4852 misdescribe the prompt. Omit the flag or send true"
4853 .into(),
4854 );
4855 }
4856 }
4858 other => {
4859 return Err(format!(
4860 "chat_template_kwargs.{other} is not supported by this server's \
4861 template renderer (it would change nothing about the prompt); the only \
4862 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4863 refuses in both directions — see its own message)"
4864 ));
4865 }
4866 }
4867 }
4868 Ok(switch)
4869}
4870
4871fn resolve_vllm_think_switch(
4875 enable_thinking: Option<bool>,
4876 kwargs: &Option<serde_json::Value>,
4877) -> Result<Option<bool>, String> {
4878 let from_kwargs = parse_template_kwargs(kwargs)?;
4879 match (enable_thinking, from_kwargs) {
4880 (Some(a), Some(b)) if a != b => Err(format!(
4881 "contradictory reasoning switches: enable_thinking={a} and \
4882 chat_template_kwargs.enable_thinking={b} — send one"
4883 )),
4884 (Some(a), _) => Ok(Some(a)),
4885 (None, b) => Ok(b),
4886 }
4887}
4888
4889pub(crate) fn canonical_effort_for(value: &str, max_tier: bool) -> Option<&'static str> {
4909 match value {
4910 "none" => Some("none"),
4911 "minimal" => Some("minimal"),
4912 "low" => Some("low"),
4913 "medium" => Some("medium"),
4914 "high" => Some("high"),
4915 "xhigh" | "max" | "ultra" => Some(if max_tier { "max" } else { "high" }),
4921 _ => None,
4922 }
4923}
4924
4925pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4928 canonical_effort_for(value, false)
4929}
4930
4931fn json_to_val(v: &serde_json::Value) -> chat::Val {
4934 match v {
4935 serde_json::Value::Null => chat::Val::Null,
4936 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4937 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4938 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4939 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4940 serde_json::Value::Object(o) => chat::Val::Obj(
4943 o.iter()
4944 .map(|(k, val)| (k.clone(), json_to_val(val)))
4945 .collect(),
4946 ),
4947 }
4948}
4949
4950#[allow(clippy::type_complexity)]
4954fn prepare_tools(
4955 tools: &[serde_json::Value],
4956) -> Result<
4957 (
4958 Vec<String>,
4959 Vec<chat::Val>,
4960 HashMap<String, HashMap<String, String>>,
4961 ),
4962 String,
4963> {
4964 let mut tools_json = Vec::with_capacity(tools.len());
4965 let mut tools_struct = Vec::with_capacity(tools.len());
4966 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4967 for t in tools {
4968 let f = t
4969 .get("function")
4970 .ok_or("each tool needs a function object")?;
4971 let name = f
4972 .get("name")
4973 .and_then(|n| n.as_str())
4974 .ok_or("each tool needs function.name")?;
4975 let mut params: HashMap<String, String> = HashMap::new();
4976 if let Some(props) = f
4977 .get("parameters")
4978 .and_then(|p| p.get("properties"))
4979 .and_then(|p| p.as_object())
4980 {
4981 for (p, def) in props {
4982 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4983 params.insert(p.clone(), ty.to_string());
4984 }
4985 }
4986 }
4987 schemas.insert(name.to_string(), params);
4988 tools_json.push(pyjson_str(t));
4989 tools_struct.push(json_to_val(f));
4991 }
4992 Ok((tools_json, tools_struct, schemas))
4993}
4994
4995fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
5000 let parsed: serde_json::Value = match &tc.function.arguments {
5001 serde_json::Value::Null => json!({}),
5002 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
5003 serde_json::Value::String(s) => serde_json::from_str(s)
5004 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
5005 v @ serde_json::Value::Object(_) => v.clone(),
5006 _ => return Err("tool_calls arguments must be a JSON object".into()),
5007 };
5008 let obj = parsed
5009 .as_object()
5010 .ok_or("tool_calls arguments must decode to a JSON object")?;
5011 let params = obj
5012 .iter()
5013 .map(|(k, v)| {
5014 let rendered = match v {
5015 serde_json::Value::String(s) => s.clone(),
5016 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
5017 scalar => scalar.to_string(),
5018 };
5019 (k.clone(), rendered)
5020 })
5021 .collect();
5022 let args = obj
5025 .iter()
5026 .map(|(k, v)| (k.clone(), json_to_val(v)))
5027 .collect();
5028 Ok(TmplToolCall {
5029 name: tc.function.name.clone(),
5030 params,
5031 args,
5032 id: tc.id.clone(),
5033 })
5034}
5035
5036fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
5038 json!({ "id": c.id, "type": "function",
5039 "function": { "name": c.name, "arguments": c.arguments } })
5040}
5041
5042async fn serve_bounded_http_with_limits<F>(
5046 listener: tokio::net::TcpListener,
5047 app: Router,
5048 shutdown: F,
5049 header_read_timeout: std::time::Duration,
5050 max_connections: usize,
5051 connection_max_lifetime: std::time::Duration,
5052) -> std::io::Result<()>
5053where
5054 F: std::future::Future<Output = ()> + Send,
5055{
5056 let connections = Arc::new(tokio::sync::Semaphore::new(max_connections));
5057 let (connection_shutdown, _) = tokio::sync::watch::channel(false);
5058 let mut connection_tasks = tokio::task::JoinSet::new();
5059 let mut shutdown = Box::pin(shutdown);
5060
5061 loop {
5062 tokio::select! {
5063 _ = &mut shutdown => break,
5064 joined = connection_tasks.join_next(), if !connection_tasks.is_empty() => {
5065 if let Some(Err(error)) = joined {
5066 eprintln!("[server] connection task failed: {error}");
5067 }
5068 }
5069 accepted = listener.accept() => {
5070 let (stream, _) = match accepted {
5071 Ok(connection) => connection,
5072 Err(error) => {
5073 eprintln!("[server] accept failed: {error}");
5074 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
5075 continue;
5076 }
5077 };
5078 let permit = match connections.clone().try_acquire_owned() {
5079 Ok(permit) => permit,
5080 Err(_) => {
5081 drop(stream);
5082 continue;
5083 }
5084 };
5085 let service = app.clone().map_request(
5086 |request: hyper::Request<hyper::body::Incoming>| request.map(Body::new),
5087 );
5088 let service = hyper_util::service::TowerToHyperService::new(service);
5089 let io = hyper_util::rt::TokioIo::new(stream);
5090 let mut builder = hyper_util::server::conn::auto::Builder::new(
5091 hyper_util::rt::TokioExecutor::new(),
5092 );
5093 builder
5094 .http1()
5095 .timer(hyper_util::rt::TokioTimer::new())
5096 .header_read_timeout(header_read_timeout)
5097 .max_headers(64);
5098 builder
5099 .http2()
5100 .timer(hyper_util::rt::TokioTimer::new())
5101 .max_concurrent_streams(MAX_HTTP2_STREAMS_PER_CONNECTION)
5102 .keep_alive_interval(Some(std::time::Duration::from_secs(30)))
5103 .keep_alive_timeout(std::time::Duration::from_secs(10));
5104 let mut connection = Box::pin(builder
5105 .serve_connection_with_upgrades(io, service)
5106 .into_owned());
5107 let mut shutdown_rx = connection_shutdown.subscribe();
5108 connection_tasks.spawn(async move {
5109 let _permit = permit;
5110 tokio::select! {
5111 result = connection.as_mut() => {
5112 let _ = result;
5113 }
5114 _ = tokio::time::sleep(connection_max_lifetime) => {
5115 connection.as_mut().graceful_shutdown();
5120 let _ = connection.await;
5121 }
5122 _ = shutdown_rx.changed() => {
5123 connection.as_mut().graceful_shutdown();
5124 let _ = connection.await;
5125 }
5126 }
5127 });
5128 }
5129 }
5130 }
5131 drop(listener);
5132 let _ = connection_shutdown.send(true);
5133 let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
5134 while connection_tasks.join_next().await.is_some() {}
5135 })
5136 .await;
5137 if drained.is_err() {
5138 connection_tasks.abort_all();
5139 eprintln!("[server] WARN: HTTP connections exceeded the 5s graceful close deadline");
5140 }
5141 Ok(())
5142}
5143
5144async fn serve_bounded_http<F>(
5145 listener: tokio::net::TcpListener,
5146 app: Router,
5147 shutdown: F,
5148) -> std::io::Result<()>
5149where
5150 F: std::future::Future<Output = ()> + Send,
5151{
5152 serve_bounded_http_with_limits(
5153 listener,
5154 app,
5155 shutdown,
5156 HTTP1_HEADER_READ_TIMEOUT,
5157 MAX_HTTP_CONNECTIONS,
5158 HTTP_CONNECTION_MAX_LIFETIME,
5159 )
5160 .await
5161}
5162
5163#[tokio::main]
5164pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
5165 serve_with(ServerWiring::stock()).await
5166}
5167
5168enum MeteringWiring {
5170 Stock,
5174 Custom(metering::MeteringFactory),
5179}
5180
5181pub struct ServerWiring {
5185 metering: MeteringWiring,
5186 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
5189 claimed_env: Vec<&'static str>,
5193}
5194
5195impl ServerWiring {
5196 pub fn stock() -> Self {
5198 ServerWiring {
5199 metering: MeteringWiring::Stock,
5200 on_ready: None,
5201 claimed_env: Vec::new(),
5202 }
5203 }
5204
5205 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
5208 ServerWiring {
5209 metering: MeteringWiring::Custom(factory),
5210 on_ready: None,
5211 claimed_env: Vec::new(),
5212 }
5213 }
5214
5215 pub fn claiming(mut self, var: &'static str) -> Self {
5219 self.claimed_env.push(var);
5220 self
5221 }
5222
5223 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
5224 self.on_ready = Some(Box::new(hook));
5225 self
5226 }
5227}
5228
5229pub struct RuntimeHandles {
5232 pub trim: TrimHandle,
5233 pub metadata_reload: MetadataReloadHandle,
5237 pub purge: PurgeHandle,
5240 pub kv_handoff: HostHandoffHandle,
5246 pub shutdown: tokio::sync::watch::Receiver<bool>,
5253}
5254
5255#[derive(Clone)]
5258pub struct TrimHandle {
5259 cmd_tx: Sender<Cmd>,
5260}
5261
5262impl TrimHandle {
5263 pub async fn trim(&self) -> Result<serde_json::Value, String> {
5265 let (tx, rx) = tokio::sync::oneshot::channel();
5266 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
5267 return Err("worker is down".into());
5268 }
5269 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5270 Ok(Ok(report)) => Ok(json!(report)),
5271 _ => Err("worker did not answer the trim within 30s".into()),
5272 }
5273 }
5274}
5275
5276#[derive(Clone)]
5285pub struct PurgeHandle {
5286 cmd_tx: Sender<Cmd>,
5287}
5288
5289impl PurgeHandle {
5290 pub async fn purge_tenant(&self, tenant: &str) -> Result<serde_json::Value, String> {
5292 let (tx, rx) = tokio::sync::oneshot::channel();
5293 let cmd = Cmd::PurgeTenantHost {
5294 tenant: tenant.to_string(),
5295 tx,
5296 };
5297 if self.cmd_tx.send(cmd).is_err() {
5298 return Err("worker is down".into());
5299 }
5300 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5301 Ok(Ok(report)) => Ok(json!(report)),
5302 _ => Err("worker did not answer the purge within 30s".into()),
5303 }
5304 }
5305}
5306
5307#[derive(Clone)]
5316pub struct HostHandoffHandle {
5317 cmd_tx: Sender<Cmd>,
5318}
5319
5320impl HostHandoffHandle {
5321 pub async fn export(&self, force: bool) -> Result<serde_json::Value, String> {
5324 let (tx, rx) = tokio::sync::oneshot::channel();
5325 if self
5326 .cmd_tx
5327 .send(Cmd::ExportHostHandoff { force, tx })
5328 .is_err()
5329 {
5330 return Err("worker is down".into());
5331 }
5332 match tokio::time::timeout(std::time::Duration::from_secs(900), rx).await {
5333 Ok(Ok(Ok(report))) => Ok(json!(report)),
5334 Ok(Ok(Err(refused))) => Err(refused),
5335 _ => Err("worker did not answer the export within 900s".into()),
5336 }
5337 }
5338
5339 pub async fn import(&self) -> Result<serde_json::Value, String> {
5342 let (tx, rx) = tokio::sync::oneshot::channel();
5343 if self.cmd_tx.send(Cmd::ImportHostHandoff { tx }).is_err() {
5344 return Err("worker is down".into());
5345 }
5346 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5347 Ok(Ok(Ok(start))) => Ok(json!(start)),
5348 Ok(Ok(Err(refused))) => Err(refused),
5349 _ => Err("worker did not answer the import within 30s".into()),
5350 }
5351 }
5352}
5353
5354pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
5355 let args: Vec<String> = std::env::args().skip(1).collect();
5358 if args.iter().any(|a| a == "--version" || a == "-V") {
5363 println!("memra-server {}", env!("CARGO_PKG_VERSION"));
5364 println!("system_fingerprint {SYSTEM_FINGERPRINT}");
5365 println!("build_id_src {BUILD_ID_SRC}");
5366 println!("git_sha {BUILD_GIT_SHA}");
5367 if !BUILD_ID_NOTE.is_empty() {
5368 println!("degraded {BUILD_ID_NOTE}");
5369 }
5370 return Ok(());
5371 }
5372 validate_stream_prefill_config().map_err(|message| {
5373 std::io::Error::new(
5374 std::io::ErrorKind::InvalidInput,
5375 format!("streaming prefill config: {message}"),
5376 )
5377 })?;
5378 if let Some(code) = auth::run_cli(&args) {
5379 std::process::exit(code);
5380 }
5381 eprintln!("{}", build_identity_line());
5385 if BUILD_ID_SRC != build_id::BUILD_ID_SRC_TREE {
5386 eprintln!(
5387 "[server] WARNING: build identity is DEGRADED: {BUILD_ID_NOTE}. \
5388 system_fingerprint {SYSTEM_FINGERPRINT} carries a version-only id, so it does \
5389 NOT identify the source this binary was compiled from and published \
5390 performance pins cannot be verified against it (darklanes \
5391 tools/check-claim-builds.mjs --live). Rebuild where the workspace source tree \
5392 is readable."
5393 );
5394 }
5395 auth::init_from_env();
5398 let api_auth = match ApiAuth::from_env() {
5399 Ok(auth) => auth,
5400 Err(err) => {
5401 eprintln!("[server] FATAL: {err}");
5402 std::process::exit(1);
5403 }
5404 };
5405 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
5406 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
5407 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
5408 Ok(resolved) => resolved,
5409 Err(err) => {
5410 eprintln!("[server] FATAL: {err}");
5411 std::process::exit(1);
5412 }
5413 };
5414 if let Err(message) = validate_bind_security(&addr, api_auth.configured(), allow_open_bind) {
5419 eprintln!("[server] FATAL: {message}");
5420 std::process::exit(1);
5421 }
5422 if !bind_loopback && !api_auth.configured() {
5423 eprintln!(
5424 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
5425 metrics remain bearer-protected"
5426 );
5427 }
5428 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
5429 Ok(token) if token.is_empty() => {
5430 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
5431 std::process::exit(1);
5432 }
5433 Ok(token) => Some(token),
5434 Err(std::env::VarError::NotPresent) => None,
5435 Err(std::env::VarError::NotUnicode(_)) => {
5436 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
5437 std::process::exit(1);
5438 }
5439 };
5440 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
5441
5442 let models = parse_models_config();
5443 let metadata_models: Arc<Vec<(String, String, Option<String>)>> = Arc::new(models.clone());
5446 let metadata_set = match load_openrouter_metadata(&models) {
5447 Ok(loaded) => loaded,
5448 Err(err) => {
5449 eprintln!("[server] FATAL: {err}");
5450 std::process::exit(1);
5451 }
5452 };
5453 let metadata_path = std::env::var("MEMRA_MODEL_METADATA")
5456 .ok()
5457 .map(std::path::PathBuf::from);
5458 let metering_obj: Option<Arc<dyn metering::Metering>> = {
5464 let factory = match wiring.metering {
5465 MeteringWiring::Stock => None,
5466 MeteringWiring::Custom(factory) => Some(factory),
5467 };
5468 for deployment_only in [
5469 "MEMRA_REQUEST_LEDGER",
5470 "MEMRA_TENANT_BUDGETS",
5471 "MEMRA_ADMIN_ADDR",
5472 "MEMRA_ADMIN_TOKEN_FILE",
5473 "MEMRA_CAPTURE_DIR",
5474 ] {
5475 if std::env::var_os(deployment_only).is_some()
5476 && !wiring.claimed_env.contains(&deployment_only)
5477 {
5478 eprintln!(
5479 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
5480 build ships no accounting/admin/capture. Wire a Metering implementation \
5481 through ServerWiring and claim the vars it consumes."
5482 );
5483 std::process::exit(1);
5484 }
5485 }
5486 match factory {
5487 None => None,
5488 Some(factory) => {
5489 let model_ids: Vec<String> =
5490 models.iter().map(|(name, _, _)| name.clone()).collect();
5491 match factory(&metering::MeteringInit { models: &model_ids }) {
5492 Ok(metering_obj) => metering_obj,
5493 Err(err) => {
5494 eprintln!("[server] FATAL: metering wiring: {err}");
5495 std::process::exit(1);
5496 }
5497 }
5498 }
5499 }
5500 };
5501 let budget_tokenizers = if metering_obj
5502 .as_ref()
5503 .is_some_and(|manager| manager.enforces_limits())
5504 {
5505 match load_budget_tokenizers(&models) {
5506 Ok(tokenizers) => Some(tokenizers),
5507 Err(err) => {
5508 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
5509 std::process::exit(1);
5510 }
5511 }
5512 } else {
5513 None
5514 };
5515 eprintln!("[server] starting; models config = {models:?}");
5516
5517 let health_state = health::WorkerHealth::new();
5522 health::spawn_gpu_watch(health_state.clone());
5526 health::spawn_sd_watchdog(health_state.clone());
5527
5528 let (cmd_tx, model_names, caps, metrics, worker_thread) =
5530 match worker::spawn(models, health_state.clone()) {
5531 Ok(v) => v,
5532 Err(err) => {
5533 eprintln!("[server] FATAL: worker init failed: {err}");
5534 health_state.mark_dead(format!("worker init failed: {err}"));
5535 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
5536 std::process::exit(1);
5537 }
5538 };
5539 eprintln!("[server] worker ready; serving models: {model_names:?}");
5540
5541 let metadata_cell = Arc::new(RwLock::new(Arc::new(metadata_set)));
5544 let metadata_reload = MetadataReloadHandle {
5545 cell: metadata_cell.clone(),
5546 models: metadata_models,
5547 path: metadata_path,
5548 };
5549
5550 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
5557 if let Some(on_ready) = wiring.on_ready {
5558 on_ready(RuntimeHandles {
5559 trim: TrimHandle {
5560 cmd_tx: cmd_tx.clone(),
5561 },
5562 metadata_reload,
5563 purge: PurgeHandle {
5564 cmd_tx: cmd_tx.clone(),
5565 },
5566 kv_handoff: HostHandoffHandle {
5567 cmd_tx: cmd_tx.clone(),
5568 },
5569 shutdown: drain_shutdown_rx.clone(),
5570 });
5571 }
5572
5573 let bg_handle = darklane::spawn_from_env(health_state.clone());
5576 let bg_state = bg_handle.as_ref().map(|h| {
5577 let mode = darklane::BgConfig::from_env()
5578 .map(|c| c.yield_mode.as_str())
5579 .unwrap_or("stop");
5580 (h.state.clone(), mode)
5581 });
5582
5583 let state = AppState {
5584 cmd_tx,
5585 models: model_names,
5586 caps,
5587 openrouter_metadata: metadata_cell,
5588 metering: metering_obj,
5589 budget_tokenizers,
5590 api_auth,
5591 metrics_auth,
5592 metrics,
5593 inflight: Arc::new(Default::default()),
5594 tenant_inflight: Arc::new(Default::default()),
5595 health: health_state.clone(),
5596 bg: bg_state,
5597 };
5598 let inflight_handle = state.inflight.clone();
5599 let drain_metering = state.metering.clone();
5602 worker::register_http_inflight(state.inflight.clone());
5607 let app = Router::new()
5608 .route("/health", get(health_live))
5613 .route("/livez", get(health_live))
5614 .route("/readyz", get(health_ready))
5615 .route("/models", get(list_models))
5616 .route("/v1/models", get(list_models_v1))
5617 .route("/v1/models/*id", get(retrieve_model_v1))
5626 .route("/v1/auth/check", get(auth_check))
5627 .route("/v1/completions", post(completions_admitted))
5628 .route("/v1/embeddings", post(embed_api::embeddings_admitted))
5629 .route("/v1/rerank", post(embed_api::rerank_admitted))
5630 .route("/v1/chat/completions", post(chat_completions_admitted))
5631 .route("/v1/messages", post(anthropic::messages_admitted))
5635 .route("/v1/responses", post(responses_api::responses_admitted))
5636 .route("/metrics", get(get_metrics))
5637 .route("/yield/metrics", get(yield_metrics))
5638 .with_state(state.clone());
5639 let app = apply_body_limit(app);
5642 let app = app.layer(middleware::from_fn_with_state(
5645 state,
5646 authenticate_inference_before_body,
5647 ));
5648 let app = if ttft::enabled() {
5649 app.layer(middleware::from_fn(ttft_request_start))
5650 } else {
5651 app
5652 };
5653
5654 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
5655 eprintln!("[server] listening on http://{bind_addr}");
5656 drop(drain_shutdown_rx);
5657 health::sd_notify("READY=1\nSTATUS=serving");
5661 let inflight = inflight_handle;
5668 let signal_admin_shutdown = drain_shutdown_tx.clone();
5669 let serve_result = serve_bounded_http(listener, app, async move {
5670 let mut sigterm =
5671 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
5672 Ok(s) => s,
5673 Err(err) => {
5674 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
5675 std::future::pending::<()>().await;
5676 unreachable!()
5677 }
5678 };
5679 sigterm.recv().await;
5680 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
5681 let _ = signal_admin_shutdown.send(true);
5682 health::sd_notify(&format!(
5686 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
5687 (drain_deadline_s() + 5) * 1_000_000
5688 ));
5689 let n: usize = inflight
5690 .iter()
5691 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5692 .sum();
5693 eprintln!(
5694 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
5695 drain_deadline_s()
5696 );
5697 let deadline = std::time::Duration::from_secs(drain_deadline_s());
5698 let t0 = std::time::Instant::now();
5699 loop {
5700 let n: usize = inflight
5701 .iter()
5702 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5703 .sum();
5704 if n == 0 {
5705 eprintln!(
5706 "[server] drain complete in {:.1}s; exiting",
5707 t0.elapsed().as_secs_f64()
5708 );
5709 break;
5710 }
5711 if t0.elapsed() >= deadline {
5712 eprintln!(
5713 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
5714 drain_deadline_s()
5715 );
5716 if let Some(metering) = drain_metering.as_ref() {
5723 metering.drain_kill();
5724 }
5725 break;
5726 }
5727 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
5728 }
5729 })
5730 .await;
5731 let _ = drain_shutdown_tx.send(true);
5734 serve_result?;
5735 if let Some(h) = bg_handle {
5739 h.shutdown();
5740 }
5741 worker_thread.join().map_err(|_| {
5750 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
5751 })?;
5752 eprintln!("[server] GPU worker shutdown complete");
5753 Ok(())
5754}
5755
5756fn validate_model_path(path: &str) -> Result<(), String> {
5762 let p = std::path::Path::new(path);
5763 if !p.exists() {
5764 return Err(format!("model path {path:?} does not exist"));
5765 }
5766 if p.is_file() {
5767 return Ok(()); }
5769 if p.join("manifest.json").exists() {
5770 return Ok(()); }
5772 let has_st =
5773 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
5774 if !has_st {
5775 return Err(format!(
5776 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
5777 model.safetensors.index.json + config.json (HF safetensors dir), or \
5778 manifest.json (memra repack dir)"
5779 ));
5780 }
5781 if !p.join("config.json").exists() {
5782 return Err(format!(
5783 "model dir {path:?} has safetensors weights but no config.json"
5784 ));
5785 }
5786 Ok(())
5787}
5788
5789fn parse_models_config() -> Vec<(String, String, Option<String>)> {
5797 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
5798 let mut out = Vec::new();
5799 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
5800 if let Some((name, path)) = entry.split_once('=') {
5801 let (mpath, dpath) = match path.trim().split_once('+') {
5804 Some((m, d)) => (m.trim(), Some(d.trim())),
5805 None => (path.trim(), None),
5806 };
5807 let resolve = |p: &str| {
5808 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
5809 eprintln!("[server] FATAL: model {name:?}: {err}");
5810 std::process::exit(1);
5811 })
5812 };
5813 let mpath = resolve(mpath);
5814 if let Err(err) = validate_model_path(&mpath) {
5815 eprintln!("[server] FATAL: model {name:?}: {err}");
5816 std::process::exit(1);
5817 }
5818 let dpath = dpath.map(|d| {
5828 let d = resolve(d);
5829 let p = std::path::Path::new(&d);
5830 if !p.exists() {
5831 eprintln!(
5832 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
5833 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
5834 rather than serving plain decode under a config that asked \
5835 for speculative decoding."
5836 );
5837 std::process::exit(1);
5838 }
5839 if !p.is_file() {
5840 eprintln!(
5841 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
5842 file — a '+draft' attach must be a NextN/MTP GGUF file."
5843 );
5844 std::process::exit(1);
5845 }
5846 d
5847 });
5848 out.push((name.trim().to_string(), mpath, dpath));
5849 } else {
5850 eprintln!(
5851 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
5852 );
5853 }
5854 }
5855 if !out.is_empty() {
5856 return out;
5857 }
5858 }
5859 vec![
5861 (
5862 "main".into(),
5863 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
5864 None,
5865 ),
5866 (
5867 "judge".into(),
5868 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
5869 None,
5870 ),
5871 ]
5872}
5873
5874fn load_budget_tokenizers(
5875 models: &[(String, String, Option<String>)],
5876) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
5877 let mut tokenizers = HashMap::new();
5878 for (alias, path, _) in models {
5879 let path = std::path::Path::new(path);
5880 let tokenizer = if path.is_dir() {
5881 let tokenizer_dir = if path.join("manifest.json").exists() {
5882 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
5883 format!("model {alias:?}: open repack tokenizer source: {err}")
5884 })?;
5885 repack
5886 .source_dir()
5887 .filter(|source| source.join("tokenizer.json").exists())
5888 .unwrap_or(path)
5889 .to_path_buf()
5890 } else {
5891 path.to_path_buf()
5892 };
5893 Tokenizer::from_hf_dir(&tokenizer_dir)
5894 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5895 } else {
5896 let gguf = memra_gguf::GgufFile::open(path)
5897 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
5898 Tokenizer::from_gguf(&gguf)
5899 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5900 };
5901 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
5902 }
5903 Ok(Arc::new(tokenizers))
5904}
5905
5906fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5908 let s = st.health.snapshot();
5909 let mut v = json!({
5910 "status": status,
5911 "models": *st.models,
5912 "worker": {
5913 "phase": health::phase_name(s.phase),
5914 "beat_age_ms": s.beat_age_ms,
5915 "tick_max_ms": s.tick_max_ms,
5916 "stall_threshold_ms": s.stall_threshold_ms,
5917 "forward_progress_age_ms": s.forward_progress_age_ms,
5921 "prime_progress": s.progress.map(|p| json!({
5922 "rows": p.rows,
5923 "chunks": p.events,
5924 "age_ms": p.age_ms,
5925 })),
5926 "generation": s.generation,
5927 "xid_warnings": s.xid_warns,
5928 },
5929 });
5930 if let Some(d) = detail {
5931 v["detail"] = json!(d);
5932 }
5933 v
5934}
5935
5936fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5940 let mut v = health_payload(st, status, detail);
5941 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
5942 v
5943}
5944
5945async fn auth_check() -> impl IntoResponse {
5949 StatusCode::NO_CONTENT
5950}
5951
5952async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
5968 if draining() {
5969 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
5972 }
5973 match st.health.live() {
5974 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
5975 Err(why) => retry_contract_response(
5976 (
5977 StatusCode::SERVICE_UNAVAILABLE,
5978 Json(health_payload(&st, "unhealthy", Some(&why))),
5979 )
5980 .into_response(),
5981 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
5982 ),
5983 }
5984}
5985
5986async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
5997 let is_draining = draining();
5998 match st.health.ready(is_draining) {
5999 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
6000 Err(why) => retry_contract_response(
6001 (
6002 StatusCode::SERVICE_UNAVAILABLE,
6003 Json(readiness_payload(&st, "not_ready", Some(&why))),
6004 )
6005 .into_response(),
6006 Some(if is_draining {
6007 drain_deadline_s()
6008 } else {
6009 worker::WORKER_RESPAWN_BACKOFF_BASE_S
6010 }),
6011 ),
6012 }
6013}
6014
6015#[derive(Clone, Copy)]
6016struct DualPpMetricsSnapshot {
6017 stage_ns: [u64; 4],
6018 stage_samples: [usize; 4],
6019 dropped_timing_samples: usize,
6020 overlaps: usize,
6021 slot_pairs: usize,
6022 slot_uses: [usize; 2],
6023 slot_collisions: usize,
6024}
6025
6026impl DualPpMetricsSnapshot {
6027 fn current() -> Self {
6028 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
6029 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
6030 Self {
6031 stage_ns,
6032 stage_samples,
6033 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
6034 overlaps: memra_engine::pp::dual_pp_overlaps(),
6035 slot_pairs,
6036 slot_uses,
6037 slot_collisions,
6038 }
6039 }
6040
6041 fn populated(self) -> bool {
6042 self.stage_samples.iter().any(|&n| n > 0)
6043 || self.dropped_timing_samples > 0
6044 || self.slot_pairs > 0
6045 || self.slot_collisions > 0
6046 }
6047}
6048
6049fn insert_dual_pp_metrics(
6050 body: &mut serde_json::Value,
6051 metrics_scope: &MetricsScope,
6052 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
6053) {
6054 if !metrics_scope.operator() {
6057 return;
6058 }
6059 let snapshot = snapshot();
6060 if !snapshot.populated() {
6061 return;
6062 }
6063 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
6064 .iter()
6065 .enumerate()
6066 .map(|(i, name)| {
6067 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
6068 (
6069 name.to_string(),
6070 json!({
6071 "samples": snapshot.stage_samples[i],
6072 "total_ms": total_ms,
6073 "mean_ms": if snapshot.stage_samples[i] > 0 {
6074 total_ms / snapshot.stage_samples[i] as f64
6075 } else { 0.0 },
6076 }),
6077 )
6078 })
6079 .collect();
6080 body["dual_pp"] = json!({
6081 "overlaps": snapshot.overlaps,
6082 "slot_pairs": snapshot.slot_pairs,
6083 "slot_uses": snapshot.slot_uses,
6084 "slot_collisions": snapshot.slot_collisions,
6085 "cuda_event_spans": timings,
6086 "dropped_timing_samples": snapshot.dropped_timing_samples,
6087 });
6088}
6089
6090#[derive(Clone, Copy)]
6091struct PpWaveMetricsSnapshot {
6092 ticks: usize,
6093 cells: usize,
6094 overlaps: usize,
6095}
6096
6097impl PpWaveMetricsSnapshot {
6098 fn current() -> Self {
6099 let (ticks, cells, overlaps) = memra_engine::pp::pp_wave_snapshot();
6100 Self {
6101 ticks,
6102 cells,
6103 overlaps,
6104 }
6105 }
6106}
6107
6108fn insert_pp_wave_metrics(
6109 body: &mut serde_json::Value,
6110 metrics_scope: &MetricsScope,
6111 snapshot: impl FnOnce() -> PpWaveMetricsSnapshot,
6112) {
6113 if !metrics_scope.operator() {
6114 return;
6115 }
6116 let snapshot = snapshot();
6117 if snapshot.ticks == 0 && snapshot.cells == 0 {
6118 return;
6119 }
6120 body["pp_wave"] = json!({
6121 "ticks": snapshot.ticks,
6122 "cells": snapshot.cells,
6123 "overlaps": snapshot.overlaps,
6124 });
6125}
6126
6127fn insert_spec_acceptance_metrics(
6128 body: &mut serde_json::Value,
6129 metrics_scope: &MetricsScope,
6130 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
6131) {
6132 if !metrics_scope.operator() {
6135 return;
6136 }
6137 let snapshot = snapshot();
6138 if snapshot.is_empty() {
6139 return;
6140 }
6141
6142 let mut tau = serde_json::Map::new();
6143 let mut by_position = serde_json::Map::new();
6144 for (model, telemetry) in snapshot {
6145 if telemetry.rounds == 0 {
6146 continue;
6147 }
6148 let n_pos = telemetry
6149 .pos_drafted
6150 .iter()
6151 .rposition(|&n| n > 0)
6152 .map_or(0, |position| position + 1);
6153 tau.insert(model.clone(), json!(telemetry.tau()));
6154 by_position.insert(
6155 model,
6156 json!({
6157 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
6158 "rounds": telemetry.rounds,
6159 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
6160 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
6161 "accept_rate": (0..n_pos).map(|position| {
6162 let offered = telemetry.pos_drafted[position];
6163 if offered > 0 {
6164 telemetry.pos_accepted[position] as f64 / offered as f64
6165 } else {
6166 0.0
6167 }
6168 }).collect::<Vec<f64>>(),
6169 }),
6170 );
6171 }
6172 if !tau.is_empty() {
6173 body["spec_tau"] = serde_json::Value::Object(tau);
6174 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
6175 }
6176}
6177
6178fn insert_peer_probe_metrics(
6179 body: &mut serde_json::Value,
6180 metrics_scope: &MetricsScope,
6181 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
6182) {
6183 if !metrics_scope.operator() {
6186 return;
6187 }
6188 let snapshot = snapshot();
6189 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
6190 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
6191 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
6192 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
6193 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
6194 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
6195 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
6196}
6197
6198async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
6200 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
6201 Ok(scope) => scope,
6202 Err(response) => return response,
6203 };
6204 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
6205 let mut body = if metrics_scope.process_wide() {
6209 json!({
6210 "admitted": m.admitted,
6211 "completed": m.completed,
6212 "tokens_out": m.tokens_out,
6213 "step_p50_ms": m.step_p50_ms,
6214 "step_p99_ms": m.step_p99_ms,
6215 "prompt_tokens_in": m.prompt_tokens_in,
6217 "cached_tokens_in": m.cached_tokens_in,
6218 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
6221 "admission_session_defers": m.admission_session_defers,
6224 "admission_vram_defers": m.admission_vram_defers,
6225 "step_oom_parks": m.step_oom_parks,
6226 "continuation_pool_hits": m.continuation_pool_hits,
6227 "continuation_pool_evictions": m.continuation_pool_evictions,
6228 "plain_affinity_rewinds": m.plain_affinity_rewinds,
6229 "served_dspark": m.served_dspark,
6230 "served_spec": m.served_spec,
6231 "served_plain": m.served_plain,
6232 "spec_pool_hits": m.spec_pool_hits,
6233 "spec_pool_misses": m.spec_pool_misses,
6234 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
6235 "spec_pool_evictions": m.spec_pool_evictions,
6236 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
6239 })
6240 } else {
6241 json!({})
6242 };
6243 if metrics_scope.operator() {
6247 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
6248 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
6249 body["budget_source_reload_consecutive"] =
6250 json!(budget_health.source_reload_consecutive);
6251 body["budget_source_available"] = json!(budget_health.source_available);
6252 }
6253 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
6255 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
6256 } else {
6257 0.0
6258 });
6259 body["prefix_cache_hits"] = json!(m.prefix_hits);
6260 body["prefix_cache_misses"] = json!(m.prefix_misses);
6261 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
6262 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
6263 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
6264 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
6265 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
6266 body["prefix_host_entries"] = json!(m.prefix_host_entries);
6270 body["prefix_host_bytes"] = json!(m.prefix_host_bytes);
6271 body["prefix_host_demotions"] = json!(m.prefix_host_demotions);
6272 body["prefix_host_promotions"] = json!(m.prefix_host_promotions);
6273 body["prefix_host_demote_ms"] = json!(m.prefix_host_demote_ms);
6274 body["prefix_host_promote_ms"] = json!(m.prefix_host_promote_ms);
6275 body["prefix_host_rejected_allocs"] = json!(m.prefix_host_rejected_allocs);
6276 body["prefix_host_purges"] = json!(m.prefix_host_purges);
6277 body["prefix_host_purged_entries"] = json!(m.prefix_host_purged_entries);
6278 body["prefix_host_purged_bytes"] = json!(m.prefix_host_purged_bytes);
6279 body["prefix_host_tenant_rejects"] = json!(m.prefix_host_tenant_rejects);
6280 body["prefix_host_pause_demotes"] = json!(m.prefix_host_pause_demotes);
6284 body["prefix_host_pause_cancels"] = json!(m.prefix_host_pause_cancels);
6285 body["prefix_host_handoff_exports"] = json!(m.prefix_host_handoff_exports);
6286 body["prefix_host_handoff_imported_entries"] =
6287 json!(m.prefix_host_handoff_imported_entries);
6288 body["prefix_host_handoff_imported_bytes"] = json!(m.prefix_host_handoff_imported_bytes);
6289 body["prefix_host_handoff_skips"] = json!(m.prefix_host_handoff_skips);
6290 body["kv_flex_borrowed_bytes"] = json!(m.kv_flex_borrowed_bytes);
6295 body["kv_flex_sheds"] = json!(m.kv_flex_sheds);
6296 body["kv_flex_shed_ms"] = json!(m.kv_flex_shed_ms);
6297 body["lcp_histogram"] = json!({
6300 "edges": worker::LCP_HIST_EDGES.to_vec(),
6301 "counts": m.lcp_hist.to_vec(),
6302 });
6303 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
6307 body["prefix_cache_entries"] = json!(m.prefix_entries);
6308 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
6309 body["active_sessions"] = json!(m.active_sessions);
6310 body["queued_requests"] = json!(m.queued_requests);
6311 body["admission_inflight"] = json!(m.admission_inflight);
6315 body["admission_booked_bytes"] = json!(m.admission_booked_bytes);
6316 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
6317 body["spec_pool_entries"] = json!(m.spec_pool_entries);
6318 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
6319 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
6320 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
6321 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
6322 if !m.constraint_compiler_fail_closed.is_empty() {
6323 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
6324 m.constraint_compiler_fail_closed
6325 .iter()
6326 .map(|(model, gauge)| {
6327 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
6328 (model.clone(), json!(value))
6329 })
6330 .collect(),
6331 );
6332 }
6333 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
6334 }
6335 if !m.ns_tokens.is_empty() {
6340 let tenants: serde_json::Map<String, serde_json::Value> = m
6341 .ns_tokens
6342 .iter()
6343 .filter(|(ns, _)| metrics_scope.includes(ns))
6344 .map(|(ns, [p, c])| {
6345 (
6346 ns.clone(),
6347 json!({
6348 "prompt_tokens_in": p,
6349 "cached_tokens_in": c,
6350 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
6351 }),
6352 )
6353 })
6354 .collect();
6355 if !tenants.is_empty() {
6356 body["tenants"] = serde_json::Value::Object(tenants);
6357 }
6358 }
6359 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
6360 .adsd_suspect_total
6361 .iter()
6362 .filter(|(tenant, _)| metrics_scope.includes(tenant))
6363 .map(|(tenant, total)| (tenant.clone(), json!(total)))
6364 .collect();
6365 if !adsd_suspect_total.is_empty() {
6366 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
6367 }
6368 if metrics_scope.operator()
6370 && let Some((bg, mode)) = &st.bg
6371 {
6372 body["bg"] = bg.to_json(mode);
6373 }
6374 if metrics_scope.operator() {
6381 let spec: serde_json::Map<String, serde_json::Value> = m
6382 .spec
6383 .iter()
6384 .map(|(model, t)| {
6385 let n_pos = t
6386 .pos_drafted
6387 .iter()
6388 .rposition(|&d| d > 0)
6389 .map_or(0, |p| p + 1);
6390 (
6391 model.clone(),
6392 json!({
6393 "rounds": t.rounds,
6394 "drafted": t.drafted,
6395 "accepted": t.accepted,
6396 "acceptance_rate": if t.drafted > 0 {
6397 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
6398 "tokens_per_round": if t.rounds > 0 {
6399 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
6400 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
6401 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
6402 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
6403 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
6404 .collect::<Vec<f64>>(),
6405 }),
6406 )
6407 })
6408 .collect();
6409 if !spec.is_empty() {
6410 body["spec"] = serde_json::Value::Object(spec);
6411 }
6412 }
6413 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
6414 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
6415 insert_pp_wave_metrics(&mut body, &metrics_scope, PpWaveMetricsSnapshot::current);
6416 insert_peer_probe_metrics(
6417 &mut body,
6418 &metrics_scope,
6419 memra_engine::pp::peer_probe_metrics,
6420 );
6421 Json(body).into_response()
6422}
6423
6424#[derive(Debug, Default, Deserialize)]
6425struct ModelsQuery {
6426 #[serde(default)]
6427 schema: Option<String>,
6428}
6429
6430fn models_openai_body(models: &[String]) -> serde_json::Value {
6431 let data: Vec<_> = models
6432 .iter()
6433 .map(|m| json!({ "id": m, "object": "model" }))
6434 .collect();
6435 json!({ "object": "list", "data": data })
6436}
6437
6438fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
6443 match metadata.and_then(|m| m.surface.as_deref()) {
6444 Some("embedding") => "embedding",
6445 Some("rerank") => "rerank",
6446 _ => "chat",
6447 }
6448}
6449
6450fn openrouter_supported_parameters(
6451 caps: Option<&ModelCaps>,
6452 max_output_length: Option<u64>,
6453 is_chat: bool,
6454) -> serde_json::Value {
6455 let mut parameters = serde_json::Map::new();
6456 if !is_chat {
6464 return serde_json::Value::Object(parameters);
6465 }
6466 for name in [
6467 "temperature",
6468 "top_p",
6469 "min_p",
6470 "frequency_penalty",
6471 "presence_penalty",
6472 "repetition_penalty",
6473 "stop",
6474 ] {
6475 parameters.insert(name.into(), json!({ "type": "unknown" }));
6476 }
6477 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
6478 parameters.insert(
6479 "seed".into(),
6480 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
6481 );
6482 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
6483 if let Some(max) = max_output_length {
6484 max_tokens["max"] = json!(max);
6485 }
6486 parameters.insert("max_tokens".into(), max_tokens);
6487 if is_chat
6500 && caps.is_some_and(|c| {
6501 !c.dsv4 && !(c.qwen_think && !c.think_switch && c.think_close.is_empty())
6502 })
6503 {
6504 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
6505 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
6506 }
6507 if is_chat && caps.is_some_and(|c| c.tools_branch) {
6508 parameters.insert("tools".into(), json!({ "type": "boolean" }));
6509 parameters.insert(
6510 "tool_choice".into(),
6511 json!({ "type": "enum", "values": ["auto", "none"] }),
6512 );
6513 }
6514 if is_chat
6525 && caps.is_some_and(|c| {
6526 (c.qwen_think || c.effort_levels || c.gemma_think)
6527 && !(c.qwen_think && !c.think_switch && !c.dsv4)
6528 })
6529 {
6530 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
6531 }
6532 if is_chat && caps.is_some_and(|c| c.glm5) {
6542 parameters.insert(
6543 "reasoning_effort".into(),
6544 json!({ "type": "enum", "values": ["low", "high", "max"] }),
6545 );
6546 }
6547 if is_chat && caps.is_some_and(|c| c.qwen_effort && !c.glm5) {
6553 parameters.insert(
6554 "reasoning_effort".into(),
6555 json!({ "type": "enum", "values": ["xhigh", "medium", "low"] }),
6556 );
6557 }
6558 serde_json::Value::Object(parameters)
6559}
6560
6561fn published_context_length(
6577 caps: Option<&ModelCaps>,
6578 metadata: Option<&OpenRouterModelMetadata>,
6579) -> Option<u64> {
6580 let trained = caps
6581 .map(|c| c.context_length as u64)
6582 .filter(|&value| value > 0)?;
6583 let envelope = metadata.and_then(|m| {
6584 let prompt = m.max_prompt_length?;
6585 let output = m.max_output_length?;
6586 prompt.checked_add(output)
6587 });
6588 Some(envelope.map_or(trained, |envelope| trained.min(envelope)))
6589}
6590
6591fn model_entry_openrouter(
6592 name: &str,
6593 caps: Option<&ModelCaps>,
6594 metadata: Option<&OpenRouterModelMetadata>,
6595) -> serde_json::Value {
6596 let empty = OpenRouterModelMetadata::default();
6597 let metadata = metadata.unwrap_or(&empty);
6598 let context_length =
6599 published_context_length(caps, Some(metadata)).filter(|&v| v <= JSON_SAFE_INTEGER_MAX);
6600 let tokenizer = caps
6601 .map(|c| c.tokenizer.as_str())
6602 .filter(|tokenizer| !tokenizer.is_empty());
6603
6604 let mut input = serde_json::Map::new();
6605 input.insert("type".into(), json!("text"));
6606 let mut supported_inputs = serde_json::Map::new();
6607 if let Some(value) = context_length {
6608 supported_inputs.insert(
6609 "max_context_length".into(),
6610 json!({ "value": value, "unit": "token" }),
6611 );
6612 }
6613 if let Some(value) = metadata.max_prompt_length {
6614 supported_inputs.insert(
6615 "max_prompt_length".into(),
6616 json!({ "value": value, "unit": "token" }),
6617 );
6618 }
6619 if !supported_inputs.is_empty() {
6620 input.insert(
6621 "supported_inputs".into(),
6622 serde_json::Value::Object(supported_inputs),
6623 );
6624 }
6625 let mut input_pricing = Vec::new();
6626 for (kind, cost) in [
6627 ("prompt", metadata.pricing.prompt.as_deref()),
6628 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
6629 ("cache_write", metadata.pricing.cache_write.as_deref()),
6630 ] {
6631 if let Some(cost) = cost {
6632 input_pricing.push(json!({
6633 "type": kind,
6634 "unit": "token",
6635 "cost_usd": cost,
6636 }));
6637 }
6638 }
6639 if !input_pricing.is_empty() {
6640 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
6641 }
6642 let mut input_capacity = Vec::new();
6643 for (kind, value) in [
6644 ("prompt", metadata.capacity.prompt_tpm),
6645 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
6646 ] {
6647 if let Some(value) = value {
6648 input_capacity.push(json!({
6649 "type": kind,
6650 "unit": "token",
6651 "per": "minute",
6652 "value": value,
6653 }));
6654 }
6655 }
6656 if !input_capacity.is_empty() {
6657 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
6658 }
6659
6660 let or_surface = declared_surface(Some(metadata));
6661 let or_is_chat = or_surface == "chat";
6662 let mut output = serde_json::Map::new();
6663 output.insert(
6671 "type".into(),
6672 json!(match or_surface {
6673 "embedding" => "embeddings",
6674 "rerank" => "rerank",
6675 _ => "text",
6676 }),
6677 );
6678 output.insert(
6679 "supported_parameters".into(),
6680 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
6681 );
6682 if or_is_chat {
6686 output.insert("streaming".into(), json!(true));
6687 }
6688 if let Some(value) = metadata.max_output_length
6691 && or_is_chat
6692 {
6693 output.insert(
6694 "max_length".into(),
6695 json!({ "value": value, "unit": "token" }),
6696 );
6697 }
6698 let mut output_pricing = Vec::new();
6699 for (kind, cost) in [
6700 ("completion", metadata.pricing.completion.as_deref()),
6701 (
6702 "internal_reasoning",
6703 metadata.pricing.internal_reasoning.as_deref(),
6704 ),
6705 ] {
6706 if let Some(cost) = cost {
6707 output_pricing.push(json!({
6708 "type": kind,
6709 "unit": "token",
6710 "cost_usd": cost,
6711 }));
6712 }
6713 }
6714 if !output_pricing.is_empty() {
6715 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
6716 }
6717 let mut output_capacity = Vec::new();
6718 if let Some(value) = metadata.capacity.completion_tpm {
6719 output_capacity.push(json!({
6720 "type": "completion",
6721 "unit": "token",
6722 "per": "minute",
6723 "value": value,
6724 }));
6725 }
6726 if let Some(value) = metadata.capacity.concurrency {
6727 output_capacity.push(json!({
6728 "type": "concurrency",
6729 "unit": "request",
6730 "value": value,
6731 }));
6732 }
6733 if !output_capacity.is_empty() {
6734 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
6735 }
6736
6737 let mut entry = serde_json::Map::new();
6738 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
6739 entry.insert("id".into(), json!(name));
6740 entry.insert("name".into(), json!(name));
6741 if let Some(value) = metadata.hugging_face_id.as_deref() {
6742 entry.insert("hugging_face_id".into(), json!(value));
6743 }
6744 if let Some(value) = metadata.created {
6745 entry.insert("created".into(), json!(value));
6746 }
6747 if let Some(value) = metadata.quantization.as_deref() {
6748 entry.insert("quantization".into(), json!(value));
6749 }
6750 if let Some(value) = tokenizer {
6751 entry.insert("tokenizer".into(), json!(value));
6752 }
6753 if let Some(value) = metadata.description.as_deref() {
6754 entry.insert("description".into(), json!(value));
6755 }
6756 let mut input_modalities = vec![serde_json::Value::Object(input)];
6757 for m in &metadata.input_modalities {
6758 let mut extra = serde_json::Map::new();
6759 extra.insert("type".into(), json!(m));
6760 if let Some(cost) = metadata.pricing.prompt.as_deref() {
6761 extra.insert(
6763 "pricing".into(),
6764 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
6765 );
6766 }
6767 input_modalities.push(serde_json::Value::Object(extra));
6768 }
6769 entry.insert(
6770 "input_modalities".into(),
6771 serde_json::Value::Array(input_modalities),
6772 );
6773 entry.insert(
6774 "output_modalities".into(),
6775 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
6776 );
6777 if let Some(cost) = metadata.pricing.request.as_deref() {
6778 entry.insert(
6779 "pricing".into(),
6780 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
6781 );
6782 }
6783 if let Some(value) = metadata.capacity.request_rpm {
6784 entry.insert(
6785 "capacity".into(),
6786 json!([{
6787 "type": "request",
6788 "unit": "request",
6789 "per": "minute",
6790 "value": value,
6791 }]),
6792 );
6793 }
6794 if let Some(value) = metadata.is_ready {
6795 entry.insert("is_ready".into(), json!(value));
6796 }
6797 if let Some(value) = metadata.is_free {
6798 entry.insert("is_free".into(), json!(value));
6799 }
6800 if let Some(value) = metadata.discount_to_user {
6801 entry.insert("discount_to_user".into(), json!(value));
6802 }
6803 if let Some(value) = metadata.openrouter_slug.as_deref() {
6804 entry.insert("openrouter".into(), json!({ "slug": value }));
6805 }
6806 if !metadata.datacenters.is_empty() {
6807 entry.insert("datacenters".into(), json!(metadata.datacenters));
6808 }
6809 let mut compliance = serde_json::Map::new();
6810 if let Some(value) = metadata.zdr {
6811 compliance.insert("zdr".into(), json!(value));
6812 }
6813 if let Some(value) = metadata.hipaa {
6814 compliance.insert("hipaa".into(), json!(value));
6815 }
6816 if !compliance.is_empty() {
6817 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
6818 }
6819 serde_json::Value::Object(entry)
6820}
6821
6822fn models_openrouter_body(st: &AppState) -> serde_json::Value {
6823 let md = st.metadata();
6824 let data: Vec<_> = st
6825 .models
6826 .iter()
6827 .map(|model| model_entry_openrouter(model, st.caps.get(model), md.models.get(model)))
6828 .collect();
6829 json!({ "data": data })
6830}
6831
6832fn model_entry_openmodels(
6833 name: &str,
6834 caps: Option<&ModelCaps>,
6835 metadata: Option<&OpenRouterModelMetadata>,
6836) -> Result<serde_json::Value, String> {
6837 let metadata = metadata.ok_or_else(|| {
6838 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
6839 })?;
6840 let context_length = published_context_length(caps, Some(metadata))
6841 .filter(|&value| value <= JSON_SAFE_INTEGER_MAX)
6842 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
6843 let created = metadata
6844 .created
6845 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
6846 let max_output_length = metadata
6847 .max_output_length
6848 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
6849 let prompt = metadata
6850 .pricing
6851 .prompt
6852 .as_deref()
6853 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
6854 let completion =
6855 metadata.pricing.completion.as_deref().ok_or_else(|| {
6856 format!("OpenModels feed requires pricing.completion for model {name:?}")
6857 })?;
6858 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
6859 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
6860 })?;
6861 let is_ready = metadata
6862 .is_ready
6863 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
6864 let is_free = metadata
6865 .is_free
6866 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
6867 let discount_to_user = metadata
6868 .discount_to_user
6869 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
6870
6871 let mut pricing = serde_json::Map::new();
6872 pricing.insert("prompt".into(), json!(prompt));
6873 pricing.insert("completion".into(), json!(completion));
6874 pricing.insert("input_cache_read".into(), json!(input_cache_read));
6875 if let Some(value) = metadata.pricing.request.as_deref() {
6876 pricing.insert("request".into(), json!(value));
6877 }
6878
6879 let om_surface = declared_surface(Some(metadata));
6880 let om_is_chat = om_surface == "chat";
6881 let mut supported_features = Vec::new();
6882 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
6883 supported_features.push("tool_calling");
6884 }
6885 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
6886 supported_features.push("reasoning");
6887 }
6888
6889 let mut entry = serde_json::Map::new();
6890 entry.insert("id".into(), json!(name));
6891 entry.insert("name".into(), json!(name));
6892 entry.insert("created".into(), json!(created));
6893 entry.insert("input_modalities".into(), json!(["text"]));
6894 entry.insert(
6895 "output_modalities".into(),
6896 json!(match om_surface {
6897 "embedding" => ["embeddings"],
6898 "rerank" => ["rerank"],
6899 _ => ["text"],
6900 }),
6901 );
6902 entry.insert("context_length".into(), json!(context_length));
6903 entry.insert("max_output_length".into(), json!(max_output_length));
6904 entry.insert("currency".into(), json!("USD"));
6907 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
6908 entry.insert("supported_features".into(), json!(supported_features));
6909 entry.insert("is_ready".into(), json!(is_ready));
6910 entry.insert("is_free".into(), json!(is_free));
6911 entry.insert("discount_to_user".into(), json!(discount_to_user));
6912 Ok(serde_json::Value::Object(entry))
6913}
6914
6915fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
6916 let md = st.metadata();
6917 let data: Result<Vec<_>, _> = st
6918 .models
6919 .iter()
6920 .map(|model| model_entry_openmodels(model, st.caps.get(model), md.models.get(model)))
6921 .collect();
6922 Ok(json!({ "data": data? }))
6923}
6924
6925async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
6926 match query.schema.as_deref() {
6927 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
6928 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
6929 Some("openmodels") => match models_openmodels_body(&st) {
6930 Ok(body) => Json(body).into_response(),
6931 Err(error) => bad_request(&error, Some("schema")),
6932 },
6933 Some(schema) => bad_request(
6934 &format!(
6935 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
6936 ),
6937 Some("schema"),
6938 ),
6939 }
6940}
6941
6942fn model_entry_v1(
6950 name: &str,
6951 caps: Option<&ModelCaps>,
6952 metadata: Option<&OpenRouterModelMetadata>,
6953) -> serde_json::Value {
6954 let ctx = published_context_length(caps, metadata);
6955 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
6959 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
6962 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
6963 Some(p) => json!(p),
6964 None => serde_json::Value::Null,
6965 };
6966 let owned_by = metadata
6967 .and_then(|m| m.owned_by.as_deref())
6968 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
6969 let mut input_modalities = vec!["text"];
6970 if let Some(meta) = metadata {
6971 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
6972 }
6973 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
6974 let reliability = metadata.and_then(|m| m.reliability.as_ref());
6975 let surface = declared_surface(metadata);
6981 let (model_type, endpoints, output_modalities) = match surface {
6982 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
6986 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
6987 _ => ("chat", vec!["chat/completions"], vec!["text"]),
6988 };
6989 let is_chat = surface == "chat";
6990 json!({
6991 "id": name,
6992 "name": name,
6993 "object": "model",
6994 "owned_by": owned_by,
6995 "type": model_type,
6996 "context_length": ctx,
6997 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
7000 "endpoints": endpoints,
7001 "input_modalities": input_modalities,
7002 "output_modalities": output_modalities,
7003 "capabilities": {
7004 "streaming": is_chat,
7007 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
7008 "structured_output": is_chat
7016 && !is_dsv4
7017 && !caps.is_some_and(|c| c.qwen_think && !c.think_switch && c.think_close.is_empty()),
7018 "reasoning": is_chat && thinking,
7019 "prompt_caching": is_chat && !is_dsv4,
7020 },
7021 "pricing": {
7022 "currency": "USD",
7023 "unit": "per_1m_tokens",
7024 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
7025 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
7026 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
7027 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
7028 "minimum_request": metadata
7030 .and_then(|m| m.pricing.request.as_deref())
7031 .unwrap_or("0"),
7032 },
7033 "lifecycle": {
7034 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
7035 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
7036 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
7037 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
7038 },
7039 "reliability": {
7040 "first_token_timeout_seconds":
7041 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
7042 "completion_timeout_seconds":
7043 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
7044 "stream_idle_timeout_seconds":
7045 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
7046 "capacity_scope":
7047 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
7048 },
7049 })
7050}
7051
7052async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
7055 let md = st.metadata();
7058 let data: Vec<_> = st
7059 .models
7060 .iter()
7061 .map(|m| model_entry_v1(m, st.caps.get(m), md.models.get(m)))
7062 .collect();
7063 let mut body = json!({
7064 "object": "list",
7065 "contract_version": "2.0",
7066 "data": data,
7067 });
7068 if let Some(provider) = md.provider.as_ref() {
7073 body["provider"] = json!({
7074 "id": provider.id,
7075 "status_url": provider.status_url,
7076 "support_contact": provider.support_contact,
7077 "incident_contact": provider.incident_contact,
7078 "regions": provider.regions,
7079 "request_id_header": "x-request-id",
7080 "error_contract": {
7081 "rate_limit_status": 429,
7082 "overload_status": 503,
7083 "retry_after_header": "Retry-After",
7084 "account_quota_error_codes": ["insufficient_balance"],
7085 },
7086 });
7087 }
7088 Json(body)
7089}
7090
7091fn retrieve_model_row(
7097 models: &[String],
7098 caps: &HashMap<String, ModelCaps>,
7099 md: &ModelMetadataSet,
7100 id: &str,
7101) -> Option<serde_json::Value> {
7102 models
7103 .iter()
7104 .find(|m| m.as_str() == id)
7105 .map(|name| model_entry_v1(name, caps.get(name), md.models.get(name)))
7106}
7107
7108async fn retrieve_model_v1(State(st): State<AppState>, Path(id): Path<String>) -> Response {
7125 let md = st.metadata();
7126 match retrieve_model_row(&st.models, &st.caps, &md, &id) {
7127 Some(row) => Json(row).into_response(),
7128 None => error_response_coded(
7129 StatusCode::NOT_FOUND,
7130 &format!("model {id:?} not found"),
7131 "invalid_request_error",
7132 Some("id"),
7133 Some("model_not_found"),
7134 ),
7135 }
7136}
7137
7138async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
7141 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
7142 Ok(scope) => scope,
7143 Err(response) => return response,
7144 };
7145 if !metrics_scope.process_wide() {
7146 return error_response(
7147 StatusCode::FORBIDDEN,
7148 "completion api keys do not authorize process-wide yield metrics; configure \
7149 MEMRA_METRICS_TOKEN",
7150 "authentication_error",
7151 None,
7152 );
7153 }
7154 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
7155 let lane = |i: usize| {
7156 json!({
7157 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
7158 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
7159 })
7160 };
7161 let mut body = json!({
7162 "lanes": {
7163 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
7164 },
7165 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
7166 });
7167 if metrics_scope.operator() {
7168 body["batch_size_last"] = json!(m.batch_size_last);
7169 }
7170 Json(body).into_response()
7171}
7172
7173async fn peek_admission(
7186 mut rx: worker::EventReceiver,
7187) -> Result<worker::EventReceiver, (Response, &'static str)> {
7188 match rx.recv().await {
7189 Some(Event::Error(e)) => {
7194 let error_code = engine_error_code(e.class);
7195 Err((engine_error_response(&e), error_code))
7196 }
7197 first => {
7198 let (tx2, rx2) = worker::event_channel();
7199 if let Some(ev) = first {
7200 let _ = tx2.send(ev);
7201 }
7202 tokio::spawn(forward_events(rx, tx2));
7203 Ok(rx2)
7204 }
7205 }
7206}
7207
7208async fn forward_events(mut rx: worker::EventReceiver, tx2: worker::EventSender) {
7219 loop {
7220 tokio::select! {
7221 biased;
7222 () = tx2.closed() => break,
7223 ev = rx.recv() => match ev {
7224 Some(ev) => {
7225 if tx2.send(ev).is_err() {
7226 break;
7227 }
7228 }
7229 None => break,
7230 },
7231 }
7232 }
7233}
7234
7235async fn forward_prefill_until_first_delivery(
7240 mut rx: worker::EventReceiver,
7241 tx2: worker::EventSender,
7242 deadline: RequestDeadline,
7243 mut receipt: Option<prefill_receipt::SharedReceipt>,
7244) {
7245 loop {
7246 tokio::select! {
7247 biased;
7248 () = tokio::time::sleep_until(deadline.at) => {
7249 if let Some(receipt) = receipt.as_mut() {
7250 use metering::Receipt;
7251 if let Err(error) = receipt.settle_unbilled("deadline_exceeded", 408, "deadline_exceeded") {
7252 eprintln!("[ledger] ERROR: prefill deadline settlement failed: {error}");
7253 }
7254 }
7255 let _ = tx2.send(Event::DeadlineExceeded { ms: deadline.ms });
7256 break;
7257 }
7258 () = tx2.closed() => break,
7259 ev = rx.recv() => match ev {
7260 Some(ev) => {
7261 let first_delivery = matches!(
7262 ev,
7263 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
7264 );
7265 if tx2.send(ev).is_err() {
7266 break;
7267 }
7268 if first_delivery {
7269 drop(receipt);
7270 forward_events(rx, tx2).await;
7271 break;
7272 }
7273 }
7274 None => break,
7275 },
7276 }
7277 }
7278}
7279
7280fn sse_prefill_commit_ms() -> Option<u64> {
7283 static V: std::sync::OnceLock<Option<u64>> = std::sync::OnceLock::new();
7284 *V.get_or_init(|| {
7285 std::env::var("MEMRA_SSE_PREFILL_COMMIT_MS")
7286 .ok()
7287 .and_then(|s| s.parse::<u64>().ok())
7288 .filter(|&ms| ms > 0)
7289 })
7290}
7291
7292fn validate_stream_prefill_config_values(
7293 stream_max: Option<u64>,
7294 commit_ms: Option<u64>,
7295) -> Result<(), String> {
7296 let Some(stream_max) = stream_max else {
7297 return Ok(()); };
7299 if stream_max < TIMEOUT_MS_MIN {
7300 return Err(format!(
7301 "MEMRA_STREAM_TTFT_MS_MAX must be at least {TIMEOUT_MS_MIN}, got {stream_max}"
7302 ));
7303 }
7304 if stream_max <= TIMEOUT_MS_MAX {
7305 return Ok(());
7306 }
7307 let Some(commit_ms) = commit_ms.filter(|ms| *ms > 0) else {
7308 return Err(format!(
7309 "MEMRA_STREAM_TTFT_MS_MAX={stream_max} extends streaming TTFT past the shipped \
7310 {TIMEOUT_MS_MAX} ms pre-header window, but MEMRA_SSE_PREFILL_COMMIT_MS is unset \
7311 or zero; refusing a customer promise that cannot start SSE keepalives"
7312 ));
7313 };
7314 if commit_ms >= stream_max {
7315 return Err(format!(
7316 "MEMRA_SSE_PREFILL_COMMIT_MS={commit_ms} must be below the extended streaming \
7317 deadline {stream_max}"
7318 ));
7319 }
7320 Ok(())
7321}
7322
7323fn validate_stream_prefill_config() -> Result<(), String> {
7324 let stream_max = match std::env::var("MEMRA_STREAM_TTFT_MS_MAX") {
7325 Ok(raw) => Some(
7326 raw.parse::<u64>()
7327 .map_err(|_| format!("MEMRA_STREAM_TTFT_MS_MAX must be an integer, got {raw:?}"))?,
7328 ),
7329 Err(std::env::VarError::NotPresent) => None,
7330 Err(err) => return Err(format!("MEMRA_STREAM_TTFT_MS_MAX is unreadable: {err}")),
7331 };
7332 let commit_ms =
7333 match std::env::var("MEMRA_SSE_PREFILL_COMMIT_MS") {
7334 Ok(raw) => Some(raw.parse::<u64>().map_err(|_| {
7335 format!("MEMRA_SSE_PREFILL_COMMIT_MS must be an integer, got {raw:?}")
7336 })?),
7337 Err(std::env::VarError::NotPresent) => None,
7338 Err(err) => return Err(format!("MEMRA_SSE_PREFILL_COMMIT_MS is unreadable: {err}")),
7339 };
7340 validate_stream_prefill_config_values(stream_max, commit_ms)
7341}
7342
7343async fn peek_first_token(
7358 rx: worker::EventReceiver,
7359 deadline: RequestDeadline,
7360 receipt: &mut Option<Box<dyn metering::Receipt>>,
7361) -> Result<worker::EventReceiver, ()> {
7362 let commit_after = (deadline.ms > TIMEOUT_MS_MAX)
7363 .then(sse_prefill_commit_ms)
7364 .flatten()
7365 .map(std::time::Duration::from_millis);
7366 let shared = commit_after.and_then(|_| prefill_receipt::SharedReceipt::wrap(receipt));
7367 peek_first_token_with_receipt(rx, deadline, commit_after, shared).await
7368}
7369
7370#[cfg(test)]
7371async fn peek_first_token_with_commit(
7372 rx: worker::EventReceiver,
7373 deadline: RequestDeadline,
7374 commit_after: Option<std::time::Duration>,
7375) -> Result<worker::EventReceiver, ()> {
7376 peek_first_token_with_receipt(rx, deadline, commit_after, None).await
7377}
7378
7379async fn peek_first_token_with_receipt(
7380 mut rx: worker::EventReceiver,
7381 deadline: RequestDeadline,
7382 commit_after: Option<std::time::Duration>,
7383 receipt: Option<prefill_receipt::SharedReceipt>,
7384) -> Result<worker::EventReceiver, ()> {
7385 let mut buffered: Vec<Event> = Vec::new();
7386 let request_started_at = deadline
7387 .at
7388 .checked_sub(std::time::Duration::from_millis(deadline.ms))
7389 .unwrap_or_else(tokio::time::Instant::now);
7390 let commit_at = commit_after.map(|after| request_started_at + after);
7391 loop {
7392 let wait_until = commit_at.map_or(deadline.at, |at| at.min(deadline.at));
7393 match tokio::time::timeout_at(wait_until, rx.recv()).await {
7394 Err(_) if wait_until == deadline.at => {
7395 return Err(()); }
7397 Err(_) => {
7398 let (tx2, rx2) = worker::event_channel();
7399 for ev in buffered {
7400 let _ = tx2.send(ev);
7401 }
7402 tokio::spawn(forward_prefill_until_first_delivery(
7403 rx, tx2, deadline, receipt,
7404 ));
7405 return Ok(rx2);
7406 }
7407 Ok(None) => break, Ok(Some(ev)) => {
7409 let first_delivery = matches!(
7410 ev,
7411 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
7412 );
7413 buffered.push(ev);
7414 if first_delivery {
7415 break;
7416 }
7417 }
7418 }
7419 }
7420 let (tx2, rx2) = worker::event_channel();
7421 for ev in buffered {
7422 let _ = tx2.send(ev);
7423 }
7424 tokio::spawn(forward_events(rx, tx2));
7425 Ok(rx2)
7426}
7427
7428#[cfg(test)]
7430fn build_request(
7434 req: &CompletionReq,
7435 tx: worker::EventSender,
7436 lane: lanes::Lane,
7437 affinity: Option<String>,
7438) -> Request {
7439 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
7440}
7441
7442fn build_request_with_trace(
7443 req: &CompletionReq,
7444 tx: worker::EventSender,
7445 lane: lanes::Lane,
7446 affinity: Option<String>,
7447 ttft: Option<Arc<ttft::Trace>>,
7448 sampling_defaults: &SamplingDefaults,
7449) -> Request {
7450 let params = GenParams {
7451 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
7452 max_ctx: req.max_ctx,
7453 eos: Vec::new(), };
7455 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
7460 Request {
7461 model: req.model.clone(),
7462 prompt_ids: req.prompt_ids.clone(),
7463 prompt_text: req.prompt.clone(),
7464 chat: req.chat,
7465 chat_turns: Vec::new(),
7466 tools_json: Vec::new(),
7467 tools_struct: Vec::new(),
7468 think: ThinkMode::Default,
7469 reasoning_effort: None, params,
7471 sampler_cfg,
7472 stop_strings: req.stop.clone().into_vec(),
7473 trace_id: req.trace_id.clone(),
7474 request_id: String::new(),
7477 admit_predict_logged: false,
7478 max_prompt_tokens: None,
7479 cache_ns: cache_namespace(&req.cache_salt),
7480 affinity,
7481 lane,
7482 grammar: None, prepared_constraint: None,
7484 constraint_ready: None,
7485 oom_retries: 0, spec_k_replay: None,
7487 prepared_prompt: None,
7488 capture: None, images: Vec::new(), gemma_images: Vec::new(),
7491 glm5_images: Vec::new(),
7492 step_images: Vec::new(),
7493 vision_memory: None,
7494 wire_deadline: None, ttft,
7496 tx,
7497 }
7498}
7499
7500struct ChatPlan {
7503 request: Request,
7504 parser: Option<ToolStreamParser>,
7507 pending_images: Vec<PendingVisionUnit>,
7510 pending_gemma: Vec<PendingGemmaImage>,
7511 pending_glm5: Vec<PendingGlm5Image>,
7512 pending_step: Vec<PendingStepImage>,
7513 vision_memory: Option<VisionMemoryPermit>,
7517}
7518
7519pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
7520 req.messages.iter().any(|message| {
7521 message.content.as_array().is_some_and(|parts| {
7522 parts.iter().any(|part| {
7523 matches!(
7524 part.get("type").and_then(serde_json::Value::as_str),
7525 Some("image_url" | "video_url")
7526 )
7527 })
7528 })
7529 })
7530}
7531
7532fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
7533 let mut total = 0usize;
7534 let mut add = |bytes: usize| {
7535 total = total.checked_add(bytes).ok_or_else(|| {
7536 "vision patch memory reservation overflowed while planning".to_string()
7537 })?;
7538 Ok::<(), String>(())
7539 };
7540 for unit in &plan.pending_images {
7541 let bytes = match unit {
7542 PendingVisionUnit::Still { gh, gw, .. } => gh
7543 .checked_mul(*gw)
7544 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7545 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7546 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
7547 PendingVisionUnit::Video { groups, .. } => {
7548 groups.iter().try_fold(0usize, |total, group| {
7549 let bytes = group
7550 .gh
7551 .checked_mul(group.gw)
7552 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7553 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7554 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7555 total.checked_add(bytes).ok_or_else(|| {
7556 "vision patch memory reservation overflowed while planning".to_string()
7557 })
7558 })?
7559 }
7560 };
7561 add(bytes)?;
7562 }
7563 for unit in &plan.pending_gemma {
7564 let bytes = unit
7565 .gw
7566 .checked_mul(unit.gh)
7567 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
7568 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7569 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7570 add(bytes)?;
7571 }
7572 for unit in &plan.pending_glm5 {
7573 let bytes = unit
7574 .gh
7575 .checked_mul(unit.gw)
7576 .and_then(|n| n.checked_mul(memra_engine::vision_glm5::G5V_PATCH_IN))
7577 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7578 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7579 add(bytes)?;
7580 }
7581 for unit in &plan.pending_step {
7582 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
7583 let patches = unit
7585 .plan
7586 .n_tiles
7587 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
7588 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
7589 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7590 let bytes = patches
7591 .checked_mul(SV_PATCH_IN)
7592 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7593 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7594 add(bytes)?;
7595 }
7596 Ok(total)
7597}
7598
7599pub(crate) fn reserve_vision_memory(
7600 plan: &ChatPlan,
7601) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
7602 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
7603 try_reserve_vision_memory(bytes)
7604}
7605
7606#[cfg(test)]
7607fn build_chat_request(
7608 req: ChatCompletionReq,
7609 caps: Option<&ModelCaps>,
7610 tx: worker::EventSender,
7611 lane: lanes::Lane,
7612 affinity: Option<String>,
7613) -> Result<ChatPlan, String> {
7614 let defaults = ModelSamplingDefaults::resolve(None, caps);
7617 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
7618}
7619
7620#[allow(clippy::too_many_arguments)]
7634fn build_chat_request_with_trace(
7635 req: ChatCompletionReq,
7636 caps: Option<&ModelCaps>,
7637 tx: worker::EventSender,
7638 lane: lanes::Lane,
7639 affinity: Option<String>,
7640 ttft: Option<Arc<ttft::Trace>>,
7641 default_effort: Option<&str>,
7642 sampling_defaults: &ModelSamplingDefaults,
7643) -> Result<ChatPlan, String> {
7644 req.stop.validate()?;
7645 let client_sampling: ClientSampling = (&req).into();
7648 let tool_choice = parse_tool_choice(&req.tool_choice)?;
7649 if let Some(c) = caps
7654 && !c.chat_ok
7655 {
7656 return Err(format!(
7657 "model {:?} has no chat template (checkpoint carries neither \
7658 tokenizer_config.json chat_template nor chat_template.jinja) — \
7659 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
7660 req.model
7661 ));
7662 }
7663 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
7664 let (mut think, effort_level, think_client_explicit) = parse_think(
7665 &req.reasoning_effort,
7666 &req.reasoning,
7667 vllm_switch,
7668 req.include_reasoning,
7669 default_effort,
7670 caps.is_some_and(|c| c.dsv4 || c.glm5),
7674 )?;
7675 let level_template = caps
7679 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort || c.glm5)
7680 .unwrap_or(false);
7681 if think_client_explicit
7702 && think == ThinkMode::NoThink
7703 && let Some(c) = caps
7704 && c.qwen_think
7705 && !c.think_switch
7706 && !c.dsv4
7707 {
7708 return Err(format!(
7709 "model {:?} cannot disable reasoning: its chat template opens a think \
7710 tail unconditionally and carries no enable_thinking switch, so \
7711 reasoning_effort/enable_thinking cannot turn it off on this model",
7712 req.model
7713 ));
7714 }
7715 let reasoning_effort = if level_template { effort_level } else { None };
7748 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
7751 if grammar.is_some()
7765 && let Some(c) = caps
7766 && c.qwen_think
7767 && think != ThinkMode::NoThink
7768 {
7769 if c.think_switch {
7770 think = ThinkMode::NoThink;
7771 } else if c.think_close.is_empty() {
7772 return Err(
7773 "response_format requires the model's think channel to close \
7774 before the grammar can engage, but this chat template has \
7775 neither an enable_thinking switch nor a recognizable \
7776 think-close token sequence"
7777 .into(),
7778 );
7779 }
7780 }
7784
7785 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
7792
7793 let (tools_json, tools_struct, schemas) =
7796 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
7797 prepare_tools(&req.tools)?
7798 } else {
7799 (Vec::new(), Vec::new(), HashMap::new())
7800 };
7801
7802 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
7803 let mut images: Vec<PendingVisionUnit> = Vec::new();
7804 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
7805 let mut glm5_images: Vec<PendingGlm5Image> = Vec::new();
7806 let mut step_images: Vec<PendingStepImage> = Vec::new();
7807 let mut next_video = 0usize;
7808 for msg in &req.messages {
7809 let content = content_to_text_vision(
7810 &msg.content,
7811 &mut images,
7812 &mut gemma_images,
7813 &mut glm5_images,
7814 &mut step_images,
7815 &mut next_video,
7816 )
7817 .map_err(|e| format!("{} message: {e}", msg.role))?;
7818 let tool_calls = msg
7819 .tool_calls
7820 .iter()
7821 .map(render_req_tool_call)
7822 .collect::<Result<Vec<_>, _>>()?;
7823 if !tool_calls.is_empty() && msg.role != "assistant" {
7824 return Err("tool_calls are only valid on assistant messages".into());
7825 }
7826 let role = if msg.role == "developer" {
7829 "system".to_string()
7830 } else {
7831 msg.role.clone()
7832 };
7833 turns.push(TmplTurn {
7834 role,
7835 content,
7836 tool_calls,
7837 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
7839 tool_call_id: msg.tool_call_id.clone(),
7840 tool_name: msg.name.clone(),
7841 tool_responses: Vec::new(),
7842 task: None,
7846 tools: Vec::new(),
7847 });
7848 }
7849
7850 let has_tool_features = !tools_json.is_empty()
7853 || turns
7854 .iter()
7855 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
7856 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
7857 return Err(format!(
7858 "model {:?} chat template has no tools branch",
7859 req.model
7860 ));
7861 }
7862
7863 let think_open = caps
7866 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
7867 .unwrap_or(false);
7868 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
7886 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
7892 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
7893 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
7894 let glm5 = caps.map(|c| c.glm5).unwrap_or(false);
7900 let is_hy3 = caps.map(|c| c.hy3).unwrap_or(false);
7903 let hy3_think_open = is_hy3 && think == ThinkMode::Think;
7904 let hy3_tools = is_hy3 && !tools_json.is_empty();
7905 let parser = if glm5 {
7906 Some(ToolStreamParser::glm5(think_open, schemas))
7907 } else if is_hy3 && (hy3_tools || hy3_think_open) {
7908 Some(ToolStreamParser::hy3(schemas, hy3_think_open))
7909 } else if is_dsv4 && (dsv4_tools || dsv4_think_open) {
7910 Some(ToolStreamParser::dsv4(dsv4_think_open))
7911 } else if gemma_tools {
7912 Some(ToolStreamParser::gemma_tools())
7913 } else if !tools_json.is_empty() {
7914 Some(ToolStreamParser::new(schemas, think_open))
7915 } else if think_open {
7916 Some(ToolStreamParser::reasoning_only())
7917 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
7918 Some(ToolStreamParser::gemma_thought())
7926 } else {
7927 None
7928 };
7929
7930 Ok(ChatPlan {
7931 request: Request {
7932 model: req.model,
7933 prompt_ids: Vec::new(),
7934 prompt_text: String::new(),
7935 chat: false,
7936 chat_turns: turns,
7937 tools_json,
7938 tools_struct,
7939 think,
7940 reasoning_effort,
7941 params: GenParams {
7942 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
7943 max_ctx: req.max_ctx,
7944 eos: Vec::new(),
7945 },
7946 sampler_cfg,
7947 stop_strings: {
7948 let mut stops = req.stop.into_vec();
7953 if gemma_tools {
7954 stops.push("<tool_call|>".to_string());
7955 }
7956 if dsv4_tools {
7961 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
7962 }
7963 if hy3_tools {
7966 stops.push("</tool_calls:opensource>".to_string());
7967 }
7968 stops
7969 },
7970 trace_id: None,
7971 request_id: String::new(),
7974 admit_predict_logged: false,
7975 max_prompt_tokens: None,
7976 cache_ns: cache_namespace(&req.cache_salt),
7977 affinity,
7978 lane,
7979 grammar,
7980 prepared_constraint: None,
7981 constraint_ready: None,
7982 oom_retries: 0, spec_k_replay: None,
7984 prepared_prompt: None,
7985 images: Vec::new(),
7990 gemma_images: Vec::new(),
7991 glm5_images: Vec::new(),
7992 step_images: Vec::new(),
7993 capture: None, vision_memory: None,
7995 wire_deadline: None, ttft,
7997 tx,
7998 },
7999 parser,
8000 pending_images: images,
8001 pending_gemma: gemma_images,
8002 pending_glm5: glm5_images,
8003 pending_step: step_images,
8004 vision_memory: None,
8005 })
8006}
8007
8008fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
8014 for (i, unit) in plan.pending_images.drain(..).enumerate() {
8015 match unit {
8016 PendingVisionUnit::Still { bytes, gh, gw } => {
8017 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
8018 .map_err(|e| format!("image {}: {e}", i + 1))?;
8019 if (prep.gh, prep.gw) != (gh, gw) {
8020 return Err(format!(
8021 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
8022 i + 1,
8023 prep.gh,
8024 prep.gw
8025 ));
8026 }
8027 plan.request
8028 .images
8029 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
8030 }
8031 PendingVisionUnit::Video {
8032 bytes,
8033 groups,
8034 video,
8035 } => {
8036 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
8037 .map_err(|e| format!("video {}: {e}", i + 1))?;
8038 if prepared.groups.len() != groups.len() {
8039 return Err(format!(
8040 "video {}: decoded {} groups differ from its header-planned {} groups",
8041 i + 1,
8042 prepared.groups.len(),
8043 groups.len()
8044 ));
8045 }
8046 for ((group, prep), timestamp) in
8047 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
8048 {
8049 if (prep.gh, prep.gw) != (group.gh, group.gw) {
8050 return Err(format!(
8051 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
8052 i + 1,
8053 prep.gh,
8054 prep.gw,
8055 group.gh,
8056 group.gw
8057 ));
8058 }
8059 if (timestamp - group.timestamp).abs() > 0.001 {
8060 return Err(format!(
8061 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
8062 i + 1,
8063 group.timestamp
8064 ));
8065 }
8066 plan.request
8067 .images
8068 .push(memra_engine::vision_pre::VisionUnit {
8069 prep,
8070 video: Some(video),
8071 });
8072 }
8073 }
8074 }
8075 }
8076 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
8077 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
8078 .map_err(|e| format!("image {}: {e}", i + 1))?;
8079 if (gw, gh) != (unit.gw, unit.gh) {
8080 return Err(format!(
8081 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
8082 i + 1,
8083 unit.gw,
8084 unit.gh
8085 ));
8086 }
8087 plan.request
8088 .gemma_images
8089 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
8090 }
8091 for (i, unit) in plan.pending_glm5.drain(..).enumerate() {
8092 let (patches, gh, gw) = memra_engine::vision_glm5::glm5_prep_image(&unit.bytes)
8093 .map_err(|e| format!("image {}: {e}", i + 1))?;
8094 if (gh, gw) != (unit.gh, unit.gw) {
8095 return Err(format!(
8096 "image {}: decoded grid {gh}x{gw} differs from its header-planned grid {}x{} — refusing (placeholder runs already rendered)",
8097 i + 1,
8098 unit.gh,
8099 unit.gw
8100 ));
8101 }
8102 plan.request
8103 .glm5_images
8104 .push(memra_engine::vision_glm5::Glm5VisionUnit { patches, gh, gw });
8105 }
8106 for (i, unit) in plan.pending_step.drain(..).enumerate() {
8107 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
8108 .map_err(|e| format!("image {}: {e}", i + 1))?;
8109 if prepped.tiles.len() != unit.plan.n_tiles
8110 || prepped.newline_mask != unit.plan.newline_mask
8111 {
8112 return Err(format!(
8113 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
8114 ({} tiles) — refusing (pad runs already rendered)",
8115 i + 1,
8116 prepped.tiles.len(),
8117 unit.plan.n_tiles
8118 ));
8119 }
8120 plan.request.step_images.push(prepped);
8121 }
8122 Ok(())
8123}
8124
8125fn bearer_token(headers: &HeaderMap) -> Option<&str> {
8133 headers
8134 .get("authorization")
8135 .and_then(|value| value.to_str().ok())
8136 .and_then(|value| value.strip_prefix("Bearer "))
8137}
8138
8139fn authentication_error(why: auth::AuthDenied) -> Response {
8140 match why {
8141 auth::AuthDenied::Unknown => error_response(
8142 StatusCode::UNAUTHORIZED,
8143 "invalid api key",
8144 "authentication_error",
8145 None,
8146 ),
8147 auth::AuthDenied::Disabled => error_response(
8148 StatusCode::FORBIDDEN,
8149 "api key is disabled",
8150 "authentication_error",
8151 None,
8152 ),
8153 }
8154}
8155
8156#[allow(clippy::result_large_err)] fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
8158 auth::authenticate_with(
8159 api_auth.keyring,
8160 api_auth.single_key.as_deref(),
8161 bearer_token(headers),
8162 )
8163 .map_err(authentication_error)
8164}
8165
8166#[derive(Debug, Clone, PartialEq, Eq)]
8167enum MetricsScope {
8168 All,
8169 CompletionDomain,
8170 Tenant(String),
8171}
8172
8173impl MetricsScope {
8174 fn operator(&self) -> bool {
8175 matches!(self, MetricsScope::All)
8176 }
8177
8178 fn process_wide(&self) -> bool {
8179 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
8180 }
8181
8182 fn includes(&self, tenant_row: &str) -> bool {
8183 match self {
8184 MetricsScope::All | MetricsScope::CompletionDomain => true,
8185 MetricsScope::Tenant(tenant) => tenant == tenant_row,
8186 }
8187 }
8188}
8189
8190#[allow(clippy::result_large_err)] fn authorize_metrics(
8192 api_auth: &ApiAuth,
8193 metrics_auth: &MetricsAuth,
8194 headers: &HeaderMap,
8195) -> Result<MetricsScope, Response> {
8196 if !metrics_auth.required {
8197 return Ok(MetricsScope::All);
8198 }
8199 let Some(candidate) = bearer_token(headers) else {
8200 return Err(authentication_error(auth::AuthDenied::Unknown));
8201 };
8202 if let Some(token) = metrics_auth.token.as_deref() {
8203 if auth::constant_time_secret_eq(token, candidate) {
8204 return Ok(MetricsScope::All);
8205 }
8206 if api_auth.configured() {
8207 return match auth::authenticate_with(
8208 api_auth.keyring,
8209 api_auth.single_key.as_deref(),
8210 Some(candidate),
8211 ) {
8212 Ok(_) => Err(error_response(
8213 StatusCode::FORBIDDEN,
8214 "completion api keys do not authorize metrics while \
8215 MEMRA_METRICS_TOKEN is configured",
8216 "authentication_error",
8217 None,
8218 )),
8219 Err(why) => Err(authentication_error(why)),
8220 };
8221 }
8222 return Err(authentication_error(auth::AuthDenied::Unknown));
8223 }
8224 if api_auth.configured() {
8225 let tenant = authenticate(api_auth, headers)?;
8226 return Ok(if api_auth.keyring.is_some() {
8227 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
8228 } else {
8229 MetricsScope::CompletionDomain
8233 });
8234 }
8235 Err(authentication_error(auth::AuthDenied::Unknown))
8236}
8237
8238#[allow(clippy::result_large_err)] fn lane_for_tenant(
8245 headers: &axum::http::HeaderMap,
8246 tenant: &auth::TenantCtx,
8247) -> Result<lanes::Lane, Response> {
8248 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
8249 None => None,
8250 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
8255 error_response_coded(
8256 StatusCode::BAD_REQUEST,
8257 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
8258 "invalid_request_error",
8259 Some("x-lane"),
8260 Some("invalid_lane"),
8261 )
8262 })?),
8263 };
8264 match tenant.lane_class {
8265 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
8266 auth::LaneClass::Batch => match requested {
8267 None => Ok(lanes::Lane::Harvest),
8268 Some(lanes::Lane::Interactive) => Err(error_response(
8269 StatusCode::FORBIDDEN,
8270 "this api key is batch-class: x-lane interactive is not permitted \
8271 (use judge or harvest)",
8272 "authentication_error",
8273 Some("x-lane"),
8274 )),
8275 Some(l) => Ok(l),
8276 },
8277 }
8278}
8279
8280fn tenant_namespace(
8284 tenant: &auth::TenantCtx,
8285 cache_salt: &Option<String>,
8286) -> Result<String, &'static str> {
8287 let keyring_configured = auth::global().is_some();
8288 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
8289 if keyring_configured {
8290 Ok(auth::scope_namespace(&tenant.tenant, &raw))
8291 } else {
8292 Ok(raw)
8293 }
8294}
8295
8296fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
8301 eprintln!(
8302 "[meter] admit id={} tenant={} lane={} model={:?}",
8303 env.id,
8304 tenant.tenant,
8305 lane.as_str(),
8306 model
8307 );
8308}
8309
8310fn apply_model_request_limits(
8311 request: &mut Request,
8312 metadata: Option<&OpenRouterModelMetadata>,
8313 caps: Option<&ModelCaps>,
8314) -> Result<(), (String, &'static str)> {
8315 let Some(metadata) = metadata else {
8316 return Ok(());
8317 };
8318 let max_prompt = metadata
8319 .max_prompt_length
8320 .map(usize::try_from)
8321 .transpose()
8322 .map_err(|_| {
8323 (
8324 "configured model prompt limit does not fit this platform".into(),
8325 "model",
8326 )
8327 })?;
8328 let max_output = metadata
8329 .max_output_length
8330 .map(usize::try_from)
8331 .transpose()
8332 .map_err(|_| {
8333 (
8334 "configured model output limit does not fit this platform".into(),
8335 "model",
8336 )
8337 })?;
8338
8339 request.max_prompt_tokens = max_prompt;
8340 if let Some(max_output) = max_output {
8341 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
8342 request.params.max_new = metadata
8343 .default_output_length
8344 .map(usize::try_from)
8345 .transpose()
8346 .map_err(|_| {
8347 (
8348 "configured default output length does not fit this platform".into(),
8349 "model",
8350 )
8351 })?
8352 .unwrap_or(max_output);
8353 } else if request.params.max_new > max_output {
8354 return Err((
8355 format!(
8356 "max_tokens {} exceeds configured model maximum {max_output}",
8357 request.params.max_new
8358 ),
8359 "max_tokens",
8360 ));
8361 }
8362 }
8363
8364 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
8368 (max_prompt, max_output, request.params.max_ctx)
8369 {
8370 let operational_ctx = max_prompt
8371 .checked_add(max_output)
8372 .and_then(|value| value.checked_add(8))
8373 .ok_or_else(|| {
8374 (
8375 "configured model context envelope overflowed".into(),
8376 "model",
8377 )
8378 })?;
8379 let operational_ctx = caps
8380 .map(|caps| caps.context_length)
8381 .filter(|&context| context > 0)
8382 .map_or(operational_ctx, |context| operational_ctx.min(context));
8383 if requested_ctx > operational_ctx {
8384 return Err((
8385 format!(
8386 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
8387 ),
8388 "max_ctx",
8389 ));
8390 }
8391 }
8392 Ok(())
8393}
8394
8395fn effective_max_tokens(request: &worker::Request) -> Option<u64> {
8399 (request.params.max_new != worker::MAX_NEW_CTX_BOUNDED).then_some(request.params.max_new as u64)
8400}
8401
8402#[allow(clippy::too_many_arguments)]
8403fn start_request_receipt(
8404 st: &AppState,
8405 env: &Envelope,
8406 tenant: &auth::TenantCtx,
8407 model: &str,
8408 route: &'static str,
8409 lane: lanes::Lane,
8410 stream: bool,
8411 max_tokens: Option<u64>,
8412 reserved_ctx: Option<u64>,
8413 budget_permit: Option<metering::Permit>,
8414) -> Option<Box<dyn metering::Receipt>> {
8415 st.metering.as_ref().map(|accounting| {
8416 accounting.open(
8417 &metering::RequestMeta {
8418 request_id: &env.id,
8419 tenant: &tenant.tenant,
8420 principal: tenant.key_prefix.as_deref(),
8421 model,
8422 route,
8423 lane: lane.as_str(),
8424 stream,
8425 max_tokens,
8426 reserved_ctx,
8427 },
8428 budget_permit,
8429 )
8430 })
8431}
8432
8433fn arm_capture(
8439 mut receipt: Option<Box<dyn metering::Receipt>>,
8440 prompt: impl FnOnce() -> serde_json::Value,
8441) -> Option<Box<dyn metering::Receipt>> {
8442 if let Some(receipt) = receipt.as_mut()
8443 && receipt.wants_capture()
8444 {
8445 receipt.arm_capture(prompt());
8446 }
8447 receipt
8448}
8449
8450fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
8454 serde_json::Value::Array(
8455 messages
8456 .iter()
8457 .map(|message| {
8458 let mut row = json!({ "role": message.role, "content": message.content });
8459 if !message.tool_calls.is_empty() {
8460 row["tool_calls"] = serde_json::Value::Array(
8461 message
8462 .tool_calls
8463 .iter()
8464 .map(|call| {
8465 json!({
8466 "id": call.id,
8467 "function": {
8468 "name": call.function.name,
8469 "arguments": call.function.arguments,
8470 },
8471 })
8472 })
8473 .collect(),
8474 );
8475 }
8476 row
8477 })
8478 .collect(),
8479 )
8480}
8481
8482enum BudgetRejection {
8483 Invalid(String),
8484 Insufficient,
8485 Unenrolled,
8486 PrincipalCapped,
8489 Unavailable(String),
8490}
8491
8492impl BudgetRejection {
8493 fn into_response(self) -> (Response, &'static str) {
8494 match self {
8495 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
8496 Self::Insufficient => (
8497 error_response_coded(
8498 StatusCode::PAYMENT_REQUIRED,
8499 "tenant prepaid balance is insufficient for this request",
8500 "insufficient_balance",
8501 None,
8502 Some("insufficient_balance"),
8503 ),
8504 "insufficient_balance",
8505 ),
8506 Self::Unenrolled => (
8507 error_response_coded(
8508 StatusCode::PAYMENT_REQUIRED,
8509 "tenant is not enrolled for prepaid billing",
8510 "tenant_not_enrolled",
8511 None,
8512 Some("tenant_not_enrolled"),
8513 ),
8514 "tenant_not_enrolled",
8515 ),
8516 Self::PrincipalCapped => (
8517 error_response_coded(
8518 StatusCode::PAYMENT_REQUIRED,
8519 "this API key's spend cap is reached; raise or clear the key's cap to continue",
8520 "key_spend_cap_reached",
8521 None,
8522 Some("key_spend_cap_reached"),
8523 ),
8524 "key_spend_cap_reached",
8525 ),
8526 Self::Unavailable(err) => {
8527 eprintln!("[budget] ERROR: admission unavailable: {err}");
8528 (
8529 error_response_coded(
8530 StatusCode::SERVICE_UNAVAILABLE,
8531 "tenant budget accounting is unavailable",
8532 "server_error",
8533 None,
8534 Some("tenant_budget_unavailable"),
8535 ),
8536 "tenant_budget_unavailable",
8537 )
8538 }
8539 }
8540 }
8541}
8542
8543fn prepare_budget_prompt(
8544 request: &mut Request,
8545 tokenizer: Option<&Tokenizer>,
8546) -> Result<usize, String> {
8547 if let Some(error) = worker::prompt_source_limit_error(request) {
8548 return Err(error);
8549 }
8550 if request.prepared_prompt.is_none() {
8551 if let Some(trace) = request.ttft.as_ref() {
8552 trace.mark_tokenize_start();
8553 }
8554 let prompt = if !request.prompt_ids.is_empty() {
8555 request.prompt_ids.clone()
8556 } else if !request.chat_turns.is_empty() {
8557 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8558 let plain = worker::plain_chat_render_path(
8565 &request.tools_json,
8566 &request.think,
8567 request.reasoning_effort.as_deref(),
8568 &request.chat_turns,
8569 tokenizer.has_qwen_effort_ladder(),
8570 );
8571 let rendered = if plain {
8572 let messages: Vec<_> = request
8573 .chat_turns
8574 .iter()
8575 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
8576 .collect();
8577 tokenizer.apply_chat_template(&messages, true)
8578 } else {
8579 tokenizer
8580 .apply_chat_template_tools_ex(
8581 &request.chat_turns,
8582 true,
8583 &request.tools_json,
8584 &request.tools_struct,
8585 request.think,
8586 request.reasoning_effort.as_deref(),
8587 )
8588 .map_err(|err| format!("chat template: {err}"))?
8589 };
8590 tokenizer.encode(&rendered, true)
8591 } else if request.chat {
8592 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8593 let rendered =
8594 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
8595 tokenizer.encode(&rendered, true)
8596 } else {
8597 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8598 tokenizer.encode(&request.prompt_text, true)
8599 };
8600 if prompt.is_empty() {
8601 return Err("empty prompt after tokenization".into());
8602 }
8603 if let Some(trace) = request.ttft.as_ref() {
8604 trace.mark_tokenize_end(prompt.len());
8605 }
8606 request.prepared_prompt = Some(prompt);
8607 }
8608 let prompt_tokens = request
8609 .prepared_prompt
8610 .as_ref()
8611 .expect("budget prompt was prepared")
8612 .len();
8613 if let Some(limit) = request.max_prompt_tokens
8614 && prompt_tokens > limit
8615 {
8616 return Err(format!(
8617 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
8618 ));
8619 }
8620 Ok(prompt_tokens)
8621}
8622
8623fn budget_completion_bound(
8624 request: &Request,
8625 prompt_tokens: usize,
8626 caps: Option<&ModelCaps>,
8627) -> Result<usize, String> {
8628 let max_new = request.params.max_new;
8629 let requested_ctx = match (request.params.max_ctx, max_new) {
8630 (Some(cap), _) => cap,
8631 (None, worker::MAX_NEW_CTX_BOUNDED) => {
8632 let server_ctx = std::env::var("MEMRA_CTX")
8633 .ok()
8634 .and_then(|value| value.parse().ok())
8635 .unwrap_or(8192usize);
8636 let mut cap = server_ctx;
8637 if prompt_tokens.saturating_add(16) > cap {
8638 cap = prompt_tokens.saturating_add(server_ctx);
8639 }
8640 cap
8641 }
8642 (None, max_new) => prompt_tokens
8643 .checked_add(max_new)
8644 .and_then(|value| value.checked_add(8))
8645 .ok_or_else(|| "request context bound overflowed".to_string())?,
8646 };
8647 let ctx_cap = caps
8648 .map(|caps| caps.context_length)
8649 .filter(|&context| context > 0)
8650 .map_or(requested_ctx, |context| requested_ctx.min(context));
8651 if prompt_tokens >= ctx_cap {
8652 return Err(format!(
8653 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
8654 ));
8655 }
8656 Ok(max_new.min(ctx_cap - prompt_tokens))
8657}
8658
8659struct BudgetAdmission {
8664 permit: Option<metering::Permit>,
8665 reserved_ctx: Option<u64>,
8666}
8667
8668impl std::fmt::Debug for BudgetAdmission {
8670 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
8671 f.debug_struct("BudgetAdmission")
8672 .field("permit", &self.permit.is_some())
8673 .field("reserved_ctx", &self.reserved_ctx)
8674 .finish()
8675 }
8676}
8677
8678fn admit_tenant_budget(
8679 st: &AppState,
8680 tenant: &auth::TenantCtx,
8681 request: &mut Request,
8682) -> Result<BudgetAdmission, BudgetRejection> {
8683 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
8684 return Ok(BudgetAdmission {
8685 permit: None,
8686 reserved_ctx: None,
8687 });
8688 };
8689 match accounting.is_limited(&tenant.tenant) {
8690 Ok(false) => return Err(BudgetRejection::Unenrolled),
8691 Ok(true) => {}
8692 Err(metering::AdmitError::Unavailable(err)) => {
8693 return Err(BudgetRejection::Unavailable(err));
8694 }
8695 Err(other) => {
8696 return Err(BudgetRejection::Unavailable(format!(
8697 "unexpected budget enrollment result: {other:?}"
8698 )));
8699 }
8700 }
8701 let tokenizer = st
8702 .budget_tokenizers
8703 .as_ref()
8704 .and_then(|tokenizers| tokenizers.get(&request.model))
8705 .map(Arc::as_ref);
8706 if request.prompt_ids.is_empty() && tokenizer.is_none() {
8707 return Err(BudgetRejection::Unavailable(format!(
8708 "no reservation tokenizer for model {:?}",
8709 request.model
8710 )));
8711 }
8712 let prompt_tokens =
8713 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
8714 let completion_tokens =
8715 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
8716 .map_err(BudgetRejection::Invalid)?;
8717 let prompt_tokens = u64::try_from(prompt_tokens)
8718 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
8719 let completion_tokens = u64::try_from(completion_tokens)
8720 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
8721 match accounting.reserve(
8722 &tenant.tenant,
8723 tenant.key_prefix.as_deref(),
8724 &request.model,
8725 prompt_tokens,
8726 completion_tokens,
8727 ) {
8728 Ok(permit) => Ok(BudgetAdmission {
8729 permit,
8730 reserved_ctx: Some(prompt_tokens.saturating_add(completion_tokens)),
8731 }),
8732 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
8733 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
8734 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
8738 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
8739 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
8740 }
8741}
8742
8743fn request_ledger_error_response() -> Response {
8744 error_response_coded(
8745 StatusCode::INTERNAL_SERVER_ERROR,
8746 "request completion could not be committed to the billing ledger",
8747 "server_error",
8748 None,
8749 Some("request_ledger_unavailable"),
8750 )
8751}
8752
8753fn request_ledger_error_body() -> serde_json::Value {
8754 error_body(
8755 "request completion could not be committed to the billing ledger",
8756 "server_error",
8757 None,
8758 Some("request_ledger_unavailable"),
8759 )
8760}
8761
8762fn ledger_rejected(
8763 mut receipt: Option<Box<dyn metering::Receipt>>,
8764 response: Response,
8765 error_code: &str,
8766 request_id: &str,
8767) -> Response {
8768 let status = response.status().as_u16();
8769 if let Some(receipt) = receipt.as_mut()
8770 && let Err(err) = receipt.reject(status, error_code)
8771 {
8772 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
8773 return with_request_id(request_id, request_ledger_error_response());
8774 }
8775 with_request_id(request_id, response)
8776}
8777
8778fn ledger_unbilled(
8783 mut receipt: Option<Box<dyn metering::Receipt>>,
8784 response: Response,
8785 outcome: &'static str,
8786 error_code: &str,
8787 request_id: &str,
8788) -> Response {
8789 let status = response.status().as_u16();
8790 if let Some(receipt) = receipt.as_mut()
8791 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
8792 {
8793 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
8794 return with_request_id(request_id, request_ledger_error_response());
8795 }
8796 with_request_id(request_id, response)
8797}
8798
8799fn engine_error_code(class: worker::ErrClass) -> &'static str {
8800 use worker::ErrClass as C;
8801 match class {
8802 C::InvalidRequest => "invalid_request",
8803 C::ContextLength => "context_length_exceeded",
8804 C::ModelNotFound => "model_not_found",
8805 C::RateLimit => "rate_limit_exceeded",
8806 C::Overloaded => "overloaded",
8807 C::Engine => "engine_error",
8808 }
8809}
8810
8811fn model_not_found_response(models: &[String], requested: &str) -> Response {
8831 error_response_coded(
8832 StatusCode::BAD_REQUEST,
8833 &format!("unknown model {requested:?}; loaded: {models:?}"),
8834 "invalid_request_error",
8835 Some("model"),
8836 Some("model_not_found"),
8837 )
8838}
8839
8840fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
8849 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
8850 return Ok(());
8851 };
8852 if let Some((pos, &id)) = ids
8853 .iter()
8854 .enumerate()
8855 .find(|&(_, &id)| id as usize >= n_vocab)
8856 {
8857 return Err(format!(
8858 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
8859 ));
8860 }
8861 Ok(())
8862}
8863
8864#[cfg(test)]
8865mod prompt_ids_tests {
8866 use super::*;
8867
8868 #[test]
8869 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
8870 let caps = ModelCaps {
8871 n_vocab: 8,
8872 ..Default::default()
8873 };
8874 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
8876 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
8877 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
8879 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
8880 assert!(err.contains("vocab size 8"), "{err}");
8881 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
8882 assert!(err.contains("4294967295"), "{err}");
8883 let unknown = ModelCaps::default();
8885 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
8886 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
8887 }
8888}
8889
8890fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
8891 if models.iter().any(|m| m == requested) {
8892 return Some(requested.to_string());
8893 }
8894 if requested.is_empty() || requested.contains('/') {
8895 return None;
8896 }
8897 let mut matches = models.iter().filter(|m| {
8898 m.rsplit('/')
8899 .next()
8900 .is_some_and(|suffix| suffix == requested)
8901 });
8902 match (matches.next(), matches.next()) {
8903 (Some(only), None) => Some(only.clone()),
8904 _ => None,
8905 }
8906}
8907
8908async fn completions_admitted(
8909 state: State<AppState>,
8910 headers: axum::http::HeaderMap,
8911 trace: Option<Extension<TtftRequestTrace>>,
8912 AdmittedJson(req, admission): AdmittedJson<CompletionReq>,
8913) -> Response {
8914 completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
8915}
8916
8917#[cfg(test)]
8918async fn completions(
8919 State(st): State<AppState>,
8920 headers: axum::http::HeaderMap,
8921 trace: Option<Extension<TtftRequestTrace>>,
8922 request: Json<CompletionReq>,
8923) -> Response {
8924 completions_with_admission(State(st), headers, trace, request, None).await
8925}
8926
8927async fn completions_with_admission(
8928 State(st): State<AppState>,
8929 headers: axum::http::HeaderMap,
8930 trace: Option<Extension<TtftRequestTrace>>,
8931 Json(mut req): Json<CompletionReq>,
8932 mut body_admission: Option<BodyAdmissionLease>,
8933) -> Response {
8934 let env = Envelope::new(false);
8935 if let Err(msg) = req.stop.validate() {
8936 return with_request_id(&env.id, bad_request(&msg, Some("stop")));
8937 }
8938 if let Err(msg) = validate_client_identifier(req.trace_id.as_deref(), "trace_id") {
8939 return with_request_id(&env.id, bad_request(&msg, Some("trace_id")));
8940 }
8941 match canonical_model_id(&st.models, &req.model) {
8942 Some(canonical) => req.model = canonical,
8943 None => {
8944 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
8945 }
8946 }
8947 let ttft = trace.and_then(|Extension(trace)| trace.0);
8950 if let Some(trace) = ttft.as_ref() {
8951 trace.mark_parsed();
8952 trace.bind_request(&env.id, &req.model);
8953 }
8954 let tenant = match authenticate(&st.api_auth, &headers) {
8955 Ok(t) => t,
8956 Err(resp) => return with_request_id(&env.id, resp),
8957 };
8958 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
8959 Ok(ns) => ns,
8960 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
8961 };
8962 if let Err((msg, param)) = reject_unsupported(&[
8964 (
8965 "logit_bias",
8966 req.logit_bias.is_some(),
8967 " (device-side sampling has no bias hook yet)",
8968 ),
8969 ("logprobs", req.logprobs.is_some(), ""),
8970 (
8971 "n",
8972 req.n.is_some_and(|n| n != 1),
8973 " for n != 1 (single choice only)",
8974 ),
8975 (
8976 "best_of",
8977 req.best_of.is_some_and(|n| n != 1),
8978 " (single choice only)",
8979 ),
8980 ]) {
8981 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8982 }
8983 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
8986 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
8987 }
8988 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref(), req.stream) {
8991 Ok(ms) => RequestDeadline::starting_now(ms),
8992 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8993 };
8994 let lane = match lane_for_tenant(&headers, &tenant) {
8995 Ok(l) => l,
8996 Err(resp) => return resp,
8997 };
8998 let (tx, rx) = worker::event_channel();
8999 let model = req.model.clone();
9000 let stream = req.stream;
9001 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
9002 Ok(affinity) => affinity,
9003 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
9004 };
9005 let md = st.metadata();
9008 let mut request = build_request_with_trace(
9009 &req,
9010 tx,
9011 lane,
9012 affinity,
9013 ttft.clone(),
9014 AppState::sampling_defaults_in(&md, &st.caps, &model).for_mode(ThinkMode::Default),
9018 );
9019 request.cache_ns = cache_ns;
9020 request.request_id = env.id.clone();
9021 request.wire_deadline = Some(deadline.at.into_std());
9024 if let Err((message, param)) =
9025 apply_model_request_limits(&mut request, md.models.get(&model), st.caps.get(&model))
9026 {
9027 return with_request_id(&env.id, bad_request(&message, Some(param)));
9028 }
9029 if let Err(msg) = nonstream_deadline_gate(
9034 &request,
9035 req.stream,
9036 deadline,
9037 req.max_tokens.is_some(),
9038 st.budget_tokenizers
9039 .as_ref()
9040 .and_then(|t| t.get(&req.model))
9041 .map(Arc::as_ref),
9042 ) {
9043 return with_request_id(
9044 &env.id,
9045 error_response_coded(
9046 StatusCode::BAD_REQUEST,
9047 &msg,
9048 "invalid_request_error",
9049 Some("max_tokens"),
9050 Some("nonstream_deadline_infeasible"),
9051 ),
9052 );
9053 }
9054 if draining() {
9057 let receipt = start_request_receipt(
9058 &st,
9059 &env,
9060 &tenant,
9061 &req.model,
9062 "/v1/completions",
9063 lane,
9064 req.stream,
9065 effective_max_tokens(&request),
9066 None,
9067 None,
9068 );
9069 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
9070 }
9071 let budget = match admit_tenant_budget(&st, &tenant, &mut request) {
9072 Ok(budget) => budget,
9073 Err(rejection) => {
9074 let (response, error_code) = rejection.into_response();
9075 let receipt = start_request_receipt(
9076 &st,
9077 &env,
9078 &tenant,
9079 &req.model,
9080 "/v1/completions",
9081 lane,
9082 req.stream,
9083 effective_max_tokens(&request),
9084 None,
9085 None,
9086 );
9087 return ledger_rejected(receipt, response, error_code, &env.id);
9088 }
9089 };
9090 let receipt = start_request_receipt(
9091 &st,
9092 &env,
9093 &tenant,
9094 &req.model,
9095 "/v1/completions",
9096 lane,
9097 req.stream,
9098 effective_max_tokens(&request),
9099 budget.reserved_ctx,
9100 budget.permit,
9101 );
9102 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
9103 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
9106 Ok(slot) => slot,
9107 Err(resp) => {
9108 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
9109 }
9110 };
9111 if let Some(admission) = body_admission.as_mut() {
9112 admission.release();
9113 }
9114 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline.preheader(stream)) {
9117 Ok(guard) => guard,
9118 Err((resp, outcome)) => {
9119 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
9120 }
9121 };
9122 meter_admit(&env, &tenant, &model, lane);
9123 let stop_strings = request.stop_strings.clone();
9124
9125 if let Some(trace) = ttft.as_ref() {
9130 trace.mark_submitted();
9131 }
9132 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
9133 drop(pending_admit);
9134 return ledger_rejected(
9135 receipt,
9136 rl.attach(worker_unavailable_response()),
9137 "worker_unavailable",
9138 &env.id,
9139 );
9140 }
9141 pending_admit.commit();
9142 let rx = match tokio::time::timeout_at(deadline.preheader(stream).at, peek_admission(rx)).await
9146 {
9147 Ok(Ok(rx)) => rx,
9148 Ok(Err((resp, error_code))) => {
9149 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
9150 }
9151 Err(_) => {
9152 return ledger_unbilled(
9153 receipt,
9154 rl.attach(admission_deadline_response(deadline, stream)),
9155 "deadline_exceeded",
9156 "deadline_exceeded",
9157 &env.id,
9158 );
9159 }
9160 };
9161
9162 let resp = if stream {
9163 let mut receipt = receipt;
9167 let rx = match peek_first_token(rx, deadline, &mut receipt).await {
9168 Ok(rx) => rx,
9169 Err(()) => {
9170 return ledger_unbilled(
9171 receipt,
9172 rl.attach(deadline_exceeded_response(deadline.ms, true)),
9173 "deadline_exceeded",
9174 "deadline_exceeded",
9175 &env.id,
9176 );
9177 }
9178 };
9179 sse_response_with_receipt(
9180 rx,
9181 model,
9182 false,
9183 None,
9184 env.clone(),
9185 stop_strings,
9186 Some(guard),
9187 receipt,
9188 )
9189 .into_response()
9190 } else {
9191 let mut receipt = receipt;
9197 let resp = blocking_response_with_receipt(
9198 rx,
9199 model,
9200 false,
9201 stop_strings,
9202 None,
9203 env.clone(),
9204 &mut receipt,
9205 Some(deadline),
9206 )
9207 .await;
9208 drop(guard); resp.into_response()
9210 };
9211 rl.attach(with_request_id(&env.id, resp))
9212}
9213
9214async fn chat_completions_admitted(
9215 state: State<AppState>,
9216 headers: axum::http::HeaderMap,
9217 trace: Option<Extension<TtftRequestTrace>>,
9218 AdmittedJson(req, admission): AdmittedJson<ChatCompletionReq>,
9219) -> Response {
9220 chat_completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
9221}
9222
9223#[cfg(test)]
9224async fn chat_completions(
9225 State(st): State<AppState>,
9226 headers: axum::http::HeaderMap,
9227 trace: Option<Extension<TtftRequestTrace>>,
9228 request: Json<ChatCompletionReq>,
9229) -> Response {
9230 chat_completions_with_admission(State(st), headers, trace, request, None).await
9231}
9232
9233async fn chat_completions_with_admission(
9234 State(st): State<AppState>,
9235 headers: axum::http::HeaderMap,
9236 trace: Option<Extension<TtftRequestTrace>>,
9237 Json(mut req): Json<ChatCompletionReq>,
9238 mut body_admission: Option<BodyAdmissionLease>,
9239) -> Response {
9240 let env = Envelope::new(true);
9241 match canonical_model_id(&st.models, &req.model) {
9246 Some(canonical) => req.model = canonical,
9247 None => {
9248 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
9249 }
9250 }
9251 let ttft = trace.and_then(|Extension(trace)| trace.0);
9252 if let Some(trace) = ttft.as_ref() {
9253 trace.mark_parsed();
9254 trace.bind_request(&env.id, &req.model);
9255 }
9256 let tenant = match authenticate(&st.api_auth, &headers) {
9257 Ok(t) => t,
9258 Err(resp) => return with_request_id(&env.id, resp),
9259 };
9260 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
9261 Ok(ns) => ns,
9262 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
9263 };
9264 if req.messages.is_empty()
9265 || req.messages.iter().any(|message| {
9266 !matches!(
9267 message.role.as_str(),
9268 "system" | "developer" | "user" | "assistant" | "tool"
9269 )
9270 })
9271 {
9272 return with_request_id(
9273 &env.id,
9274 bad_request(
9275 "messages must use system/developer/user/assistant/tool roles",
9276 Some("messages"),
9277 ),
9278 );
9279 }
9280 if let Err((msg, param)) = reject_unsupported(&[
9285 (
9286 "logit_bias",
9287 req.logit_bias.is_some(),
9288 " (device-side sampling has no bias hook yet)",
9289 ),
9290 (
9291 "logprobs",
9292 req.logprobs
9293 .as_ref()
9294 .is_some_and(|v| v.as_bool() != Some(false)),
9295 "",
9296 ),
9297 ("top_logprobs", req.top_logprobs.is_some(), ""),
9298 (
9299 "n",
9300 req.n.is_some_and(|n| n != 1),
9301 " for n != 1 (single choice only)",
9302 ),
9303 ]) {
9304 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
9305 }
9306 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref(), req.stream) {
9309 Ok(ms) => RequestDeadline::starting_now(ms),
9310 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
9311 };
9312 let lane = match lane_for_tenant(&headers, &tenant) {
9313 Ok(l) => l,
9314 Err(resp) => return resp,
9315 };
9316 let model = req.model.clone();
9317 let stream = req.stream;
9318 let capture_prompt = st
9321 .metering
9322 .as_ref()
9323 .filter(|m| m.captures(&tenant.tenant))
9324 .map(|_| capture_chat_messages(&req.messages));
9325 let declared_max_tokens = req.max_tokens.is_some();
9329 let vision_preprocess_permit = match try_vision_preprocess(request_has_vision(&req)) {
9333 Ok(permit) => permit,
9334 Err(response) => return with_request_id(&env.id, response),
9335 };
9336 let (tx, rx) = worker::event_channel();
9337 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
9338 Ok(affinity) => affinity,
9339 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
9340 };
9341 let md = st.metadata();
9344 let mut plan = match build_chat_request_with_trace(
9345 req,
9346 st.caps.get(&model),
9347 tx,
9348 lane,
9349 affinity,
9350 ttft.clone(),
9351 md.models
9352 .get(&model)
9353 .and_then(|m| m.default_reasoning_effort.as_deref()),
9354 &AppState::sampling_defaults_in(&md, &st.caps, &model),
9355 ) {
9356 Ok(plan) => plan,
9357 Err(err) => {
9358 return with_request_id(&env.id, bad_request(&err, None));
9359 }
9360 };
9361 plan.request.cache_ns = cache_ns;
9362 plan.request.request_id = env.id.clone();
9363 plan.request.wire_deadline = Some(deadline.at.into_std());
9364 if let Err((message, param)) = apply_model_request_limits(
9365 &mut plan.request,
9366 md.models.get(&model),
9367 st.caps.get(&model),
9368 ) {
9369 return with_request_id(&env.id, bad_request(&message, Some(param)));
9370 }
9371 if let Err(msg) = nonstream_deadline_gate(
9374 &plan.request,
9375 stream,
9376 deadline,
9377 declared_max_tokens,
9378 st.budget_tokenizers
9379 .as_ref()
9380 .and_then(|t| t.get(&model))
9381 .map(Arc::as_ref),
9382 ) {
9383 return with_request_id(
9384 &env.id,
9385 error_response_coded(
9386 StatusCode::BAD_REQUEST,
9387 &msg,
9388 "invalid_request_error",
9389 Some("max_tokens"),
9390 Some("nonstream_deadline_infeasible"),
9391 ),
9392 );
9393 }
9394 plan.vision_memory = match reserve_vision_memory(&plan) {
9395 Ok(permit) => permit,
9396 Err(err) => {
9397 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
9398 }
9399 };
9400 if draining() {
9403 let receipt = start_request_receipt(
9404 &st,
9405 &env,
9406 &tenant,
9407 &model,
9408 "/v1/chat/completions",
9409 lane,
9410 stream,
9411 effective_max_tokens(&plan.request),
9412 None,
9413 None,
9414 );
9415 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
9416 }
9417 let budget = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
9418 Ok(budget) => budget,
9419 Err(rejection) => {
9420 let (response, error_code) = rejection.into_response();
9421 let receipt = start_request_receipt(
9422 &st,
9423 &env,
9424 &tenant,
9425 &model,
9426 "/v1/chat/completions",
9427 lane,
9428 stream,
9429 effective_max_tokens(&plan.request),
9430 None,
9431 None,
9432 );
9433 return ledger_rejected(receipt, response, error_code, &env.id);
9434 }
9435 };
9436 let receipt = start_request_receipt(
9437 &st,
9438 &env,
9439 &tenant,
9440 &model,
9441 "/v1/chat/completions",
9442 lane,
9443 stream,
9444 effective_max_tokens(&plan.request),
9445 budget.reserved_ctx,
9446 budget.permit,
9447 );
9448 let receipt = if let Some(prompt) = capture_prompt {
9449 arm_capture(receipt, move || prompt)
9450 } else {
9451 receipt
9452 };
9453 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
9457 Ok(slot) => slot,
9458 Err(resp) => {
9459 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
9460 }
9461 };
9462 if let Some(admission) = body_admission.as_mut() {
9463 admission.release();
9464 }
9465 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline.preheader(stream)) {
9468 Ok(guard) => guard,
9469 Err((resp, outcome)) => {
9470 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
9471 }
9472 };
9473 if let Err(err) = decode_pending_vision(&mut plan) {
9478 return ledger_rejected(
9479 receipt,
9480 rl.attach(bad_request(&err, Some("messages"))),
9481 "invalid_request_error",
9482 &env.id,
9483 );
9484 }
9485 plan.request.vision_memory = plan.vision_memory.take();
9486 drop(vision_preprocess_permit);
9487 let constraint_ready = if plan.request.grammar.is_some() {
9488 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
9489 plan.request.constraint_ready = Some(ready_tx);
9490 Some(ready_rx)
9491 } else {
9492 None
9493 };
9494 meter_admit(&env, &tenant, &model, lane);
9495 let stop_strings = plan.request.stop_strings.clone();
9496 if let Some(trace) = ttft.as_ref() {
9498 trace.mark_submitted();
9499 }
9500 if st
9501 .cmd_tx
9502 .send(Cmd::Generate(Box::new(plan.request)))
9503 .is_err()
9504 {
9505 drop(pending_admit);
9506 return ledger_rejected(
9507 receipt,
9508 rl.attach(worker_unavailable_response()),
9509 "worker_unavailable",
9510 &env.id,
9511 );
9512 }
9513 pending_admit.commit();
9514 if let Some(ready) = constraint_ready {
9520 let bound =
9521 constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.preheader(stream).remaining());
9522 match tokio::time::timeout(bound, ready).await {
9523 Ok(Ok(Ok(()))) => {}
9524 Ok(Ok(Err(err))) => {
9525 return ledger_rejected(
9526 receipt,
9527 rl.attach(engine_error_response(&err)),
9528 engine_error_code(err.class),
9529 &env.id,
9530 );
9531 }
9532 Ok(Err(_)) => {
9533 return ledger_rejected(
9534 receipt,
9535 rl.attach(worker_unavailable_response()),
9536 "worker_unavailable",
9537 &env.id,
9538 );
9539 }
9540 Err(_) if deadline.preheader(stream).remaining().is_zero() => {
9541 return ledger_unbilled(
9542 receipt,
9543 rl.attach(admission_deadline_response(deadline, stream)),
9544 "deadline_exceeded",
9545 "deadline_exceeded",
9546 &env.id,
9547 );
9548 }
9549 Err(_) => {
9550 return ledger_rejected(
9551 receipt,
9552 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
9553 "constraint_compile_timeout",
9554 &env.id,
9555 );
9556 }
9557 }
9558 }
9559 let rx = match tokio::time::timeout_at(deadline.preheader(stream).at, peek_admission(rx)).await
9561 {
9562 Ok(Ok(rx)) => rx,
9563 Ok(Err((resp, error_code))) => {
9564 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
9565 }
9566 Err(_) => {
9567 return ledger_unbilled(
9568 receipt,
9569 rl.attach(admission_deadline_response(deadline, stream)),
9570 "deadline_exceeded",
9571 "deadline_exceeded",
9572 &env.id,
9573 );
9574 }
9575 };
9576 let resp = if stream {
9577 let mut receipt = receipt;
9579 let rx = match peek_first_token(rx, deadline, &mut receipt).await {
9580 Ok(rx) => rx,
9581 Err(()) => {
9582 return ledger_unbilled(
9583 receipt,
9584 rl.attach(deadline_exceeded_response(deadline.ms, true)),
9585 "deadline_exceeded",
9586 "deadline_exceeded",
9587 &env.id,
9588 );
9589 }
9590 };
9591 sse_response_with_receipt(
9592 rx,
9593 model,
9594 true,
9595 plan.parser,
9596 env.clone(),
9597 stop_strings,
9598 Some(guard),
9599 receipt,
9600 )
9601 .into_response()
9602 } else {
9603 let mut receipt = receipt;
9606 let resp = blocking_response_with_receipt(
9607 rx,
9608 model,
9609 true,
9610 stop_strings,
9611 plan.parser,
9612 env.clone(),
9613 &mut receipt,
9614 Some(deadline),
9615 )
9616 .await;
9617 drop(guard); resp.into_response()
9619 };
9620 rl.attach(with_request_id(&env.id, resp))
9621}
9622
9623#[cfg(test)]
9632fn sse_response(
9633 rx: worker::EventReceiver,
9634 model: String,
9635 chat: bool,
9636 parser: Option<ToolStreamParser>,
9637 env: Envelope,
9638 stop_strings: Vec<String>,
9639 guard: Option<InflightGuard>,
9640) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9641 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
9642}
9643
9644#[allow(clippy::too_many_arguments)] fn sse_response_with_receipt(
9646 mut rx: worker::EventReceiver,
9647 model: String,
9648 chat: bool,
9649 mut parser: Option<ToolStreamParser>,
9650 env: Envelope,
9651 stop_strings: Vec<String>,
9652 guard: Option<InflightGuard>,
9653 mut receipt: Option<Box<dyn metering::Receipt>>,
9654) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9655 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
9659 .then(|| StopScrubber::new(stop_strings));
9660 let stream = async_stream::stream! {
9661 let _guard = guard;
9664 let mut call_index: usize = 0;
9665 let mut role_sent = false;
9668 macro_rules! chat_chunk {
9669 ($delta:expr, $finish:expr) => {{
9670 let mut delta = $delta;
9671 if chat && !role_sent {
9672 role_sent = true;
9673 delta["role"] = json!("assistant");
9674 }
9675 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
9676 "choices": [{ "index": 0, "delta": delta,
9677 "finish_reason": $finish }] }))
9678 .to_string()
9679 }};
9680 }
9681 macro_rules! piece_chunks {
9683 ($piece:expr) => {{
9684 let mut payloads: Vec<String> = Vec::new();
9685 match $piece {
9686 Piece::Content(text) => {
9687 let text = match scrub.as_mut() {
9688 Some(sc) => sc.push(&text),
9689 None => text,
9690 };
9691 if !text.is_empty() {
9692 payloads.push(chat_chunk!(json!({ "content": text }),
9693 serde_json::Value::Null));
9694 }
9695 }
9696 Piece::Reasoning(text) => payloads.push(
9700 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
9701 Piece::Call(call) => {
9702 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9703 "index": call_index, "id": call.id, "type": "function",
9704 "function": { "name": call.name, "arguments": "" } }] }),
9705 serde_json::Value::Null));
9706 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9707 "index": call_index,
9708 "function": { "arguments": call.arguments } }] }),
9709 serde_json::Value::Null));
9710 call_index += 1;
9711 }
9712 }
9713 payloads
9714 }};
9715 }
9716 let mut terminal = false;
9720 while let Some(ev) = rx.recv().await {
9721 match ev {
9722 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9724 if let Some(receipt) = receipt.as_mut()
9725 && let Err(err) = receipt.record_prompt_usage(
9726 n_prompt as u64,
9727 n_cached as u64,
9728 )
9729 {
9730 eprintln!(
9731 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9732 env.id
9733 );
9734 let _ = receipt.reject(500, "request_ledger_unavailable");
9737 let payload = request_ledger_error_body().to_string();
9738 if chat || openai_compat() {
9739 yield Ok(SseEvent::default().data(payload));
9740 yield Ok(SseEvent::default().data("[DONE]"));
9741 } else {
9742 yield Ok(SseEvent::default().event("error").data(payload));
9743 }
9744 terminal = true;
9745 break;
9746 }
9747 }
9748 Event::DeadlineExceeded { ms } => {
9749 let ledger_error = if let Some(receipt) = receipt.as_mut() {
9750 receipt
9751 .settle_unbilled(
9752 "deadline_exceeded",
9753 StatusCode::REQUEST_TIMEOUT.as_u16(),
9754 "deadline_exceeded",
9755 )
9756 .err()
9757 } else {
9758 None
9759 };
9760 let payload = if ledger_error.is_some() {
9761 request_ledger_error_body().to_string()
9762 } else {
9763 deadline_exceeded_error(ms, true).to_string()
9764 };
9765 if chat || openai_compat() {
9766 yield Ok(SseEvent::default().data(payload));
9767 yield Ok(SseEvent::default().data("[DONE]"));
9768 } else {
9769 yield Ok(SseEvent::default().event("error").data(payload));
9770 }
9771 terminal = true;
9772 break;
9773 }
9774 Event::Token { id, text } => {
9775 if let Some(receipt) = receipt.as_mut()
9776 && let Err(err) = receipt.record_completion_token()
9777 {
9778 eprintln!(
9779 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9780 env.id
9781 );
9782 let _ = receipt.reject(500, "request_ledger_unavailable");
9783 let payload = request_ledger_error_body().to_string();
9784 if chat || openai_compat() {
9785 yield Ok(SseEvent::default().data(payload));
9786 yield Ok(SseEvent::default().data("[DONE]"));
9787 } else {
9788 yield Ok(SseEvent::default().event("error").data(payload));
9789 }
9790 terminal = true;
9791 break;
9792 }
9793 if let Some(receipt) = receipt.as_mut() {
9796 receipt.capture_completion_delta(&text);
9797 }
9798 if let Some(p) = parser.as_mut() {
9799 for piece in p.push(&text) {
9800 for payload in piece_chunks!(piece) {
9801 yield Ok(SseEvent::default().data(payload));
9802 }
9803 }
9804 continue;
9805 }
9806 let text = match scrub.as_mut() {
9807 Some(sc) => sc.push(&text),
9808 None => text,
9809 };
9810 if text.is_empty() && scrub.is_some() {
9811 continue; }
9813 let payload = if chat {
9814 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
9815 } else if openai_compat() {
9816 env.stamp(json!({ "object": "text_completion", "model": model,
9817 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
9818 .to_string()
9819 } else {
9820 json!({ "model": model, "id": id, "text": text }).to_string()
9821 };
9822 yield Ok(SseEvent::default().data(payload));
9823 }
9824 Event::TokenSnapshot(_) => {}
9828 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
9829 let mut finish = stop_reason_to_finish(&stop_reason);
9830 if let Some(p) = parser.as_mut() {
9831 for piece in p.finish() {
9832 for payload in piece_chunks!(piece) {
9833 yield Ok(SseEvent::default().data(payload));
9834 }
9835 }
9836 if p.n_calls() > 0 { finish = "tool_calls"; }
9837 }
9838 if let Some(sc) = scrub.as_mut() {
9840 let tail = sc.finish();
9841 if !tail.is_empty() {
9842 let payload = if chat {
9843 chat_chunk!(json!({ "content": tail }),
9844 serde_json::Value::Null)
9845 } else {
9846 env.stamp(json!({ "object": "text_completion",
9847 "model": model,
9848 "choices": [{ "index": 0, "text": tail,
9849 "finish_reason": null }] })).to_string()
9850 };
9851 yield Ok(SseEvent::default().data(payload));
9852 }
9853 }
9854 if let Some(receipt) = receipt.as_mut()
9855 && let Err(err) = receipt.complete(
9856 metering::UsageCounts {
9857 prompt_tokens: n_prompt as u64,
9858 cached_prompt_tokens: n_cached as u64,
9859 completion_tokens: n_tokens as u64,
9860 },
9861 elapsed_s,
9862 )
9863 {
9864 eprintln!(
9865 "[ledger] ERROR: request {} completion receipt failed: {err}",
9866 env.id
9867 );
9868 let _ = receipt.reject(500, "request_ledger_unavailable");
9872 let payload = request_ledger_error_body().to_string();
9873 if chat || openai_compat() {
9874 yield Ok(SseEvent::default().data(payload));
9875 yield Ok(SseEvent::default().data("[DONE]"));
9876 } else {
9877 yield Ok(SseEvent::default().event("error").data(payload));
9878 }
9879 terminal = true;
9880 break;
9881 }
9882 if chat || openai_compat() {
9883 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
9884 let fin = if chat {
9885 let mut v = env.stamp(json!({
9886 "object": "chat.completion.chunk", "model": model,
9887 "choices": [{ "index": 0, "delta": {},
9888 "finish_reason": finish }],
9889 "usage": usage }));
9890 if !role_sent {
9892 v["choices"][0]["delta"]["role"] = json!("assistant");
9893 }
9894 v
9895 } else {
9896 env.stamp(json!({ "object": "text_completion", "model": model,
9897 "choices": [{ "index": 0, "text": "",
9898 "finish_reason": finish }],
9899 "usage": usage }))
9900 }.to_string();
9901 yield Ok(SseEvent::default().data(fin));
9902 yield Ok(SseEvent::default().data("[DONE]"));
9903 } else {
9904 let payload = json!({
9905 "stop_reason": stop_reason, "n_tokens": n_tokens,
9906 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
9907 "elapsed_s": elapsed_s
9908 }).to_string();
9909 yield Ok(SseEvent::default().event("done").data(payload));
9910 }
9911 terminal = true;
9912 break;
9913 }
9914 Event::Error(err) => {
9915 let ledger_error = if let Some(receipt) = receipt.as_mut() {
9925 receipt
9926 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
9927 .err()
9928 } else {
9929 None
9930 };
9931 if let Some(ref ledger_error) = ledger_error {
9932 eprintln!(
9933 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
9934 env.id
9935 );
9936 }
9937 let payload = if ledger_error.is_some() {
9938 request_ledger_error_body().to_string()
9939 } else {
9940 engine_error_body(&err).to_string()
9941 };
9942 if chat || openai_compat() {
9943 yield Ok(SseEvent::default().data(payload));
9946 yield Ok(SseEvent::default().data("[DONE]"));
9947 } else {
9948 yield Ok(SseEvent::default().event("error").data(payload));
9951 }
9952 terminal = true;
9953 break;
9954 }
9955 }
9956 }
9957 if !terminal {
9958 let e = worker::EngineError::overloaded(
9964 "worker closed the stream without completing (worker restart in progress)",
9965 );
9966 if let Some(receipt) = receipt.as_mut()
9967 && let Err(ledger_err) = receipt.reject(
9968 class_http(e.class).0.as_u16(),
9969 engine_error_code(e.class),
9970 )
9971 {
9972 eprintln!(
9973 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9974 env.id
9975 );
9976 }
9977 let payload = engine_error_body(&e).to_string();
9978 if chat || openai_compat() {
9979 yield Ok(SseEvent::default().data(payload));
9980 yield Ok(SseEvent::default().data("[DONE]"));
9981 } else {
9982 yield Ok(SseEvent::default().event("error").data(payload));
9983 }
9984 }
9985 };
9986 Sse::new(stream).keep_alive(
9987 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
9990 )
9991}
9992
9993fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
9995 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
9996 text.truncate(offset);
9997 }
9998}
9999
10000fn partial_stop_suffix(s: &str, tag: &str) -> usize {
10003 let mut best = 0;
10004 for (k, _) in tag.char_indices().skip(1) {
10005 if k <= s.len() && s.ends_with(&tag[..k]) {
10006 best = k;
10007 }
10008 }
10009 best
10010}
10011
10012struct StopScrubber {
10018 stops: Vec<String>,
10019 buf: String,
10020 done: bool,
10021}
10022
10023impl StopScrubber {
10024 fn new(stops: Vec<String>) -> Self {
10025 Self {
10026 stops,
10027 buf: String::new(),
10028 done: false,
10029 }
10030 }
10031
10032 fn push(&mut self, text: &str) -> String {
10034 if self.done {
10035 return String::new();
10036 }
10037 self.buf.push_str(text);
10038 if let Some(i) = self
10039 .stops
10040 .iter()
10041 .filter_map(|s| self.buf.find(s.as_str()))
10042 .min()
10043 {
10044 self.done = true;
10045 let out = self.buf[..i].to_string();
10046 self.buf.clear();
10047 return out;
10048 }
10049 let keep = self
10050 .stops
10051 .iter()
10052 .map(|s| partial_stop_suffix(&self.buf, s))
10053 .max()
10054 .unwrap_or(0);
10055 let emit_to = self.buf.len() - keep;
10056 let out = self.buf[..emit_to].to_string();
10057 self.buf.drain(..emit_to);
10058 out
10059 }
10060
10061 fn finish(&mut self) -> String {
10063 if self.done {
10064 self.buf.clear();
10065 return String::new();
10066 }
10067 std::mem::take(&mut self.buf)
10068 }
10069}
10070
10071#[cfg(test)]
10072async fn blocking_response(
10073 rx: worker::EventReceiver,
10074 model: String,
10075 chat: bool,
10076 stop_strings: Vec<String>,
10077 parser: Option<ToolStreamParser>,
10078 env: Envelope,
10079) -> Response {
10080 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
10081 .await
10082}
10083
10084struct BlockingPayload<'a> {
10088 env: &'a Envelope,
10089 model: String,
10090 chat: bool,
10091 finish: &'static str,
10092 text: String,
10093 reasoning: String,
10094 calls: Vec<ParsedToolCall>,
10095 tokens: Vec<u32>,
10096 stop_reason: String,
10097 n_prompt: usize,
10098 n_tokens: usize,
10099 n_cached: usize,
10100 elapsed_s: f64,
10101 spec: Option<worker::SpecUsage>,
10102 deadline_error: Option<serde_json::Value>,
10108}
10109
10110fn blocking_payload(p: BlockingPayload<'_>) -> Response {
10111 let BlockingPayload {
10112 env,
10113 model,
10114 chat,
10115 finish,
10116 text,
10117 reasoning,
10118 calls,
10119 tokens,
10120 stop_reason,
10121 n_prompt,
10122 n_tokens,
10123 n_cached,
10124 elapsed_s,
10125 spec,
10126 deadline_error,
10127 } = p;
10128 if chat {
10129 let content = if !calls.is_empty() && text.is_empty() {
10131 serde_json::Value::Null
10132 } else {
10133 serde_json::Value::String(text)
10134 };
10135 let mut message = json!({ "role": "assistant", "content": content });
10136 if !reasoning.is_empty() {
10139 message["reasoning"] = json!(reasoning);
10140 message["reasoning_details"] = json!([{
10141 "type": "reasoning.text", "text": reasoning }]);
10142 }
10143 if !calls.is_empty() {
10144 message["tool_calls"] =
10145 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
10146 }
10147 let mut body = json!({
10148 "object": "chat.completion", "model": model,
10149 "choices": [{ "index": 0,
10150 "message": message,
10151 "finish_reason": finish }],
10152 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
10153 });
10154 if let Some(err) = deadline_error {
10155 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
10156 body["error"] = err;
10157 }
10158 return Json(env.stamp(body)).into_response();
10159 }
10160 if openai_compat() {
10161 let mut body = json!({
10162 "object": "text_completion", "model": model,
10163 "choices": [{ "index": 0, "text": text,
10164 "finish_reason": finish }],
10165 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
10166 });
10167 if let Some(err) = deadline_error {
10168 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
10169 body["error"] = err;
10170 }
10171 return Json(env.stamp(body)).into_response();
10172 }
10173 Json(CompletionResp {
10174 model,
10175 text,
10176 tokens,
10177 stop_reason,
10178 error: deadline_error,
10179 n_tokens,
10180 prompt_tokens: n_prompt,
10181 cached_tokens: n_cached,
10182 elapsed_s,
10183 })
10184 .into_response()
10185}
10186
10187#[allow(clippy::too_many_arguments)] async fn blocking_response_with_receipt(
10213 mut rx: worker::EventReceiver,
10214 model: String,
10215 chat: bool,
10216 stop_strings: Vec<String>,
10217 mut parser: Option<ToolStreamParser>,
10218 env: Envelope,
10219 receipt: &mut Option<Box<dyn metering::Receipt>>,
10220 deadline: Option<RequestDeadline>,
10221) -> Response {
10222 let mut text = String::new();
10223 let mut reasoning = String::new();
10224 let mut tokens: Vec<u32> = Vec::new();
10225 let mut calls: Vec<ParsedToolCall> = Vec::new();
10226 let consume = |pieces: Vec<Piece>,
10227 text: &mut String,
10228 reasoning: &mut String,
10229 calls: &mut Vec<ParsedToolCall>| {
10230 for piece in pieces {
10231 match piece {
10232 Piece::Content(t) => text.push_str(&t),
10233 Piece::Reasoning(t) => reasoning.push_str(&t),
10234 Piece::Call(c) => calls.push(c),
10235 }
10236 }
10237 };
10238 let started = std::time::Instant::now();
10240 let mut seen_prompt: usize = 0;
10241 let mut seen_cached: usize = 0;
10242 let mut seen_tokens: usize = 0;
10243 loop {
10244 let ev = match deadline {
10245 Some(d) => tokio::select! {
10246 biased;
10247 ev = rx.recv() => ev,
10248 () = tokio::time::sleep_until(d.at) => {
10249 drop(rx);
10252 if seen_tokens == 0 {
10253 if let Some(receipt) = receipt.as_mut()
10257 && let Err(err) = receipt.settle_unbilled(
10258 "deadline_exceeded",
10259 StatusCode::REQUEST_TIMEOUT.as_u16(),
10260 "deadline_exceeded",
10261 )
10262 {
10263 eprintln!(
10264 "[ledger] ERROR: request {} deadline receipt failed: {err}",
10265 env.id
10266 );
10267 return request_ledger_error_response();
10268 }
10269 return deadline_exceeded_response(d.ms, false);
10270 }
10271 if let Some(p) = parser.as_mut() {
10272 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
10273 }
10274 truncate_at_stop(&mut text, &stop_strings);
10275 let elapsed_s = started.elapsed().as_secs_f64();
10276 if let Some(receipt) = receipt.as_mut()
10279 && let Err(err) = receipt.complete_deadline_partial(
10280 metering::UsageCounts {
10281 prompt_tokens: seen_prompt as u64,
10282 cached_prompt_tokens: seen_cached as u64,
10283 completion_tokens: seen_tokens as u64,
10284 },
10285 elapsed_s,
10286 )
10287 {
10288 eprintln!(
10289 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
10290 env.id
10291 );
10292 let _ = receipt.reject(500, "request_ledger_unavailable");
10293 return request_ledger_error_response();
10294 }
10295 eprintln!(
10296 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
10297 {} ms deadline (prompt {}); non-streaming caller advised to stream",
10298 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
10299 );
10300 let err_obj = json!({
10301 "message": format!(
10302 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
10303 the {} tokens produced before the cut are delivered above and are \
10304 billed. Set \"stream\": true for work this long — a stream's \
10305 deadline bounds only the time to first token — or lower max_tokens.",
10306 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
10307 ),
10308 "code": "deadline_exceeded",
10309 "metadata": { "error_type": "timeout", "provider_name": "memra" }
10310 });
10311 return blocking_payload(BlockingPayload {
10312 env: &env,
10313 model,
10314 chat,
10315 finish: "error",
10316 text,
10317 reasoning,
10318 calls,
10319 tokens,
10320 stop_reason: "Deadline".to_string(),
10321 n_prompt: seen_prompt,
10322 n_tokens: seen_tokens,
10323 n_cached: seen_cached,
10324 elapsed_s,
10325 spec: None,
10326 deadline_error: Some(err_obj),
10327 });
10328 }
10329 },
10330 None => rx.recv().await,
10331 };
10332 let Some(ev) = ev else { break };
10333 match ev {
10334 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
10336 if let Some(receipt) = receipt.as_mut()
10337 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
10338 {
10339 eprintln!(
10340 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
10341 env.id
10342 );
10343 let _ = receipt.reject(500, "request_ledger_unavailable");
10346 return request_ledger_error_response();
10347 }
10348 seen_prompt = n_prompt;
10349 seen_cached = n_cached;
10350 }
10351 Event::Token { id, text: delta } => {
10352 if let Some(receipt) = receipt.as_mut()
10353 && let Err(err) = receipt.record_completion_token()
10354 {
10355 eprintln!(
10356 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
10357 env.id
10358 );
10359 let _ = receipt.reject(500, "request_ledger_unavailable");
10360 return request_ledger_error_response();
10361 }
10362 if let Some(receipt) = receipt.as_mut() {
10364 receipt.capture_completion_delta(&delta);
10365 }
10366 tokens.push(id);
10367 seen_tokens += 1;
10368 match parser.as_mut() {
10369 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
10370 None => text.push_str(&delta),
10371 }
10372 }
10373 Event::TokenSnapshot(ids) => tokens = ids,
10374 Event::DeadlineExceeded { ms } => {
10375 if let Some(receipt) = receipt.as_mut()
10376 && let Err(ledger_err) = receipt.settle_unbilled(
10377 "deadline_exceeded",
10378 StatusCode::REQUEST_TIMEOUT.as_u16(),
10379 "deadline_exceeded",
10380 )
10381 {
10382 eprintln!(
10383 "[ledger] ERROR: request {} deadline receipt failed: {ledger_err}",
10384 env.id
10385 );
10386 return request_ledger_error_response();
10387 }
10388 return deadline_exceeded_response(ms, false);
10389 }
10390 Event::Done {
10391 stop_reason,
10392 n_tokens,
10393 n_prompt,
10394 n_cached,
10395 elapsed_s,
10396 spec,
10397 } => {
10398 if let Some(p) = parser.as_mut() {
10399 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
10400 }
10401 truncate_at_stop(&mut text, &stop_strings);
10402 let finish = if calls.is_empty() {
10403 stop_reason_to_finish(&stop_reason)
10404 } else {
10405 "tool_calls"
10406 };
10407 if let Some(receipt) = receipt.as_mut()
10408 && let Err(err) = receipt.complete(
10409 metering::UsageCounts {
10410 prompt_tokens: n_prompt as u64,
10411 cached_prompt_tokens: n_cached as u64,
10412 completion_tokens: n_tokens as u64,
10413 },
10414 elapsed_s,
10415 )
10416 {
10417 eprintln!(
10418 "[ledger] ERROR: request {} completion receipt failed: {err}",
10419 env.id
10420 );
10421 let _ = receipt.reject(500, "request_ledger_unavailable");
10424 return request_ledger_error_response();
10425 }
10426 return blocking_payload(BlockingPayload {
10427 env: &env,
10428 model,
10429 chat,
10430 finish,
10431 text,
10432 reasoning,
10433 calls,
10434 tokens,
10435 stop_reason,
10436 n_prompt,
10437 n_tokens,
10438 n_cached,
10439 elapsed_s,
10440 spec,
10441 deadline_error: None,
10442 });
10443 }
10444 Event::Error(err) => {
10445 if let Some(receipt) = receipt.as_mut()
10449 && let Err(ledger_err) = receipt.reject(
10450 class_http(err.class).0.as_u16(),
10451 engine_error_code(err.class),
10452 )
10453 {
10454 eprintln!(
10455 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
10456 env.id
10457 );
10458 return request_ledger_error_response();
10459 }
10460 return engine_error_response(&err);
10461 }
10462 }
10463 }
10464 let e = worker::EngineError::overloaded(
10469 "worker closed the stream without completing (worker restart in progress)",
10470 );
10471 if let Some(receipt) = receipt.as_mut()
10472 && let Err(ledger_err) =
10473 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
10474 {
10475 eprintln!(
10476 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
10477 env.id
10478 );
10479 return request_ledger_error_response();
10480 }
10481 engine_error_response(&e)
10482}
10483
10484#[cfg(test)]
10485mod tests {
10486 use super::*;
10487
10488 #[test]
10494 fn capture_children_are_distinct_ledger_identities_under_the_parent_id() {
10495 let parent = Envelope::new(false);
10496 assert!(parent.id.starts_with("cmpl-"));
10497 let a = parent.capture_child(0);
10498 let b = parent.capture_child(1);
10499 let c = parent.capture_child(2);
10500 assert_eq!(a.id, format!("{}.0", parent.id));
10501 assert_eq!(b.id, format!("{}.1", parent.id));
10502 assert_eq!(c.id, format!("{}.2", parent.id));
10503 assert_ne!(a.id, b.id);
10504 assert_ne!(b.id, c.id);
10505 for child in [&a, &b, &c] {
10506 assert!(
10507 child.id.starts_with(&parent.id),
10508 "child nests under the parent by prefix"
10509 );
10510 assert_ne!(
10511 child.id, parent.id,
10512 "a child never reuses the parent's ledger id"
10513 );
10514 assert_eq!(child.created, parent.created);
10515 }
10516 assert_eq!(parent.capture_child(1).id, b.id);
10519 }
10520
10521 #[derive(Debug, Clone, PartialEq)]
10529 enum MeterEvent {
10530 Reserve {
10531 tenant: String,
10532 principal: Option<String>,
10533 model: String,
10534 },
10535 Open {
10536 request_id: String,
10537 tenant: String,
10538 model: String,
10539 route: &'static str,
10540 stream: bool,
10541 with_permit: bool,
10542 },
10543 PromptUsage {
10544 prompt: u64,
10545 cached: u64,
10546 },
10547 Token,
10548 CapturePrompt(serde_json::Value),
10549 CaptureDelta(String),
10550 Complete {
10551 prompt: u64,
10552 cached: u64,
10553 completion: u64,
10554 },
10555 DeadlinePartial {
10556 prompt: u64,
10557 cached: u64,
10558 completion: u64,
10559 },
10560 Reject {
10561 status: u16,
10562 code: String,
10563 },
10564 Unbilled {
10565 outcome: &'static str,
10566 status: u16,
10567 code: String,
10568 },
10569 Dropped {
10572 prompt: u64,
10573 cached: u64,
10574 completion: u64,
10575 },
10576 }
10577
10578 enum ReserveScript {
10581 Admit { with_permit: bool },
10582 Insufficient,
10583 Blocked,
10584 PrincipalCapped,
10585 }
10586
10587 struct MockMetering {
10588 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10589 limits: bool,
10590 limited: bool,
10591 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
10592 captures: bool,
10593 }
10594
10595 impl MockMetering {
10596 fn admit_all() -> Arc<Self> {
10597 Arc::new(MockMetering {
10598 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10599 limits: false,
10600 limited: true,
10601 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10602 captures: false,
10603 })
10604 }
10605
10606 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
10607 Arc::new(MockMetering {
10608 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10609 limits: true,
10610 limited: true,
10611 reserve_script: std::sync::Mutex::new(script.into()),
10612 captures: false,
10613 })
10614 }
10615
10616 fn capturing() -> Arc<Self> {
10617 Arc::new(MockMetering {
10618 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10619 limits: false,
10620 limited: true,
10621 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10622 captures: true,
10623 })
10624 }
10625
10626 fn events(&self) -> Vec<MeterEvent> {
10627 self.events.lock().unwrap().clone()
10628 }
10629 }
10630
10631 impl metering::Metering for MockMetering {
10632 fn enforces_limits(&self) -> bool {
10633 self.limits
10634 }
10635
10636 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
10637 Ok(self.limited)
10638 }
10639
10640 fn reserve(
10641 &self,
10642 tenant: &str,
10643 principal: Option<&str>,
10644 model: &str,
10645 _prompt_tokens: u64,
10646 _completion_bound: u64,
10647 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
10648 self.events.lock().unwrap().push(MeterEvent::Reserve {
10649 tenant: tenant.into(),
10650 principal: principal.map(str::to_owned),
10651 model: model.into(),
10652 });
10653 match self.reserve_script.lock().unwrap().pop_front() {
10654 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
10655 Some(ReserveScript::Admit { with_permit: true }) => {
10656 Ok(Some(Box::new(()) as metering::Permit))
10657 }
10658 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
10659 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
10660 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
10661 }
10662 }
10663
10664 fn open(
10665 &self,
10666 meta: &metering::RequestMeta<'_>,
10667 permit: Option<metering::Permit>,
10668 ) -> Box<dyn metering::Receipt> {
10669 self.events.lock().unwrap().push(MeterEvent::Open {
10670 request_id: meta.request_id.into(),
10671 tenant: meta.tenant.into(),
10672 model: meta.model.into(),
10673 route: meta.route,
10674 stream: meta.stream,
10675 with_permit: permit.is_some(),
10676 });
10677 Box::new(MockReceipt {
10678 events: self.events.clone(),
10679 wants_capture: self.captures,
10680 prompt: 0,
10681 cached: 0,
10682 completion: 0,
10683 finalized: false,
10684 })
10685 }
10686
10687 fn captures(&self, _tenant: &str) -> bool {
10688 self.captures
10689 }
10690
10691 fn limits_health(&self) -> Option<metering::LimitsHealth> {
10692 self.limits.then_some(metering::LimitsHealth {
10693 source_reload_failed: 0,
10694 source_reload_consecutive: 0,
10695 source_available: true,
10696 })
10697 }
10698 }
10699
10700 struct MockReceipt {
10701 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10702 wants_capture: bool,
10703 prompt: u64,
10704 cached: u64,
10705 completion: u64,
10706 finalized: bool,
10707 }
10708
10709 impl metering::Receipt for MockReceipt {
10710 fn wants_capture(&self) -> bool {
10711 self.wants_capture
10712 }
10713
10714 fn arm_capture(&mut self, prompt: serde_json::Value) {
10715 self.events
10716 .lock()
10717 .unwrap()
10718 .push(MeterEvent::CapturePrompt(prompt));
10719 }
10720
10721 fn capture_completion_delta(&mut self, text: &str) {
10722 if self.wants_capture {
10723 self.events
10724 .lock()
10725 .unwrap()
10726 .push(MeterEvent::CaptureDelta(text.into()));
10727 }
10728 }
10729
10730 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
10731 self.prompt = prompt;
10732 self.cached = cached;
10733 self.events
10734 .lock()
10735 .unwrap()
10736 .push(MeterEvent::PromptUsage { prompt, cached });
10737 Ok(())
10738 }
10739
10740 fn record_completion_token(&mut self) -> Result<(), String> {
10741 self.completion += 1;
10742 self.events.lock().unwrap().push(MeterEvent::Token);
10743 Ok(())
10744 }
10745
10746 fn complete(
10747 &mut self,
10748 usage: metering::UsageCounts,
10749 _worker_elapsed_s: f64,
10750 ) -> Result<(), String> {
10751 self.finalized = true;
10752 self.events.lock().unwrap().push(MeterEvent::Complete {
10753 prompt: usage.prompt_tokens,
10754 cached: usage.cached_prompt_tokens,
10755 completion: usage.completion_tokens,
10756 });
10757 Ok(())
10758 }
10759
10760 fn complete_deadline_partial(
10761 &mut self,
10762 usage: metering::UsageCounts,
10763 _worker_elapsed_s: f64,
10764 ) -> Result<(), String> {
10765 self.finalized = true;
10766 self.events
10767 .lock()
10768 .unwrap()
10769 .push(MeterEvent::DeadlinePartial {
10770 prompt: usage.prompt_tokens,
10771 cached: usage.cached_prompt_tokens,
10772 completion: usage.completion_tokens,
10773 });
10774 Ok(())
10775 }
10776
10777 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
10778 self.finalized = true;
10779 self.events.lock().unwrap().push(MeterEvent::Reject {
10780 status,
10781 code: error_code.into(),
10782 });
10783 Ok(())
10784 }
10785
10786 fn settle_unbilled(
10787 &mut self,
10788 outcome: &'static str,
10789 status: u16,
10790 error_code: &str,
10791 ) -> Result<(), String> {
10792 self.finalized = true;
10793 self.events.lock().unwrap().push(MeterEvent::Unbilled {
10794 outcome,
10795 status,
10796 code: error_code.into(),
10797 });
10798 Ok(())
10799 }
10800 }
10801
10802 impl Drop for MockReceipt {
10803 fn drop(&mut self) {
10804 if !self.finalized {
10805 self.events.lock().unwrap().push(MeterEvent::Dropped {
10806 prompt: self.prompt,
10807 cached: self.cached,
10808 completion: self.completion,
10809 });
10810 }
10811 }
10812 }
10813
10814 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
10819
10820 fn gate_env_lock() -> std::sync::MutexGuard<'static, ()> {
10827 let guard = GATE_ENV_LOCK.lock().unwrap_or_else(|poisoned| {
10828 GATE_ENV_LOCK.clear_poison();
10831 poisoned.into_inner()
10832 });
10833 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10834 guard
10835 }
10836
10837 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
10840 let req: CompletionReq = serde_json::from_value(json!({
10841 "model": "qwen/qwen3.8-27b",
10842 "prompt_ids": vec![7u32; prompt_ids],
10843 }))
10844 .unwrap();
10845 let (tx, _rx) = worker::event_channel();
10846 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
10847 request.params.max_new = max_new;
10848 request
10849 }
10850
10851 #[test]
10858 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
10859 let prompt = 30_278u64;
10860 let deadline_ms = TIMEOUT_MS_DEFAULT;
10861 let margin = |max_new: u64| {
10862 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
10863 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
10864 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
10865 };
10866 for allowed in [64u64, 2048, 4096, 5120, 6144] {
10867 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
10868 }
10869 for refused in [8192u64, 16384, 262_144] {
10870 assert!(
10871 !margin(refused),
10872 "{refused} measured as a 408 and must be refused"
10873 );
10874 }
10875 }
10876
10877 #[test]
10878 fn the_gate_names_a_max_tokens_that_actually_fits() {
10879 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
10882 assert!(
10883 fits > 0 && fits < 7_800,
10884 "advice {fits} must fit the measured ceiling"
10885 );
10886 assert_eq!(
10888 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
10889 None
10890 );
10891 }
10892
10893 #[test]
10894 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
10895 let req = gate_request(262_144, 30_000);
10896 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
10897 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
10899 assert!(
10900 err.contains("stream"),
10901 "message must name the streaming alternative: {err}"
10902 );
10903 assert!(
10904 err.contains("max_tokens"),
10905 "message must name the knob: {err}"
10906 );
10907 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
10909 let _l = gate_env_lock(); for off in ["0", "off", "false"] {
10916 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
10917 assert!(
10918 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
10919 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
10920 );
10921 }
10922 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
10923 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
10924 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10925 assert!(
10926 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
10927 "unset means ON (the documented default)"
10928 );
10929 }
10930
10931 #[test]
10938 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
10939 let strip = |src: &str| -> String {
10946 src.lines()
10947 .map(|line| match line.find("//") {
10948 Some(i) => line[..i].to_string(),
10949 None => line.to_string(),
10950 })
10951 .collect::<Vec<_>>()
10952 .join("\n")
10953 };
10954 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
10956 let start = src
10957 .find(signature)
10958 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
10959 let rest = &src[start + signature.len()..];
10960 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
10961 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
10962 &rest[..end]
10963 }
10964 let main_src = strip(include_str!("lib.rs"));
10965 let surfaces_src = strip(include_str!("surfaces.rs"));
10966 for (surface, src, signature) in [
10967 (
10968 "/v1/completions",
10969 &main_src,
10970 "async fn completions_with_admission(",
10971 ),
10972 (
10973 "/v1/chat/completions",
10974 &main_src,
10975 "async fn chat_completions_with_admission(",
10976 ),
10977 (
10978 "/v1/messages + /v1/responses (shared admission)",
10979 &surfaces_src,
10980 "pub(crate) async fn admit_translated(",
10981 ),
10982 ] {
10983 let handler = body(src, signature);
10984 assert!(
10985 handler.contains("nonstream_deadline_gate("),
10986 "{surface} must CALL the feasibility gate inside {signature}"
10987 );
10988 let limits = handler
10991 .find("apply_model_request_limits(")
10992 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
10993 let gate = handler.find("nonstream_deadline_gate(").unwrap();
10994 assert!(
10995 limits < gate,
10996 "{surface}: the gate must run after apply_model_request_limits"
10997 );
10998 }
10999 }
11000
11001 #[test]
11005 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
11006 let err = json!({"code": "deadline_exceeded",
11007 "metadata": {"error_type": "timeout"}});
11008 let cut = CompletionResp {
11009 model: "m".into(),
11010 text: "partial".into(),
11011 tokens: vec![1, 2],
11012 stop_reason: "Deadline".into(),
11013 error: Some(err.clone()),
11014 n_tokens: 2,
11015 prompt_tokens: 9,
11016 cached_tokens: 0,
11017 elapsed_s: 1.0,
11018 };
11019 let v = serde_json::to_value(&cut).unwrap();
11020 assert_eq!(v["stop_reason"], "Deadline");
11021 assert_eq!(v["error"]["code"], "deadline_exceeded");
11022 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
11023 let whole = CompletionResp {
11025 error: None,
11026 stop_reason: "Eos".into(),
11027 ..cut
11028 };
11029 let v = serde_json::to_value(&whole).unwrap();
11030 assert!(
11031 v.get("error").is_none(),
11032 "a complete response must not grow an error key: {v}"
11033 );
11034 }
11035
11036 #[test]
11037 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
11038 let _l = gate_env_lock();
11039 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
11043 assert!(
11044 nonstream_deadline_gate(
11045 &req,
11046 false,
11047 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
11048 false,
11049 None,
11050 )
11051 .is_ok(),
11052 "an omitted max_tokens is never gated — context is its only limit"
11053 );
11054 let resolved = gate_request(32_768, 30_000);
11059 assert!(
11060 nonstream_deadline_gate(
11061 &resolved,
11062 false,
11063 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
11064 false,
11065 None,
11066 )
11067 .is_ok(),
11068 "a resolved-but-undeclared cap is not the caller's number to be refused over"
11069 );
11070 assert!(
11072 nonstream_deadline_gate(
11073 &resolved,
11074 false,
11075 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
11076 true,
11077 None,
11078 )
11079 .is_err()
11080 );
11081 }
11082
11083 #[test]
11084 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
11085 let req = gate_request(64, 1234);
11086 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
11087 let mut text = gate_request(64, 0);
11088 text.prompt_ids.clear();
11089 text.prompt_text = "x".repeat(6_000);
11090 assert_eq!(
11091 prompt_tokens_estimate(&text, None),
11092 1_000,
11093 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
11094 that would have succeeded"
11095 );
11096 }
11097
11098 #[test]
11099 fn vision_memory_reservation_is_bounded_and_released() {
11100 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
11101 let Err(capacity) = try_reserve_vision_memory(1) else {
11102 panic!("a full process vision budget admitted another request");
11103 };
11104 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
11105 let response = vision_memory_error_response(capacity, Some("messages"));
11106 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
11107 assert_eq!(response.headers()["retry-after"], "5");
11108 assert_eq!(response.headers()["retry-after-ms"], "5000");
11109 drop(permit);
11110 assert!(try_reserve_vision_memory(1).is_ok());
11111 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
11112 panic!("an over-limit vision request was admitted");
11113 };
11114 assert!(matches!(request, VisionMemoryError::Request(_)));
11115 let response = vision_memory_error_response(request, Some("messages"));
11116 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
11117 assert_eq!(response.headers()["x-should-retry"], "false");
11118 let _ = try_reserve_vision_memory(1);
11119 }
11120
11121 #[test]
11122 fn header_auth_gate_covers_only_inference_dialects() {
11123 for path in [
11124 "/v1/auth/check",
11125 "/v1/completions",
11126 "/v1/chat/completions",
11127 "/v1/messages",
11128 "/v1/responses",
11129 "/v1/embeddings",
11130 "/v1/rerank",
11131 ] {
11132 assert!(protected_inference_path(path), "{path}");
11133 }
11134 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
11135 assert!(!protected_inference_path(path), "{path}");
11136 }
11137 }
11138 #[tokio::test]
11145 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
11146 use crate::metering::Metering as _;
11147 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
11148
11149 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
11150 let (tx, rx) = worker::event_channel();
11151 tx.send(Event::PromptUsage {
11152 n_prompt: 7,
11153 n_cached: 0,
11154 })
11155 .unwrap();
11156 tx.send(Event::Token {
11157 id: 1,
11158 text: "Hel".into(),
11159 })
11160 .unwrap();
11161 tx.send(Event::Token {
11162 id: 2,
11163 text: "lo".into(),
11164 })
11165 .unwrap();
11166 tx.send(Event::Done {
11167 stop_reason: "eos".into(),
11168 n_tokens: 2,
11169 n_prompt: 7,
11170 n_cached: 0,
11171 elapsed_s: 0.05,
11172 spec: None,
11173 })
11174 .unwrap();
11175 drop(tx);
11176 let mut receipt = receipt;
11177 blocking_response_with_receipt(
11178 rx,
11179 "m".into(),
11180 true,
11181 Vec::new(),
11182 None,
11183 Envelope::new(true),
11184 &mut receipt,
11185 None,
11186 )
11187 .await
11188 };
11189
11190 let plain = MockMetering::admit_all();
11192 let receipt = plain.open(
11193 &metering::RequestMeta {
11194 request_id: "cap-unmarked",
11195 tenant: "unmarked",
11196 principal: None,
11197 model: "m",
11198 route: "/v1/chat/completions",
11199 lane: "interactive",
11200 stream: false,
11201 max_tokens: None,
11202 reserved_ctx: None,
11203 },
11204 None,
11205 );
11206 let response = drive(Some(receipt)).await;
11207 assert_eq!(response.status(), StatusCode::OK);
11208 assert!(
11209 !plain.events().iter().any(|e| matches!(
11210 e,
11211 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
11212 )),
11213 "an unarmed receipt must see no capture traffic: {:?}",
11214 plain.events()
11215 );
11216
11217 let capturing = MockMetering::capturing();
11220 let mut receipt = capturing.open(
11221 &metering::RequestMeta {
11222 request_id: "cap-marked",
11223 tenant: "marked",
11224 principal: None,
11225 model: "m",
11226 route: "/v1/chat/completions",
11227 lane: "interactive",
11228 stream: false,
11229 max_tokens: None,
11230 reserved_ctx: None,
11231 },
11232 None,
11233 );
11234 assert!(receipt.wants_capture());
11235 receipt.arm_capture(prompt.clone());
11236 let response = drive(Some(receipt)).await;
11237 assert_eq!(response.status(), StatusCode::OK);
11238 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
11239 .await
11240 .unwrap();
11241 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
11242 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
11243
11244 let events = capturing.events();
11245 assert!(
11246 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
11247 "prompt must arm byte-exact: {events:?}"
11248 );
11249 let completion: String = events
11250 .iter()
11251 .filter_map(|e| match e {
11252 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
11253 _ => None,
11254 })
11255 .collect();
11256 assert_eq!(
11257 completion, "Hello",
11258 "the deltas must reassemble the served completion byte-exact: {events:?}"
11259 );
11260 assert!(
11261 events.contains(&MeterEvent::Complete {
11262 prompt: 7,
11263 cached: 0,
11264 completion: 2,
11265 }),
11266 "worker-truth usage settles alongside the capture: {events:?}"
11267 );
11268 }
11269
11270 fn tool_caps() -> ModelCaps {
11271 ModelCaps {
11272 tools_branch: true,
11273 qwen_think: true,
11274 think_switch: true,
11275 chat_ok: true,
11276 ..Default::default()
11277 }
11278 }
11279
11280 fn ladder_caps() -> ModelCaps {
11285 ModelCaps {
11286 qwen_effort: true,
11287 ..tool_caps()
11288 }
11289 }
11290
11291 fn gemma_tool_caps() -> ModelCaps {
11292 ModelCaps {
11293 tools_branch: true,
11294 gemma_think: true,
11295 chat_ok: true,
11296 instruct_type: Some("gemma".into()),
11297 ..Default::default()
11298 }
11299 }
11300
11301 fn hy3_tool_caps() -> ModelCaps {
11302 ModelCaps {
11303 tools_branch: true,
11304 hy3: true,
11305 chat_ok: true,
11306 effort_levels: true,
11307 instruct_type: Some("hy3".into()),
11308 ..Default::default()
11309 }
11310 }
11311
11312 fn gemma_template(kind: &str) -> String {
11313 let file = match kind {
11314 "qat" => "qat-trunk-template.jinja",
11315 _ => "official-tooluse-template.jinja",
11316 };
11317 let path = format!(
11318 "{}/../../research/gemma4-tools-20260817/{file}",
11319 env!("CARGO_MANIFEST_DIR")
11320 );
11321 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
11322 }
11323
11324 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
11329 let tools_arr = request
11330 .get("tools")
11331 .and_then(|t| t.as_array())
11332 .cloned()
11333 .unwrap_or_default();
11334 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
11335 (Vec::new(), Vec::new(), HashMap::new())
11336 } else {
11337 prepare_tools(&tools_arr).unwrap()
11338 };
11339 let effort = request
11340 .get("reasoning_effort")
11341 .and_then(|v| v.as_str())
11342 .map(String::from);
11343 let (think, _lvl, _explicit) =
11344 parse_think(&effort, &None, None, None, None, false).unwrap();
11345
11346 let mut turns: Vec<TmplTurn> = Vec::new();
11347 for msg in request["messages"].as_array().unwrap() {
11348 let role = msg["role"].as_str().unwrap();
11349 let role = if role == "developer" { "system" } else { role };
11350 let content =
11351 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
11352 let tool_calls = msg
11353 .get("tool_calls")
11354 .and_then(|a| a.as_array())
11355 .map(|a| {
11356 a.iter()
11357 .map(|tc| {
11358 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
11359 render_req_tool_call(&rtc).unwrap()
11360 })
11361 .collect()
11362 })
11363 .unwrap_or_default();
11364 let tool_responses = msg
11365 .get("tool_responses")
11366 .and_then(|a| a.as_array())
11367 .map(|a| {
11368 a.iter()
11369 .map(|tr| {
11370 (
11371 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
11372 json_to_val(&tr["response"]),
11373 )
11374 })
11375 .collect()
11376 })
11377 .unwrap_or_default();
11378 turns.push(TmplTurn {
11379 role: role.to_string(),
11380 content,
11381 tool_calls,
11382 reasoning: msg
11383 .get("reasoning")
11384 .and_then(|r| r.as_str())
11385 .map(String::from)
11386 .filter(|s| !s.is_empty()),
11387 tool_call_id: msg
11388 .get("tool_call_id")
11389 .and_then(|s| s.as_str())
11390 .map(String::from),
11391 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
11392 tool_responses,
11393 task: None,
11394 tools: Vec::new(),
11395 });
11396 }
11397 chat::apply_chat_template_tools_ex(
11398 Some(template),
11399 &turns,
11400 true,
11401 &tools_json,
11402 &tools_struct,
11403 think,
11404 None,
11405 None,
11406 )
11407 .unwrap()
11408 }
11409
11410 #[test]
11414 fn gemma4_tools_fixtures_match_the_official_jinja() {
11415 let dir = format!(
11416 "{}/../../research/gemma4-tools-20260817/fixtures",
11417 env!("CARGO_MANIFEST_DIR")
11418 );
11419 let mut entries: Vec<_> = std::fs::read_dir(&dir)
11420 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
11421 .map(|e| e.unwrap().path())
11422 .filter(|p| p.is_dir())
11423 .collect();
11424 entries.sort();
11425 assert!(
11426 entries.len() >= 14,
11427 "expected >=14 fixtures, found {}",
11428 entries.len()
11429 );
11430 let (mut official, mut qat) = (0u32, 0u32);
11431 for d in entries {
11432 let input: serde_json::Value =
11433 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11434 .unwrap();
11435 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11436 let kind = input
11437 .get("template")
11438 .and_then(|t| t.as_str())
11439 .unwrap_or("official");
11440 match kind {
11441 "qat" => qat += 1,
11442 _ => official += 1,
11443 }
11444 let tmpl = gemma_template(kind);
11445 let got = render_fixture(&input["request"], &tmpl);
11446 assert_eq!(
11447 got, expected,
11448 "fixture {:?} diverged from the jinja oracle",
11449 d
11450 );
11451 }
11452 assert!(
11453 official >= 12 && qat >= 2,
11454 "coverage: {official} official, {qat} qat"
11455 );
11456 }
11457
11458 #[test]
11464 fn gemma4_tools_flow_through_build_chat_request() {
11465 let tmpl = gemma_template("official");
11466 for name in [
11467 "01-system-tools-basic",
11468 "04-single-call-cycle",
11469 "07-multi-cycle-agentic",
11470 ] {
11471 let path = format!(
11472 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
11473 env!("CARGO_MANIFEST_DIR")
11474 );
11475 let input: serde_json::Value =
11476 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
11477 let expected_path = format!(
11478 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
11479 env!("CARGO_MANIFEST_DIR")
11480 );
11481 let expected = std::fs::read_to_string(&expected_path).unwrap();
11482 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
11483 let (tx, _rx) = worker::event_channel();
11484 let plan = build_chat_request(
11485 req,
11486 Some(&gemma_tool_caps()),
11487 tx,
11488 lanes::Lane::Interactive,
11489 None,
11490 )
11491 .unwrap();
11492 let got = chat::apply_chat_template_tools_ex(
11493 Some(&tmpl),
11494 &plan.request.chat_turns,
11495 true,
11496 &plan.request.tools_json,
11497 &plan.request.tools_struct,
11498 plan.request.think,
11499 plan.request.reasoning_effort.as_deref(),
11500 None,
11501 )
11502 .unwrap();
11503 assert_eq!(got, expected, "pipeline render diverged for {name}");
11504 }
11505 }
11506
11507 fn glm5_template() -> String {
11516 let path = format!(
11517 "{}/../../research/glm53-flash-bringup-20260827/chat_template.jinja",
11518 env!("CARGO_MANIFEST_DIR")
11519 );
11520 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
11521 }
11522
11523 fn glm5_caps() -> ModelCaps {
11528 ModelCaps {
11529 tools_branch: true,
11530 qwen_think: true,
11531 think_switch: false,
11532 chat_ok: true,
11533 context_length: 1_048_576,
11534 tokenizer: "glm4".into(),
11535 instruct_type: Some("glm".into()),
11536 effort_levels: true,
11537 glm5: true,
11538 ..Default::default()
11539 }
11540 }
11541
11542 fn glm5_render(body: serde_json::Value) -> Result<String, String> {
11544 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11545 let (tx, _rx) = worker::event_channel();
11546 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)?;
11547 chat::apply_chat_template_tools_ex(
11548 Some(&glm5_template()),
11549 &plan.request.chat_turns,
11550 true,
11551 &plan.request.tools_json,
11552 &plan.request.tools_struct,
11553 plan.request.think,
11554 plan.request.reasoning_effort.as_deref(),
11555 None,
11556 )
11557 }
11558
11559 #[test]
11564 fn glm5_fixtures_match_the_vendor_jinja() {
11565 let dir = format!(
11566 "{}/../../research/glm53-flash-bringup-20260827/surface-fixtures",
11567 env!("CARGO_MANIFEST_DIR")
11568 );
11569 let mut entries: Vec<_> = std::fs::read_dir(&dir)
11570 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
11571 .map(|e| e.unwrap().path())
11572 .filter(|p| p.is_dir())
11573 .collect();
11574 entries.sort();
11575 assert!(
11576 entries.len() >= 22,
11577 "expected >=22 fixtures, found {}",
11578 entries.len()
11579 );
11580 for d in entries {
11581 let input: serde_json::Value =
11582 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11583 .unwrap();
11584 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11585 let got = glm5_render(input["request"].clone())
11586 .unwrap_or_else(|e| panic!("fixture {d:?} refused: {e}"));
11587 assert_eq!(
11588 got, expected,
11589 "fixture {d:?} diverged from the jinja oracle"
11590 );
11591 }
11592 }
11593
11594 #[test]
11600 fn glm5_never_renders_chatml() {
11601 let tmpl = glm5_template();
11602 assert!(tmpl.contains("<think>") && tmpl.contains("add_generation_prompt"));
11604 assert!(tmpl.contains("<tools>"));
11605 assert!(chat::template_is_glm5(&tmpl));
11606 for body in [
11607 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11608 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11609 "tools": [{"type": "function", "function": {"name": "f",
11610 "parameters": {"type": "object", "properties": {}}}}]}),
11611 ] {
11612 let got = glm5_render(body).unwrap();
11613 assert!(
11614 !got.contains("<|im_start|>") && !got.contains("<|im_end|>"),
11615 "glm5 rendered ChatML frames: {got:?}"
11616 );
11617 assert!(
11618 got.starts_with("[gMASK]<sop><|system|>Reasoning Effort: "),
11619 "{got:?}"
11620 );
11621 assert!(got.ends_with("<|assistant|><think>"), "{got:?}");
11622 }
11623 }
11624
11625 #[test]
11629 fn glm5_reasoning_effort_renders_and_keeps_its_max_tier() {
11630 for (sent, line) in [
11631 (None, "Max"),
11632 (Some("low"), "Low"),
11633 (Some("medium"), "High"),
11638 (Some("high"), "High"),
11639 (Some("xhigh"), "Max"),
11640 (Some("max"), "Max"),
11641 (Some("ultra"), "Max"),
11642 ] {
11643 let mut body = json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]});
11644 if let Some(v) = sent {
11645 body["reasoning_effort"] = json!(v);
11646 }
11647 let got = glm5_render(body).unwrap();
11648 assert!(
11649 got.starts_with(&format!("[gMASK]<sop><|system|>Reasoning Effort: {line}<|")),
11650 "reasoning_effort {sent:?} should render {line:?}: {got:?}"
11651 );
11652 }
11653 let sampler_of = |v: &str| {
11656 let req: ChatCompletionReq = serde_json::from_value(
11657 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11660 "reasoning_effort": v, "seed": 7}),
11661 )
11662 .unwrap();
11663 let (tx, _rx) = worker::event_channel();
11664 let plan =
11665 build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11666 .unwrap();
11667 format!("{:?}", plan.request.sampler_cfg)
11668 };
11669 assert_eq!(sampler_of("low"), sampler_of("max"));
11670 assert_eq!(canonical_effort_for("max", true), Some("max"));
11672 assert_eq!(canonical_effort_for("xhigh", true), Some("max"));
11673 assert_eq!(canonical_effort_for("max", false), Some("high"));
11674 }
11675
11676 #[test]
11680 fn glm5_refuses_what_its_template_cannot_honour() {
11681 for (value, needle) in [
11682 ("none", "cannot disable reasoning"),
11683 ("minimal", "cannot disable reasoning"),
11684 ("bogus", "bad reasoning_effort"),
11685 ] {
11686 let err = glm5_render(json!({"model": "m",
11687 "messages": [{"role": "user", "content": "hi"}],
11688 "reasoning_effort": value}))
11689 .err()
11690 .unwrap_or_else(|| panic!("reasoning_effort {value:?} must be refused"));
11691 assert!(err.contains(needle), "{value}: {err}");
11692 }
11693 }
11694
11695 #[test]
11699 fn one_glm5_request_renders_identical_bytes_on_all_three_surfaces() {
11700 let chat = json!({
11707 "model": "m",
11708 "reasoning_effort": "high",
11709 "messages": [
11710 {"role": "user", "content": "Weather in Paris and Rome?"},
11711 {"role": "assistant", "content": null,
11712 "tool_calls": [
11713 {"id": "c1", "type": "function",
11714 "function": {"name": "get_weather",
11715 "arguments": "{\"city\": \"Paris\"}"}},
11716 {"id": "c2", "type": "function",
11717 "function": {"name": "get_weather",
11718 "arguments": "{\"city\": \"Rome\"}"}}]},
11719 {"role": "tool", "tool_call_id": "c2", "content": "rome:27"},
11720 {"role": "tool", "tool_call_id": "c1", "content": "paris:21"}
11721 ],
11722 "tools": [{"type": "function", "function": {
11723 "name": "get_weather", "description": "Get the current weather for a city",
11724 "parameters": {"type": "object",
11725 "properties": {"city": {"type": "string"}},
11726 "required": ["city"]}}}]
11727 });
11728 let responses = responses_api::translate(&json!({
11729 "model": "m",
11730 "reasoning": {"effort": "high"},
11731 "input": [
11732 {"type": "message", "role": "user",
11733 "content": [{"type": "input_text", "text": "Weather in Paris and Rome?"}]},
11734 {"type": "function_call", "call_id": "c1", "name": "get_weather",
11735 "arguments": "{\"city\": \"Paris\"}"},
11736 {"type": "function_call", "call_id": "c2", "name": "get_weather",
11737 "arguments": "{\"city\": \"Rome\"}"},
11738 {"type": "function_call_output", "call_id": "c2", "output": "rome:27"},
11739 {"type": "function_call_output", "call_id": "c1", "output": "paris:21"}
11740 ],
11741 "tools": [{"type": "function", "name": "get_weather",
11742 "description": "Get the current weather for a city",
11743 "parameters": {"type": "object",
11744 "properties": {"city": {"type": "string"}},
11745 "required": ["city"]}}]
11746 }))
11747 .expect("/v1/responses translate");
11748 let messages = anthropic::translate(&json!({
11749 "model": "m",
11750 "max_tokens": 256,
11751 "output_config": {"effort": "high"},
11752 "messages": [
11753 {"role": "user", "content": "Weather in Paris and Rome?"},
11754 {"role": "assistant", "content": [
11755 {"type": "tool_use", "id": "c1", "name": "get_weather",
11756 "input": {"city": "Paris"}},
11757 {"type": "tool_use", "id": "c2", "name": "get_weather",
11758 "input": {"city": "Rome"}}]},
11759 {"role": "user", "content": [
11760 {"type": "tool_result", "tool_use_id": "c2", "content": "rome:27"},
11761 {"type": "tool_result", "tool_use_id": "c1", "content": "paris:21"}]}
11762 ],
11763 "tools": [{"name": "get_weather",
11764 "description": "Get the current weather for a city",
11765 "input_schema": {"type": "object",
11766 "properties": {"city": {"type": "string"}},
11767 "required": ["city"]}}]
11768 }))
11769 .expect("/v1/messages translate");
11770 let want = glm5_render(chat).expect("chat");
11771 assert!(
11773 want.contains(
11774 "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value>\
11775 </tool_call><tool_call>get_weather<arg_key>city</arg_key>\
11776 <arg_value>Rome</arg_value></tool_call>"
11777 ),
11778 "{want:?}"
11779 );
11780 assert!(
11784 want.contains(
11785 "<|observation|><tool_response>paris:21</tool_response>\
11786 <tool_response>rome:27</tool_response>"
11787 ),
11788 "{want:?}"
11789 );
11790 assert!(
11791 want.contains("<|system|>Reasoning Effort: High"),
11792 "{want:?}"
11793 );
11794 for (surface, body) in [
11795 ("/v1/responses", responses),
11796 ("/v1/messages", messages.clone()),
11797 ] {
11798 let got = glm5_render(body).unwrap_or_else(|e| panic!("{surface}: {e}"));
11799 assert_eq!(
11800 got, want,
11801 "{surface} rendered DIFFERENT glm5 prompt bytes than /v1/chat/completions"
11802 );
11803 }
11804 let mut idless = messages;
11810 for m in idless["messages"].as_array_mut().unwrap() {
11811 if m["role"] == "tool" {
11812 m.as_object_mut().unwrap().remove("tool_call_id");
11813 }
11814 }
11815 let got = glm5_render(idless).expect("id-less render");
11816 assert_ne!(
11817 got, want,
11818 "dropping tool_call_id must change the rendered order — this test cannot detect \
11819 a surface that loses ids otherwise"
11820 );
11821 assert!(
11822 got.contains(
11823 "<|observation|><tool_response>rome:27</tool_response>\
11824 <tool_response>paris:21</tool_response>"
11825 ),
11826 "{got:?}"
11827 );
11828 }
11829
11830 #[test]
11833 fn glm5_chat_arms_the_native_tool_parser() {
11834 let req: ChatCompletionReq = serde_json::from_value(json!({
11835 "model": "m", "messages": [{"role": "user", "content": "weather?"}],
11836 "tools": [{"type": "function", "function": {"name": "get_weather",
11837 "parameters": {"type": "object",
11838 "properties": {"city": {"type": "string"}}}}}]}))
11839 .unwrap();
11840 let (tx, _rx) = worker::event_channel();
11841 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11842 .unwrap();
11843 let mut parser = plan.parser.expect("glm5 tools request must carry a parser");
11844 let pieces = parser.push(
11845 "reasoning here</think><tool_call>get_weather<arg_key>city</arg_key>\
11846 <arg_value>Paris</arg_value></tool_call>",
11847 );
11848 let calls: Vec<_> = pieces
11849 .iter()
11850 .filter_map(|p| match p {
11851 toolcall::Piece::Call(c) => Some((c.name.as_str(), c.arguments.as_str())),
11852 _ => None,
11853 })
11854 .collect();
11855 assert_eq!(
11856 calls,
11857 vec![("get_weather", r#"{"city":"Paris"}"#)],
11858 "{pieces:?}"
11859 );
11860 assert!(
11861 pieces
11862 .iter()
11863 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "reasoning here")),
11864 "{pieces:?}"
11865 );
11866 assert!(
11868 !pieces
11869 .iter()
11870 .any(|p| matches!(p, toolcall::Piece::Content(_))),
11871 "{pieces:?}"
11872 );
11873 let req: ChatCompletionReq = serde_json::from_value(
11877 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11878 )
11879 .unwrap();
11880 let (tx, _rx) = worker::event_channel();
11881 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11882 .unwrap();
11883 let mut parser = plan
11884 .parser
11885 .expect("glm5 non-tools request must still split reasoning");
11886 let pieces = parser.push("weighing it</think>The answer.");
11887 assert!(
11888 pieces
11889 .iter()
11890 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "weighing it")),
11891 "{pieces:?}"
11892 );
11893 assert!(
11894 pieces
11895 .iter()
11896 .any(|p| matches!(p, toolcall::Piece::Content(c) if c == "The answer.")),
11897 "{pieces:?}"
11898 );
11899 }
11900
11901 #[test]
11908 fn glm5_plain_fast_path_never_drops_replayed_reasoning() {
11909 let with_reasoning = vec![
11910 chat::Turn {
11911 role: "user".into(),
11912 content: "a".into(),
11913 ..Default::default()
11914 },
11915 chat::Turn {
11916 role: "assistant".into(),
11917 content: "A".into(),
11918 reasoning: Some("I considered a.".into()),
11919 ..Default::default()
11920 },
11921 chat::Turn {
11922 role: "user".into(),
11923 content: "b".into(),
11924 ..Default::default()
11925 },
11926 ];
11927 assert!(!worker::plain_chat_render_path(
11929 &[],
11930 &chat::ThinkMode::Default,
11931 None,
11932 &with_reasoning,
11933 false,
11934 ));
11935 let plain_turns: Vec<chat::Turn> = with_reasoning
11938 .iter()
11939 .cloned()
11940 .map(|mut t| {
11941 t.reasoning = None;
11942 t
11943 })
11944 .collect();
11945 assert!(worker::plain_chat_render_path(
11946 &[],
11947 &chat::ThinkMode::Default,
11948 None,
11949 &plain_turns,
11950 false,
11951 ));
11952 let tmpl = glm5_template();
11955 let via_tools = chat::apply_chat_template_tools_ex(
11956 Some(&tmpl),
11957 &with_reasoning,
11958 true,
11959 &[],
11960 &[],
11961 chat::ThinkMode::Default,
11962 None,
11963 None,
11964 )
11965 .unwrap();
11966 let msgs: Vec<(&str, &str)> = with_reasoning
11967 .iter()
11968 .map(|t| (t.role.as_str(), t.content.as_str()))
11969 .collect();
11970 let via_plain = chat::apply_chat_template_str(Some(&tmpl), &msgs, true);
11971 assert!(
11972 via_tools.contains("<think>I considered a.</think>"),
11973 "{via_tools:?}"
11974 );
11975 assert_ne!(via_tools, via_plain);
11976 let plain_msgs: Vec<(&str, &str)> = plain_turns
11979 .iter()
11980 .map(|t| (t.role.as_str(), t.content.as_str()))
11981 .collect();
11982 assert_eq!(
11983 chat::apply_chat_template_tools_ex(
11984 Some(&tmpl),
11985 &plain_turns,
11986 true,
11987 &[],
11988 &[],
11989 chat::ThinkMode::Default,
11990 None,
11991 None,
11992 )
11993 .unwrap(),
11994 chat::apply_chat_template_str(Some(&tmpl), &plain_msgs, true)
11995 );
11996 }
11997
11998 #[test]
12002 fn glm5_model_row_does_not_claim_structured_output() {
12003 let caps = glm5_caps();
12004 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), None);
12005 assert_eq!(row["capabilities"]["structured_output"], json!(false));
12006 assert_eq!(row["capabilities"]["tools"], json!(true));
12007 assert_eq!(row["capabilities"]["reasoning"], json!(true));
12008 let err = glm5_render(json!({"model": "m",
12010 "messages": [{"role": "user", "content": "hi"}],
12011 "response_format": {"type": "json_object"}}))
12012 .expect_err("response_format must be refused on a switchless think template");
12013 assert!(
12017 err.contains("neither an enable_thinking switch nor a recognizable"),
12018 "{err}"
12019 );
12020 let switchable = model_entry_v1("q", Some(&tool_caps()), None);
12022 assert_eq!(switchable["capabilities"]["structured_output"], json!(true));
12023 let glm_params = openrouter_supported_parameters(Some(&caps), None, true);
12026 assert!(
12027 glm_params.get("structured_outputs").is_none(),
12028 "{glm_params}"
12029 );
12030 let step_like = ModelCaps {
12036 chat_ok: true,
12037 qwen_think: true,
12038 think_switch: false,
12039 think_close: vec![128799],
12040 ..caps.clone()
12041 };
12042 let step_row = model_entry_v1("stepfun/step-3.7-flash", Some(&step_like), None);
12043 assert_eq!(step_row["capabilities"]["structured_output"], json!(true));
12044 let step_params = openrouter_supported_parameters(Some(&step_like), None, true);
12045 assert!(
12046 step_params.get("structured_outputs").is_some(),
12047 "{step_params}"
12048 );
12049 assert!(glm_params.get("json_mode").is_none(), "{glm_params}");
12050 assert!(glm_params.get("tools").is_some(), "{glm_params}");
12051 assert_eq!(
12054 glm_params.get("reasoning_effort"),
12055 Some(&json!({ "type": "enum", "values": ["low", "high", "max"] })),
12056 "{glm_params}"
12057 );
12058 let qwen_params = openrouter_supported_parameters(Some(&tool_caps()), None, true);
12059 assert!(
12060 qwen_params.get("structured_outputs").is_some(),
12061 "{qwen_params}"
12062 );
12063 assert!(qwen_params.get("json_mode").is_some(), "{qwen_params}");
12064 assert!(
12065 qwen_params.get("reasoning_effort").is_none(),
12066 "{qwen_params}"
12067 );
12068 let ladder_params = openrouter_supported_parameters(Some(&ladder_caps()), None, true);
12073 assert!(glm_params.get("reasoning").is_none(), "{glm_params}");
12077 assert!(step_params.get("reasoning").is_none(), "{step_params}");
12078 assert!(qwen_params.get("reasoning").is_some(), "{qwen_params}");
12079 assert!(ladder_params.get("reasoning").is_some(), "{ladder_params}");
12080 assert_eq!(
12081 ladder_params.get("reasoning_effort"),
12082 Some(&json!({ "type": "enum", "values": ["xhigh", "medium", "low"] })),
12083 "{ladder_params}"
12084 );
12085 }
12086
12087 #[test]
12095 fn catalog_context_claim_is_capped_by_the_deployment_envelope() {
12096 let caps = glm5_caps();
12097 assert_eq!(caps.context_length, 1_048_576);
12098 let metadata = OpenRouterModelMetadata {
12099 max_prompt_length: Some(126_976),
12100 max_output_length: Some(4_096),
12101 ..Default::default()
12102 };
12103 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
12105 assert_eq!(row["context_length"], json!(131_072));
12106 let or_row = model_entry_openrouter("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
12107 assert_eq!(
12108 or_row["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
12109 json!(131_072)
12110 );
12111 assert_eq!(
12112 published_context_length(Some(&caps), Some(&metadata)),
12113 Some(131_072)
12114 );
12115 assert_eq!(published_context_length(Some(&caps), None), Some(1_048_576));
12117 let half = OpenRouterModelMetadata {
12118 max_output_length: Some(4_096),
12119 ..Default::default()
12120 };
12121 assert_eq!(
12122 published_context_length(Some(&caps), Some(&half)),
12123 Some(1_048_576)
12124 );
12125 let wide = OpenRouterModelMetadata {
12127 max_prompt_length: Some(2_000_000),
12128 max_output_length: Some(2_000_000),
12129 ..Default::default()
12130 };
12131 assert_eq!(
12132 published_context_length(Some(&caps), Some(&wide)),
12133 Some(1_048_576)
12134 );
12135 }
12136
12137 #[test]
12147 fn retrieve_model_row_matches_the_list_row_for_a_slash_name() {
12148 let name = "zai/glm-5.3-flash";
12149 let caps_map = HashMap::from([(name.to_string(), glm5_caps())]);
12150 let models = vec![name.to_string()];
12151 let md = ModelMetadataSet::default();
12152 let got = retrieve_model_row(&models, &caps_map, &md, name).expect("known id resolves");
12153 assert_eq!(got["id"], json!(name));
12154 assert_eq!(
12155 got,
12156 model_entry_v1(name, Some(&glm5_caps()), None),
12157 "the retrieve row must equal the list row for the same model",
12158 );
12159 assert!(retrieve_model_row(&models, &caps_map, &md, "no/such-model").is_none());
12167 assert!(retrieve_model_row(&models, &caps_map, &md, "").is_none());
12168 }
12169
12170 #[tokio::test]
12177 async fn retrieve_route_delivers_both_id_spellings_to_one_handler() {
12178 async fn echo(Path(id): Path<String>) -> String {
12179 id
12180 }
12181 let app = Router::new().route("/v1/models/*id", get(echo));
12182 for url in [
12183 "/v1/models/zai/glm-5.3-flash",
12184 "/v1/models/zai%2Fglm-5.3-flash",
12185 "/v1/models/zai%2fglm-5.3-flash",
12186 ] {
12187 let resp = app
12188 .clone()
12189 .oneshot(axum::http::Request::get(url).body(Body::empty()).unwrap())
12190 .await
12191 .unwrap();
12192 assert_eq!(resp.status(), StatusCode::OK, "{url}");
12193 let raw = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
12194 let got = String::from_utf8(raw.to_vec()).unwrap();
12195 assert_eq!(
12196 got, "zai/glm-5.3-flash",
12197 "wildcard extraction for {url}: axum 0.7 percent-decodes the tail, so the raw-slash and %2F spellings arrive identical",
12198 );
12199 }
12200 }
12201
12202 fn dsv4_sentinel() -> String {
12210 let path = format!(
12211 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
12212 env!("CARGO_MANIFEST_DIR")
12213 );
12214 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
12215 }
12216
12217 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
12222 let role = msg["role"].as_str().unwrap().to_string();
12223 let content =
12224 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
12225 let reasoning = msg
12226 .get("reasoning")
12227 .or_else(|| msg.get("reasoning_content"))
12228 .and_then(|r| r.as_str())
12229 .map(String::from)
12230 .filter(|s| !s.is_empty());
12231 let tool_calls = msg
12232 .get("tool_calls")
12233 .and_then(|a| a.as_array())
12234 .map(|a| {
12235 a.iter()
12236 .map(|tc| {
12237 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
12238 render_req_tool_call(&rtc).unwrap()
12239 })
12240 .collect()
12241 })
12242 .unwrap_or_default();
12243 let tools = msg
12244 .get("tools")
12245 .and_then(|a| a.as_array())
12246 .map(|a| {
12247 a.iter()
12248 .filter_map(|t| t.get("function").map(json_to_val))
12249 .collect()
12250 })
12251 .unwrap_or_default();
12252 TmplTurn {
12253 role,
12254 content,
12255 tool_calls,
12256 reasoning,
12257 tool_call_id: msg
12258 .get("tool_call_id")
12259 .and_then(|s| s.as_str())
12260 .map(String::from),
12261 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
12262 tool_responses: Vec::new(),
12263 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
12264 tools,
12265 }
12266 }
12267
12268 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
12269 v.and_then(|t| t.as_array())
12270 .map(|a| {
12271 a.iter()
12272 .filter_map(|t| t.get("function").map(json_to_val))
12273 .collect()
12274 })
12275 .unwrap_or_default()
12276 }
12277
12278 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
12282 let dir = format!(
12283 "{}/../../research/dsv4-template-20260818/{subdir}",
12284 env!("CARGO_MANIFEST_DIR")
12285 );
12286 let tmpl = dsv4_sentinel();
12287 let mut entries: Vec<_> = std::fs::read_dir(&dir)
12288 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
12289 .map(|e| e.unwrap().path())
12290 .filter(|p| p.is_dir())
12291 .collect();
12292 entries.sort();
12293 assert!(
12294 entries.len() >= min_fixtures,
12295 "expected >={min_fixtures} fixtures, found {}",
12296 entries.len()
12297 );
12298 for d in &entries {
12299 let input: serde_json::Value =
12300 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
12301 .unwrap();
12302 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
12303 let turns: Vec<TmplTurn> = input["turns"]
12304 .as_array()
12305 .unwrap()
12306 .iter()
12307 .map(dsv4_turn)
12308 .collect();
12309 let think = match input["think"].as_str().unwrap() {
12310 "chat" => ThinkMode::NoThink,
12311 _ => ThinkMode::Think,
12312 };
12313 let effort = input
12314 .get("reasoning_effort")
12315 .and_then(|v| v.as_str())
12316 .map(String::from);
12317 let req_tools = dsv4_req_tools(input.get("req_tools"));
12318 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
12319 let got = chat::apply_chat_template_tools_ex(
12320 Some(&tmpl),
12321 &turns,
12322 agp,
12323 &[],
12324 &req_tools,
12325 think,
12326 effort.as_deref(),
12327 Some(encoding),
12328 )
12329 .unwrap();
12330 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
12331 }
12332 }
12333
12334 #[test]
12335 fn dsv4_template_fixtures_match_the_oracle() {
12336 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
12337 }
12338
12339 #[test]
12345 fn dsv4_0731_fixtures_match_the_oracle() {
12346 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
12347 }
12348
12349 #[test]
12350 fn dsv4_artifact_fixtures_are_byte_identical() {
12351 let base = format!(
12355 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
12356 env!("CARGO_MANIFEST_DIR")
12357 );
12358 let tmpl = dsv4_sentinel();
12359 for (n, think) in [
12360 (1u32, ThinkMode::Think),
12361 (2, ThinkMode::Think),
12362 (3, ThinkMode::Think),
12363 (4, ThinkMode::NoThink),
12364 ] {
12365 let td: serde_json::Value = serde_json::from_str(
12366 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
12367 )
12368 .unwrap();
12369 let (messages, tools) = if td.is_object() {
12370 (td["messages"].clone(), td.get("tools").cloned())
12371 } else {
12372 (td.clone(), None)
12373 };
12374 let mut turns: Vec<TmplTurn> = Vec::new();
12375 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
12376 let mut t = dsv4_turn(msg);
12377 if i == 0
12378 && let Some(tl) = &tools
12379 {
12380 t.tools = tl
12381 .as_array()
12382 .unwrap()
12383 .iter()
12384 .filter_map(|x| x.get("function").map(json_to_val))
12385 .collect();
12386 }
12387 turns.push(t);
12388 }
12389 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
12390 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
12394 let got = chat::apply_chat_template_tools_ex(
12395 Some(&tmpl),
12396 &turns,
12397 true,
12398 &[],
12399 &[],
12400 think,
12401 None,
12402 Some(encoding),
12403 )
12404 .unwrap();
12405 assert_eq!(
12406 got, expected,
12407 "artifact fixture {n} diverged from the oracle under {encoding:?}"
12408 );
12409 }
12410 }
12411 }
12412
12413 #[test]
12414 fn dsv4_default_thinkmode_renders_thinking() {
12415 let tmpl = dsv4_sentinel();
12418 let turns = vec![TmplTurn {
12419 role: "user".into(),
12420 content: "Hi".into(),
12421 ..Default::default()
12422 }];
12423 let dflt = chat::apply_chat_template_tools_ex(
12424 Some(&tmpl),
12425 &turns,
12426 true,
12427 &[],
12428 &[],
12429 ThinkMode::Default,
12430 None,
12431 None,
12432 )
12433 .unwrap();
12434 let think = chat::apply_chat_template_tools_ex(
12435 Some(&tmpl),
12436 &turns,
12437 true,
12438 &[],
12439 &[],
12440 ThinkMode::Think,
12441 None,
12442 None,
12443 )
12444 .unwrap();
12445 assert_eq!(dflt, think);
12446 assert!(
12447 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
12448 "{dflt:?}"
12449 );
12450 let chat_mode = chat::apply_chat_template_tools_ex(
12451 Some(&tmpl),
12452 &turns,
12453 true,
12454 &[],
12455 &[],
12456 ThinkMode::NoThink,
12457 None,
12458 None,
12459 )
12460 .unwrap();
12461 assert!(
12462 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
12463 "{chat_mode:?}"
12464 );
12465 }
12466
12467 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
12472 let base = format!(
12473 "{}/../../research/dsv4-template-20260818",
12474 env!("CARGO_MANIFEST_DIR")
12475 );
12476 let refdir = std::path::Path::new(&base).join("ref");
12477 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
12478 .expect("load dsv4 tokenizer from ref dir");
12479 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
12480 let banked: serde_json::Value = serde_json::from_str(
12481 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
12482 .unwrap(),
12483 )
12484 .unwrap();
12485 let obj = banked.as_object().unwrap();
12486 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
12487 for (name, ids_v) in obj {
12488 let rendered =
12489 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
12490 let want: Vec<u32> = ids_v
12491 .as_array()
12492 .unwrap()
12493 .iter()
12494 .map(|v| v.as_u64().unwrap() as u32)
12495 .collect();
12496 let got = tok.encode(&rendered, true);
12497 assert_eq!(got, want, "tokenization diverged for {name}");
12498 }
12499 }
12500
12501 #[test]
12502 fn dsv4_tokenization_crosscheck_matches_official_ids() {
12503 dsv4_run_tokenization_crosscheck("fixtures");
12504 }
12505
12506 #[test]
12510 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
12511 dsv4_run_tokenization_crosscheck("fixtures-0731");
12512 }
12513
12514 #[test]
12515 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
12516 let base = format!(
12526 "{}/../../research/dsv4-template-20260818",
12527 env!("CARGO_MANIFEST_DIR")
12528 );
12529 let refdir = std::path::Path::new(&base).join("ref");
12530 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
12531 .expect("load dsv4 tokenizer from ref dir");
12532 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
12533 let tmpl = dsv4_sentinel();
12534 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
12535 {"type": "function", "function": {
12536 "name": "get_data",
12537 "description": "Fetch a blob",
12538 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
12539 "required": ["key"]}
12540 }}
12541 ])));
12542
12543 let cases: Vec<(&str, String)> = vec![
12544 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
12546 ("space-131k", " ".repeat(131_072)),
12547 ("digit-131k", "7".repeat(131_072)),
12548 (
12549 "mixed-runs",
12550 format!(
12551 "{}{}{}{}",
12552 "Z".repeat(65_536),
12553 " ".repeat(65_536),
12554 "7".repeat(65_536),
12555 "\n".repeat(65_536)
12556 ),
12557 ),
12558 ("cjk-64k", "中".repeat(65_536)),
12559 ("accented-letter-64k", "é".repeat(65_536)),
12560 ];
12561 for (name, blob) in &cases {
12562 let msgs = serde_json::json!([
12563 {"role": "system", "content": "You are a tool-using assistant."},
12564 {"role": "user", "content": "Fetch the blob."},
12565 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
12566 "tool_calls": [{"id": "call_001", "type": "function",
12567 "function": {"name": "get_data",
12568 "arguments": "{\"key\": \"blob\"}"}}]},
12569 {"role": "tool", "tool_call_id": "call_001", "content": blob}
12570 ]);
12571 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
12572 let rendered = chat::apply_chat_template_tools_ex(
12573 Some(&tmpl),
12574 &turns,
12575 true,
12576 &[],
12577 &req_tools,
12578 ThinkMode::Think,
12579 None,
12580 None,
12581 )
12582 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
12583 assert!(
12584 rendered.contains(blob.as_str()),
12585 "{name}: tool result missing from render"
12586 );
12587 let t0 = std::time::Instant::now();
12588 let ids = tok.encode(&rendered, true);
12589 let encode_dt = t0.elapsed();
12590 assert!(!ids.is_empty(), "{name}: empty encode");
12591 let back = tok.decode(&ids);
12592 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
12593 assert!(
12597 encode_dt < std::time::Duration::from_secs(60),
12598 "{name}: encode took {encode_dt:?}"
12599 );
12600 if *name == "ascii-letter-131k"
12604 && let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR")
12605 {
12606 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
12607 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
12608 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
12609 }
12610 }
12611 }
12612
12613 #[test]
12614 fn models_v1_entry_advertises_thinking_support() {
12615 let step_caps = ModelCaps {
12618 effort_levels: true,
12619 ..tool_caps()
12620 };
12621 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
12622 assert_eq!(entry["capabilities"]["reasoning"], true);
12623 assert_eq!(entry["capabilities"]["tools"], true);
12624
12625 let plain = ModelCaps {
12627 chat_ok: true,
12628 ..Default::default()
12629 };
12630 let entry = model_entry_v1("plain", Some(&plain), None);
12631 assert_eq!(entry["capabilities"]["reasoning"], false);
12632 assert_eq!(entry["capabilities"]["tools"], false);
12633 let entry = model_entry_v1("unknown", None, None);
12635 assert_eq!(entry["capabilities"]["reasoning"], false);
12636 assert_eq!(entry["capabilities"]["streaming"], true);
12637 }
12638
12639 #[test]
12640 fn chat_request_preserves_turns_and_openai_stop_forms() {
12641 let payload = serde_json::json!({
12642 "model": "plain_quant",
12643 "messages": [
12644 {"role": "system", "content": "rules"},
12645 {"role": "developer", "content": "dev rules"},
12646 {"role": "user", "content": "task"},
12647 {"role": "assistant", "content": "work"}
12648 ],
12649 "max_tokens": 64,
12650 "temperature": 0.0,
12651 "stop": "<stop>"
12652 });
12653 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12654 let (tx, _rx) = worker::event_channel();
12655 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
12656 let request = plan.request;
12657 assert!(
12658 plan.parser.is_none(),
12659 "no tools -> no parser (isolation contract)"
12660 );
12661 assert!(request.tools_json.is_empty());
12662 assert_eq!(request.think, ThinkMode::Default);
12663 assert_eq!(request.model, "plain_quant");
12664 assert_eq!(request.params.max_new, 64);
12665 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12667 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
12668 }))
12669 .unwrap();
12670 let (tx, _rx) = worker::event_channel();
12671 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
12672 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
12673 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12675 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12676 "max_completion_tokens": 7
12677 }))
12678 .unwrap();
12679 let (tx, _rx) = worker::event_channel();
12680 assert_eq!(
12681 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12682 .unwrap()
12683 .request
12684 .params
12685 .max_new,
12686 7
12687 );
12688 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12690 "model": "plain_quant", "prompt": "task"
12691 }))
12692 .unwrap();
12693 let (tx, _rx) = worker::event_channel();
12694 assert_eq!(
12695 build_request(&req, tx, lanes::Lane::Interactive, None)
12696 .params
12697 .max_new,
12698 worker::MAX_NEW_CTX_BOUNDED
12699 );
12700 let turns: Vec<(String, String)> = request
12701 .chat_turns
12702 .iter()
12703 .map(|t| (t.role.clone(), t.content.clone()))
12704 .collect();
12705 assert_eq!(
12706 turns,
12707 vec![
12708 ("system".into(), "rules".into()),
12709 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
12711 ("assistant".into(), "work".into()),
12712 ]
12713 );
12714 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
12715 assert_eq!(request.stop_strings, vec!["<stop>"]);
12716
12717 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12718 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12719 "stop": ["a", "b"]
12720 }))
12721 .unwrap();
12722 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
12723
12724 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12728 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12729 "stop": ["", "real", ""]
12730 }))
12731 .unwrap();
12732 assert_eq!(req.stop.into_vec(), vec!["real"]);
12733 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12734 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12735 "stop": ""
12736 }))
12737 .unwrap();
12738 assert!(req.stop.into_vec().is_empty());
12739
12740 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12741 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12742 "stop": null
12743 }))
12744 .unwrap();
12745 assert!(req.stop.into_vec().is_empty());
12746 }
12747
12748 #[test]
12749 fn stop_sequence_limits_bound_count_individual_and_aggregate_work() {
12750 let at_limit = StopSequences::Many(vec!["x".repeat(256); MAX_STOP_SEQUENCES]);
12751 assert!(at_limit.validate().is_ok());
12752 assert!(
12753 StopSequences::Many(vec![String::new(); MAX_STOP_SEQUENCES + 1])
12754 .validate()
12755 .unwrap_err()
12756 .contains("at most")
12757 );
12758 assert!(
12759 StopSequences::One("x".repeat(MAX_STOP_SEQUENCE_BYTES + 1))
12760 .validate()
12761 .unwrap_err()
12762 .contains("each stop")
12763 );
12764 assert!(
12765 StopSequences::Many(vec!["x".repeat(300); MAX_STOP_SEQUENCES])
12766 .validate()
12767 .unwrap_err()
12768 .contains("total at most")
12769 );
12770 }
12771
12772 #[tokio::test]
12773 async fn chat_response_has_openai_message_shape() {
12774 let (tx, rx) = worker::event_channel();
12775 tx.send(Event::Token {
12776 id: 1,
12777 text: "hello".into(),
12778 })
12779 .unwrap();
12780 tx.send(Event::Done {
12781 stop_reason: "Eos".into(),
12782 n_tokens: 1,
12783 n_prompt: 42,
12784 n_cached: 30,
12785 elapsed_s: 0.5,
12786 spec: None,
12787 })
12788 .unwrap();
12789 drop(tx);
12790 let response = blocking_response(
12791 rx,
12792 "plain_quant".into(),
12793 true,
12794 Vec::new(),
12795 None,
12796 Envelope::new(true),
12797 )
12798 .await;
12799 assert_eq!(response.status(), StatusCode::OK);
12800 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12801 .await
12802 .unwrap();
12803 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12804 assert_eq!(payload["object"], "chat.completion");
12805 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
12807 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
12808 let fingerprint = payload["system_fingerprint"].as_str().unwrap();
12812 assert!(
12813 build_id::fingerprint_is_well_formed(fingerprint),
12814 "system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
12815 );
12816 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
12817 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
12818 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
12819 assert_eq!(payload["usage"]["prompt_tokens"], 42);
12821 assert_eq!(payload["usage"]["completion_tokens"], 1);
12822 assert_eq!(payload["usage"]["total_tokens"], 43);
12823 assert_eq!(
12824 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12825 30
12826 );
12827 assert!(payload["usage"].get("spec").is_none());
12830 }
12831
12832 #[tokio::test]
12833 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
12834 let (tx, rx) = worker::event_channel();
12835 tx.send(Event::Token {
12837 id: 4,
12838 text: "hello".into(),
12839 })
12840 .unwrap();
12841 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
12842 tx.send(Event::Done {
12843 stop_reason: "MaxNew".into(),
12844 n_tokens: 4,
12845 n_prompt: 2,
12846 n_cached: 0,
12847 elapsed_s: 0.5,
12848 spec: None,
12849 })
12850 .unwrap();
12851 drop(tx);
12852
12853 let response = blocking_response(
12854 rx,
12855 "plain_quant".into(),
12856 false,
12857 Vec::new(),
12858 None,
12859 Envelope::new(false),
12860 )
12861 .await;
12862 assert_eq!(response.status(), StatusCode::OK);
12863 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12864 .await
12865 .unwrap();
12866 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12867 assert_eq!(payload["text"], "hello");
12868 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
12869 assert_eq!(payload["n_tokens"], 4);
12870 }
12871
12872 #[tokio::test]
12875 async fn chat_usage_carries_spec_acceptance_summary() {
12876 let (tx, rx) = worker::event_channel();
12877 tx.send(Event::Token {
12878 id: 1,
12879 text: "hello".into(),
12880 })
12881 .unwrap();
12882 tx.send(Event::Done {
12883 stop_reason: "Eos".into(),
12884 n_tokens: 1,
12885 n_prompt: 42,
12886 n_cached: 0,
12887 elapsed_s: 0.5,
12888 spec: Some(worker::SpecUsage {
12889 rounds: 10,
12890 drafted: 30,
12891 accepted: 21,
12892 }),
12893 })
12894 .unwrap();
12895 drop(tx);
12896 let response = blocking_response(
12897 rx,
12898 "plain_quant".into(),
12899 true,
12900 Vec::new(),
12901 None,
12902 Envelope::new(true),
12903 )
12904 .await;
12905 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12906 .await
12907 .unwrap();
12908 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12909 let sp = &payload["usage"]["spec"];
12910 assert_eq!(sp["rounds"], 10);
12911 assert_eq!(sp["drafted"], 30);
12912 assert_eq!(sp["accepted"], 21);
12913 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
12914 assert_eq!(payload["usage"]["total_tokens"], 43);
12916 }
12917
12918 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
12919 let mut payload = serde_json::json!({
12920 "model": "m",
12921 "messages": [{"role": "user", "content": "Weather in Paris?"}],
12922 "tools": [{"type": "function", "function": {
12923 "name": "get_weather",
12924 "description": "Get current weather",
12925 "parameters": {"type": "object",
12926 "properties": {"city": {"type": "string"},
12927 "days": {"type": "integer"}},
12928 "required": ["city"]}}}],
12929 });
12930 if let Some(obj) = extra.as_object() {
12931 for (k, v) in obj {
12932 payload[k] = v.clone();
12933 }
12934 }
12935 serde_json::from_value(payload).unwrap()
12936 }
12937
12938 #[test]
12942 fn glm5_vision_decode_is_deferred_and_grid_pinned() {
12943 let (tx, _rx) = worker::event_channel();
12944 let req: ChatCompletionReq = serde_json::from_value(json!({
12945 "model": "m", "messages": [{"role": "user", "content": "hi"}],
12946 }))
12947 .unwrap();
12948 let mut plan = build_chat_request(
12949 req,
12950 Some(&ModelCaps {
12951 chat_ok: true,
12952 ..Default::default()
12953 }),
12954 tx,
12955 lanes::Lane::Interactive,
12956 None,
12957 )
12958 .unwrap();
12959 let bmp = |w: u32, h: u32| -> Vec<u8> {
12962 let row = (w * 3).div_ceil(4) * 4;
12963 let size = 54 + row * h;
12964 let mut b = vec![0x42u8, 0x4d];
12965 b.extend_from_slice(&size.to_le_bytes());
12966 b.extend_from_slice(&[0; 4]);
12967 b.extend_from_slice(&54u32.to_le_bytes());
12968 b.extend_from_slice(&40u32.to_le_bytes());
12969 b.extend_from_slice(&w.to_le_bytes());
12970 b.extend_from_slice(&h.to_le_bytes());
12971 b.extend_from_slice(&1u16.to_le_bytes());
12972 b.extend_from_slice(&24u16.to_le_bytes());
12973 b.extend_from_slice(&[0u8; 24]);
12974 b.extend(std::iter::repeat_n(0x7fu8, (row * h) as usize));
12975 b
12976 };
12977 let bytes = bmp(112, 112);
12978 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes).unwrap();
12979 assert_eq!((gh, gw), (8, 8), "identity resize grid");
12980 assert_eq!(memra_engine::vision_glm5::n_merged_for_grid(gh, gw), 16);
12981 plan.pending_glm5.push(PendingGlm5Image {
12982 bytes: bytes.clone(),
12983 gh,
12984 gw,
12985 });
12986 decode_pending_vision(&mut plan).unwrap();
12987 assert_eq!(plan.request.glm5_images.len(), 1);
12988 let unit = &plan.request.glm5_images[0];
12989 assert_eq!((unit.gh, unit.gw), (gh, gw));
12990 assert_eq!(
12991 unit.patches.len(),
12992 gh * gw * memra_engine::vision_glm5::G5V_PATCH_IN
12993 );
12994 plan.request.glm5_images.clear();
12996 plan.pending_glm5.push(PendingGlm5Image {
12997 bytes,
12998 gh: gh + 2,
12999 gw,
13000 });
13001 let err = decode_pending_vision(&mut plan).unwrap_err();
13002 assert!(err.contains("header-planned"), "got: {err}");
13003 }
13004
13005 #[test]
13006 fn vision_decode_is_deferred_and_grid_pinned() {
13007 let (tx, _rx) = worker::event_channel();
13012 let req: ChatCompletionReq = serde_json::from_value(json!({
13013 "model": "m", "messages": [{"role": "user", "content": "hi"}],
13014 }))
13015 .unwrap();
13016 let mut plan = build_chat_request(
13017 req,
13018 Some(&ModelCaps {
13019 chat_ok: true,
13020 ..Default::default()
13021 }),
13022 tx,
13023 lanes::Lane::Interactive,
13024 None,
13025 )
13026 .unwrap();
13027 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
13031 let mut b = Vec::new();
13032 b.extend_from_slice(b"BM");
13033 b.extend_from_slice(&54u32.to_le_bytes());
13034 b.extend_from_slice(&0u32.to_le_bytes());
13035 b.extend_from_slice(&54u32.to_le_bytes());
13036 b.extend_from_slice(&40u32.to_le_bytes());
13037 b.extend_from_slice(&w.to_le_bytes());
13038 b.extend_from_slice(&h.to_le_bytes());
13039 b.extend_from_slice(&1u16.to_le_bytes());
13040 b.extend_from_slice(&24u16.to_le_bytes());
13041 b.extend_from_slice(&[0u8; 24]);
13042 if with_pixels {
13043 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
13044 }
13045 b
13046 };
13047 let bytes = bmp(64, 64, true);
13048 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
13049 plan.pending_images.push(PendingVisionUnit::Still {
13050 bytes: bytes.clone(),
13051 gh,
13052 gw,
13053 });
13054 decode_pending_vision(&mut plan).unwrap();
13055 assert_eq!(plan.request.images.len(), 1);
13056 assert_eq!(
13057 (
13058 plan.request.images[0].prep.gh,
13059 plan.request.images[0].prep.gw
13060 ),
13061 (gh, gw),
13062 "decoded grid must equal the header-planned grid the pad run was rendered from"
13063 );
13064 plan.request.images.clear();
13066 plan.pending_images.push(PendingVisionUnit::Still {
13067 bytes,
13068 gh: gh + 2,
13069 gw,
13070 });
13071 let err = decode_pending_vision(&mut plan).unwrap_err();
13072 assert!(err.contains("header-planned"), "got: {err}");
13073 let bomb = bmp(16_000, 16_000, false);
13076 plan.pending_images.clear();
13077 plan.pending_images.push(PendingVisionUnit::Still {
13078 bytes: bomb,
13079 gh: 2,
13080 gw: 2,
13081 });
13082 let err = decode_pending_vision(&mut plan).unwrap_err();
13083 assert!(err.contains("exceeds the decode budget"), "got: {err}");
13084 }
13085
13086 #[test]
13087 fn tools_request_renders_client_key_order_and_arms_parser() {
13088 let (tx, _rx) = worker::event_channel();
13089 let plan = build_chat_request(
13090 weather_request(json!({})),
13091 Some(&tool_caps()),
13092 tx,
13093 lanes::Lane::Interactive,
13094 None,
13095 )
13096 .unwrap();
13097 assert!(plan.parser.is_some());
13098 assert_eq!(plan.request.tools_json.len(), 1);
13099 assert_eq!(
13101 plan.request.tools_json[0],
13102 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
13103 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
13104 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
13105 \"integer\"}}, \"required\": [\"city\"]}}}"
13106 );
13107 }
13108
13109 #[test]
13110 fn hy3_tools_and_reasoning_flow_through_the_real_chat_plan() {
13111 let (tx, _rx) = worker::event_channel();
13112 let plan = build_chat_request(
13113 weather_request(json!({"reasoning_effort": "high"})),
13114 Some(&hy3_tool_caps()),
13115 tx,
13116 lanes::Lane::Interactive,
13117 None,
13118 )
13119 .unwrap();
13120 assert_eq!(plan.request.think, ThinkMode::Think);
13121 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
13122 assert!(
13123 plan.request
13124 .stop_strings
13125 .iter()
13126 .any(|stop| stop == "</tool_calls:opensource>")
13127 );
13128 let rendered = chat::apply_chat_template_tools_ex(
13129 Some("... hy_User ... <tools> ..."),
13130 &plan.request.chat_turns,
13131 true,
13132 &plan.request.tools_json,
13133 &plan.request.tools_struct,
13134 plan.request.think,
13135 plan.request.reasoning_effort.as_deref(),
13136 None,
13137 )
13138 .unwrap();
13139 assert!(rendered.contains("<tool_calls:opensource>"));
13140 assert!(rendered.ends_with("<think:opensource>"));
13141
13142 let mut parser = plan.parser.expect("HY3 tools arm its native parser");
13143 let pieces = parser.push(concat!(
13144 "Need weather.</think:opensource>",
13145 "<tool_calls:opensource><tool_call:opensource>get_weather",
13146 "<tool_sep:opensource>\n<arg_key:opensource>city</arg_key:opensource>\n",
13147 "<arg_value:opensource>Paris</arg_value:opensource>\n",
13148 "</tool_call:opensource></tool_calls:opensource>",
13149 ));
13150 assert!(pieces.contains(&Piece::Reasoning("Need weather.".into())));
13151 assert!(pieces.iter().any(|piece| matches!(piece, Piece::Call(call)
13152 if call.name == "get_weather" && call.arguments == r#"{"city":"Paris"}"#)));
13153 }
13154
13155 #[test]
13156 fn tool_choice_none_strips_tools_and_parser() {
13157 let (tx, _rx) = worker::event_channel();
13158 let plan = build_chat_request(
13159 weather_request(json!({"tool_choice": "none"})),
13160 Some(&tool_caps()),
13161 tx,
13162 lanes::Lane::Interactive,
13163 None,
13164 )
13165 .unwrap();
13166 let mut p = plan
13169 .parser
13170 .expect("think-open chat arms the reasoning splitter");
13171 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
13172 assert_eq!(
13173 pieces,
13174 vec![
13175 Piece::Reasoning("x".into()),
13176 Piece::Content("<tool_call> stays prose".into()),
13177 ]
13178 );
13179 assert!(plan.request.tools_json.is_empty());
13180 let (tx, _rx) = worker::event_channel();
13182 assert!(
13183 build_chat_request(
13184 weather_request(json!({"tool_choice": "required"})),
13185 Some(&tool_caps()),
13186 tx,
13187 lanes::Lane::Interactive,
13188 None
13189 )
13190 .is_err()
13191 );
13192 let (tx, _rx) = worker::event_channel();
13193 assert!(
13194 build_chat_request(
13195 weather_request(json!({"tool_choice":
13196 {"type": "function", "function": {"name": "get_weather"}}})),
13197 Some(&tool_caps()),
13198 tx,
13199 lanes::Lane::Interactive,
13200 None
13201 )
13202 .is_err()
13203 );
13204 }
13205
13206 #[test]
13207 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
13208 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
13209 let _ = std::fs::remove_dir_all(&root);
13210
13211 let st = root.join("st_single");
13213 std::fs::create_dir_all(&st).unwrap();
13214 std::fs::write(st.join("config.json"), "{}").unwrap();
13215 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
13216 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
13217
13218 let sh = root.join("st_sharded");
13220 std::fs::create_dir_all(&sh).unwrap();
13221 std::fs::write(sh.join("config.json"), "{}").unwrap();
13222 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
13223 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
13224
13225 let rp = root.join("repack");
13227 std::fs::create_dir_all(&rp).unwrap();
13228 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
13229 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
13230
13231 let bogus = root.join("bogus");
13233 std::fs::create_dir_all(&bogus).unwrap();
13234 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
13235 assert!(
13236 err.contains("model.safetensors"),
13237 "error should say what is missing: {err}"
13238 );
13239 assert!(
13240 err.contains("manifest.json"),
13241 "error should mention the repack form: {err}"
13242 );
13243
13244 let nc = root.join("no_config");
13246 std::fs::create_dir_all(&nc).unwrap();
13247 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
13248 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
13249 assert!(
13250 err.contains("config.json"),
13251 "error should name config.json: {err}"
13252 );
13253
13254 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
13256 assert!(err.contains("does not exist"), "{err}");
13257
13258 let f = root.join("model.gguf");
13260 std::fs::write(&f, b"g").unwrap();
13261 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
13262
13263 let _ = std::fs::remove_dir_all(&root);
13264 }
13265
13266 #[test]
13267 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
13268 let caps = ModelCaps {
13271 tools_branch: false,
13272 qwen_think: false,
13273 think_switch: false,
13274 chat_ok: false,
13275 ..Default::default()
13276 };
13277 let payload = serde_json::json!({
13278 "model": "st_model",
13279 "messages": [{"role": "user", "content": "hello"}],
13280 });
13281 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13282 let (tx, _rx) = worker::event_channel();
13283 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
13284 Err(e) => e,
13285 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
13286 };
13287 assert!(
13288 err.contains("no chat template"),
13289 "message should name the cause: {err}"
13290 );
13291 assert!(
13292 err.contains("/v1/completions"),
13293 "message should point at the raw-prompt escape hatch: {err}"
13294 );
13295 }
13296
13297 #[test]
13298 fn tools_on_model_without_tools_branch_is_rejected() {
13299 let (tx, _rx) = worker::event_channel();
13300 let caps = ModelCaps {
13301 chat_ok: true,
13302 ..Default::default()
13303 };
13304 assert!(
13305 build_chat_request(
13306 weather_request(json!({})),
13307 Some(&caps),
13308 tx,
13309 lanes::Lane::Interactive,
13310 None
13311 )
13312 .is_err()
13313 );
13314 let (tx, _rx) = worker::event_channel();
13315 assert!(
13316 build_chat_request(
13317 weather_request(json!({})),
13318 None,
13319 tx,
13320 lanes::Lane::Interactive,
13321 None
13322 )
13323 .is_err()
13324 );
13325 }
13326
13327 #[test]
13328 fn reasoning_effort_maps_to_think_switch() {
13329 for (extra, want) in [
13335 (json!({}), ThinkMode::Default),
13336 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
13337 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
13338 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
13339 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
13340 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
13341 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
13342 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
13343 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
13344 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
13348 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
13349 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
13350 (
13354 json!({"reasoning": {"enabled": true, "effort": "none"}}),
13355 ThinkMode::Think,
13356 ),
13357 (
13358 json!({"reasoning": {"enabled": false, "effort": "high"}}),
13359 ThinkMode::NoThink,
13360 ),
13361 ] {
13362 let (tx, _rx) = worker::event_channel();
13363 let plan = build_chat_request(
13364 weather_request(extra.clone()),
13365 Some(&ladder_caps()),
13370 tx,
13371 lanes::Lane::Interactive,
13372 None,
13373 )
13374 .unwrap();
13375 assert_eq!(plan.request.think, want, "extra={extra}");
13376 }
13377 for extra in [
13381 json!({"reasoning_effort": "extreme"}),
13382 json!({"reasoning": {"effort": "banana"}}),
13383 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
13384 json!({"reasoning": {"enabled": true, "effort": ""}}),
13385 ] {
13386 let (tx, _rx) = worker::event_channel();
13387 assert!(
13388 build_chat_request(
13389 weather_request(extra.clone()),
13390 Some(&tool_caps()),
13391 tx,
13392 lanes::Lane::Interactive,
13393 None
13394 )
13395 .is_err(),
13396 "extra={extra} must be rejected by the one allowlist"
13397 );
13398 }
13399 for (raw, want) in [
13402 ("none", Some("none")),
13403 ("minimal", Some("minimal")),
13404 ("low", Some("low")),
13405 ("medium", Some("medium")),
13406 ("high", Some("high")),
13407 ("xhigh", Some("high")),
13408 ("max", Some("high")),
13409 ("ultra", Some("high")),
13410 ("banana", None),
13411 ("", None),
13412 ("HIGH", None),
13413 ] {
13414 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
13415 }
13416 for (raw, want) in [
13419 ("none", Some("none")),
13420 ("minimal", Some("minimal")),
13421 ("low", Some("low")),
13422 ("medium", Some("medium")),
13423 ("high", Some("high")),
13424 ("xhigh", Some("max")),
13425 ("max", Some("max")),
13426 ("ultra", Some("max")),
13427 ("banana", None),
13428 ("", None),
13429 ("MAX", None),
13430 ] {
13431 assert_eq!(
13432 canonical_effort_for(raw, true),
13433 want,
13434 "canonical_effort_for({raw:?}, dsv4)"
13435 );
13436 }
13437 }
13438
13439 #[test]
13440 fn dsv4_reasoning_effort_max_survives_canonicalization() {
13441 let dsv4_caps = ModelCaps {
13447 chat_ok: true,
13448 dsv4: true,
13449 ..Default::default()
13450 };
13451 let build = |caps: &ModelCaps, effort: &str| {
13452 let (tx, _rx) = worker::event_channel();
13453 let req: ChatCompletionReq = serde_json::from_value(json!({
13454 "model": "m",
13455 "messages": [{"role": "user", "content": "hi"}],
13456 "reasoning_effort": effort,
13457 }))
13458 .unwrap();
13459 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
13460 };
13461 for raw in ["max", "xhigh", "ultra"] {
13462 let plan = build(&dsv4_caps, raw).unwrap();
13463 assert_eq!(
13464 plan.request.reasoning_effort.as_deref(),
13465 Some("max"),
13466 "dsv4 {raw:?} must reach the renderer as the max rung"
13467 );
13468 assert_eq!(plan.request.think, chat::ThinkMode::Think);
13469 }
13470 let plan = build(&dsv4_caps, "high").unwrap();
13472 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
13473 let step_caps = ModelCaps {
13475 chat_ok: true,
13476 effort_levels: true,
13477 ..Default::default()
13478 };
13479 let plan = build(&step_caps, "max").unwrap();
13480 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
13481 }
13482
13483 #[test]
13484 fn default_reasoning_effort_flips_only_the_unset_request() {
13485 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
13492 let (tx, _rx) = worker::event_channel();
13493 build_chat_request_with_trace(
13494 weather_request(extra),
13495 Some(&ladder_caps()),
13496 tx,
13497 lanes::Lane::Interactive,
13498 None,
13499 None,
13500 default_effort,
13501 &ModelSamplingDefaults::default(),
13502 )
13503 .unwrap()
13504 };
13505 for (extra, want) in [
13506 (json!({}), ThinkMode::Think),
13508 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
13512 (json!({"include_reasoning": false}), ThinkMode::NoThink),
13513 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
13515 (json!({"include_reasoning": true}), ThinkMode::Think),
13516 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
13518 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
13519 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
13520 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
13522 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
13523 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
13524 ] {
13525 let plan = build(extra.clone(), Some("high"));
13526 assert_eq!(plan.request.think, want, "extra={extra}");
13527 }
13528 assert_eq!(
13530 build(json!({}), Some("none")).request.think,
13531 ThinkMode::NoThink
13532 );
13533 assert_eq!(
13534 build(json!({"reasoning_effort": "high"}), Some("none"))
13535 .request
13536 .think,
13537 ThinkMode::Think
13538 );
13539 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
13543 }
13544
13545 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
13551 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
13552 {%- else %}'<think>\\n'{%- endif %}";
13553
13554 #[test]
13555 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
13556 let build = |extra: serde_json::Value| {
13562 let (tx, _rx) = worker::event_channel();
13563 build_chat_request(
13564 weather_request(extra),
13565 Some(&tool_caps()),
13566 tx,
13567 lanes::Lane::Interactive,
13568 None,
13569 )
13570 };
13571 for (extra, want) in [
13572 (json!({"enable_thinking": false}), ThinkMode::NoThink),
13573 (json!({"enable_thinking": true}), ThinkMode::Think),
13574 (
13575 json!({"chat_template_kwargs": {"enable_thinking": false}}),
13576 ThinkMode::NoThink,
13577 ),
13578 (
13579 json!({"chat_template_kwargs": {"enable_thinking": true}}),
13580 ThinkMode::Think,
13581 ),
13582 (
13585 json!({"enable_thinking": false, "reasoning_effort": "high"}),
13586 ThinkMode::NoThink,
13587 ),
13588 (
13590 json!({"enable_thinking": false,
13591 "chat_template_kwargs": {"enable_thinking": false}}),
13592 ThinkMode::NoThink,
13593 ),
13594 ] {
13595 let plan = build(extra.clone()).unwrap_or_else(|e| {
13596 panic!("{extra} must be accepted and honored, got 400: {e}");
13597 });
13598 assert_eq!(
13599 plan.request.think, want,
13600 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
13601 );
13602 }
13603 let render = |extra: serde_json::Value| -> String {
13606 let plan = build(extra).unwrap();
13607 chat::apply_chat_template_tools_ex(
13608 Some(SWITCHED_QWEN_TMPL),
13609 &plan.request.chat_turns,
13610 true,
13611 &plan.request.tools_json,
13612 &plan.request.tools_struct,
13613 plan.request.think,
13614 plan.request.reasoning_effort.as_deref(),
13615 None,
13616 )
13617 .unwrap()
13618 };
13619 let off = render(json!({"enable_thinking": false}));
13620 assert!(
13621 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
13622 "enable_thinking:false must render the CLOSED think pair: {off:?}"
13623 );
13624 let on = render(json!({}));
13625 assert!(
13626 on.ends_with("<|im_start|>assistant\n<think>\n"),
13627 "an unset request must still render the template's OPEN think tail: {on:?}"
13628 );
13629 assert_eq!(
13630 off,
13631 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
13632 "both vLLM spellings must render byte-identically"
13633 );
13634 assert_eq!(
13635 off,
13636 render(json!({"reasoning_effort": "none"})),
13637 "the vLLM spelling must render byte-identically to the OpenAI spelling"
13638 );
13639 }
13640
13641 #[test]
13642 fn unknown_chat_template_kwarg_refuses_by_name() {
13643 let build = |extra: serde_json::Value| {
13646 let (tx, _rx) = worker::event_channel();
13647 build_chat_request(
13648 weather_request(extra),
13649 Some(&tool_caps()),
13650 tx,
13651 lanes::Lane::Interactive,
13652 None,
13653 )
13654 };
13655 let refusal = |extra: serde_json::Value, why: &str| -> String {
13656 build(extra).err().unwrap_or_else(|| panic!("{why}"))
13657 };
13658 let err = refusal(
13659 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
13660 "an unimplementable template kwarg must not be accepted",
13661 );
13662 assert!(
13663 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
13664 "the refusal must name the offending key AND the supported one: {err}"
13665 );
13666 let err = refusal(
13667 json!({"chat_template_kwargs": "enable_thinking=false"}),
13668 "a non-object chat_template_kwargs must not be accepted",
13669 );
13670 assert!(
13671 err.contains("must be an object"),
13672 "refusal must say what shape is expected: {err}"
13673 );
13674 let err = refusal(
13675 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
13676 "a stringly-typed switch must not be accepted",
13677 );
13678 assert!(
13679 err.contains("true or false"),
13680 "refusal must name the expected type: {err}"
13681 );
13682 let plan = build(json!({"chat_template_kwargs": null}))
13684 .expect("null chat_template_kwargs is the unset case");
13685 assert_eq!(plan.request.think, ThinkMode::Default);
13686 }
13687
13688 const Q38_TMPL: &str =
13703 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
13704
13705 fn render_with(
13708 tmpl: &str,
13709 caps: &ModelCaps,
13710 extra: serde_json::Value,
13711 default_effort: Option<&str>,
13712 ) -> Result<String, String> {
13713 let mut payload = serde_json::json!({
13714 "model": "m",
13715 "messages": [{"role": "user", "content": "hi"}],
13716 });
13717 if let Some(obj) = extra.as_object() {
13718 for (k, v) in obj {
13719 payload[k] = v.clone();
13720 }
13721 }
13722 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13723 let (tx, _rx) = worker::event_channel();
13724 let plan = build_chat_request_with_trace(
13725 req,
13726 Some(caps),
13727 tx,
13728 lanes::Lane::Interactive,
13729 None,
13730 None,
13731 default_effort,
13732 &ModelSamplingDefaults::default(),
13733 )?;
13734 Ok(chat::apply_chat_template_tools_ex(
13735 Some(tmpl),
13736 &plan.request.chat_turns,
13737 true,
13738 &plan.request.tools_json,
13739 &plan.request.tools_struct,
13740 plan.request.think,
13741 plan.request.reasoning_effort.as_deref(),
13742 None,
13743 )
13744 .unwrap())
13745 }
13746
13747 #[test]
13748 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
13749 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
13756 let xhigh = "Reasoning effort is set to xhigh.";
13757 let low = "Reasoning effort is set to low.";
13758 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
13760 assert!(
13761 r(json!({"reasoning_effort": "high"}))
13762 .unwrap()
13763 .contains(xhigh)
13764 );
13765 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
13768 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
13769 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
13771 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
13772 assert_ne!(low_p, high_p);
13773 assert_ne!(low_p, medium);
13774 assert_ne!(high_p, medium);
13775 for alias in ["xhigh", "max", "ultra"] {
13778 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
13779 }
13780 assert_eq!(r(json!({})).unwrap(), high_p);
13783 assert_eq!(
13786 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
13787 medium
13788 );
13789 let off = r(json!({"reasoning_effort": "none"})).unwrap();
13792 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
13793 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
13794 }
13795
13796 #[test]
13797 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
13798 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
13811focused, moving directly to the conclusion without unnecessary elaboration.";
13812 let expected = format!(
13813 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
13814 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
13815 );
13816 let after_fix = render_with(
13818 Q38_TMPL,
13819 &ladder_caps(),
13820 json!({"reasoning_effort": "low"}),
13821 None,
13822 )
13823 .unwrap();
13824 assert_eq!(
13825 after_fix, expected,
13826 "the shipped prompt for reasoning_effort:\"low\""
13827 );
13828 const ORNITH_TMPL: &str = include_str!(
13831 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
13832 );
13833 let on_deployed_binary = render_with(
13834 ORNITH_TMPL,
13835 &tool_caps(),
13836 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
13837 {"role": "user", "content": "hi"}]}),
13838 None,
13839 )
13840 .unwrap();
13841 assert_eq!(
13842 on_deployed_binary, expected,
13843 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
13844 level, or its reasoning-volume numbers do not describe the shipped prompt"
13845 );
13846 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
13849 assert!(
13850 !ladderless_unset.contains("Reasoning effort is set to"),
13851 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
13852 );
13853 assert_eq!(
13854 ladderless_unset,
13855 render_with(
13856 Q38_TMPL,
13857 &ladder_caps(),
13858 json!({"reasoning_effort": "medium"}),
13859 None
13860 )
13861 .unwrap(),
13862 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
13863 );
13864 }
13865
13866 #[test]
13867 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
13868 let off = render_with(
13875 Q38_TMPL,
13876 &ladder_caps(),
13877 json!({"reasoning_effort": "none"}),
13878 None,
13879 )
13880 .unwrap();
13881 for extra in [
13882 json!({"include_reasoning": false}),
13883 json!({"reasoning": {"exclude": true}}),
13884 ] {
13885 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
13886 assert!(
13887 got.ends_with("<think>\n\n</think>\n\n"),
13888 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
13889 );
13890 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
13891 }
13892 for extra in [
13897 json!({"enable_thinking": true, "include_reasoning": false}),
13898 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
13899 json!({"reasoning": {"enabled": true, "exclude": true}}),
13900 ] {
13901 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
13902 .err()
13903 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
13904 assert!(e.contains("contradictory"), "{extra}: {e}");
13905 assert!(
13906 e.contains("include_reasoning") || e.contains("exclude"),
13907 "{extra}: the refusal must name the suppression field the caller sent: {e}"
13908 );
13909 }
13910 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
13913 for extra in [
13914 json!({"include_reasoning": true}),
13915 json!({"reasoning": {"exclude": false}}),
13916 ] {
13917 assert_eq!(
13918 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
13919 dflt,
13920 "{extra} must not perturb the model's default"
13921 );
13922 }
13923 let switchless = ModelCaps {
13927 think_switch: false,
13928 ..tool_caps()
13929 };
13930 let err = render_with(
13931 Q38_TMPL,
13932 &switchless,
13933 json!({"include_reasoning": false}),
13934 None,
13935 )
13936 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
13937 assert!(err.contains("cannot disable reasoning"), "{err}");
13938 }
13939
13940 #[test]
13941 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
13942 let build = |extra: serde_json::Value| {
13943 let (tx, _rx) = worker::event_channel();
13944 build_chat_request(
13945 weather_request(extra),
13946 Some(&ladder_caps()),
13947 tx,
13948 lanes::Lane::Interactive,
13949 None,
13950 )
13951 };
13952 let err = |extra: serde_json::Value, why: &str| -> String {
13953 build(extra).err().unwrap_or_else(|| panic!("{why}"))
13954 };
13955 let e = err(
13959 json!({"reasoning": {"max_tokens": 1024}}),
13960 "reasoning.max_tokens must not be accepted-and-ignored",
13961 );
13962 assert!(e.contains("reasoning.max_tokens"), "{e}");
13963 assert!(e.contains("ONE output budget"), "{e}");
13964 for extra in [
13969 json!({"reasoning": {"max_tokens": null}}),
13970 json!({"reasoning": {"banana": null}}),
13971 ] {
13972 let e = err(
13973 extra.clone(),
13974 "a null-valued unhonourable key must still refuse",
13975 );
13976 assert!(
13977 e.contains("max_tokens") || e.contains("banana"),
13978 "{extra}: {e}"
13979 );
13980 }
13981 let e = err(
13983 json!({"reasoning": {"budget": 5}}),
13984 "an unknown reasoning key must not be accepted",
13985 );
13986 assert!(
13987 e.contains("reasoning.budget") && e.contains("enabled"),
13988 "{e}"
13989 );
13990 for (extra, want) in [
13994 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
13995 (json!({"reasoning": {"exclude": 1}}), "true or false"),
13996 (json!({"reasoning": {"effort": 3}}), "must be a string"),
13997 ] {
13998 let e = err(
13999 extra.clone(),
14000 "a wrong-typed reasoning key must not be ignored",
14001 );
14002 assert!(e.contains(want), "{extra}: {e}");
14003 }
14004 for extra in [
14009 json!({"reasoning": {"enabled": true}}),
14010 json!({"reasoning": {"effort": "low"}}),
14011 json!({"reasoning": {"exclude": false}}),
14012 json!({"reasoning": null}),
14013 json!({"reasoning": {"effort": null}}),
14014 json!({"reasoning": {"enabled": null, "exclude": null}}),
14015 ] {
14016 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
14017 }
14018 }
14019
14020 #[test]
14021 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
14022 const ORNITH_TMPL: &str = include_str!(
14030 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
14031 );
14032 let explicit_on = render_with(
14036 ORNITH_TMPL,
14037 &tool_caps(),
14038 json!({"reasoning": {"enabled": true}}),
14039 None,
14040 )
14041 .unwrap();
14042 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
14043 for extra in [
14044 json!({"reasoning_effort": "low"}),
14045 json!({"reasoning_effort": "medium"}),
14046 json!({"reasoning_effort": "high"}),
14047 json!({"reasoning_effort": "xhigh"}),
14049 json!({"reasoning": {"effort": "xhigh"}}),
14050 ] {
14051 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
14052 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
14053 assert_eq!(
14054 got, explicit_on,
14055 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
14056 documented translation, not a decorative accept"
14057 );
14058 }
14059 for extra in [
14061 json!({}),
14062 json!({"reasoning_effort": "none"}),
14063 json!({"reasoning_effort": "minimal"}),
14064 json!({"enable_thinking": false}),
14065 ] {
14066 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
14067 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
14068 }
14069 let minimal = render_with(
14072 ORNITH_TMPL,
14073 &tool_caps(),
14074 json!({"reasoning_effort": "minimal"}),
14075 None,
14076 )
14077 .unwrap();
14078 assert!(
14079 minimal.ends_with("<think>\n\n</think>\n\n"),
14080 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
14081 );
14082 let ladder_low = render_with(
14085 Q38_TMPL,
14086 &ladder_caps(),
14087 json!({"reasoning_effort": "low"}),
14088 None,
14089 )
14090 .unwrap();
14091 assert!(
14092 ladder_low.contains("Reasoning effort is set to low."),
14093 "{ladder_low:?}"
14094 );
14095 assert_ne!(
14096 ladder_low,
14097 render_with(
14098 Q38_TMPL,
14099 &ladder_caps(),
14100 json!({"reasoning_effort": "high"}),
14101 None
14102 )
14103 .unwrap(),
14104 "the ladder model's rungs stay distinct prompts"
14105 );
14106 }
14107
14108 #[test]
14109 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
14110 let render_chat = |body: serde_json::Value| -> Result<String, String> {
14121 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14122 let (tx, _rx) = worker::event_channel();
14123 let plan = build_chat_request(
14124 req,
14125 Some(&ladder_caps()),
14126 tx,
14127 lanes::Lane::Interactive,
14128 None,
14129 )?;
14130 Ok(chat::apply_chat_template_tools_ex(
14131 Some(Q38_TMPL),
14132 &plan.request.chat_turns,
14133 true,
14134 &plan.request.tools_json,
14135 &plan.request.tools_struct,
14136 plan.request.think,
14137 plan.request.reasoning_effort.as_deref(),
14138 None,
14139 )
14140 .unwrap())
14141 };
14142 for (intent, chat_body, responses_body, messages_body) in [
14147 (
14148 "reasoning OFF",
14149 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
14150 "reasoning_effort": "none"}),
14151 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
14152 json!({"model": "m", "max_tokens": 16,
14153 "messages": [{"role": "user", "content": "hi"}],
14154 "thinking": {"type": "disabled"}}),
14155 ),
14156 (
14157 "reasoning ON at the top rung",
14158 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
14159 "reasoning_effort": "xhigh"}),
14160 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
14161 json!({"model": "m", "max_tokens": 16,
14162 "messages": [{"role": "user", "content": "hi"}],
14163 "output_config": {"effort": "xhigh"}}),
14164 ),
14165 (
14166 "reasoning ON at the bottom rung",
14167 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
14168 "reasoning_effort": "low"}),
14169 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
14170 json!({"model": "m", "max_tokens": 16,
14171 "messages": [{"role": "user", "content": "hi"}],
14172 "output_config": {"effort": "low"}}),
14173 ),
14174 (
14175 "the model's own default",
14176 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
14177 json!({"model": "m", "input": "hi"}),
14178 json!({"model": "m", "max_tokens": 16,
14179 "messages": [{"role": "user", "content": "hi"}]}),
14180 ),
14181 ] {
14182 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
14183 let via_responses = responses_api::translate(&responses_body)
14184 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
14185 let via_messages = anthropic::translate(&messages_body)
14186 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
14187 for (surface, translated) in [
14188 ("/v1/responses", via_responses),
14189 ("/v1/messages", via_messages),
14190 ] {
14191 let got = render_chat(translated)
14192 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
14193 assert_eq!(
14194 got, chat,
14195 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
14196 /v1/chat/completions — the parameter is honoured on one format and not \
14197 the other"
14198 );
14199 }
14200 }
14201 let switchless = ModelCaps {
14204 think_switch: false,
14205 ..ladder_caps()
14206 };
14207 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
14208 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14209 let (tx, _rx) = worker::event_channel();
14210 let plan =
14211 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
14212 Ok(format!("{:?}", plan.request.think))
14213 };
14214 for (surface, body) in [
14215 (
14216 "/v1/responses",
14217 responses_api::translate(&json!({
14218 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
14219 .unwrap(),
14220 ),
14221 (
14222 "/v1/messages",
14223 anthropic::translate(&json!({
14224 "model": "m", "max_tokens": 16,
14225 "messages": [{"role": "user", "content": "hi"}],
14226 "thinking": {"type": "disabled"}}))
14227 .unwrap(),
14228 ),
14229 ] {
14230 let err = render_switchless(body)
14231 .err()
14232 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
14233 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
14234 }
14235 }
14236
14237 #[test]
14238 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
14239 let build = |extra: serde_json::Value| {
14247 let (tx, _rx) = worker::event_channel();
14248 build_chat_request(
14249 weather_request(extra),
14250 Some(&ladder_caps()),
14251 tx,
14252 lanes::Lane::Interactive,
14253 None,
14254 )
14255 };
14256 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
14257 .expect("preserve_thinking:true is the vendor default the renderer implements");
14258 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
14259 .err()
14260 .expect("preserve_thinking:false (the strip arm) must refuse");
14261 assert!(e.contains("preserve_thinking"), "{e}");
14262 assert!(e.contains("strip"), "{e}");
14263 assert_eq!(
14266 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
14267 .unwrap()
14268 .request
14269 .think,
14270 ThinkMode::NoThink
14271 );
14272 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
14274 .err()
14275 .expect("a stringly-typed preserve_thinking must not be accepted");
14276 assert!(e.contains("true or false"), "{e}");
14277 }
14278
14279 #[test]
14280 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
14281 let dsv4_caps = ModelCaps {
14287 qwen_think: true,
14288 think_switch: false,
14289 dsv4: true,
14290 ..tool_caps()
14291 };
14292 for extra in [
14293 json!({"reasoning_effort": "none"}),
14294 json!({"reasoning": {"enabled": false}}),
14295 json!({"enable_thinking": false}),
14296 json!({"include_reasoning": false}),
14297 ] {
14298 let (tx, _rx) = worker::event_channel();
14299 let plan = build_chat_request(
14300 weather_request(extra.clone()),
14301 Some(&dsv4_caps),
14302 tx,
14303 lanes::Lane::Interactive,
14304 None,
14305 )
14306 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
14307 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
14308 }
14309 }
14310
14311 #[test]
14312 fn contradictory_think_switches_refuse_instead_of_picking_one() {
14313 let build = |extra: serde_json::Value| {
14316 let (tx, _rx) = worker::event_channel();
14317 build_chat_request(
14318 weather_request(extra),
14319 Some(&tool_caps()),
14320 tx,
14321 lanes::Lane::Interactive,
14322 None,
14323 )
14324 };
14325 for extra in [
14326 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
14327 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
14328 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
14329 ] {
14330 match build(extra.clone()) {
14331 Err(err) => assert!(
14332 err.contains("contradictory"),
14333 "the refusal must say the switches contradict: {err}"
14334 ),
14335 Ok(plan) => panic!(
14336 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
14337 plan.request.think
14338 ),
14339 }
14340 }
14341 for extra in [
14343 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
14344 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
14345 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
14346 ] {
14347 build(extra.clone())
14348 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
14349 }
14350 }
14351
14352 #[test]
14353 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
14354 let switchless = ModelCaps {
14359 tools_branch: true,
14360 qwen_think: true,
14361 think_switch: false,
14362 chat_ok: true,
14363 ..Default::default()
14364 };
14365 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
14366 let (tx, _rx) = worker::event_channel();
14367 build_chat_request_with_trace(
14368 weather_request(extra),
14369 Some(caps),
14370 tx,
14371 lanes::Lane::Interactive,
14372 None,
14373 None,
14374 default_effort,
14375 &ModelSamplingDefaults::default(),
14376 )
14377 };
14378 for extra in [
14379 json!({"reasoning_effort": "none"}),
14380 json!({"reasoning_effort": "minimal"}),
14381 json!({"reasoning": {"enabled": false}}),
14382 json!({"enable_thinking": false}),
14383 json!({"chat_template_kwargs": {"enable_thinking": false}}),
14384 ] {
14385 let err = build(extra.clone(), &switchless, None)
14386 .err()
14387 .unwrap_or_else(|| {
14388 panic!(
14389 "{extra} on a switchless think template must not be accepted-and-ignored"
14390 )
14391 });
14392 assert!(
14393 err.contains("cannot disable reasoning"),
14394 "the refusal must say the model cannot disable reasoning: {err}"
14395 );
14396 }
14397 for (extra, default_effort) in [
14401 (json!({}), None),
14402 (json!({"reasoning_effort": "high"}), None),
14405 (json!({"reasoning": {"enabled": true}}), None),
14406 (json!({"enable_thinking": true}), None),
14407 (json!({}), Some("none")),
14408 (json!({}), Some("minimal")),
14409 (json!({}), Some("high")),
14410 ] {
14411 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
14412 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
14413 });
14414 }
14415 assert_eq!(
14418 build(json!({"enable_thinking": false}), &tool_caps(), None)
14419 .unwrap()
14420 .request
14421 .think,
14422 ThinkMode::NoThink
14423 );
14424 }
14425
14426 #[test]
14427 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
14428 let gemma_caps = ModelCaps {
14435 tools_branch: true,
14436 chat_ok: true,
14437 gemma_think: true,
14438 instruct_type: Some("gemma".into()),
14439 ..Default::default()
14440 };
14441 let render =
14442 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
14443 let mut payload = serde_json::json!({
14444 "model": "google/gemma-4-31b-it",
14445 "messages": [{"role": "user", "content": "Weather in Paris?"}],
14446 });
14447 if let Some(obj) = extra.as_object() {
14448 for (k, v) in obj {
14449 payload[k] = v.clone();
14450 }
14451 }
14452 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
14453 let (tx, _rx) = worker::event_channel();
14454 let plan = build_chat_request_with_trace(
14455 req,
14456 Some(&gemma_caps),
14457 tx,
14458 lanes::Lane::Interactive,
14459 None,
14460 None,
14461 default_effort,
14462 &ModelSamplingDefaults::default(),
14463 )
14464 .unwrap();
14465 chat::apply_chat_template_tools_ex(
14466 Some(tmpl),
14467 &plan.request.chat_turns,
14468 true,
14469 &plan.request.tools_json,
14470 &plan.request.tools_struct,
14471 plan.request.think,
14472 plan.request.reasoning_effort.as_deref(),
14473 None, )
14475 .unwrap()
14476 };
14477 let official = gemma_template("official");
14478 let unset_with_knob = render(&official, json!({}), Some("high"));
14479 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
14480 assert_eq!(
14481 unset_with_knob, explicit_on,
14482 "knob render must be byte-identical to the explicit think-on render"
14483 );
14484 assert!(
14485 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
14486 "think-on injects the <|think|> system token: {unset_with_knob:?}"
14487 );
14488 assert!(
14489 unset_with_knob.ends_with("<|turn>model\n"),
14490 "think-on generation turn is OPEN: {unset_with_knob:?}"
14491 );
14492 let explicit_off_with_knob =
14496 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
14497 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
14498 assert_eq!(explicit_off_with_knob, explicit_off);
14499 assert!(
14500 !explicit_off_with_knob.contains("<|think|>")
14501 && explicit_off_with_knob.ends_with("<|turn>model\n"),
14502 "explicit off keeps the official template's thinking-off bytes: \
14503 {explicit_off_with_knob:?}"
14504 );
14505 let unset_no_knob = render(&official, json!({}), None);
14507 assert_eq!(
14508 unset_no_knob, explicit_off,
14509 "knobless unset stays the template's own thinking-off default"
14510 );
14511 assert_ne!(unset_no_knob, unset_with_knob);
14512 let qat = gemma_template("qat");
14515 assert!(
14516 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
14517 "QAT knobless unset keeps the closed-channel default"
14518 );
14519 assert_eq!(
14520 render(&qat, json!({}), Some("high")),
14521 render(&qat, json!({"reasoning_effort": "high"}), None),
14522 "QAT knob render must equal the explicit think-on render"
14523 );
14524 }
14525
14526 #[test]
14527 fn default_reasoning_effort_is_validated_at_metadata_load() {
14528 let parsed = OpenRouterMetadataFile::from_toml(
14530 r#"
14531[models.g]
14532default_reasoning_effort = "high"
14533"#,
14534 )
14535 .unwrap();
14536 assert_eq!(
14537 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
14538 Some("high")
14539 );
14540 let err = OpenRouterMetadataFile::from_toml(
14541 r#"
14542[models.g]
14543default_reasoning_effort = "always"
14544"#,
14545 )
14546 .unwrap_err();
14547 assert!(err.contains("default_reasoning_effort"), "{err}");
14548 }
14549
14550 #[test]
14551 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
14552 let effort_caps = ModelCaps {
14564 effort_levels: true,
14565 think_switch: false,
14566 ..tool_caps()
14567 };
14568 for (extra, want) in [
14569 (json!({}), None),
14570 (json!({"reasoning_effort": "low"}), Some("low")),
14571 (json!({"reasoning_effort": "medium"}), Some("medium")),
14572 (json!({"reasoning_effort": "high"}), Some("high")),
14573 (json!({"reasoning": {"effort": "high"}}), Some("high")),
14574 (json!({"reasoning_effort": "xhigh"}), Some("high")),
14576 (json!({"reasoning": {"effort": "max"}}), Some("high")),
14577 ] {
14578 let (tx, _rx) = worker::event_channel();
14579 let plan = build_chat_request(
14580 weather_request(extra.clone()),
14581 Some(&effort_caps),
14582 tx,
14583 lanes::Lane::Interactive,
14584 None,
14585 )
14586 .unwrap();
14587 assert_eq!(
14588 plan.request.reasoning_effort.as_deref(),
14589 want,
14590 "extra={extra}"
14591 );
14592 }
14593 for extra in [
14600 json!({"reasoning_effort": "none"}),
14601 json!({"reasoning_effort": "minimal"}),
14602 json!({"reasoning": {"enabled": false}}),
14603 json!({"enable_thinking": false}),
14604 json!({"include_reasoning": false}),
14605 ] {
14606 let (tx, _rx) = worker::event_channel();
14607 let err = build_chat_request(
14608 weather_request(extra.clone()),
14609 Some(&effort_caps),
14610 tx,
14611 lanes::Lane::Interactive,
14612 None,
14613 )
14614 .err()
14615 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
14616 assert!(
14617 err.contains("cannot disable reasoning"),
14618 "extra={extra}: {err}"
14619 );
14620 }
14621 for extra in [
14628 json!({"reasoning_effort": "high"}),
14629 json!({"reasoning": {"effort": "low"}}),
14630 ] {
14631 let (tx, _rx) = worker::event_channel();
14632 let plan = build_chat_request(
14633 weather_request(extra.clone()),
14634 Some(&tool_caps()),
14635 tx,
14636 lanes::Lane::Interactive,
14637 None,
14638 )
14639 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
14640 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
14641 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
14642 }
14643 let (tx, _rx) = worker::event_channel();
14645 let plan = build_chat_request(
14646 weather_request(json!({})),
14647 Some(&tool_caps()),
14648 tx,
14649 lanes::Lane::Interactive,
14650 None,
14651 )
14652 .unwrap();
14653 assert_eq!(plan.request.reasoning_effort, None);
14654 }
14655
14656 #[test]
14657 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
14658 let payload = serde_json::json!({
14659 "model": "m",
14660 "messages": [
14661 {"role": "user", "content": "Weather in Paris?"},
14662 {"role": "assistant", "content": null, "tool_calls": [
14663 {"id": "call_x", "type": "function", "function": {
14664 "name": "get_weather",
14665 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
14666 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
14667 ],
14668 });
14669 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
14670 let (tx, _rx) = worker::event_channel();
14671 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
14672 .unwrap();
14673 let turns = &plan.request.chat_turns;
14674 assert_eq!(turns[1].tool_calls.len(), 1);
14675 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
14676 assert_eq!(
14677 turns[1].tool_calls[0].params,
14678 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
14679 );
14680 assert_eq!(turns[2].role, "tool");
14681 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
14682 let mut p = plan
14685 .parser
14686 .expect("think-open chat arms the reasoning splitter");
14687 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
14688 assert_eq!(
14689 pieces,
14690 vec![
14691 Piece::Reasoning("thought".into()),
14692 Piece::Content("answer <tool_call> is prose here".into()),
14693 ]
14694 );
14695 }
14696
14697 #[tokio::test]
14698 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
14699 let (tx, rx) = worker::event_channel();
14700 tx.send(Event::Token {
14701 id: 1,
14702 text: "plan</think>\n\n".into(),
14703 })
14704 .unwrap();
14705 tx.send(Event::Token {
14706 id: 2,
14707 text: "<tool_call>\n<function=get_weather>\n\
14708<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
14709 .into(),
14710 })
14711 .unwrap();
14712 tx.send(Event::Done {
14713 stop_reason: "Eos".into(),
14714 n_tokens: 2,
14715 n_prompt: 40,
14716 n_cached: 0,
14717 elapsed_s: 0.5,
14718 spec: None,
14719 })
14720 .unwrap();
14721 drop(tx);
14722 let parser = ToolStreamParser::new(HashMap::new(), true);
14723 let response = blocking_response(
14724 rx,
14725 "m".into(),
14726 true,
14727 Vec::new(),
14728 Some(parser),
14729 Envelope::new(true),
14730 )
14731 .await;
14732 assert_eq!(response.status(), StatusCode::OK);
14733 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
14734 .await
14735 .unwrap();
14736 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14737 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
14738 assert_eq!(
14741 payload["choices"][0]["message"]["content"],
14742 serde_json::Value::Null
14743 );
14744 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
14745 assert_eq!(
14746 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
14747 "plan"
14748 );
14749 let call = &payload["choices"][0]["message"]["tool_calls"][0];
14750 assert_eq!(call["type"], "function");
14751 assert_eq!(call["function"]["name"], "get_weather");
14752 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
14753 assert_eq!(payload["usage"]["prompt_tokens"], 40);
14756 assert_eq!(payload["usage"]["completion_tokens"], 2);
14757 assert_eq!(payload["usage"]["total_tokens"], 42);
14758 assert_eq!(
14759 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
14760 0
14761 );
14762 }
14763
14764 #[test]
14765 fn cache_salt_plumbs_to_the_worker_namespace() {
14766 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14768 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
14769 }))
14770 .unwrap();
14771 let (tx, _rx) = worker::event_channel();
14772 assert_eq!(
14773 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14774 "tenant-a"
14775 );
14776
14777 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14778 "model": "m", "messages": [{"role": "user", "content": "task"}],
14779 "cache_salt": "tenant-b"
14780 }))
14781 .unwrap();
14782 let (tx, _rx) = worker::event_channel();
14783 assert_eq!(
14784 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14785 .unwrap()
14786 .request
14787 .cache_ns,
14788 "tenant-b"
14789 );
14790
14791 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14793 "model": "m", "prompt": "task"
14794 }))
14795 .unwrap();
14796 let (tx, _rx) = worker::event_channel();
14797 assert_eq!(
14798 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14799 ""
14800 );
14801 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14802 "model": "m", "messages": [{"role": "user", "content": "task"}]
14803 }))
14804 .unwrap();
14805 let (tx, _rx) = worker::event_channel();
14806 assert_eq!(
14807 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14808 .unwrap()
14809 .request
14810 .cache_ns,
14811 ""
14812 );
14813 }
14814
14815 #[test]
14816 fn cache_salt_validation_rejects_oversized_value() {
14817 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
14818 assert_eq!(
14819 validate_cache_namespace(&salt, false),
14820 Err("cache_salt must be at most 64 bytes")
14821 );
14822 }
14823
14824 #[test]
14825 fn cache_salt_validation_rejects_reserved_open_namespace() {
14826 let salt = Some("t:acme\u{1f}private".to_string());
14827 assert_eq!(
14828 validate_cache_namespace(&salt, false),
14829 Err("cache_salt must not use the reserved t: prefix without a keyring")
14830 );
14831 }
14832
14833 #[test]
14834 fn cache_salt_validation_accepts_normal_value() {
14835 let raw = "tenant-A_7.c2VjcmV0LXNjb3Bl+/=";
14836 let salt = Some(raw.to_string());
14837 assert_eq!(validate_cache_namespace(&salt, false).unwrap(), raw);
14838 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
14839 let max_raw = "a".repeat(CACHE_SALT_MAX_BYTES);
14840 let max = Some(max_raw.clone());
14841 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max_raw);
14842 }
14843
14844 #[test]
14845 fn cache_salt_validation_rejects_unsupported_characters() {
14846 let salt = Some("tenant salt".to_string());
14847 assert_eq!(
14848 validate_cache_namespace(&salt, false),
14849 Err("cache_salt contains unsupported characters")
14850 );
14851 }
14852
14853 #[test]
14854 fn affinity_key_honors_both_client_conventions_in_priority_order() {
14855 use axum::http::HeaderMap;
14856 let hdr = |v: &str| {
14857 let mut h = HeaderMap::new();
14858 h.insert("x-session-id", v.parse().unwrap());
14859 h
14860 };
14861 let empty = HeaderMap::new();
14862 let s = |v: &str| Some(v.to_string());
14863 assert_eq!(
14865 affinity_key(&s("explicit"), &None, &empty).unwrap(),
14866 s("explicit")
14867 );
14868 assert_eq!(
14869 affinity_key(&None, &s("openai-user"), &empty).unwrap(),
14870 s("openai-user")
14871 );
14872 assert_eq!(
14873 affinity_key(&None, &None, &hdr("hdr-id")).unwrap(),
14874 s("hdr-id")
14875 );
14876 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")).unwrap(), s("a"));
14879 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")).unwrap(), s("b"));
14880 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")).unwrap(), None);
14883 assert_eq!(affinity_key(&s(""), &s("real"), &empty).unwrap(), s("real"));
14884 assert_eq!(
14886 affinity_key(&s(" padded "), &None, &empty).unwrap(),
14887 s("padded")
14888 );
14889 assert_eq!(affinity_key(&None, &None, &empty).unwrap(), None);
14891 assert!(
14892 affinity_key(
14893 &s(&"x".repeat(MAX_CLIENT_IDENTIFIER_BYTES + 1)),
14894 &None,
14895 &empty,
14896 )
14897 .unwrap_err()
14898 .contains("at most")
14899 );
14900 assert!(
14901 affinity_key(&s("forged\nlog"), &None, &empty)
14902 .unwrap_err()
14903 .contains("control")
14904 );
14905 }
14906
14907 #[test]
14908 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
14909 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14910 "model": "m", "prompt": "task", "session_id": "conv-1"
14911 }))
14912 .unwrap();
14913 let (tx, _rx) = worker::event_channel();
14914 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14915 assert_eq!(
14916 build_request(&req, tx, lanes::Lane::Interactive, key)
14917 .affinity
14918 .as_deref(),
14919 Some("conv-1")
14920 );
14921 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14923 "model": "m", "messages": [{"role": "user", "content": "task"}],
14924 "user": "conv-2"
14925 }))
14926 .unwrap();
14927 let (tx, _rx) = worker::event_channel();
14928 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14929 assert_eq!(
14930 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
14931 .unwrap()
14932 .request
14933 .affinity
14934 .as_deref(),
14935 Some("conv-2")
14936 );
14937 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14939 "model": "m", "prompt": "task"
14940 }))
14941 .unwrap();
14942 let (tx, _rx) = worker::event_channel();
14943 assert!(
14944 build_request(&req, tx, lanes::Lane::Interactive, None)
14945 .affinity
14946 .is_none()
14947 );
14948 }
14949
14950 async fn sse_data_lines(resp: Response) -> Vec<String> {
14952 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14953 .await
14954 .unwrap();
14955 String::from_utf8(bytes.to_vec())
14956 .unwrap()
14957 .lines()
14958 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
14959 .collect()
14960 }
14961
14962 #[tokio::test]
14963 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
14964 let feed = |think: bool| {
14970 let (tx, rx) = worker::event_channel();
14971 let body = if think {
14972 "a plan</think>\n\nanswer"
14973 } else {
14974 "answer"
14975 };
14976 tx.send(Event::Token {
14977 id: 1,
14978 text: body.into(),
14979 })
14980 .unwrap();
14981 tx.send(Event::Done {
14982 stop_reason: "Eos".into(),
14983 n_tokens: 3,
14984 n_prompt: 10,
14985 n_cached: 0,
14986 elapsed_s: 0.1,
14987 spec: None,
14988 })
14989 .unwrap();
14990 drop(tx);
14991 rx
14992 };
14993 let resp = blocking_response(
14995 feed(true),
14996 "m".into(),
14997 true,
14998 Vec::new(),
14999 Some(ToolStreamParser::reasoning_only()),
15000 Envelope::new(true),
15001 )
15002 .await;
15003 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15004 .await
15005 .unwrap();
15006 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15007 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
15008 assert_eq!(
15009 v["choices"][0]["message"]["reasoning_details"][0]["text"],
15010 "a plan"
15011 );
15012 assert_eq!(v["choices"][0]["message"]["content"], "answer");
15013 let resp = blocking_response(
15016 feed(false),
15017 "m".into(),
15018 true,
15019 Vec::new(),
15020 None,
15021 Envelope::new(true),
15022 )
15023 .await;
15024 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15025 .await
15026 .unwrap();
15027 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15028 assert!(
15029 v["choices"][0]["message"].get("reasoning").is_none(),
15030 "a reasoning-off response must carry no reasoning field: {v}"
15031 );
15032 assert_eq!(v["choices"][0]["message"]["content"], "answer");
15033 let resp = sse_response(
15035 feed(true),
15036 "m".into(),
15037 true,
15038 Some(ToolStreamParser::reasoning_only()),
15039 Envelope::new(true),
15040 Vec::new(),
15041 None,
15042 )
15043 .into_response();
15044 let lines = sse_data_lines(resp).await;
15045 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
15046 .iter()
15047 .map(|l| serde_json::from_str(l).unwrap())
15048 .collect();
15049 let reasoning: String = chunks
15050 .iter()
15051 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
15052 .collect();
15053 assert_eq!(
15054 reasoning, "a plan",
15055 "think text must stream as delta.reasoning"
15056 );
15057 let content: String = chunks
15058 .iter()
15059 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
15060 .collect();
15061 assert_eq!(content, "answer", "content must exclude the think segment");
15062 let resp = sse_response(
15064 feed(false),
15065 "m".into(),
15066 true,
15067 None,
15068 Envelope::new(true),
15069 Vec::new(),
15070 None,
15071 )
15072 .into_response();
15073 let lines = sse_data_lines(resp).await;
15074 for l in &lines[..lines.len() - 1] {
15075 let c: serde_json::Value = serde_json::from_str(l).unwrap();
15076 assert!(
15077 c["choices"][0]["delta"].get("reasoning").is_none(),
15078 "a reasoning-off stream must carry no reasoning deltas: {c}"
15079 );
15080 }
15081 }
15082
15083 #[tokio::test]
15084 async fn stream_chunks_carry_envelope_and_first_delta_role() {
15085 let (tx, rx) = worker::event_channel();
15086 tx.send(Event::Token {
15087 id: 1,
15088 text: "he".into(),
15089 })
15090 .unwrap();
15091 tx.send(Event::Token {
15092 id: 2,
15093 text: "llo".into(),
15094 })
15095 .unwrap();
15096 tx.send(Event::Done {
15097 stop_reason: "Eos".into(),
15098 n_tokens: 2,
15099 n_prompt: 10,
15100 n_cached: 0,
15101 elapsed_s: 0.1,
15102 spec: None,
15103 })
15104 .unwrap();
15105 drop(tx);
15106 let resp = sse_response(
15107 rx,
15108 "m".into(),
15109 true,
15110 None,
15111 Envelope::new(true),
15112 Vec::new(),
15113 None,
15114 )
15115 .into_response();
15116 let lines = sse_data_lines(resp).await;
15117 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
15118 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
15119 .iter()
15120 .map(|l| serde_json::from_str(l).unwrap())
15121 .collect();
15122 let id = chunks[0]["id"].as_str().unwrap().to_string();
15124 assert!(id.starts_with("chatcmpl-"));
15125 for c in &chunks {
15126 assert_eq!(c["id"], id.as_str());
15127 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
15128 let fingerprint = c["system_fingerprint"].as_str().unwrap();
15129 assert!(
15130 build_id::fingerprint_is_well_formed(fingerprint),
15131 "chunk system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
15132 );
15133 assert_eq!(c["object"], "chat.completion.chunk");
15134 }
15135 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
15137 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
15138 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
15139 let fin = chunks.last().unwrap();
15141 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
15142 assert_eq!(fin["usage"]["prompt_tokens"], 10);
15143 }
15144
15145 #[tokio::test]
15146 async fn stream_token_events_equal_usage_on_every_finish_path() {
15147 for (stop_reason, expected_finish) in [
15148 ("Eos", "stop"),
15149 ("Callback", "stop"),
15150 ("MaxNew", "length"),
15151 ("ContextFull", "length"),
15152 ] {
15153 let (tx, rx) = worker::event_channel();
15154 tx.send(Event::Token {
15157 id: 248_046,
15158 text: String::new(),
15159 })
15160 .unwrap();
15161 tx.send(Event::Done {
15162 stop_reason: stop_reason.into(),
15163 n_tokens: 1,
15164 n_prompt: 8,
15165 n_cached: 8,
15166 elapsed_s: 0.1,
15167 spec: None,
15168 })
15169 .unwrap();
15170 drop(tx);
15171
15172 let resp = sse_response(
15173 rx,
15174 "m".into(),
15175 true,
15176 None,
15177 Envelope::new(true),
15178 Vec::new(),
15179 None,
15180 )
15181 .into_response();
15182 let lines = sse_data_lines(resp).await;
15183 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
15184 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
15185 .iter()
15186 .map(|line| serde_json::from_str(line).unwrap())
15187 .collect();
15188 let token_events = chunks
15189 .iter()
15190 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
15191 .count();
15192 let terminal = chunks.last().unwrap();
15193 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
15194 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
15195 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
15196 }
15197 }
15198
15199 #[tokio::test]
15200 async fn stream_excludes_stop_text_like_non_stream_does() {
15201 let (tx, rx) = worker::event_channel();
15205 tx.send(Event::Token {
15206 id: 1,
15207 text: "answer\nPro".into(),
15208 })
15209 .unwrap();
15210 tx.send(Event::Token {
15211 id: 2,
15212 text: "blem: leaked prompt".into(),
15213 })
15214 .unwrap();
15215 tx.send(Event::Done {
15216 stop_reason: "Callback".into(),
15217 n_tokens: 2,
15218 n_prompt: 8,
15219 n_cached: 0,
15220 elapsed_s: 0.1,
15221 spec: None,
15222 })
15223 .unwrap();
15224 drop(tx);
15225 let resp = sse_response(
15226 rx,
15227 "m".into(),
15228 true,
15229 None,
15230 Envelope::new(true),
15231 vec!["Problem:".into()],
15232 None,
15233 )
15234 .into_response();
15235 let lines = sse_data_lines(resp).await;
15236 let content: String = lines
15237 .iter()
15238 .filter(|l| *l != "[DONE]")
15239 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
15240 .filter_map(|c| {
15241 c["choices"][0]["delta"]["content"]
15242 .as_str()
15243 .map(str::to_string)
15244 })
15245 .collect();
15246 assert_eq!(content, "answer\n");
15247
15248 let (tx, rx) = worker::event_channel();
15250 tx.send(Event::Token {
15251 id: 1,
15252 text: "ends in Pro".into(),
15253 })
15254 .unwrap();
15255 tx.send(Event::Done {
15256 stop_reason: "Eos".into(),
15257 n_tokens: 1,
15258 n_prompt: 8,
15259 n_cached: 0,
15260 elapsed_s: 0.1,
15261 spec: None,
15262 })
15263 .unwrap();
15264 drop(tx);
15265 let resp = sse_response(
15266 rx,
15267 "m".into(),
15268 true,
15269 None,
15270 Envelope::new(true),
15271 vec!["Problem:".into()],
15272 None,
15273 )
15274 .into_response();
15275 let lines = sse_data_lines(resp).await;
15276 let content: String = lines
15277 .iter()
15278 .filter(|l| *l != "[DONE]")
15279 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
15280 .filter_map(|c| {
15281 c["choices"][0]["delta"]["content"]
15282 .as_str()
15283 .map(str::to_string)
15284 })
15285 .collect();
15286 assert_eq!(content, "ends in Pro");
15287 }
15288
15289 #[tokio::test]
15290 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
15291 let (tx, rx) = worker::event_channel();
15292 tx.send(Event::Error(worker::EngineError::engine("boom")))
15293 .unwrap();
15294 drop(tx);
15295 let resp = sse_response(
15296 rx,
15297 "m".into(),
15298 true,
15299 None,
15300 Envelope::new(true),
15301 Vec::new(),
15302 None,
15303 )
15304 .into_response();
15305 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15306 .await
15307 .unwrap();
15308 let body = String::from_utf8(bytes.to_vec()).unwrap();
15309 assert!(
15311 !body.contains("event: error"),
15312 "named SSE event leaked: {body}"
15313 );
15314 let lines: Vec<&str> = body
15315 .lines()
15316 .filter_map(|l| l.strip_prefix("data: "))
15317 .collect();
15318 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
15319 assert_eq!(
15322 err["error"]["message"],
15323 worker::engine_client_message(worker::ErrClass::Engine)
15324 );
15325 assert!(
15326 !body.contains("boom"),
15327 "producer text leaked into the stream: {body}"
15328 );
15329 assert_eq!(err["error"]["type"], "server_error");
15330 assert_eq!(err["error"]["code"], "engine_error");
15331 assert_eq!(lines.last(), Some(&"[DONE]"));
15332 }
15333
15334 #[test]
15335 fn ttft_sse_marker_ignores_keepalive_comments() {
15336 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
15337 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
15338 assert!(is_sse_data_frame(
15339 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
15340 ));
15341 }
15342
15343 #[tokio::test]
15344 async fn error_bodies_use_the_openai_object_shape() {
15345 let (tx, rx) = worker::event_channel();
15346 tx.send(Event::Error(worker::EngineError::model_not_found(
15347 "unknown model \"x\"",
15348 )))
15349 .unwrap();
15350 drop(tx);
15351 let response =
15352 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
15353 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15354 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15355 .await
15356 .unwrap();
15357 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15358 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
15360 assert_eq!(payload["error"]["type"], "invalid_request_error");
15361 assert_eq!(payload["error"]["param"], "model");
15362 assert_eq!(payload["error"]["code"], "model_not_found");
15363 }
15364
15365 fn retry_after(resp: &Response) -> Option<String> {
15372 resp.headers()
15373 .get(axum::http::header::RETRY_AFTER)
15374 .and_then(|v| v.to_str().ok())
15375 .map(str::to_string)
15376 }
15377
15378 async fn body_value(resp: Response) -> serde_json::Value {
15381 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15382 .await
15383 .expect("body");
15384 serde_json::from_slice(&bytes).expect("json body")
15385 }
15386
15387 async fn chat_completion_admitted(st: &AppState, req: serde_json::Value) -> Response {
15403 let mut last_shed = serde_json::Value::Null;
15404 for _ in 0..50 {
15405 let resp = chat_completions(
15406 State(st.clone()),
15407 HeaderMap::new(),
15408 None,
15409 Json(serde_json::from_value(req.clone()).unwrap()),
15410 )
15411 .await;
15412 if resp.status() != StatusCode::TOO_MANY_REQUESTS {
15413 return resp;
15414 }
15415 let body = body_value(resp).await;
15416 let code = body["error"]["code"].as_str().unwrap_or_default();
15417 assert!(
15418 code.starts_with("shed_"),
15419 "only a contention shed may be retried; any other 429 is a finding: {body}"
15420 );
15421 last_shed = body;
15422 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
15423 }
15424 panic!(
15427 "still shed after 50 attempts — either load the retry budget cannot absorb \
15428 or a shed that no longer clears; last refusal: {last_shed}"
15429 );
15430 }
15431
15432 #[test]
15433 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
15434 assert_eq!(parse_timeout_ms(None, false).unwrap(), TIMEOUT_MS_DEFAULT);
15436 assert_eq!(
15437 parse_timeout_ms(Some(&serde_json::Value::Null), false).unwrap(),
15438 TIMEOUT_MS_DEFAULT
15439 );
15440 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
15444 assert_eq!(parse_timeout_ms(Some(&json!(ms)), false).unwrap(), ms);
15445 }
15446 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
15448 let err =
15449 parse_timeout_ms(Some(&json!(bad)), false).expect_err("out of range must refuse");
15450 assert!(err.contains("timeout_ms"), "{err}");
15451 assert!(
15452 err.contains(&TIMEOUT_MS_MIN.to_string())
15453 && err.contains(&TIMEOUT_MS_MAX.to_string()),
15454 "the message must state the range: {err}"
15455 );
15456 assert!(
15457 err.contains("stream"),
15458 "the message must point at streaming for longer work: {err}"
15459 );
15460 }
15461 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
15463 let err = parse_timeout_ms(Some(&bad), false).expect_err("bad type must refuse");
15464 assert!(
15465 err.contains("timeout_ms") && err.contains("stream"),
15466 "{err}"
15467 );
15468 }
15469 assert!(parse_timeout_ms(Some(&json!(-1)), false).is_err());
15471 }
15472
15473 #[tokio::test]
15476 #[allow(clippy::await_holding_lock)] async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
15478 let _l = drain_lock();
15479 let st = fake_worker_state();
15480
15481 let comp = completions(
15482 State(st.clone()),
15483 HeaderMap::new(),
15484 None,
15485 Json(
15486 serde_json::from_value(json!({
15487 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
15488 .unwrap(),
15489 ),
15490 )
15491 .await;
15492 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
15493 let chat = chat_completions(
15494 State(st.clone()),
15495 HeaderMap::new(),
15496 None,
15497 Json(
15498 serde_json::from_value(json!({
15499 "model": "m", "messages": [{"role": "user", "content": "t"}],
15500 "timeout_ms": 90_001}))
15501 .unwrap(),
15502 ),
15503 )
15504 .await;
15505 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
15506 let resp_api = responses_api::responses(
15507 State(st.clone()),
15508 HeaderMap::new(),
15509 None,
15510 axum::body::Bytes::from(
15511 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
15512 ),
15513 )
15514 .await;
15515 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
15516 let msgs = anthropic::messages(
15517 State(st.clone()),
15518 HeaderMap::new(),
15519 None,
15520 axum::body::Bytes::from(
15521 json!({"model": "m", "max_tokens": 16,
15522 "messages": [{"role": "user", "content": "t"}],
15523 "timeout_ms": 90_001})
15524 .to_string(),
15525 ),
15526 )
15527 .await;
15528 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
15529
15530 for (surface, resp) in [
15532 ("/v1/completions", comp),
15533 ("/v1/chat/completions", chat),
15534 ("/v1/responses", resp_api),
15535 ] {
15536 let body = body_value(resp).await;
15537 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
15538 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
15539 let m = body["error"]["message"].as_str().unwrap();
15540 assert!(
15541 m.contains("90000") && m.contains("stream"),
15542 "{surface}: {m}"
15543 );
15544 }
15545 let body = body_value(msgs).await;
15547 assert_eq!(body["error"]["type"], "invalid_request_error");
15548 let m = body["error"]["message"].as_str().unwrap();
15549 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
15550 }
15551
15552 #[tokio::test]
15555 #[allow(clippy::await_holding_lock)] async fn a_non_integer_timeout_ms_is_a_named_400() {
15557 let _l = drain_lock();
15558 let st = fake_worker_state();
15559 let resp = chat_completions(
15560 State(st),
15561 HeaderMap::new(),
15562 None,
15563 Json(
15564 serde_json::from_value(json!({
15565 "model": "m", "messages": [{"role": "user", "content": "t"}],
15566 "timeout_ms": "30s"}))
15567 .unwrap(),
15568 ),
15569 )
15570 .await;
15571 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
15572 let body = body_value(resp).await;
15573 assert_eq!(body["error"]["param"], "timeout_ms");
15574 }
15575
15576 #[tokio::test]
15582 #[allow(clippy::await_holding_lock)] async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
15584 let _l = drain_lock();
15585 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15590 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
15591 let worker_cancel = cancel_seen.clone();
15592 let health = health::WorkerHealth::new();
15593 let h = health.clone();
15594 std::thread::spawn(move || {
15595 h.mark_ready();
15596 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15597 worker::release_pending_admit();
15598 worker::release_admission_reservation(req.lane);
15599 let _ = req.tx.send(Event::PromptUsage {
15600 n_prompt: 1,
15601 n_cached: 0,
15602 });
15603 let _ = req.tx.send(Event::Token {
15604 id: 1,
15605 text: "partial".into(),
15606 });
15607 for _ in 0..5_000 {
15610 if req.tx.is_closed() {
15611 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
15612 break;
15613 }
15614 std::thread::sleep(std::time::Duration::from_millis(1));
15615 }
15616 }
15617 });
15618 for _ in 0..2_000 {
15619 if health.live().is_ok() {
15620 break;
15621 }
15622 std::thread::sleep(std::time::Duration::from_millis(1));
15623 }
15624 let mut st = fake_worker_state();
15625 st.cmd_tx = cmd_tx;
15626 st.health = health;
15627 let mock = MockMetering::admit_all();
15628 st.metering = Some(mock.clone());
15629
15630 let resp = chat_completion_admitted(
15631 &st,
15632 json!({
15633 "model": "m", "messages": [{"role": "user", "content": "t"}],
15634 "timeout_ms": 1_000}),
15635 )
15636 .await;
15637
15638 assert_eq!(resp.status(), StatusCode::OK);
15643 let body = body_value(resp).await;
15644 assert!(
15645 body["choices"][0]["message"]["content"]
15646 .as_str()
15647 .unwrap()
15648 .contains("partial"),
15649 "the tokens generated before the cut must be delivered: {body}"
15650 );
15651 assert_eq!(body["choices"][0]["finish_reason"], "error");
15655 assert_eq!(
15656 body["choices"][0]["native_finish_reason"],
15657 "deadline_exceeded"
15658 );
15659 assert_eq!(body["error"]["code"], "deadline_exceeded");
15660 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
15661 let message = body["error"]["message"].as_str().unwrap();
15662 assert!(
15663 message.contains("1000") && message.contains("stream"),
15664 "the partial must name the deadline and the streaming alternative: {message}"
15665 );
15666 assert_eq!(body["usage"]["completion_tokens"], 1);
15667
15668 let mut cancelled = false;
15673 for _ in 0..500 {
15674 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
15675 cancelled = true;
15676 break;
15677 }
15678 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
15679 }
15680 assert!(
15681 cancelled,
15682 "the deadline must CANCEL generation (worker's event channel closed)"
15683 );
15684
15685 let events = mock.events();
15690 assert!(
15691 events.contains(&MeterEvent::DeadlinePartial {
15692 prompt: 1,
15693 cached: 0,
15694 completion: 1,
15695 }),
15696 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
15697 );
15698 assert!(
15699 !events
15700 .iter()
15701 .any(|e| matches!(e, MeterEvent::Complete { .. })),
15702 "a deadline cut must stay distinguishable from a full answer: {events:?}"
15703 );
15704 }
15705
15706 #[tokio::test]
15710 #[allow(clippy::await_holding_lock)] async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
15712 let _l = drain_lock();
15713 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15714 let health = health::WorkerHealth::new();
15715 let h = health.clone();
15716 std::thread::spawn(move || {
15717 h.mark_ready();
15718 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15721 worker::release_pending_admit();
15722 worker::release_admission_reservation(req.lane);
15723 let _ = req.tx.send(Event::PromptUsage {
15724 n_prompt: 1,
15725 n_cached: 0,
15726 });
15727 for _ in 0..5_000 {
15728 if req.tx.is_closed() {
15729 break;
15730 }
15731 std::thread::sleep(std::time::Duration::from_millis(1));
15732 }
15733 }
15734 });
15735 for _ in 0..2_000 {
15736 if health.live().is_ok() {
15737 break;
15738 }
15739 std::thread::sleep(std::time::Duration::from_millis(1));
15740 }
15741 let mut st = fake_worker_state();
15742 st.cmd_tx = cmd_tx;
15743 st.health = health;
15744 let mock = MockMetering::admit_all();
15745 st.metering = Some(mock.clone());
15746 let resp = chat_completion_admitted(
15747 &st,
15748 json!({
15749 "model": "m", "messages": [{"role": "user", "content": "t"}],
15750 "timeout_ms": 1_000}),
15751 )
15752 .await;
15753 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15754 assert!(resp.headers().get("x-should-retry").is_none());
15756 assert_eq!(retry_after(&resp), None);
15757 let body = body_value(resp).await;
15758 assert_eq!(body["error"]["code"], "deadline_exceeded");
15759 assert!(
15760 body["error"]["message"]
15761 .as_str()
15762 .unwrap()
15763 .contains("not billed"),
15764 "the zero-token 408 keeps the billing promise: {body}"
15765 );
15766 let events = mock.events();
15767 assert!(
15768 events.contains(&MeterEvent::Unbilled {
15769 outcome: "deadline_exceeded",
15770 status: 408,
15771 code: "deadline_exceeded".into(),
15772 }),
15773 "the named zero-debit census outcome, not the generic reject — every sibling \
15774 deadline path settles this one: {events:?}"
15775 );
15776 }
15777
15778 #[tokio::test]
15781 #[allow(clippy::await_holding_lock)] async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
15783 let _l = drain_lock();
15784 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15786 let health = health::WorkerHealth::new();
15787 let h = health.clone();
15788 std::thread::spawn(move || {
15789 h.mark_ready();
15790 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15791 worker::release_pending_admit();
15792 worker::release_admission_reservation(req.lane);
15793 let _ = req.tx.send(Event::PromptUsage {
15794 n_prompt: 1,
15795 n_cached: 0,
15796 });
15797 while !req.tx.is_closed() {
15798 std::thread::sleep(std::time::Duration::from_millis(1));
15799 }
15800 }
15801 });
15802 for _ in 0..2_000 {
15803 if health.live().is_ok() {
15804 break;
15805 }
15806 std::thread::sleep(std::time::Duration::from_millis(1));
15807 }
15808 let mut st = fake_worker_state();
15809 st.cmd_tx = cmd_tx;
15810 st.health = health;
15811 let mock = MockMetering::admit_all();
15812 st.metering = Some(mock.clone());
15813
15814 let resp = chat_completion_admitted(
15815 &st,
15816 json!({
15817 "model": "m", "messages": [{"role": "user", "content": "t"}],
15818 "stream": true, "timeout_ms": 1_000}),
15819 )
15820 .await;
15821 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15824 let body = body_value(resp).await;
15825 assert_eq!(body["error"]["code"], "deadline_exceeded");
15826 assert!(
15827 body["error"]["message"]
15828 .as_str()
15829 .unwrap()
15830 .contains("first token"),
15831 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
15832 );
15833 let events = mock.events();
15834 assert!(
15835 events.contains(&MeterEvent::Unbilled {
15836 outcome: "deadline_exceeded",
15837 status: 408,
15838 code: "deadline_exceeded".into(),
15839 }),
15840 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
15841 );
15842 }
15843
15844 #[tokio::test]
15845 async fn an_extended_stream_commits_prefill_then_injects_the_original_deadline() {
15846 let (tx, rx) = worker::event_channel();
15847 tx.send(Event::PromptUsage {
15848 n_prompt: 218_000,
15849 n_cached: 0,
15850 })
15851 .unwrap();
15852 let deadline = RequestDeadline::starting_now(80);
15853 let started = tokio::time::Instant::now();
15854 let mut committed =
15855 peek_first_token_with_commit(rx, deadline, Some(std::time::Duration::from_millis(10)))
15856 .await
15857 .expect("the extended stream must commit before first token");
15858 assert!(
15859 started.elapsed() < std::time::Duration::from_millis(60),
15860 "the bridge waited for the first-token deadline instead of committing"
15861 );
15862 assert!(matches!(
15863 committed.recv().await,
15864 Some(Event::PromptUsage {
15865 n_prompt: 218_000,
15866 n_cached: 0
15867 })
15868 ));
15869 assert!(matches!(
15870 tokio::time::timeout(std::time::Duration::from_millis(150), committed.recv()).await,
15871 Ok(Some(Event::DeadlineExceeded { ms: 80 }))
15872 ));
15873 for _ in 0..50 {
15874 if tx.is_closed() {
15875 break;
15876 }
15877 tokio::time::sleep(std::time::Duration::from_millis(2)).await;
15878 }
15879 assert!(
15880 tx.is_closed(),
15881 "dropping the prefill bridge at the deadline must cancel the worker request"
15882 );
15883 }
15884
15885 #[test]
15886 fn an_extended_streaming_ceiling_requires_an_effective_prefill_commit() {
15887 assert!(validate_stream_prefill_config_values(None, None).is_ok());
15888 assert!(
15889 validate_stream_prefill_config_values(Some(TIMEOUT_MS_MAX), None).is_ok(),
15890 "the shipped 90 s posture needs no early commit"
15891 );
15892 let extended = TIMEOUT_MS_MAX + 90_000;
15893 assert!(validate_stream_prefill_config_values(Some(extended), None).is_err());
15894 assert!(validate_stream_prefill_config_values(Some(extended), Some(0)).is_err());
15895 assert!(validate_stream_prefill_config_values(Some(extended), Some(extended)).is_err());
15896 assert!(validate_stream_prefill_config_values(Some(extended), Some(10_000)).is_ok());
15897 }
15898
15899 #[test]
15900 fn streaming_prefill_config_is_wired_in_a_fresh_process() {
15901 if std::env::var_os("PREFILL_CONFIG_TEST_CHILD").is_none() {
15902 let result = std::process::Command::new(std::env::current_exe().unwrap())
15903 .args([
15904 "--exact",
15905 "tests::streaming_prefill_config_is_wired_in_a_fresh_process",
15906 "--nocapture",
15907 ])
15908 .env("PREFILL_CONFIG_TEST_CHILD", "1")
15909 .env_remove("MEMRA_TIMEOUT_MS_MAX")
15910 .env("MEMRA_STREAM_TTFT_MS_MAX", "300000")
15911 .env("MEMRA_SSE_PREFILL_COMMIT_MS", "100")
15912 .output()
15913 .unwrap();
15914 assert!(
15915 result.status.success(),
15916 "{}",
15917 String::from_utf8_lossy(&result.stderr)
15918 );
15919 return;
15920 }
15921 validate_stream_prefill_config().unwrap();
15922 assert_eq!(parse_timeout_ms(None, false).unwrap(), 90_000);
15923 assert_eq!(parse_timeout_ms(None, true).unwrap(), 300_000);
15924 assert!(parse_timeout_ms(Some(&json!(300_000)), false).is_err());
15925 tokio::runtime::Builder::new_current_thread()
15926 .enable_time()
15927 .build()
15928 .unwrap()
15929 .block_on(async {
15930 let (tx, rx) = worker::event_channel();
15931 tx.send(Event::PromptUsage {
15932 n_prompt: 1,
15933 n_cached: 0,
15934 })
15935 .unwrap();
15936 let rx = tokio::time::timeout(
15937 std::time::Duration::from_secs(2),
15938 peek_first_token(rx, RequestDeadline::starting_now(300_000), &mut None),
15939 )
15940 .await
15941 .expect("the configured production peek must commit before first token")
15942 .unwrap();
15943 drop(rx);
15944 tokio::time::timeout(std::time::Duration::from_secs(1), tx.closed())
15945 .await
15946 .unwrap();
15947 use tower::ServiceExt;
15950 for (path, payload) in [
15951 ("/v1/completions", json!({"model":"m","prompt":"t","stream":true,"max_tokens":16})),
15952 ("/v1/chat/completions", json!({"model":"m","messages":[{"role":"user","content":"t"}],"stream":true,"max_tokens":16})),
15953 ("/v1/messages", json!({"model":"m","messages":[{"role":"user","content":"t"}],"stream":true,"max_tokens":16})),
15954 ("/v1/responses", json!({"model":"m","input":"t","stream":true,"max_output_tokens":16})),
15955 ("/v1/chat/completions", json!({"model":"m","messages":[{"role":"user","content":"t"}],"stream":true,"max_tokens":16,"response_format":{"type":"json_object"}})),
15956 ("/v1/responses", json!({"model":"m","input":"t","stream":true,"max_output_tokens":16,"text":{"format":{"type":"json_object"}}})),
15957 ] {
15958 let mut st = fake_worker_state();
15959 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel();
15960 st.cmd_tx = cmd_tx;
15961 let mock = MockMetering::admit_all();
15962 st.metering = Some(mock.clone());
15963 let app = Router::new()
15964 .route("/v1/completions", post(completions_admitted))
15965 .route("/v1/chat/completions", post(chat_completions_admitted))
15966 .route("/v1/messages", post(anthropic::messages_admitted))
15967 .route("/v1/responses", post(responses_api::responses_admitted))
15968 .with_state(st);
15969 let response = tokio::time::timeout(std::time::Duration::from_secs(2),
15970 app.oneshot(axum::http::Request::builder().method("POST").uri(path)
15971 .header("content-type", "application/json")
15972 .body(axum::body::Body::from(payload.to_string())).unwrap()))
15973 .await.expect("admission must expire before the proxy ceiling").unwrap();
15974 assert_eq!(response.status(), StatusCode::REQUEST_TIMEOUT, "{path}");
15975 let body = body_value(response).await;
15976 assert!(body["error"]["message"].as_str().unwrap().contains("pre-header budget"), "{path}: {body}");
15977 let Cmd::Generate(req) = cmd_rx.try_recv().expect("request reached admission") else { panic!("unexpected command") };
15978 assert!(req.tx.is_closed(), "{path}: queued work must be cancelled");
15979 worker::release_pending_admit();
15980 worker::release_admission_reservation(req.lane);
15981 assert!(mock.events().contains(&MeterEvent::Unbilled {
15982 outcome: "deadline_exceeded", status: 408, code: "deadline_exceeded".into(),
15983 }));
15984 }
15985 });
15986 }
15987
15988 #[tokio::test]
15989 async fn committed_prefill_sends_a_real_keepalive_before_any_generated_token() {
15990 let (tx, rx) = worker::event_channel();
15991 tx.send(Event::PromptUsage {
15992 n_prompt: 218_000,
15993 n_cached: 0,
15994 })
15995 .unwrap();
15996 let rx = peek_first_token_with_commit(
15997 rx,
15998 RequestDeadline::starting_now(30_000),
15999 Some(std::time::Duration::from_millis(1)),
16000 )
16001 .await
16002 .unwrap();
16003 let response = sse_response(
16004 rx,
16005 "m".into(),
16006 true,
16007 None,
16008 Envelope::new(true),
16009 Vec::new(),
16010 None,
16011 )
16012 .into_response();
16013 let mut body = Box::pin(response.into_body().into_data_stream());
16014 let frame = tokio::time::timeout(
16015 std::time::Duration::from_secs(7),
16016 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)),
16017 )
16018 .await
16019 .expect("prefill must send the five-second heartbeat")
16020 .expect("stream must remain open")
16021 .unwrap();
16022 assert!(frame.starts_with(b":"), "expected SSE comment: {frame:?}");
16023 assert!(!is_sse_data_frame(&frame));
16024 drop(body);
16025 tokio::time::timeout(std::time::Duration::from_secs(1), tx.closed())
16026 .await
16027 .expect("disconnect must cancel silent prefill immediately");
16028 }
16029
16030 #[tokio::test]
16032 #[ignore = "manual loopback fixture for external proxy qualification"]
16033 async fn prefill_proxy_fixture() {
16034 async fn response(early: bool) -> Response {
16035 let (tx, rx) = worker::event_channel();
16036 tx.send(Event::PromptUsage {
16037 n_prompt: 1,
16038 n_cached: 0,
16039 })
16040 .unwrap();
16041 tokio::spawn(async move {
16042 tokio::select! {
16043 () = tx.closed() => return,
16044 () = tokio::time::sleep(std::time::Duration::from_secs(150)) => {}
16045 }
16046 let _ = tx.send(Event::Token {
16047 id: 1,
16048 text: "synthetic proxy gate".into(),
16049 });
16050 let _ = tx.send(Event::Done {
16051 stop_reason: "Eos".into(),
16052 n_tokens: 1,
16053 n_prompt: 1,
16054 n_cached: 0,
16055 elapsed_s: 150.0,
16056 spec: None,
16057 });
16058 });
16059 let rx = peek_first_token_with_commit(
16060 rx,
16061 RequestDeadline::starting_now(180_000),
16062 early.then(|| std::time::Duration::from_secs(10)),
16063 )
16064 .await
16065 .unwrap();
16066 sse_response(
16067 rx,
16068 "synthetic".into(),
16069 true,
16070 None,
16071 Envelope::new(true),
16072 Vec::new(),
16073 None,
16074 )
16075 .into_response()
16076 }
16077 let app = Router::new()
16078 .route("/early", get(|| response(true)))
16079 .route("/held", get(|| response(false)));
16080 let listener = tokio::net::TcpListener::bind("127.0.0.1:18995")
16081 .await
16082 .unwrap();
16083 eprintln!("prefill proxy fixture listening on 127.0.0.1:18995");
16084 axum::serve(listener, app)
16085 .with_graceful_shutdown(async {
16086 tokio::time::sleep(std::time::Duration::from_secs(720)).await;
16087 })
16088 .await
16089 .unwrap();
16090 }
16091
16092 #[tokio::test]
16093 async fn a_committed_prefill_timeout_is_a_named_sse_error_and_zero_debit() {
16094 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
16095 let receipt: Box<dyn metering::Receipt> = Box::new(MockReceipt {
16096 events: events.clone(),
16097 wants_capture: false,
16098 prompt: 0,
16099 cached: 0,
16100 completion: 0,
16101 finalized: false,
16102 });
16103 let (tx, rx) = worker::event_channel();
16104 tx.send(Event::PromptUsage {
16105 n_prompt: 218_000,
16106 n_cached: 0,
16107 })
16108 .unwrap();
16109 tx.send(Event::DeadlineExceeded { ms: 180_000 }).unwrap();
16110 drop(tx);
16111 let resp = sse_response_with_receipt(
16112 rx,
16113 "m".into(),
16114 true,
16115 None,
16116 Envelope::new(true),
16117 Vec::new(),
16118 None,
16119 Some(receipt),
16120 )
16121 .into_response();
16122 let lines = sse_data_lines(resp).await;
16123 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
16124 let error: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
16125 assert_eq!(error["error"]["type"], "timeout");
16126 assert_eq!(error["error"]["code"], "deadline_exceeded");
16127 assert!(
16128 error["error"]["message"]
16129 .as_str()
16130 .unwrap()
16131 .contains("not billed")
16132 );
16133 assert!(events.lock().unwrap().contains(&MeterEvent::Unbilled {
16134 outcome: "deadline_exceeded",
16135 status: 408,
16136 code: "deadline_exceeded".into(),
16137 }));
16138 }
16139
16140 #[tokio::test]
16141 async fn expired_prefill_settles_before_an_unpolled_body_is_dropped() {
16142 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
16143 let mut receipt: Option<Box<dyn metering::Receipt>> = Some(Box::new(MockReceipt {
16144 events: events.clone(),
16145 wants_capture: false,
16146 prompt: 0,
16147 cached: 0,
16148 completion: 0,
16149 finalized: false,
16150 }));
16151 let shared = prefill_receipt::SharedReceipt::wrap(&mut receipt);
16152 let (tx, rx) = worker::event_channel();
16153 tx.send(Event::PromptUsage {
16154 n_prompt: 218_000,
16155 n_cached: 0,
16156 })
16157 .unwrap();
16158 let rx = peek_first_token_with_receipt(
16159 rx,
16160 RequestDeadline::starting_now(100),
16161 Some(std::time::Duration::from_millis(1)),
16162 shared,
16163 )
16164 .await
16165 .unwrap();
16166 let response = sse_response_with_receipt(
16167 rx,
16168 "m".into(),
16169 true,
16170 None,
16171 Envelope::new(true),
16172 Vec::new(),
16173 None,
16174 receipt,
16175 )
16176 .into_response();
16177 let mut body = Box::pin(response.into_body().into_data_stream());
16178 std::future::poll_fn(|cx| {
16179 assert!(body.as_mut().poll_next(cx).is_pending());
16180 std::task::Poll::Ready(())
16181 })
16182 .await;
16183 tokio::time::timeout(std::time::Duration::from_secs(1), tx.closed())
16184 .await
16185 .unwrap();
16186 drop(body); let events = events.lock().unwrap();
16188 assert!(events.contains(&MeterEvent::PromptUsage {
16189 prompt: 218_000,
16190 cached: 0
16191 }));
16192 assert!(events.contains(&MeterEvent::Unbilled {
16193 outcome: "deadline_exceeded",
16194 status: 408,
16195 code: "deadline_exceeded".into(),
16196 }));
16197 assert!(
16198 !events
16199 .iter()
16200 .any(|e| matches!(e, MeterEvent::Dropped { .. } | MeterEvent::Complete { .. }))
16201 );
16202 }
16203
16204 #[tokio::test]
16205 async fn committed_prefill_disarms_deadline_after_first_token_and_bills_completion() {
16206 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
16207 let mut receipt: Option<Box<dyn metering::Receipt>> = Some(Box::new(MockReceipt {
16208 events: events.clone(),
16209 wants_capture: false,
16210 prompt: 0,
16211 cached: 0,
16212 completion: 0,
16213 finalized: false,
16214 }));
16215 let shared = prefill_receipt::SharedReceipt::wrap(&mut receipt);
16216 let (tx, rx) = worker::event_channel();
16217 tx.send(Event::PromptUsage {
16218 n_prompt: 1,
16219 n_cached: 0,
16220 })
16221 .unwrap();
16222 let rx = peek_first_token_with_receipt(
16223 rx,
16224 RequestDeadline::starting_now(200),
16225 Some(std::time::Duration::from_millis(1)),
16226 shared,
16227 )
16228 .await
16229 .unwrap();
16230 tokio::spawn(async move {
16231 let _ = tx.send(Event::Token {
16232 id: 1,
16233 text: "first".into(),
16234 });
16235 tokio::time::sleep(std::time::Duration::from_millis(300)).await;
16236 let _ = tx.send(Event::Token {
16237 id: 2,
16238 text: "second".into(),
16239 });
16240 let _ = tx.send(Event::Done {
16241 stop_reason: "Eos".into(),
16242 n_tokens: 2,
16243 n_prompt: 1,
16244 n_cached: 0,
16245 elapsed_s: 0.3,
16246 spec: None,
16247 });
16248 });
16249 let response = sse_response_with_receipt(
16250 rx,
16251 "m".into(),
16252 true,
16253 None,
16254 Envelope::new(true),
16255 Vec::new(),
16256 None,
16257 receipt,
16258 )
16259 .into_response();
16260 let lines = sse_data_lines(response).await;
16261 assert!(lines[0].contains("first") && lines[1].contains("second"));
16262 assert_eq!(lines.last().unwrap(), "[DONE]");
16263 let events = events.lock().unwrap();
16264 assert!(events.contains(&MeterEvent::Complete {
16265 prompt: 1,
16266 cached: 0,
16267 completion: 2
16268 }));
16269 assert!(
16270 !events
16271 .iter()
16272 .any(|e| matches!(e, MeterEvent::Unbilled { .. } | MeterEvent::Dropped { .. }))
16273 );
16274 }
16275
16276 #[tokio::test]
16280 #[allow(clippy::await_holding_lock)] async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
16282 let _l = drain_lock();
16283 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
16286 let mock = MockMetering::admit_all();
16287 st.metering = Some(mock.clone());
16288 let resp = chat_completion_admitted(
16289 &st,
16290 json!({
16291 "model": "m", "messages": [{"role": "user", "content": "t"}],
16292 "stream": true, "timeout_ms": 1_000}),
16293 )
16294 .await;
16295 assert_eq!(
16296 resp.status(),
16297 StatusCode::OK,
16298 "TTFT was met — 200 is correct"
16299 );
16300 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16301 .await
16302 .expect("the stream must run to completion past the deadline");
16303 let text = String::from_utf8(bytes.to_vec()).unwrap();
16304 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
16305 let events = mock.events();
16306 assert!(
16307 events
16308 .iter()
16309 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
16310 "a stream past its deadline after first token still settles as COMPLETE with \
16311 all four tokens: {events:?}"
16312 );
16313 }
16314
16315 fn admission_counters_guard() -> std::sync::MutexGuard<'static, ()> {
16322 static COUNTERS: std::sync::Mutex<()> = std::sync::Mutex::new(());
16323 COUNTERS
16324 .lock()
16325 .unwrap_or_else(|poisoned| poisoned.into_inner())
16326 }
16327
16328 struct CounterRestore<'a>(&'a std::sync::atomic::AtomicUsize, usize);
16334 impl Drop for CounterRestore<'_> {
16335 fn drop(&mut self) {
16336 self.0.store(self.1, std::sync::atomic::Ordering::Release);
16337 }
16338 }
16339
16340 #[allow(clippy::result_large_err)] fn reserve_interactive_through_contention(
16351 st: &AppState,
16352 rl: &RateLimit,
16353 deadline_ms: u64,
16354 ) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
16355 let reserve = || {
16356 reserve_pending_admit(
16357 st,
16358 lanes::Lane::Interactive,
16359 rl,
16360 RequestDeadline::starting_now(deadline_ms),
16361 )
16362 };
16363 let mut g = reserve();
16364 for _ in 0..50 {
16365 match &g {
16366 Ok(_) => break,
16367 Err((_, "shed_deadline")) => {
16368 std::thread::sleep(std::time::Duration::from_millis(10));
16369 g = reserve();
16370 }
16371 Err((_, outcome)) => panic!("unexpected refusal: {outcome}"),
16372 }
16373 }
16374 g
16375 }
16376
16377 #[test]
16380 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
16381 let _counters = admission_counters_guard();
16382 let st = fake_worker_state();
16383 let lane = lanes::Lane::Interactive;
16384 let cap = lane_cap(lane);
16385 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16386 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
16387 let _restore = CounterRestore(counter, prev);
16388 let rl = RateLimit {
16389 limit: cap,
16390 remaining: 0,
16391 reset_s: 1,
16392 };
16393 let (resp, outcome) = reserve_pending_admit(
16394 &st,
16395 lane,
16396 &rl,
16397 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16398 )
16399 .map(|_| ())
16400 .expect_err("a backlog at the bound must shed");
16401 assert_eq!(outcome, "shed_queue");
16402 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16403 assert!(
16404 retry_after(&resp).is_some(),
16405 "a shed must carry Retry-After so the router's spill can act on it"
16406 );
16407 let stamped = rl.attach(resp);
16409 for h in [
16410 "x-ratelimit-limit",
16411 "x-ratelimit-remaining",
16412 "x-ratelimit-reset",
16413 ] {
16414 assert!(stamped.headers().get(h).is_some(), "missing {h}");
16415 }
16416 }
16417
16418 #[test]
16422 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
16423 let _counters = admission_counters_guard();
16424 let st = fake_worker_state();
16425 let lane = lanes::Lane::Interactive;
16426 let cap = lane_cap(lane);
16427 {
16428 let mut m = st.metrics.lock().unwrap();
16429 m.completed = 10;
16430 m.tokens_out = 1_000;
16431 m.step_p50_ms = 10.0; }
16433 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16434 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
16436 let rl = RateLimit {
16437 limit: cap,
16438 remaining: 0,
16439 reset_s: 1,
16440 };
16441 let admitted = reserve_pending_admit(
16443 &st,
16444 lane,
16445 &rl,
16446 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16447 );
16448 assert!(
16449 admitted.is_ok(),
16450 "a request whose deadline covers the estimate must be admitted"
16451 );
16452 drop(admitted); let (resp, outcome) = reserve_pending_admit(
16455 &st,
16456 lane,
16457 &rl,
16458 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
16459 )
16460 .map(|_| ())
16461 .expect_err("a deadline shorter than the estimated wait must shed");
16462 assert_eq!(outcome, "shed_deadline");
16463 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16464 assert!(retry_after(&resp).is_some());
16465 }
16466
16467 #[test]
16470 fn deadline_shed_is_interactive_only_and_silent_with_free_slots() {
16471 let _counters = admission_counters_guard();
16472 let st = fake_worker_state();
16473 let cap = lane_cap(lanes::Lane::Interactive);
16474 {
16475 let mut m = st.metrics.lock().unwrap();
16476 m.completed = 10;
16477 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
16479 }
16480 let free = RateLimit {
16482 limit: cap,
16483 remaining: 1,
16484 reset_s: 0,
16485 };
16486 let g = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
16491 assert!(
16492 g.is_ok(),
16493 "free capacity must admit regardless of the estimate"
16494 );
16495 drop(g);
16496 let full = RateLimit {
16499 limit: cap,
16500 remaining: 0,
16501 reset_s: 5,
16502 };
16503 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
16504 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16505 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
16507 let g = reserve_pending_admit(
16508 &st,
16509 lane,
16510 &full,
16511 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
16512 );
16513 assert!(
16514 g.is_ok(),
16515 "{lane:?} must not be deadline-shed by the interactive gate"
16516 );
16517 drop(g);
16518 }
16519 }
16520
16521 #[test]
16528 fn a_saturated_queue_with_free_http_slots_queues_silently_without_a_ceiling() {
16529 let _counters = admission_counters_guard();
16530 let st = fake_worker_state();
16531 let lane = lanes::Lane::Interactive;
16532 let cap = lane_cap(lane);
16533 {
16534 let mut m = st.metrics.lock().unwrap();
16535 m.completed = 10;
16536 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
16539 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16540 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
16542 let free = RateLimit {
16545 limit: cap,
16546 remaining: 1,
16547 reset_s: 0,
16548 };
16549 let g = reserve_pending_admit(
16550 &st,
16551 lane,
16552 &free,
16553 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16554 );
16555 assert!(
16556 g.is_ok(),
16557 "flag off: a ~20 s projected wait whose deadline can absorb it queues \
16558 silently (no 429) - the darklanes#5 defect shape, preserved by default"
16559 );
16560 drop(g);
16561 }
16562
16563 #[test]
16568 fn the_queue_wait_ceiling_sheds_with_429_retry_after_and_the_ratelimit_trio() {
16569 let _counters = admission_counters_guard();
16570 let st = fake_worker_state();
16571 let lane = lanes::Lane::Interactive;
16572 let cap = lane_cap(lane);
16573 {
16574 let mut m = st.metrics.lock().unwrap();
16575 m.completed = 10;
16576 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
16579 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16580 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
16582 let free = RateLimit {
16583 limit: cap,
16584 remaining: 1,
16585 reset_s: 0,
16586 };
16587 let (resp, outcome) = reserve_pending_admit_with_ceiling(
16588 &st,
16589 lane,
16590 &free,
16591 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16592 5, )
16594 .map(|_| ())
16595 .expect_err("a projected wait past the ceiling must shed");
16596 assert_eq!(outcome, "shed_queue_wait");
16597 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16598 assert_eq!(
16599 retry_after(&resp).as_deref(),
16600 Some("20"),
16601 "Retry-After must carry the estimate (~10 s/wave x 2 waves)"
16602 );
16603 assert_eq!(
16604 resp.headers()
16605 .get("retry-after-ms")
16606 .and_then(|v| v.to_str().ok()),
16607 Some("20000"),
16608 "the ms twin must match"
16609 );
16610 let stamped = free.attach(resp);
16611 for h in [
16612 "x-ratelimit-limit",
16613 "x-ratelimit-remaining",
16614 "x-ratelimit-reset",
16615 ] {
16616 assert!(stamped.headers().get(h).is_some(), "missing {h}");
16617 }
16618 }
16619
16620 #[test]
16624 fn the_queue_wait_ceiling_admits_under_it_and_never_touches_dark_lanes() {
16625 let _counters = admission_counters_guard();
16626 let st = fake_worker_state();
16627 let lane = lanes::Lane::Interactive;
16628 let cap = lane_cap(lane);
16629 {
16630 let mut m = st.metrics.lock().unwrap();
16631 m.completed = 10;
16632 m.tokens_out = 1_000;
16633 m.step_p50_ms = 100.0; }
16635 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16636 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel);
16637 let _restore = CounterRestore(counter, prev);
16638 let free = RateLimit {
16639 limit: cap,
16640 remaining: 1,
16641 reset_s: 0,
16642 };
16643 let g = reserve_pending_admit_with_ceiling(
16644 &st,
16645 lane,
16646 &free,
16647 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16648 60, );
16650 assert!(
16651 g.is_ok(),
16652 "an estimate under the ceiling must admit and queue as before"
16653 );
16654 drop(g);
16655 let full = RateLimit {
16657 limit: cap,
16658 remaining: 0,
16659 reset_s: 5,
16660 };
16661 for dark in [lanes::Lane::Judge, lanes::Lane::Harvest] {
16662 let counter = &worker::ADMISSION_RESERVATIONS[dark.idx()];
16663 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
16665 let g = reserve_pending_admit_with_ceiling(
16666 &st,
16667 dark,
16668 &full,
16669 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16670 1,
16671 );
16672 assert!(
16673 g.is_ok(),
16674 "{dark:?} must not be shed by the interactive queue-wait ceiling"
16675 );
16676 drop(g);
16677 }
16678 }
16679
16680 #[test]
16684 fn the_queue_wait_ceiling_leaves_the_existing_shed_arms_first_and_unchanged() {
16685 let _counters = admission_counters_guard();
16686 let st = fake_worker_state();
16687 let lane = lanes::Lane::Interactive;
16688 let cap = lane_cap(lane);
16689 {
16690 let mut m = st.metrics.lock().unwrap();
16691 m.completed = 10;
16692 m.tokens_out = 1_000;
16693 m.step_p50_ms = 100.0;
16694 }
16695 let rl = RateLimit {
16696 limit: cap,
16697 remaining: 0,
16698 reset_s: 1,
16699 };
16700 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
16701 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
16703 let _restore = CounterRestore(counter, prev);
16704 assert!(matches!(
16705 reserve_pending_admit_with_ceiling(
16706 &st,
16707 lane,
16708 &rl,
16709 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16710 1,
16711 ),
16712 Err((_, "shed_queue"))
16713 ));
16714 counter.store(cap, std::sync::atomic::Ordering::Release);
16716 assert!(matches!(
16717 reserve_pending_admit_with_ceiling(
16718 &st,
16719 lane,
16720 &rl,
16721 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
16722 1,
16723 ),
16724 Err((_, "shed_deadline"))
16725 ));
16726 }
16727
16728 #[test]
16733 fn the_queue_wait_ceiling_is_wired_through_the_production_wrapper() {
16734 let src = include_str!("lib.rs");
16735 let code: String = src
16736 .lines()
16737 .map(|l| l.split("//").next().unwrap_or(""))
16738 .collect::<Vec<_>>()
16739 .join("\n");
16740 let start = code
16741 .find("pub(crate) fn reserve_pending_admit(")
16742 .expect("the production wrapper exists");
16743 let rest = &code[start..];
16744 let end = rest.find("\nfn ").unwrap_or(rest.len());
16745 let wrapper = &rest[..end];
16746 assert!(
16747 wrapper.contains(
16748 "reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())"
16749 ),
16750 "every production ingress must judge the ceiling the env read armed"
16751 );
16752 }
16753
16754 #[test]
16755 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
16756 let _counters = admission_counters_guard();
16757 let st = fake_worker_state();
16758 let cap = lane_cap(lanes::Lane::Interactive);
16759 let bound = max_queue_depth(cap);
16760 assert!(bound > 0, "the queue bound must admit at least one request");
16761 let rl = RateLimit {
16762 limit: cap,
16763 remaining: 0,
16764 reset_s: 1,
16765 };
16766 let _ = worker::PENDING_ADMITS.fetch_update(
16767 std::sync::atomic::Ordering::AcqRel,
16768 std::sync::atomic::Ordering::Acquire,
16769 |_| Some(0),
16770 );
16771 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
16772 let _restore = CounterRestore(counter, 0);
16773 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
16774 let guard = reserve_pending_admit(
16775 &st,
16776 lanes::Lane::Interactive,
16777 &rl,
16778 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16779 )
16780 .expect("the final queue slot should be reservable");
16781 assert_eq!(
16782 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
16783 1
16784 );
16785 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
16786 drop(guard);
16787 assert_eq!(
16788 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
16789 0
16790 );
16791 assert_eq!(
16792 counter.load(std::sync::atomic::Ordering::Acquire),
16793 bound - 1
16794 );
16795
16796 counter.store(bound, std::sync::atomic::Ordering::Release);
16797 let rejected = reserve_pending_admit(
16798 &st,
16799 lanes::Lane::Interactive,
16800 &rl,
16801 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16802 );
16803 assert!(matches!(rejected, Err((_, "shed_queue"))));
16804 }
16805
16806 #[test]
16807 fn admission_reservations_are_lane_scoped() {
16808 let _counters = admission_counters_guard();
16809 let st = fake_worker_state();
16810 let harvest = lanes::Lane::Harvest;
16811 let interactive = lanes::Lane::Interactive;
16812 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
16813 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
16814 let _restore = CounterRestore(harvest_counter, 0);
16815 harvest_counter.store(
16816 max_queue_depth(lane_cap(harvest)),
16817 std::sync::atomic::Ordering::Release,
16818 );
16819 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
16820 let free = RateLimit {
16821 limit: lane_cap(interactive),
16822 remaining: 1,
16823 reset_s: 0,
16824 };
16825 let guard = reserve_pending_admit(
16836 &st,
16837 interactive,
16838 &free,
16839 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
16840 )
16841 .expect("a full harvest queue must not consume interactive capacity");
16842 drop(guard);
16843 let tight = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
16844 assert!(
16845 tight.is_ok(),
16846 "a full harvest queue must not deadline-shed a tight interactive request \
16847 (a backlog that outlasts the retry budget here is a cross-lane leak, not \
16848 contention)"
16849 );
16850 drop(tight);
16851 let harvest_rl = RateLimit {
16852 limit: lane_cap(harvest),
16853 remaining: 0,
16854 reset_s: 1,
16855 };
16856 assert!(matches!(
16857 reserve_pending_admit(
16858 &st,
16859 harvest,
16860 &harvest_rl,
16861 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
16862 ),
16863 Err((_, "shed_queue"))
16864 ));
16865 }
16866
16867 #[test]
16868 fn taxonomy_maps_every_class_to_its_status_and_code() {
16869 use worker::{EngineError as E, ErrClass as C};
16870 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
16871 (
16872 E::invalid_param("bad json", "response_format"),
16873 StatusCode::BAD_REQUEST,
16874 "invalid_request_error",
16875 "",
16876 ),
16877 (
16878 E::context_length("prompt (9000 tok) >= context cap (8192)"),
16879 StatusCode::BAD_REQUEST,
16880 "invalid_request_error",
16881 "context_length_exceeded",
16882 ),
16883 (
16884 E::model_not_found("unknown model \"nope\""),
16885 StatusCode::BAD_REQUEST,
16886 "invalid_request_error",
16887 "model_not_found",
16888 ),
16889 (
16890 E::rate_limit("lane judge is at capacity, retry"),
16891 StatusCode::TOO_MANY_REQUESTS,
16892 "rate_limit_error",
16893 "rate_limit_exceeded",
16894 ),
16895 (
16896 E::overloaded("no VRAM for a new session"),
16897 StatusCode::SERVICE_UNAVAILABLE,
16898 "server_error",
16899 "overloaded",
16900 ),
16901 (
16902 E::engine("graph step failed: launch error"),
16903 StatusCode::INTERNAL_SERVER_ERROR,
16904 "server_error",
16905 "engine_error",
16906 ),
16907 ];
16908 for (err, want_status, want_type, want_code) in cases {
16909 let (status, etype, code) = class_http(err.class);
16910 assert_eq!(status, want_status, "{:?}", err);
16911 assert_eq!(etype, want_type, "{:?}", err);
16912 if !want_code.is_empty() {
16913 assert_eq!(code, Some(want_code), "{:?}", err);
16914 }
16915 let body = engine_error_body(&err);
16917 assert_eq!(body["error"]["message"], err.message);
16918 assert_eq!(body["error"]["type"], want_type);
16919 }
16920 for c in [
16922 C::InvalidRequest,
16923 C::ContextLength,
16924 C::ModelNotFound,
16925 C::RateLimit,
16926 C::Overloaded,
16927 C::Engine,
16928 ] {
16929 let (s, t, _) = class_http(c);
16930 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
16931 assert!(!t.is_empty());
16932 }
16933 }
16934
16935 #[test]
16936 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
16937 let e = worker::EngineError::engine(
16942 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
16943 );
16944 let resp = engine_error_response(&e);
16945 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16946 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
16947 }
16948
16949 #[test]
16950 fn retry_headers_follow_the_sdk_contract() {
16951 for e in [
16955 worker::EngineError::rate_limit("shed"),
16956 worker::EngineError::overloaded("no VRAM"),
16957 ] {
16958 let resp = engine_error_response(&e);
16959 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
16960 let secs: u64 = ra
16961 .parse()
16962 .expect("Retry-After must be integer delay-seconds");
16963 assert!(
16964 secs > 0 && secs <= 60,
16965 "Retry-After {secs}s outside the honored window"
16966 );
16967 let ms = resp
16968 .headers()
16969 .get("retry-after-ms")
16970 .unwrap()
16971 .to_str()
16972 .unwrap();
16973 assert_eq!(
16974 ms.parse::<u64>().unwrap(),
16975 secs * 1000,
16976 "the two headers disagree"
16977 );
16978 assert!(
16979 resp.headers().get("x-should-retry").is_none(),
16980 "a retryable class must not say x-should-retry: false"
16981 );
16982 }
16983 }
16984
16985 #[tokio::test]
16992 async fn admit_predict_reject_matches_shed_contract() {
16993 let shed = retry_contract_response(
16995 (
16996 StatusCode::TOO_MANY_REQUESTS,
16997 Json(error_body(
16998 "interactive queue is at its bound",
16999 "rate_limit_error",
17000 None,
17001 Some("shed_queue"),
17002 )),
17003 )
17004 .into_response(),
17005 Some(7),
17006 );
17007 let predict = engine_error_response(&worker::EngineError::rate_limit_after(
17010 "predicted KV-to-completion exceeds the box budget; retry",
17011 7,
17012 ));
17013 assert_eq!(shed.status(), predict.status());
17014 for header in ["retry-after", "retry-after-ms"] {
17015 assert_eq!(
17016 shed.headers().get(header),
17017 predict.headers().get(header),
17018 "header {header} must be byte-identical to the shed contract"
17019 );
17020 }
17021 let shed_body: serde_json::Value = serde_json::from_slice(
17022 &axum::body::to_bytes(shed.into_body(), usize::MAX)
17023 .await
17024 .unwrap(),
17025 )
17026 .unwrap();
17027 let predict_body: serde_json::Value = serde_json::from_slice(
17028 &axum::body::to_bytes(predict.into_body(), usize::MAX)
17029 .await
17030 .unwrap(),
17031 )
17032 .unwrap();
17033 assert_eq!(shed_body["error"]["type"], predict_body["error"]["type"]);
17034 assert_eq!(predict_body["error"]["type"], "rate_limit_error");
17035 let shed_keys: Vec<&String> = shed_body["error"].as_object().unwrap().keys().collect();
17036 let predict_keys: Vec<&String> =
17037 predict_body["error"].as_object().unwrap().keys().collect();
17038 assert_eq!(shed_keys, predict_keys, "same body schema, key for key");
17039 assert_eq!(predict_body["error"]["code"], "rate_limit_exceeded");
17040
17041 let clamped = engine_error_response(&worker::EngineError::rate_limit_after("m", 400));
17043 assert_eq!(retry_after(&clamped).as_deref(), Some("60"));
17044 let plain = engine_error_response(&worker::EngineError::rate_limit("m"));
17046 assert_eq!(retry_after(&plain).as_deref(), Some("2"));
17047 }
17048
17049 #[tokio::test]
17050 #[allow(clippy::await_holding_lock)] async fn command_send_failure_obeys_the_retry_contract() {
17052 let _l = drain_lock();
17053 let mut st = fake_worker_state();
17054 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
17055 drop(cmd_rx);
17056 st.cmd_tx = cmd_tx;
17057
17058 let completion = completions(
17059 State(st.clone()),
17060 axum::http::HeaderMap::new(),
17061 None,
17062 Json(
17063 serde_json::from_value(serde_json::json!({
17064 "model": "m", "prompt": "test"
17065 }))
17066 .unwrap(),
17067 ),
17068 )
17069 .await;
17070 let chat = chat_completions(
17071 State(st),
17072 axum::http::HeaderMap::new(),
17073 None,
17074 Json(
17075 serde_json::from_value(serde_json::json!({
17076 "model": "m", "messages": [{"role": "user", "content": "test"}]
17077 }))
17078 .unwrap(),
17079 ),
17080 )
17081 .await;
17082
17083 for resp in [completion, chat] {
17084 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17085 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
17086 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
17087 assert_ne!(
17088 resp.headers()
17089 .get("x-should-retry")
17090 .and_then(|v| v.to_str().ok()),
17091 Some("false")
17092 );
17093 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17094 .await
17095 .unwrap();
17096 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17097 assert_eq!(payload["error"]["type"], "server_error");
17098 assert_eq!(payload["error"]["code"], "overloaded");
17099 }
17100 }
17101
17102 #[test]
17103 fn unfixable_client_errors_say_x_should_retry_false() {
17104 for e in [
17107 worker::EngineError::model_not_found("unknown model \"x\""),
17108 worker::EngineError::context_length("prompt too long"),
17109 worker::EngineError::invalid_param("bad", "messages"),
17110 ] {
17111 let resp = engine_error_response(&e);
17112 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
17113 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
17114 assert!(
17115 retry_after(&resp).is_none(),
17116 "a 400 must not promise a retry window"
17117 );
17118 }
17119 }
17120
17121 #[tokio::test]
17122 async fn a_closed_worker_channel_is_503_not_500() {
17123 let (tx, rx) = worker::event_channel();
17127 drop(tx);
17128 let resp =
17129 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
17130 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17131 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
17132 }
17133
17134 #[tokio::test]
17135 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
17136 let (tx, rx) = worker::event_channel();
17139 tx.send(Event::Error(worker::EngineError::rate_limit(
17140 "lane judge shed: interactive p99 over budget, retry",
17141 )))
17142 .unwrap();
17143 let (resp, error_code) = peek_admission(rx)
17144 .await
17145 .expect_err("a shed must not be forwarded into the stream");
17146 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
17147 assert_eq!(error_code, "rate_limit_exceeded");
17148 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
17149 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17150 .await
17151 .unwrap();
17152 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17153 assert!(
17154 payload["error"].is_object(),
17155 "bare-string error body: {payload}"
17156 );
17157 assert_eq!(payload["error"]["type"], "rate_limit_error");
17158 assert!(
17159 payload["error"]["message"]
17160 .as_str()
17161 .unwrap()
17162 .contains("shed")
17163 );
17164 }
17165
17166 #[tokio::test]
17167 async fn interactive_admission_error_is_a_preheader_429() {
17168 let (tx, rx) = worker::event_channel();
17171 tx.send(Event::Error(worker::EngineError::rate_limit(
17172 "KV capacity unavailable",
17173 )))
17174 .unwrap();
17175 let (resp, error_code) = peek_admission(rx)
17176 .await
17177 .expect_err("admission error must stay pre-header");
17178 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
17179 assert_eq!(error_code, "rate_limit_exceeded");
17180 }
17181
17182 #[tokio::test]
17183 async fn admission_peek_preserves_context_error_for_the_ledger() {
17184 let (tx, rx) = worker::event_channel();
17185 tx.send(Event::Error(worker::EngineError::context_length(
17186 "prompt exceeds configured model maximum",
17187 )))
17188 .unwrap();
17189 let (resp, error_code) = peek_admission(rx)
17190 .await
17191 .expect_err("context rejection must stay pre-header");
17192 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
17193 assert_eq!(error_code, "context_length_exceeded");
17194 }
17195
17196 #[tokio::test]
17197 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
17198 let (tx, rx) = worker::event_channel();
17199 tx.send(Event::PromptUsage {
17200 n_prompt: 262_143,
17201 n_cached: 0,
17202 })
17203 .unwrap();
17204 let mut replay = peek_admission(rx).await.expect("successful admission");
17205 assert!(matches!(
17206 replay.recv().await,
17207 Some(Event::PromptUsage {
17208 n_prompt: 262_143,
17209 n_cached: 0
17210 }),
17211 ));
17212 }
17213
17214 #[test]
17215 fn penalties_plumb_from_http_to_sampler_config() {
17216 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17219 "model": "m", "messages": [{"role": "user", "content": "task"}],
17220 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
17221 }))
17222 .unwrap();
17223 let (tx, _rx) = worker::event_channel();
17224 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17225 .unwrap()
17226 .request
17227 .sampler_cfg;
17228 assert_eq!(cfg.penalty_freq, 0.5);
17229 assert_eq!(cfg.penalty_present, 0.25);
17230 assert_eq!(cfg.penalty_repeat, 1.1);
17231 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
17232
17233 let req: CompletionReq = serde_json::from_value(serde_json::json!({
17234 "model": "m", "prompt": "task", "frequency_penalty": 1.5
17235 }))
17236 .unwrap();
17237 let (tx, _rx) = worker::event_channel();
17238 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
17239 assert_eq!(cfg.penalty_freq, 1.5);
17240 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
17241
17242 let req: CompletionReq = serde_json::from_value(serde_json::json!({
17244 "model": "m", "prompt": "task"
17245 }))
17246 .unwrap();
17247 let (tx, _rx) = worker::event_channel();
17248 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
17249 assert_eq!(cfg.penalty_last_n, 0);
17250 assert_eq!(cfg.penalty_repeat, 1.0);
17251 }
17252
17253 #[test]
17254 fn omitted_temperature_is_openai_default_not_greedy() {
17255 let chat_temp = |body: serde_json::Value| {
17271 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17272 let (tx, _rx) = worker::event_channel();
17273 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17274 .unwrap()
17275 .request
17276 .sampler_cfg
17277 .temperature
17278 };
17279 let comp_temp = |body: serde_json::Value| {
17280 let req: CompletionReq = serde_json::from_value(body).unwrap();
17281 let (tx, _rx) = worker::event_channel();
17282 build_request(&req, tx, lanes::Lane::Interactive, None)
17283 .sampler_cfg
17284 .temperature
17285 };
17286
17287 assert_eq!(
17289 chat_temp(serde_json::json!({
17290 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
17291 1.0,
17292 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
17293 );
17294 assert_eq!(
17295 comp_temp(serde_json::json!({
17296 "model": "m", "prompt": "t"})),
17297 1.0,
17298 "omitted completions temperature must be the OpenAI 1.0 default"
17299 );
17300
17301 assert_eq!(
17303 chat_temp(serde_json::json!({
17304 "model": "m", "messages": [{"role": "user", "content": "t"}],
17305 "temperature": 0.0})),
17306 0.0,
17307 "explicit temperature 0 must stay greedy"
17308 );
17309 assert_eq!(
17310 comp_temp(serde_json::json!({
17311 "model": "m", "prompt": "t", "temperature": 0})),
17312 0.0,
17313 "explicit temperature 0 must stay greedy"
17314 );
17315 assert!(
17317 memra_engine::sampler::Sampler::new(sampler_config(
17318 0.0,
17319 0,
17320 1.0,
17321 0.0,
17322 0.0,
17323 0.0,
17324 1.0,
17325 Some(0)
17326 ))
17327 .is_greedy()
17328 );
17329 assert!(
17330 !memra_engine::sampler::Sampler::new(sampler_config(
17331 1.0,
17332 0,
17333 1.0,
17334 0.0,
17335 0.0,
17336 0.0,
17337 1.0,
17338 Some(0)
17339 ))
17340 .is_greedy()
17341 );
17342
17343 assert_eq!(
17345 chat_temp(serde_json::json!({
17346 "model": "m", "messages": [{"role": "user", "content": "t"}],
17347 "temperature": 0.7})),
17348 0.7
17349 );
17350
17351 let req: CompletionReq = serde_json::from_value(serde_json::json!({
17355 "model": "m", "prompt": "t"}))
17356 .unwrap();
17357 let (tx, _rx) = worker::event_channel();
17358 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
17359 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
17360 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
17361 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
17362 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
17363 assert!(
17368 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
17369 "the omitted-temperature default must ride sampled spec's pure-temp regime"
17370 );
17371 }
17372
17373 #[test]
17374 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
17375 let caps = ModelCaps {
17376 chat_temperature_default: Some(0.5),
17377 chat_top_p_default: Some(0.9),
17378 chat_ok: true,
17379 ..Default::default()
17380 };
17381 let cfg = |extra: serde_json::Value| {
17382 let mut body = serde_json::json!({
17383 "model": "step35",
17384 "messages": [{"role": "user", "content": "task"}]
17385 });
17386 body.as_object_mut()
17387 .unwrap()
17388 .extend(extra.as_object().unwrap().clone());
17389 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17390 let (tx, _rx) = worker::event_channel();
17391 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
17392 .unwrap()
17393 .request
17394 .sampler_cfg
17395 };
17396
17397 let omitted = cfg(serde_json::json!({}));
17398 assert_eq!(omitted.temperature, 0.5);
17399 assert_eq!(omitted.top_p, 0.9);
17400
17401 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
17402 assert_eq!(explicit_temp.temperature, 0.7);
17403 assert_eq!(
17404 explicit_temp.top_p, 0.9,
17405 "omitting top_p must retain StepFun's nucleus default"
17406 );
17407
17408 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
17409 assert_eq!(
17410 explicit.temperature, 0.0,
17411 "explicit greedy must remain authoritative"
17412 );
17413 assert_eq!(
17414 explicit.top_p, 1.0,
17415 "explicit untruncated sampling must remain authoritative"
17416 );
17417 }
17418
17419 fn qwen38_vendor_defaults() -> SamplingDefaults {
17423 SamplingDefaults {
17424 temperature: Some(1.0),
17425 top_p: Some(0.95),
17426 top_k: Some(20),
17427 min_p: Some(0.0),
17428 presence_penalty: Some(0.0),
17429 repetition_penalty: Some(1.0),
17430 frequency_penalty: None,
17431 }
17432 }
17433
17434 fn gemma4_vendor_defaults() -> SamplingDefaults {
17439 SamplingDefaults {
17440 temperature: Some(1.0),
17441 top_p: Some(0.95),
17442 top_k: Some(64),
17443 ..Default::default()
17444 }
17445 }
17446
17447 #[test]
17448 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
17449 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
17454 let chat = |extra: serde_json::Value| {
17455 let mut body = serde_json::json!({
17456 "model": "google/gemma-4-31b-it",
17457 "messages": [{"role": "user", "content": "task"}],
17458 "seed": 7
17460 });
17461 body.as_object_mut()
17462 .unwrap()
17463 .extend(extra.as_object().unwrap().clone());
17464 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17465 let (tx, _rx) = worker::event_channel();
17466 build_chat_request_with_trace(
17467 req,
17468 Some(&ModelCaps {
17469 chat_ok: true,
17470 ..Default::default()
17471 }),
17472 tx,
17473 lanes::Lane::Interactive,
17474 None,
17475 None,
17476 None,
17477 &d,
17478 )
17479 .unwrap()
17480 .request
17481 .sampler_cfg
17482 };
17483
17484 let omitted = chat(serde_json::json!({}));
17486 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
17487 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
17488 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
17489 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
17491 assert_eq!(omitted.penalty_repeat, 1.0);
17492 assert_eq!(omitted.penalty_freq, 0.0);
17493 assert_eq!(omitted.penalty_present, 0.0);
17494 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
17495 assert!(
17496 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
17497 "the vendor default must NOT be greedy — that is the whole point of the lane"
17498 );
17499
17500 let greedy = chat(serde_json::json!({"temperature": 0}));
17503 assert_eq!(
17504 greedy.temperature, 0.0,
17505 "explicit temperature 0 stays greedy"
17506 );
17507 assert!(
17508 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
17509 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
17510 spec/graph exactness arms"
17511 );
17512
17513 let one_field = chat(serde_json::json!({"top_k": 3}));
17515 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
17516 assert_eq!(
17517 one_field.temperature, 1.0,
17518 "omitting temperature still takes the vendor value"
17519 );
17520 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
17521
17522 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
17525 assert_eq!(
17526 disabled.top_k, 0,
17527 "an explicit top_k 0 means KEEP ALL, not 'unset'"
17528 );
17529 assert_eq!(
17530 disabled.top_p, 1.0,
17531 "an explicit top_p 1.0 means untruncated"
17532 );
17533
17534 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
17536 assert_eq!(penal.penalty_present, 1.5);
17537 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
17538 }
17539
17540 #[test]
17541 fn vendor_sampling_defaults_are_identical_on_every_surface() {
17542 let d = qwen38_vendor_defaults();
17554 let md = ModelSamplingDefaults::single(d);
17555 let comp = |extra: serde_json::Value| {
17556 let mut body = serde_json::json!({
17557 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
17558 body.as_object_mut()
17559 .unwrap()
17560 .extend(extra.as_object().unwrap().clone());
17561 let req: CompletionReq = serde_json::from_value(body).unwrap();
17562 let (tx, _rx) = worker::event_channel();
17563 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
17564 };
17565 let chat = |extra: serde_json::Value| {
17566 let mut body = serde_json::json!({
17567 "model": "qwen/qwen3.8-27b",
17568 "messages": [{"role": "user", "content": "task"}],
17569 "seed": 11 });
17570 body.as_object_mut()
17571 .unwrap()
17572 .extend(extra.as_object().unwrap().clone());
17573 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17574 let (tx, _rx) = worker::event_channel();
17575 build_chat_request_with_trace(
17576 req,
17577 Some(&ModelCaps {
17578 chat_ok: true,
17579 ..Default::default()
17580 }),
17581 tx,
17582 lanes::Lane::Interactive,
17583 None,
17584 None,
17585 None,
17586 &md,
17587 )
17588 .unwrap()
17589 .request
17590 .sampler_cfg
17591 };
17592
17593 for extra in [
17594 serde_json::json!({}),
17595 serde_json::json!({"temperature": 0}),
17596 serde_json::json!({"temperature": 0.0}),
17597 serde_json::json!({"temperature": 0.7}),
17598 serde_json::json!({"top_p": 1.0}),
17599 serde_json::json!({"top_k": 0}),
17600 serde_json::json!({"min_p": 0.05}),
17601 serde_json::json!({"repetition_penalty": 1.1}),
17602 serde_json::json!({"frequency_penalty": 0.5}),
17603 serde_json::json!({"presence_penalty": 1.5}),
17604 serde_json::json!({
17605 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
17606 "frequency_penalty": 0.1, "presence_penalty": 0.2,
17607 "repetition_penalty": 1.05 }),
17608 ] {
17609 let c = comp(extra.clone());
17610 let h = chat(extra.clone());
17611 assert_eq!(
17612 (
17613 c.temperature,
17614 c.top_p,
17615 c.top_k,
17616 c.min_p,
17617 c.penalty_repeat,
17618 c.penalty_freq,
17619 c.penalty_present,
17620 c.penalty_last_n,
17621 c.seed
17622 ),
17623 (
17624 h.temperature,
17625 h.top_p,
17626 h.top_k,
17627 h.min_p,
17628 h.penalty_repeat,
17629 h.penalty_freq,
17630 h.penalty_present,
17631 h.penalty_last_n,
17632 h.seed
17633 ),
17634 "/v1/completions and /v1/chat/completions disagree on {extra} — \
17635 standard-surface-law violation"
17636 );
17637 }
17638
17639 let omitted = comp(serde_json::json!({}));
17641 assert_eq!(
17642 omitted.temperature, 1.0,
17643 "qwen3.8 card thinking temperature"
17644 );
17645 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
17646 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
17647 assert!(
17649 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
17650 .is_greedy()
17651 );
17652 }
17653
17654 #[tokio::test]
17667 #[allow(clippy::await_holding_lock)] async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
17669 let _l = drain_lock();
17670 let step_caps = ModelCaps {
17671 chat_ok: true,
17672 chat_temperature_default: Some(0.5),
17673 chat_top_p_default: Some(0.9),
17674 ..Default::default()
17675 };
17676 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
17677 let st = fake_worker_state_full(
17678 1,
17679 std::time::Duration::ZERO,
17680 HashMap::from([("m".to_string(), step_caps)]),
17681 Some(cfg_tx),
17682 );
17683 let fields = |saw: &WorkerSaw| {
17688 let c = &saw.sampler_cfg;
17689 (
17690 c.temperature,
17691 c.top_p,
17692 c.top_k,
17693 c.min_p,
17694 c.penalty_repeat,
17695 c.penalty_freq,
17696 c.penalty_present,
17697 c.penalty_last_n,
17698 )
17699 };
17700 let worker_saw = |surface: &str| {
17701 cfg_rx
17702 .recv_timeout(std::time::Duration::from_secs(10))
17703 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
17704 };
17705
17706 let resp = completions(
17707 State(st.clone()),
17708 axum::http::HeaderMap::new(),
17709 None,
17710 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
17711 )
17712 .await;
17713 assert_eq!(
17714 resp.status(),
17715 StatusCode::OK,
17716 "/v1/completions rejected the omitted-sampling request"
17717 );
17718 let comp = worker_saw("/v1/completions");
17719
17720 let resp = chat_completions(
17721 State(st.clone()),
17722 axum::http::HeaderMap::new(),
17723 None,
17724 Json(
17725 serde_json::from_value(serde_json::json!({
17726 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
17727 .unwrap(),
17728 ),
17729 )
17730 .await;
17731 assert_eq!(
17732 resp.status(),
17733 StatusCode::OK,
17734 "/v1/chat/completions rejected the omitted-sampling request"
17735 );
17736 let chat = worker_saw("/v1/chat/completions");
17737
17738 let resp = anthropic::messages(
17739 State(st.clone()),
17740 axum::http::HeaderMap::new(),
17741 None,
17742 axum::body::Bytes::from(
17743 serde_json::json!({
17744 "model": "m", "max_tokens": 16,
17745 "messages": [{"role": "user", "content": "t"}]})
17746 .to_string(),
17747 ),
17748 )
17749 .await;
17750 assert_eq!(
17751 resp.status(),
17752 StatusCode::OK,
17753 "/v1/messages rejected the omitted-sampling request"
17754 );
17755 let msg = worker_saw("/v1/messages");
17756
17757 let resp = responses_api::responses(
17758 State(st.clone()),
17759 axum::http::HeaderMap::new(),
17760 None,
17761 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
17762 )
17763 .await;
17764 assert_eq!(
17765 resp.status(),
17766 StatusCode::OK,
17767 "/v1/responses rejected the omitted-sampling request"
17768 );
17769 let rsp = worker_saw("/v1/responses");
17770
17771 for (surface, cfg) in [
17772 ("/v1/completions", &comp),
17773 ("/v1/messages", &msg),
17774 ("/v1/responses", &rsp),
17775 ] {
17776 assert_eq!(
17777 fields(cfg),
17778 fields(&chat),
17779 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
17780 for the same omitted-sampling request — standard-surface-law violation \
17781 (hermes d991b51699218285)"
17782 );
17783 }
17784 assert_eq!(
17787 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
17788 (0.5, 0.9),
17789 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
17790 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
17791 );
17792 }
17793
17794 #[tokio::test]
17804 #[allow(clippy::await_holding_lock)] async fn same_effort_value_resolves_identically_on_every_surface() {
17806 let _l = drain_lock();
17807 let caps = ModelCaps {
17810 chat_ok: true,
17811 effort_levels: true,
17812 ..Default::default()
17813 };
17814 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
17815 let st = fake_worker_state_full(
17816 1,
17817 std::time::Duration::ZERO,
17818 HashMap::from([("m".to_string(), caps)]),
17819 Some(saw_tx),
17820 );
17821 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
17822 match surface {
17823 "/v1/chat/completions" => {
17824 chat_completions(
17825 State(st),
17826 axum::http::HeaderMap::new(),
17827 None,
17828 Json(
17829 serde_json::from_value(serde_json::json!({
17830 "model": "m", "max_tokens": 8,
17831 "reasoning_effort": effort,
17832 "messages": [{"role": "user", "content": "t"}]}))
17833 .unwrap(),
17834 ),
17835 )
17836 .await
17837 }
17838 "/v1/responses" => {
17839 responses_api::responses(
17840 State(st),
17841 axum::http::HeaderMap::new(),
17842 None,
17843 axum::body::Bytes::from(
17844 serde_json::json!({
17845 "model": "m", "max_output_tokens": 8, "input": "t",
17846 "reasoning": {"effort": effort}})
17847 .to_string(),
17848 ),
17849 )
17850 .await
17851 }
17852 "/v1/messages" => {
17853 anthropic::messages(
17854 State(st),
17855 axum::http::HeaderMap::new(),
17856 None,
17857 axum::body::Bytes::from(
17858 serde_json::json!({
17859 "model": "m", "max_tokens": 8,
17860 "messages": [{"role": "user", "content": "t"}],
17861 "output_config": {"effort": effort}})
17862 .to_string(),
17863 ),
17864 )
17865 .await
17866 }
17867 other => panic!("unknown surface {other}"),
17868 }
17869 };
17870 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
17871
17872 for (effort, want_think, want_level) in [
17875 ("none", ThinkMode::NoThink, Some("low")),
17876 ("minimal", ThinkMode::NoThink, Some("low")),
17877 ("low", ThinkMode::Think, Some("low")),
17878 ("medium", ThinkMode::Think, Some("medium")),
17879 ("high", ThinkMode::Think, Some("high")),
17880 ("xhigh", ThinkMode::Think, Some("high")),
17882 ] {
17883 for surface in SURFACES {
17884 let resp = send(st.clone(), surface, effort).await;
17885 assert_eq!(
17886 resp.status(),
17887 StatusCode::OK,
17888 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
17889 diverged again (issue #31)"
17890 );
17891 let saw = saw_rx
17892 .recv_timeout(std::time::Duration::from_secs(10))
17893 .unwrap_or_else(|_| {
17894 panic!("{surface}: effort {effort:?} request never reached the worker")
17895 });
17896 assert_eq!(
17897 (saw.think, saw.reasoning_effort.as_deref()),
17898 (want_think, want_level),
17899 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
17900 reasoning surface — the parameter was dropped or remapped before \
17901 parse_think (issue #31 regression)"
17902 );
17903 }
17904 }
17905
17906 for effort in ["bogus", "banana", ""] {
17909 for surface in SURFACES {
17910 let resp = send(st.clone(), surface, effort).await;
17911 assert_eq!(
17912 resp.status(),
17913 StatusCode::BAD_REQUEST,
17914 "{surface} accepted effort {effort:?} — silent-accept regression \
17915 (issue #31: the value never reached parse_think's allowlist)"
17916 );
17917 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
17919 .await
17920 .unwrap();
17921 let v: serde_json::Value = serde_json::from_slice(&body)
17922 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
17923 match surface {
17924 "/v1/messages" => {
17925 assert_eq!(v["type"], "error", "{surface} error envelope");
17926 assert_eq!(
17927 v["error"]["type"], "invalid_request_error",
17928 "{surface} error type"
17929 );
17930 }
17931 _ => {
17932 assert!(
17933 v["error"]["message"].is_string(),
17934 "{surface} OpenAI-shaped error body: {v}"
17935 );
17936 }
17937 }
17938 }
17939 }
17940
17941 let resp = anthropic::messages(
17945 State(st.clone()),
17946 axum::http::HeaderMap::new(),
17947 None,
17948 axum::body::Bytes::from(
17949 serde_json::json!({
17950 "model": "m", "max_tokens": 8,
17951 "messages": [{"role": "user", "content": "t"}],
17952 "thinking": {"type": "enabled"},
17953 "output_config": {"effort": "none"}})
17954 .to_string(),
17955 ),
17956 )
17957 .await;
17958 assert_eq!(resp.status(), StatusCode::OK);
17959 let saw = saw_rx
17960 .recv_timeout(std::time::Duration::from_secs(10))
17961 .expect("thinking+effort request never reached the worker");
17962 assert_eq!(
17963 saw.think,
17964 ThinkMode::Think,
17965 "thinking.type (the documented Anthropic lever) must win the switch over \
17966 output_config.effort"
17967 );
17968 let resp = anthropic::messages(
17969 State(st.clone()),
17970 axum::http::HeaderMap::new(),
17971 None,
17972 axum::body::Bytes::from(
17973 serde_json::json!({
17974 "model": "m", "max_tokens": 8,
17975 "messages": [{"role": "user", "content": "t"}],
17976 "thinking": {"type": "enabled"},
17977 "output_config": {"effort": "banana"}})
17978 .to_string(),
17979 ),
17980 )
17981 .await;
17982 assert_eq!(
17983 resp.status(),
17984 StatusCode::BAD_REQUEST,
17985 "an invalid effort must 400 even next to an explicit thinking.type — \
17986 precedence must not re-open the silent-accept hole"
17987 );
17988 }
17989
17990 #[test]
17991 fn vendor_sampling_defaults_are_boot_validated() {
17992 let parsed = OpenRouterMetadataFile::from_toml(
17995 r#"
17996[models.g]
17997default_temperature = 1.0
17998default_top_p = 0.95
17999default_top_k = 64
18000default_min_p = 0.0
18001default_presence_penalty = 0.0
18002default_frequency_penalty = 0.0
18003default_repetition_penalty = 1.0
18004"#,
18005 )
18006 .unwrap();
18007 let g = parsed.get("g").unwrap();
18008 assert_eq!(g.default_temperature, Some(1.0));
18009 assert_eq!(g.default_top_p, Some(0.95));
18010 assert_eq!(g.default_top_k, Some(64));
18011
18012 let err = OpenRouterMetadataFile::from_toml(
18016 r#"
18017[models.g]
18018default_temperature = 0.0
18019"#,
18020 )
18021 .unwrap_err();
18022 assert!(err.contains("default_temperature"), "{err}");
18023 assert!(
18024 err.contains("greedy"),
18025 "the refusal must say WHY a zero default is refused: {err}"
18026 );
18027
18028 for bad in [
18029 "default_temperature = 2.5",
18030 "default_temperature = -1.0",
18031 "default_top_p = 0.0",
18032 "default_top_p = 1.5",
18033 "default_min_p = 1.0",
18034 "default_min_p = -0.1",
18035 "default_presence_penalty = 3.0",
18036 "default_frequency_penalty = -2.5",
18037 "default_repetition_penalty = 0.0",
18038 ] {
18039 let err =
18040 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
18041 let key = bad.split(' ').next().unwrap();
18042 assert!(err.contains(key), "{bad} must be refused by name: {err}");
18043 }
18044
18045 let err = OpenRouterMetadataFile::from_toml(
18049 r#"
18050[models.g]
18051default_temperture = 1.0
18052"#,
18053 )
18054 .unwrap_err();
18055 assert!(
18056 err.contains("unknown field"),
18057 "an unknown key must be fatal, which is what makes binary-first ordering \
18058 mandatory: {err}"
18059 );
18060 }
18061
18062 #[test]
18063 fn non_thinking_sampling_arm_is_boot_validated() {
18064 let parsed = OpenRouterMetadataFile::from_toml(
18068 r#"
18069[models.q]
18070default_temperature = 1.0
18071default_top_p = 0.95
18072default_top_k = 20
18073
18074[models.q.non_thinking_sampling]
18075temperature = 0.7
18076top_p = 0.8
18077top_k = 20
18078presence_penalty = 1.5
18079"#,
18080 )
18081 .unwrap();
18082 let arm = parsed
18083 .get("q")
18084 .unwrap()
18085 .non_thinking_sampling
18086 .as_ref()
18087 .unwrap();
18088 assert_eq!(arm.temperature, Some(0.7));
18089 assert_eq!(arm.top_p, Some(0.8));
18090 assert_eq!(arm.top_k, Some(20));
18091 assert_eq!(arm.presence_penalty, Some(1.5));
18092 assert_eq!(
18093 arm.min_p, None,
18094 "undeclared arm fields stay undeclared, never invented"
18095 );
18096
18097 let err = OpenRouterMetadataFile::from_toml(
18101 r#"
18102[models.q]
18103[models.q.non_thinking_sampling]
18104temperature = 0.0
18105"#,
18106 )
18107 .unwrap_err();
18108 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
18109 assert!(err.contains("greedy"), "{err}");
18110
18111 let err = OpenRouterMetadataFile::from_toml(
18114 r#"
18115[models.q]
18116[models.q.non_thinking_sampling]
18117"#,
18118 )
18119 .unwrap_err();
18120 assert!(err.contains("non_thinking_sampling"), "{err}");
18121 assert!(err.contains("declare"), "{err}");
18122
18123 for bad in [
18125 "temperature = 2.5",
18126 "top_p = 0.0",
18127 "top_p = 1.5",
18128 "min_p = 1.0",
18129 "presence_penalty = 3.0",
18130 "frequency_penalty = -2.5",
18131 "repetition_penalty = 0.0",
18132 ] {
18133 let err = OpenRouterMetadataFile::from_toml(&format!(
18134 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
18135 ))
18136 .unwrap_err();
18137 let key = bad.split(' ').next().unwrap();
18138 assert!(
18139 err.contains(&format!("non_thinking_sampling.{key}")),
18140 "the refusal for {bad:?} must name the nested key: {err}"
18141 );
18142 }
18143
18144 let err = OpenRouterMetadataFile::from_toml(
18148 r#"
18149[models.q]
18150[models.q.non_thinking_sampling]
18151temperture = 0.7
18152"#,
18153 )
18154 .unwrap_err();
18155 assert!(err.contains("unknown field"), "{err}");
18156 }
18157
18158 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
18163 SamplingDefaults {
18164 temperature: Some(0.7),
18165 top_p: Some(0.8),
18166 top_k: Some(20),
18167 presence_penalty: Some(1.5),
18168 ..Default::default()
18169 }
18170 }
18171
18172 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
18173 ModelSamplingDefaults {
18174 thinking: qwen38_vendor_defaults(),
18175 non_thinking: Some(qwen38_non_thinking_defaults()),
18176 }
18177 }
18178
18179 fn qwen38_caps() -> ModelCaps {
18183 ModelCaps {
18184 chat_ok: true,
18185 qwen_think: true,
18186 think_switch: true,
18187 ..Default::default()
18188 }
18189 }
18190
18191 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
18194 (
18195 c.temperature,
18196 c.top_p,
18197 c.top_k,
18198 c.min_p,
18199 c.penalty_present,
18200 c.penalty_freq,
18201 c.penalty_repeat,
18202 c.penalty_last_n,
18203 c.seed,
18204 )
18205 }
18206
18207 fn build_with_arms(
18208 defaults: &ModelSamplingDefaults,
18209 caps: &ModelCaps,
18210 default_effort: Option<&str>,
18211 extra: serde_json::Value,
18212 ) -> Request {
18213 let mut body = serde_json::json!({
18214 "model": "m",
18215 "messages": [{"role": "user", "content": "task"}],
18216 "seed": 3
18218 });
18219 body.as_object_mut()
18220 .unwrap()
18221 .extend(extra.as_object().unwrap().clone());
18222 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
18223 let (tx, _rx) = worker::event_channel();
18224 build_chat_request_with_trace(
18225 req,
18226 Some(caps),
18227 tx,
18228 lanes::Lane::Interactive,
18229 None,
18230 None,
18231 default_effort,
18232 defaults,
18233 )
18234 .unwrap()
18235 .request
18236 }
18237
18238 #[test]
18239 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
18240 let two_arm = qwen38_two_arm_defaults();
18245 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
18246 let caps = qwen38_caps();
18247
18248 let off_spellings = [
18250 serde_json::json!({"reasoning_effort": "none"}),
18251 serde_json::json!({"enable_thinking": false}),
18252 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
18253 serde_json::json!({"reasoning": {"enabled": false}}),
18254 ];
18255 for extra in &off_spellings {
18256 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
18257 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
18258 let c = &r.sampler_cfg;
18259 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
18260 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
18261 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
18262 assert_eq!(
18263 c.penalty_present, 1.5,
18264 "{extra}: non-thinking presence_penalty"
18265 );
18266 assert_eq!(
18267 c.penalty_last_n,
18268 memra_engine::spec::PEN_WINDOW_MAX,
18269 "{extra}: the arm's presence penalty uses the cross-path history window"
18270 );
18271 assert_eq!(
18272 c.min_p, 0.0,
18273 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
18274 );
18275
18276 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
18279 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
18280 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
18281 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
18282 assert_eq!(
18283 s.sampler_cfg.penalty_present, 0.0,
18284 "{extra}: single-arm model"
18285 );
18286 }
18287
18288 for extra in [
18291 serde_json::json!({}),
18292 serde_json::json!({"enable_thinking": true}),
18293 serde_json::json!({"reasoning_effort": "high"}),
18294 serde_json::json!({"reasoning": {"enabled": true}}),
18295 ] {
18296 for defaults in [&two_arm, &single_arm] {
18297 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
18298 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
18299 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
18300 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
18301 assert_eq!(
18302 c.penalty_present, 0.0,
18303 "{extra}: thinking arm has no presence"
18304 );
18305 }
18306 }
18307
18308 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
18311 assert_eq!(
18312 c.temperature, 0.7,
18313 "deployment-default off = non-thinking arm"
18314 );
18315 let c = build_with_arms(
18317 &two_arm,
18318 &caps,
18319 Some("none"),
18320 serde_json::json!({"enable_thinking": true}),
18321 )
18322 .sampler_cfg;
18323 assert_eq!(
18324 c.temperature, 1.0,
18325 "explicit ON beats the deployment default"
18326 );
18327
18328 let c = build_with_arms(
18330 &two_arm,
18331 &caps,
18332 None,
18333 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
18334 )
18335 .sampler_cfg;
18336 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
18337 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
18338 let c = build_with_arms(
18339 &two_arm,
18340 &caps,
18341 None,
18342 serde_json::json!({
18343 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
18344 )
18345 .sampler_cfg;
18346 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
18347 assert_eq!(
18348 c.penalty_present, 0.0,
18349 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
18350 is a value, not an absence"
18351 );
18352 assert_eq!(
18353 c.penalty_last_n, 0,
18354 "all penalties off => no history window"
18355 );
18356 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
18357
18358 let c = build_with_arms(
18361 &two_arm,
18362 &caps,
18363 None,
18364 serde_json::json!({"enable_thinking": false, "temperature": 0}),
18365 )
18366 .sampler_cfg;
18367 assert!(
18368 memra_engine::sampler::Sampler::new(c).is_greedy(),
18369 "explicit temperature 0 must stay greedy on the non-thinking arm"
18370 );
18371
18372 let c = build_with_arms(
18375 &single_arm,
18376 &caps,
18377 None,
18378 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
18379 )
18380 .sampler_cfg;
18381 assert_eq!(c.temperature, 0.55);
18382 assert_eq!(
18383 c.top_p, 0.95,
18384 "single-arm model: unset top_p takes its one arm"
18385 );
18386 }
18387
18388 #[test]
18389 fn sampling_arms_never_blend_field_by_field() {
18390 let parsed = OpenRouterMetadataFile::from_toml(
18395 r#"
18396[models.m]
18397default_temperature = 1.0
18398default_min_p = 0.05
18399
18400[models.m.non_thinking_sampling]
18401temperature = 0.6
18402"#,
18403 )
18404 .unwrap();
18405 let caps = ModelCaps {
18406 chat_temperature_default: Some(0.5),
18407 chat_top_p_default: Some(0.9),
18408 ..Default::default()
18409 };
18410 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
18411 let client = ClientSampling {
18412 seed: Some(1),
18413 ..Default::default()
18414 };
18415
18416 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
18417 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
18418 assert_eq!(
18419 off.min_p, 0.0,
18420 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
18421 );
18422 assert_eq!(
18423 off.top_p, 1.0,
18424 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
18425 );
18426
18427 for mode in [ThinkMode::Default, ThinkMode::Think] {
18429 let on = resolve_sampler_config(client, d.for_mode(mode));
18430 assert_eq!(on.temperature, 1.0);
18431 assert_eq!(on.min_p, 0.05);
18432 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
18433 }
18434 }
18435
18436 #[test]
18437 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
18438 let caps = qwen38_caps();
18447 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
18448 let two_arm = qwen38_two_arm_defaults();
18449
18450 let bodies = [
18451 serde_json::json!({}),
18452 serde_json::json!({"enable_thinking": true}),
18453 serde_json::json!({"reasoning_effort": "high"}),
18454 serde_json::json!({"reasoning_effort": "none"}),
18455 serde_json::json!({"enable_thinking": false}),
18456 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
18457 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
18458 serde_json::json!({"enable_thinking": false, "temperature": 0}),
18459 ];
18460 for extra in &bodies {
18461 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
18463 let mut client = ClientSampling {
18464 seed: Some(3),
18465 ..Default::default()
18466 };
18467 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
18468 client.temperature = Some(t as f32);
18469 }
18470 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
18471 client.top_p = Some(p as f32);
18472 }
18473 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
18474 assert_eq!(
18475 sampler_key(&r.sampler_cfg),
18476 sampler_key(&pre_arm),
18477 "{extra}: single-arm model diverged from the pre-arm resolution law"
18478 );
18479
18480 if r.think != ThinkMode::NoThink {
18483 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
18484 assert_eq!(t.think, r.think, "{extra}");
18485 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
18486 assert_eq!(
18487 sampler_key(&t.sampler_cfg),
18488 sampler_key(&r.sampler_cfg),
18489 "{extra}: a thinking-on request must not feel the non-thinking arm"
18490 );
18491 }
18492 }
18493 }
18494
18495 #[test]
18496 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
18497 let r = build_with_arms(
18503 &qwen38_two_arm_defaults(),
18504 &qwen38_caps(),
18505 None,
18506 serde_json::json!({"response_format": {"type": "json_object"}}),
18507 );
18508 assert_eq!(
18509 r.think,
18510 ThinkMode::NoThink,
18511 "constraint forces the switch off"
18512 );
18513 assert_eq!(
18514 r.sampler_cfg.temperature, 0.7,
18515 "and the arm follows the real mode"
18516 );
18517 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
18518 }
18519
18520 #[test]
18521 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
18522 let caps = ModelCaps {
18527 chat_temperature_default: Some(0.5),
18528 chat_top_p_default: Some(0.9),
18529 chat_ok: true,
18530 ..Default::default()
18531 };
18532 let metadata = OpenRouterModelMetadata {
18533 default_temperature: Some(1.0),
18534 default_top_p: Some(0.95),
18535 default_top_k: Some(64),
18536 ..Default::default()
18537 };
18538
18539 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
18540 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
18541 assert_eq!(caps_only.top_p, Some(0.9));
18542 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
18543
18544 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
18545 assert_eq!(
18546 both.temperature,
18547 Some(1.0),
18548 "metadata outranks the arch cap"
18549 );
18550 assert_eq!(both.top_p, Some(0.95));
18551 assert_eq!(both.top_k, Some(64));
18552
18553 let partial = SamplingDefaults::resolve(
18555 Some(&OpenRouterModelMetadata {
18556 default_temperature: Some(0.7),
18557 ..Default::default()
18558 }),
18559 Some(&caps),
18560 );
18561 assert_eq!(partial.temperature, Some(0.7));
18562 assert_eq!(
18563 partial.top_p,
18564 Some(0.9),
18565 "an undeclared metadata field must fall through to the cap, not to 1.0"
18566 );
18567
18568 assert_eq!(
18570 SamplingDefaults::resolve(None, None),
18571 SamplingDefaults::default()
18572 );
18573 }
18574
18575 #[test]
18576 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
18577 let resolved = |d: &SamplingDefaults| {
18590 resolve_sampler_config(
18591 ClientSampling {
18592 seed: Some(1),
18593 ..Default::default()
18594 },
18595 d,
18596 )
18597 };
18598
18599 assert!(
18601 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
18602 .is_spec_sampling(),
18603 "the API-standard default must stay in the fast pure-temp regime"
18604 );
18605
18606 for (name, d) in [
18607 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
18608 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
18609 ] {
18610 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
18611 assert!(
18612 !sampler.is_greedy(),
18613 "{name}: vendor default must not be greedy"
18614 );
18615 assert!(
18616 !sampler.is_spec_sampling(),
18617 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
18618 starts passing, either the vendor numbers changed or the in-graph draft \
18619 learned filters, and the perf note in docs/SERVING.md needs revisiting"
18620 );
18621 }
18622
18623 let opted_out = resolve_sampler_config(
18625 ClientSampling {
18626 top_p: Some(1.0),
18627 top_k: Some(0),
18628 seed: Some(1),
18629 ..Default::default()
18630 },
18631 &qwen38_vendor_defaults(),
18632 );
18633 assert!(
18634 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
18635 "explicitly disabling the filters must restore the pure-temp regime"
18636 );
18637 }
18638
18639 #[test]
18640 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
18641 let comp_seed = |body: serde_json::Value| {
18648 let req: CompletionReq = serde_json::from_value(body).unwrap();
18649 let (tx, _rx) = worker::event_channel();
18650 build_request(&req, tx, lanes::Lane::Interactive, None)
18651 .sampler_cfg
18652 .seed
18653 };
18654 let chat_seed = |body: serde_json::Value| {
18655 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
18656 let (tx, _rx) = worker::event_channel();
18657 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
18658 .unwrap()
18659 .request
18660 .sampler_cfg
18661 .seed
18662 };
18663
18664 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
18667 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
18668 let c = chat_seed(serde_json::json!({
18669 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
18670 assert_ne!(
18671 a, 0,
18672 "omitted seed must not be the pinned 0 that caused the loop"
18673 );
18674 assert_ne!(b, 0);
18675 assert_ne!(c, 0);
18676 assert_ne!(
18677 a, b,
18678 "two seed-omitting requests must get DIFFERENT streams"
18679 );
18680 assert_ne!(a, c);
18681
18682 assert_eq!(
18685 comp_seed(serde_json::json!({
18686 "model": "m", "prompt": "t", "seed": 0})),
18687 0,
18688 "explicit seed 0 must stay 0 — the determinism gates depend on it"
18689 );
18690 assert_eq!(
18691 comp_seed(serde_json::json!({
18692 "model": "m", "prompt": "t", "seed": 12345})),
18693 12345
18694 );
18695 assert_eq!(
18696 chat_seed(serde_json::json!({
18697 "model": "m", "messages": [{"role": "user", "content": "t"}],
18698 "seed": 777})),
18699 777
18700 );
18701 assert_eq!(
18703 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
18704 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
18705 );
18706
18707 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
18710 assert_eq!(
18711 seeds.len(),
18712 256,
18713 "fresh_seed must not collide across rapid calls"
18714 );
18715 assert!(!seeds.contains(&0));
18716 }
18717
18718 #[test]
18719 fn response_format_builds_grammar_only_when_present() {
18720 let mk = |rf: Option<serde_json::Value>| {
18724 let mut body = serde_json::json!({
18725 "model": "m", "messages": [{"role": "user", "content": "t"}]});
18726 if let Some(rf) = rf {
18727 body["response_format"] = rf;
18728 }
18729 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
18730 let (tx, _rx) = worker::event_channel();
18731 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
18732 };
18733 assert!(mk(None).unwrap().request.grammar.is_none());
18734 assert!(
18735 mk(Some(serde_json::json!({"type": "text"})))
18736 .unwrap()
18737 .request
18738 .grammar
18739 .is_none()
18740 );
18741 assert!(matches!(
18742 mk(Some(serde_json::json!({"type": "json_object"})))
18743 .unwrap()
18744 .request
18745 .grammar,
18746 Some(constrained::GrammarSpec::JsonObject)
18747 ));
18748 assert!(matches!(
18749 mk(Some(serde_json::json!({"type": "json_schema",
18750 "json_schema": {"schema": {"type": "object"}}})))
18751 .unwrap()
18752 .request
18753 .grammar,
18754 Some(constrained::GrammarSpec::JsonSchema(_))
18755 ));
18756 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
18758 }
18759
18760 #[test]
18770 fn response_format_think_table_switch_postthink_refusal() {
18771 let mk = |caps: &ModelCaps| {
18772 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
18773 "model": "m", "messages": [{"role": "user", "content": "t"}],
18774 "response_format": {"type": "json_object"}}))
18775 .unwrap();
18776 let (tx, _rx) = worker::event_channel();
18777 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
18778 };
18779 let switch = ModelCaps {
18781 chat_ok: true,
18782 qwen_think: true,
18783 think_switch: true,
18784 ..Default::default()
18785 };
18786 let plan = mk(&switch).unwrap();
18787 assert_eq!(
18788 plan.request.think,
18789 memra_tokenizer::chat::ThinkMode::NoThink,
18790 "switch-carrying template must keep the grammar-from-token-1 path"
18791 );
18792 assert!(plan.request.grammar.is_some());
18793
18794 let postthink = ModelCaps {
18796 chat_ok: true,
18797 qwen_think: true,
18798 think_switch: false,
18799 think_close: vec![128799],
18800 ..Default::default()
18801 };
18802 let plan = mk(&postthink).unwrap();
18803 assert_ne!(
18804 plan.request.think,
18805 memra_tokenizer::chat::ThinkMode::NoThink,
18806 "post-think constrained request must keep the think channel ON"
18807 );
18808 assert!(plan.request.grammar.is_some());
18809
18810 let no_contract = ModelCaps {
18812 chat_ok: true,
18813 qwen_think: true,
18814 think_switch: false,
18815 think_close: Vec::new(),
18816 ..Default::default()
18817 };
18818 let err = match mk(&no_contract) {
18819 Err(err) => err,
18820 Ok(_) => panic!("think-forced template with no close contract must refuse"),
18821 };
18822 assert!(
18823 err.contains("think-close"),
18824 "refusal must name the missing close contract: {err}"
18825 );
18826 }
18827
18828 #[test]
18829 fn unsupported_semantic_params_are_named_rejections() {
18830 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
18832 "model": "m", "messages": [{"role": "user", "content": "t"}],
18833 "response_format": {"type": "json_object"}
18834 }))
18835 .unwrap();
18836 assert!(req.response_format.is_some());
18837 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
18838 "model": "m", "messages": [{"role": "user", "content": "t"}],
18839 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
18840 "user": "u-1", "stream_options": {"include_usage": true}
18841 }))
18842 .unwrap();
18843 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
18845 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
18846 assert_eq!(req.n, Some(1));
18847 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
18849 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
18850 assert_eq!(param, "logit_bias");
18851 assert_eq!(msg, "logit_bias is not supported (why)");
18852 }
18853
18854 #[test]
18855 fn completions_accept_openai_stop_forms() {
18856 for (value, expected) in [
18857 (serde_json::json!("Problem:"), vec!["Problem:"]),
18858 (
18859 serde_json::json!(["Question:", "Problem:"]),
18860 vec!["Question:", "Problem:"],
18861 ),
18862 (serde_json::Value::Null, Vec::<&str>::new()),
18863 ] {
18864 let req: CompletionReq = serde_json::from_value(serde_json::json!({
18865 "model": "plain_quant", "prompt": "task", "stop": value
18866 }))
18867 .unwrap();
18868 assert_eq!(req.stop.into_vec(), expected);
18869 }
18870 }
18871
18872 fn fake_worker_state() -> AppState {
18879 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
18880 }
18881
18882 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
18883 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
18884 }
18885
18886 struct WorkerSaw {
18891 sampler_cfg: SamplerConfig,
18892 think: ThinkMode,
18893 reasoning_effort: Option<String>,
18894 }
18895
18896 fn fake_worker_state_full(
18903 steps: usize,
18904 step_delay: std::time::Duration,
18905 caps: HashMap<String, ModelCaps>,
18906 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
18907 ) -> AppState {
18908 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
18909 let health = health::WorkerHealth::new();
18910 let h = health.clone();
18911 std::thread::spawn(move || {
18912 h.mark_ready();
18913 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
18914 if let Some(tx) = &saw_tx {
18915 let _ = tx.send(WorkerSaw {
18916 sampler_cfg: req.sampler_cfg.clone(),
18917 think: req.think,
18918 reasoning_effort: req.reasoning_effort.clone(),
18919 });
18920 }
18921 worker::release_pending_admit();
18925 worker::release_admission_reservation(req.lane);
18926 h.beat_busy();
18927 if let Some(ready) = req.constraint_ready.take() {
18928 let _ = ready.send(Ok(()));
18929 }
18930 let _ = req.tx.send(Event::PromptUsage {
18931 n_prompt: 1,
18932 n_cached: 0,
18933 });
18934 if let Some(spec) = req.capture.as_ref() {
18939 let _ = req.tx.send(Event::PromptCapture {
18940 hidden: spec.hidden.then(|| vec![1.0, 0.0]),
18941 logits: if spec.logit_pieces.is_empty() {
18942 Vec::new()
18943 } else {
18944 vec![2.0, 0.0]
18945 },
18946 });
18947 }
18948 for step in 0..steps {
18949 h.beat_busy();
18950 let text = if steps == 1 { "ok" } else { "x" };
18951 let _ = req.tx.send(Event::Token {
18952 id: step as u32 + 1,
18953 text: text.into(),
18954 });
18955 if !step_delay.is_zero() {
18956 std::thread::sleep(step_delay);
18957 }
18958 }
18959 let _ = req.tx.send(Event::Done {
18960 stop_reason: "Eos".into(),
18961 n_tokens: steps,
18962 n_prompt: 1,
18963 n_cached: 0,
18964 elapsed_s: 0.01,
18965 spec: None,
18966 });
18967 h.set_phase(health::PHASE_IDLE);
18968 }
18969 });
18970 for _ in 0..2000 {
18973 if health.live().is_ok() {
18974 break;
18975 }
18976 std::thread::sleep(std::time::Duration::from_millis(1));
18977 }
18978 AppState {
18979 cmd_tx,
18980 models: Arc::new(vec!["m".into()]),
18981 caps: Arc::new(caps),
18982 openrouter_metadata: Arc::new(RwLock::new(Arc::new(ModelMetadataSet::default()))),
18983 metering: None,
18984
18985 budget_tokenizers: None,
18986 api_auth: ApiAuth::default(),
18987 metrics_auth: MetricsAuth::default(),
18988 metrics: SharedMetrics::default(),
18989 inflight: Arc::new(Default::default()),
18990 tenant_inflight: Arc::new(Default::default()),
18991 health,
18992 bg: None,
18993 }
18994 }
18995
18996 #[tokio::test]
18997 #[allow(clippy::await_holding_lock)] async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
18999 let _l = drain_lock();
19000 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
19001 let normal_state = st.clone();
19002 let normal = tokio::spawn(async move {
19003 chat_completions(
19004 State(normal_state),
19005 axum::http::HeaderMap::new(),
19006 None,
19007 Json(
19008 serde_json::from_value(serde_json::json!({
19009 "model": "m",
19010 "messages": [{"role": "user", "content": "keep decoding"}],
19011 }))
19012 .unwrap(),
19013 ),
19014 )
19015 .await
19016 });
19017 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
19018
19019 let mut deep = serde_json::json!({"type": "string"});
19020 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
19021 deep = serde_json::json!({"allOf": [deep]});
19022 }
19023 let bad = chat_completions(
19024 State(st.clone()),
19025 axum::http::HeaderMap::new(),
19026 None,
19027 Json(
19028 serde_json::from_value(serde_json::json!({
19029 "model": "m",
19030 "messages": [{"role": "user", "content": "bad schema"}],
19031 "response_format": {
19032 "type": "json_schema",
19033 "json_schema": {"schema": deep},
19034 },
19035 }))
19036 .unwrap(),
19037 ),
19038 )
19039 .await;
19040 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
19041 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
19042 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
19043 .await
19044 .unwrap();
19045 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19046 assert!(
19047 payload["error"]["message"]
19048 .as_str()
19049 .unwrap()
19050 .contains("maximum nesting depth")
19051 );
19052 assert!(
19053 !normal.is_finished(),
19054 "bad schema stalled or replaced the normal decode"
19055 );
19056
19057 let normal_response = normal.await.unwrap();
19058 assert_eq!(normal_response.status(), StatusCode::OK);
19059 let snapshot = st.health.snapshot();
19060 assert!(
19061 st.health.live().is_ok(),
19062 "normal decode left health stalled"
19063 );
19064 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
19065 }
19066
19067 #[tokio::test]
19068 #[allow(clippy::await_holding_lock)] async fn valid_response_format_preflight_preserves_generation() {
19070 let _l = drain_lock();
19071 let response = chat_completions(
19072 State(fake_worker_state()),
19073 axum::http::HeaderMap::new(),
19074 None,
19075 Json(
19076 serde_json::from_value(serde_json::json!({
19077 "model": "m",
19078 "messages": [{"role": "user", "content": "valid schema"}],
19079 "response_format": {"type": "json_object"},
19080 }))
19081 .unwrap(),
19082 ),
19083 )
19084 .await;
19085 assert_eq!(response.status(), StatusCode::OK);
19086 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19087 .await
19088 .unwrap();
19089 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19090 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
19091 }
19092
19093 #[tokio::test]
19094 #[allow(clippy::await_holding_lock)] async fn unknown_model_refuses_model_not_found_before_admission() {
19096 let _l = drain_lock();
19097 let response = chat_completions(
19102 State(fake_worker_state()),
19103 axum::http::HeaderMap::new(),
19104 None,
19105 Json(
19106 serde_json::from_value(serde_json::json!({
19107 "model": "qwen/qwen3.8-27b-typo",
19108 "messages": [{"role": "user", "content": "hi"}],
19109 }))
19110 .unwrap(),
19111 ),
19112 )
19113 .await;
19114 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
19115 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19116 .await
19117 .unwrap();
19118 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19119 assert_eq!(payload["error"]["code"], "model_not_found");
19120 assert_eq!(payload["error"]["type"], "invalid_request_error");
19121
19122 let response = completions(
19124 State(fake_worker_state()),
19125 axum::http::HeaderMap::new(),
19126 None,
19127 Json(
19128 serde_json::from_value(serde_json::json!({
19129 "model": "nope",
19130 "prompt": "hi",
19131 }))
19132 .unwrap(),
19133 ),
19134 )
19135 .await;
19136 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
19137 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19138 .await
19139 .unwrap();
19140 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19141 assert_eq!(payload["error"]["code"], "model_not_found");
19142 }
19143
19144 const METRICS_KEY_ACME: &str = "completion-acme-secret";
19145 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
19146
19147 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
19148 let spec = format!(
19149 "acme:{},blue:{}",
19150 auth::sha256_hex(METRICS_KEY_ACME),
19151 auth::sha256_hex(METRICS_KEY_BLUE),
19152 );
19153 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
19154 let mut st = fake_worker_state();
19155 st.api_auth.keyring = Some(keyring);
19156 st.metrics_auth = MetricsAuth::new(
19157 true,
19158 st.api_auth.configured(),
19159 metrics_token.map(str::to_string),
19160 );
19161 {
19162 let mut metrics = st.metrics.lock().unwrap();
19163 metrics.admitted = 17;
19164 metrics.prompt_tokens_in = 400;
19165 metrics.cached_tokens_in = 60;
19166 metrics.prefix_hits = 2;
19167 metrics.prefix_misses = 3;
19168 metrics.prefix_inserts = 5;
19169 metrics.prefix_evictions = 7;
19170 metrics.prefix_skips_budget = 9;
19171 metrics.prefix_skips_pinned = 10;
19172 metrics.prefix_hit_tokens = 11;
19173 metrics.lcp_hist[4] = 13;
19174 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
19175 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
19176 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
19177 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
19178 metrics.prefix_entries = 29;
19179 metrics.prefix_bytes = 31;
19180 metrics.active_sessions = 3;
19181 metrics.queued_requests = 5;
19182 metrics.admission_inflight.insert("m".into(), 4);
19183 metrics
19184 .admission_booked_bytes
19185 .insert("m".into(), 41_000_000);
19186 metrics.continuation_pool_entries = 7;
19187 metrics.spec_pool_entries = 11;
19188 metrics.cuda_driver_free_bytes = 13;
19189 metrics.cuda_pool_reserved_bytes = 17;
19190 metrics.cuda_pool_used_bytes = 19;
19191 metrics.cuda_pool_cached_bytes = 23;
19192 metrics.batch_size_last = 37;
19193 metrics.spec.insert(
19194 "m".into(),
19195 memra_engine::spec::SpecTelemetry {
19196 rounds: 2,
19197 drafted: 6,
19198 accepted: 4,
19199 ..Default::default()
19200 },
19201 );
19202 let mut spec_window = memra_engine::spec::SpecTelemetry {
19203 rounds: 4,
19204 drafted: 12,
19205 accepted: 6,
19206 ..Default::default()
19207 };
19208 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
19209 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
19210 metrics.spec_window.insert("m".into(), spec_window);
19211 metrics.constraint_compiler_fail_closed.insert(
19212 "m".into(),
19213 Arc::new(std::sync::atomic::AtomicBool::new(true)),
19214 );
19215 }
19216 st
19217 }
19218
19219 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
19220 let mut headers = HeaderMap::new();
19221 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
19222 let response = get_metrics(State(st), headers).await;
19223 assert_eq!(response.status(), StatusCode::OK);
19224 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19225 .await
19226 .unwrap();
19227 serde_json::from_slice(&bytes).unwrap()
19228 }
19229
19230 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
19231 let mut headers = HeaderMap::new();
19232 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
19233 let response = yield_metrics(State(st), headers).await;
19234 assert_eq!(response.status(), StatusCode::OK);
19235 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19236 .await
19237 .unwrap();
19238 serde_json::from_slice(&bytes).unwrap()
19239 }
19240
19241 #[test]
19242 fn exposed_open_bind_is_refused_before_server_start() {
19243 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
19244 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
19245
19246 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
19247 assert!(err.contains("refusing unauthenticated non-loopback bind"));
19248 assert!(err.contains("MEMRA_API_KEY"));
19249 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
19250 assert!(validate_bind_security("[::]:8000", false, false).is_err());
19251
19252 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
19253 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
19254 }
19255
19256 #[tokio::test]
19257 async fn keyed_metrics_require_and_accept_api_bearer() {
19258 let mut st = fake_worker_state();
19259 st.api_auth.single_key = Some(Arc::from("completion-secret"));
19260 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
19261
19262 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
19263 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
19264 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
19265 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
19266
19267 let mut headers = HeaderMap::new();
19268 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
19269 assert_eq!(
19270 get_metrics(State(st.clone()), headers.clone())
19271 .await
19272 .status(),
19273 StatusCode::OK,
19274 );
19275 let body = metrics_json(st.clone(), "completion-secret").await;
19276 assert!(
19277 body.get("admitted").is_some(),
19278 "the legacy single-key domain keeps cumulative counters",
19279 );
19280 assert!(
19281 body.get("active_sessions").is_none(),
19282 "a static completion key is not an operator metrics principal",
19283 );
19284 assert_eq!(
19285 yield_metrics(State(st), headers).await.status(),
19286 StatusCode::OK
19287 );
19288 }
19289
19290 #[tokio::test]
19291 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
19292 let st = multi_key_metrics_state(None);
19293 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
19294 assert_eq!(
19295 body.as_object().unwrap().len(),
19296 2,
19297 "completion metrics must contain only tenant-scoped rows",
19298 );
19299 let tenants = body["tenants"].as_object().unwrap();
19300 assert_eq!(tenants.len(), 1);
19301 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
19302 assert!(!tenants.contains_key("t:blue"));
19303 let adsd = body["adsd_suspect_total"].as_object().unwrap();
19304 assert_eq!(adsd.len(), 1);
19305 assert_eq!(adsd["t:acme"], 1);
19306 assert!(!adsd.contains_key("t:blue"));
19307
19308 let mut headers = HeaderMap::new();
19309 headers.insert(
19310 "authorization",
19311 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
19312 );
19313 assert_eq!(
19314 yield_metrics(State(st), headers).await.status(),
19315 StatusCode::FORBIDDEN,
19316 "the process-wide yield view requires an operator metrics token",
19317 );
19318 }
19319
19320 #[tokio::test]
19321 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
19322 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
19323 for operator_only in [
19324 "prefix_cache_entries",
19325 "prefix_cache_bytes",
19326 "prefix_cache_skips_budget",
19327 "prefix_cache_skips_pinned",
19328 "active_sessions",
19329 "queued_requests",
19330 "admission_inflight",
19331 "admission_booked_bytes",
19332 "continuation_pool_entries",
19333 "spec_pool_entries",
19334 "cuda_driver_free_bytes",
19335 "cuda_pool_reserved_bytes",
19336 "cuda_pool_used_bytes",
19337 "cuda_pool_cached_bytes",
19338 "constraint_compiler_fail_closed",
19339 "serve_idle_seconds",
19340 "spec",
19341 "spec_tau",
19342 "spec_accept_by_position",
19343 "dual_pp",
19344 "pp_wave",
19345 "peer_probe_bypassed",
19346 "peer_probe_boundary_copies",
19347 "peer_probe_runtime_reprobes",
19348 "peer_probe_runtime_failures",
19349 "peer_probe_deferred_total",
19350 "peer_probe_integrity_degraded",
19351 "peer_probe_degraded_to_host_bounce",
19352 ] {
19353 assert!(
19354 body.get(operator_only).is_none(),
19355 "tenant metrics must not expose operator field {operator_only}",
19356 );
19357 }
19358 }
19359
19360 #[test]
19361 fn populated_spec_acceptance_metrics_are_operator_only() {
19362 for scope in [
19363 MetricsScope::CompletionDomain,
19364 MetricsScope::Tenant("t:acme".into()),
19365 ] {
19366 let mut body = json!({});
19367 insert_spec_acceptance_metrics(&mut body, &scope, || {
19368 panic!("tenant scope evaluated the process-wide spec snapshot")
19369 });
19370 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
19371 assert!(
19372 body.get("spec_accept_by_position").is_none(),
19373 "{scope:?} leaked the accept histogram"
19374 );
19375 }
19376
19377 let mut telemetry = memra_engine::spec::SpecTelemetry {
19378 rounds: 4,
19379 drafted: 12,
19380 accepted: 6,
19381 ..Default::default()
19382 };
19383 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
19384 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
19385 let mut body = json!({});
19386 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
19387 HashMap::from([("model-a".to_string(), telemetry)])
19388 });
19389 assert_eq!(body["spec_tau"]["model-a"], 1.5);
19390 let histogram = &body["spec_accept_by_position"]["model-a"];
19391 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
19392 assert_eq!(histogram["rounds"], 4);
19393 assert_eq!(histogram["offered"], json!([4, 4, 4]));
19394 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
19395 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
19396 }
19397
19398 #[test]
19399 fn populated_dual_pp_metrics_are_operator_only() {
19400 let populated = DualPpMetricsSnapshot {
19401 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
19402 stage_samples: [1, 1, 1, 1],
19403 dropped_timing_samples: 0,
19404 overlaps: 17,
19405 slot_pairs: 19,
19406 slot_uses: [19, 19],
19407 slot_collisions: 0,
19408 };
19409 for scope in [
19410 MetricsScope::CompletionDomain,
19411 MetricsScope::Tenant("t:acme".into()),
19412 ] {
19413 let mut body = json!({});
19414 insert_dual_pp_metrics(&mut body, &scope, || populated);
19415 assert!(
19416 body.get("dual_pp").is_none(),
19417 "{scope:?} leaked dual PP topology"
19418 );
19419 }
19420
19421 let mut body = json!({});
19422 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
19423 assert_eq!(body["dual_pp"]["overlaps"], 17);
19424 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
19425 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
19426 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
19427 assert_eq!(
19428 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
19429 1.0
19430 );
19431 }
19432
19433 #[test]
19434 fn populated_pp_wave_metrics_are_operator_only() {
19435 let populated = PpWaveMetricsSnapshot {
19436 ticks: 11,
19437 cells: 96,
19438 overlaps: 37,
19439 };
19440 for scope in [
19441 MetricsScope::CompletionDomain,
19442 MetricsScope::Tenant("t:acme".into()),
19443 ] {
19444 let mut body = json!({});
19445 insert_pp_wave_metrics(&mut body, &scope, || populated);
19446 assert!(
19447 body.get("pp_wave").is_none(),
19448 "{scope:?} leaked PP wave topology"
19449 );
19450 }
19451
19452 let mut body = json!({});
19453 insert_pp_wave_metrics(&mut body, &MetricsScope::All, || populated);
19454 assert_eq!(body["pp_wave"]["ticks"], 11);
19455 assert_eq!(body["pp_wave"]["cells"], 96);
19456 assert_eq!(body["pp_wave"]["overlaps"], 37);
19457 }
19458
19459 #[test]
19460 fn peer_probe_metrics_are_operator_only() {
19461 let populated = memra_engine::pp::PeerProbeMetrics {
19462 bypassed: 1,
19463 boundary_copies: 8_192,
19464 runtime_probes: 1,
19465 runtime_failures: 0,
19466 deferred_total: 4,
19467 integrity_degraded: true,
19468 degraded_to_host_bounce: true,
19469 };
19470 for scope in [
19471 MetricsScope::CompletionDomain,
19472 MetricsScope::Tenant("t:acme".into()),
19473 ] {
19474 let mut body = json!({});
19475 insert_peer_probe_metrics(&mut body, &scope, || populated);
19476 assert!(body.get("peer_probe_bypassed").is_none());
19477 }
19478
19479 let mut body = json!({});
19480 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
19481 assert_eq!(body["peer_probe_bypassed"], 1);
19482 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
19483 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
19484 assert_eq!(body["peer_probe_runtime_failures"], 0);
19485 assert_eq!(body["peer_probe_deferred_total"], 4);
19486 assert_eq!(body["peer_probe_integrity_degraded"], true);
19487 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
19488 }
19489
19490 #[tokio::test]
19491 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
19492 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
19493 for operator_only in [
19494 "lcp_histogram",
19495 "cache_hit_token_ratio",
19496 "prefix_cache_hits",
19497 "prefix_cache_misses",
19498 "prefix_cache_inserts",
19499 "prefix_cache_evictions",
19500 "prefix_cache_skips_budget",
19501 "prefix_cache_skips_pinned",
19502 "prefix_cache_hit_tokens",
19503 ] {
19504 assert!(
19505 tenant_body.get(operator_only).is_none(),
19506 "tenant metrics must not expose global prefix field {operator_only}",
19507 );
19508 }
19509 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
19510 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
19511 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
19512 assert_eq!(
19513 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
19514 0.4
19515 );
19516
19517 let operator_body = metrics_json(
19518 multi_key_metrics_state(Some("scrape-secret")),
19519 "scrape-secret",
19520 )
19521 .await;
19522 assert_eq!(operator_body["prefix_cache_hits"], 2);
19523 assert_eq!(operator_body["prefix_cache_misses"], 3);
19524 assert_eq!(operator_body["prefix_cache_inserts"], 5);
19525 assert_eq!(operator_body["prefix_cache_evictions"], 7);
19526 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
19527 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
19528 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
19529 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
19530 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
19531 }
19532
19533 #[tokio::test]
19534 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
19535 let st = multi_key_metrics_state(Some("scrape-secret"));
19536 let mut completion_headers = HeaderMap::new();
19537 completion_headers.insert(
19538 "authorization",
19539 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
19540 );
19541 assert_eq!(
19542 get_metrics(State(st.clone()), completion_headers.clone())
19543 .await
19544 .status(),
19545 StatusCode::FORBIDDEN,
19546 );
19547 assert_eq!(
19548 yield_metrics(State(st.clone()), completion_headers)
19549 .await
19550 .status(),
19551 StatusCode::FORBIDDEN,
19552 );
19553
19554 let body = metrics_json(st.clone(), "scrape-secret").await;
19555 let tenants = body["tenants"].as_object().unwrap();
19556 assert_eq!(tenants.len(), 2);
19557 assert!(tenants.contains_key("t:acme"));
19558 assert!(tenants.contains_key("t:blue"));
19559 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
19560 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
19561 assert_eq!(body["active_sessions"], 3);
19562 assert_eq!(body["queued_requests"], 5);
19563 assert_eq!(body["admission_inflight"]["m"], 4);
19565 assert_eq!(body["admission_booked_bytes"]["m"], 41_000_000);
19566 assert_eq!(body["prefix_cache_bytes"], 31);
19567 assert_eq!(body["cuda_driver_free_bytes"], 13);
19568 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
19569 assert_eq!(body["spec"]["m"]["drafted"], 6);
19570 assert_eq!(body["spec_tau"]["m"], 1.5);
19571 assert_eq!(
19572 body["spec_accept_by_position"]["m"]["accepted"],
19573 json!([3, 2, 1])
19574 );
19575 let yield_body = yield_metrics_json(st, "scrape-secret").await;
19576 assert_eq!(yield_body["batch_size_last"], 37);
19577 }
19578
19579 #[tokio::test]
19580 async fn metrics_token_protects_public_override_without_api_keys() {
19581 let mut st = fake_worker_state();
19582 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
19583
19584 assert_eq!(
19585 get_metrics(State(st.clone()), HeaderMap::new())
19586 .await
19587 .status(),
19588 StatusCode::UNAUTHORIZED,
19589 );
19590 let mut headers = HeaderMap::new();
19591 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
19592 assert_eq!(
19593 get_metrics(State(st.clone()), headers.clone())
19594 .await
19595 .status(),
19596 StatusCode::OK,
19597 );
19598 assert_eq!(
19599 yield_metrics(State(st), headers).await.status(),
19600 StatusCode::OK
19601 );
19602 }
19603
19604 #[tokio::test]
19605 async fn no_key_loopback_metrics_remain_open_for_development() {
19606 let mut st = fake_worker_state();
19607 st.metrics_auth = MetricsAuth::new(true, false, None);
19608 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
19609 assert_eq!(response.status(), StatusCode::OK);
19610 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19611 .await
19612 .unwrap();
19613 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19614 assert!(
19615 body.get("active_sessions").is_some(),
19616 "no-key loopback development keeps full operator visibility",
19617 );
19618 assert_eq!(
19619 yield_metrics(State(st), HeaderMap::new()).await.status(),
19620 StatusCode::OK,
19621 );
19622 }
19623
19624 #[test]
19625 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
19626 let metrics = SharedMetrics::default();
19627 let rl = RateLimit::compute(4, 1, &metrics);
19629 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
19630 let rl = RateLimit::compute(4, 3, &metrics);
19631 assert_eq!(rl.remaining, 1);
19632 let rl = RateLimit::compute(4, 4, &metrics);
19634 assert_eq!(rl.remaining, 0);
19635 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
19636 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
19638 let m = worker::Metrics {
19640 completed: 2,
19641 tokens_out: 200,
19642 step_p50_ms: 20.0,
19643 ..Default::default()
19644 };
19645 assert_eq!(reset_estimate_s(&m), 2); }
19647
19648 #[test]
19649 fn inflight_guard_counts_up_and_frees_on_drop() {
19650 let counts: InflightCounts = Arc::new(Default::default());
19651 let tenants: TenantGauge = Arc::new(Default::default());
19652 let (g1, n1, t1) = InflightGuard::try_acquire(
19653 counts.clone(),
19654 lanes::Lane::Interactive,
19655 tenants.clone(),
19656 "acme",
19657 None,
19658 )
19659 .unwrap();
19660 let (g2, n2, t2) = InflightGuard::try_acquire(
19661 counts.clone(),
19662 lanes::Lane::Interactive,
19663 tenants.clone(),
19664 "acme",
19665 None,
19666 )
19667 .unwrap();
19668 assert_eq!((n1, n2), (1, 2));
19669 assert_eq!((t1, t2), (1, 2));
19671 let (gj, nj, tj) = InflightGuard::try_acquire(
19673 counts.clone(),
19674 lanes::Lane::Judge,
19675 tenants.clone(),
19676 "blue",
19677 None,
19678 )
19679 .unwrap();
19680 assert_eq!((nj, tj), (1, 1));
19681 drop(g1);
19682 drop(gj);
19683 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
19684 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
19685 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
19686 assert!(tenants.lock().unwrap().get("blue").is_none());
19688 drop(g2);
19689 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
19690 assert!(tenants.lock().unwrap().is_empty());
19691 }
19692
19693 #[test]
19694 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
19695 let counts: InflightCounts = Arc::new(Default::default());
19696 let tenants: TenantGauge = Arc::new(Default::default());
19697 let start = Arc::new(std::sync::Barrier::new(3));
19698 let attempted = Arc::new(std::sync::Barrier::new(3));
19699 let mut joins = Vec::new();
19700 for _ in 0..2 {
19701 let counts = counts.clone();
19702 let tenants = tenants.clone();
19703 let start = start.clone();
19704 let attempted = attempted.clone();
19705 joins.push(std::thread::spawn(move || {
19706 start.wait();
19707 let result = InflightGuard::try_acquire(
19708 counts,
19709 lanes::Lane::Interactive,
19710 tenants,
19711 "preview_001",
19712 Some(1),
19713 );
19714 let won = result.is_ok();
19715 attempted.wait(); drop(result);
19717 won
19718 }));
19719 }
19720 start.wait();
19721 attempted.wait();
19722 let wins = joins
19723 .into_iter()
19724 .map(|join| join.join().unwrap())
19725 .filter(|won| *won)
19726 .count();
19727 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
19728 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
19729 assert!(tenants.lock().unwrap().is_empty());
19730 }
19731
19732 #[tokio::test]
19733 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
19734 let st = fake_worker_state();
19735 let tenant = auth::TenantCtx {
19736 tenant: "preview_001".into(),
19737 lane_class: auth::LaneClass::Interactive,
19738 rate_limit: Some(1),
19739 key_prefix: None,
19740 };
19741 let first_env = Envelope::new(true);
19742 let (guard, first_rl) =
19743 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
19744 Ok(slot) => slot,
19745 Err(_) => panic!("the first request must acquire the tenant slot"),
19746 };
19747 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
19748
19749 let second_env = Envelope::new(true);
19750 let response =
19751 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
19752 Err(response) => response,
19753 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
19754 };
19755 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
19756 assert_eq!(response.headers()["retry-after"], "2");
19757 assert_eq!(response.headers()["retry-after-ms"], "2000");
19758 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
19759 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
19760 assert_eq!(response.headers()["x-request-id"], second_env.id);
19761 assert_eq!(
19762 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
19763 1,
19764 "rejected request must not consume a lane slot"
19765 );
19766 assert_eq!(
19767 st.tenant_inflight
19768 .lock()
19769 .unwrap()
19770 .get("preview_001")
19771 .copied(),
19772 Some(1),
19773 "rejected request must not increment the tenant gauge"
19774 );
19775 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19776 .await
19777 .unwrap();
19778 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19779 assert_eq!(payload["error"]["type"], "rate_limit_error");
19780 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
19781 assert!(
19782 payload["error"]["message"]
19783 .as_str()
19784 .unwrap()
19785 .contains("concurrent request limit")
19786 );
19787
19788 drop(guard);
19789 let _ = InflightGuard::try_acquire(
19790 st.inflight.clone(),
19791 lanes::Lane::Interactive,
19792 st.tenant_inflight.clone(),
19793 "preview_001",
19794 Some(1),
19795 )
19796 .expect("slot must reopen after the in-flight request completes");
19797 }
19798
19799 #[test]
19800 fn tenant_rate_limit_override_is_min_with_global_cap() {
19801 let metrics = SharedMetrics::default();
19802 let unlimited = auth::TenantCtx::default_tenant();
19803 let capped = auth::TenantCtx {
19804 tenant: "acme".into(),
19805 lane_class: auth::LaneClass::Interactive,
19806 rate_limit: Some(2),
19807 key_prefix: None,
19808 };
19809 let global = lane_cap(lanes::Lane::Interactive);
19810 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
19812 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
19813 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
19815 assert_eq!((rl.limit, rl.remaining), (2, 1));
19816 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
19817 assert_eq!(rl.remaining, 0);
19818 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
19819 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
19823 assert_eq!(rl.remaining, 0);
19824 let wide = auth::TenantCtx {
19825 rate_limit: Some(global + 100),
19826 ..capped.clone()
19827 };
19828 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
19829 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
19830 }
19831
19832 #[test]
19833 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
19834 let batch = auth::TenantCtx {
19835 tenant: "bulk".into(),
19836 lane_class: auth::LaneClass::Batch,
19837 rate_limit: None,
19838 key_prefix: None,
19839 };
19840 let interactive = auth::TenantCtx::default_tenant();
19841 let hdr = |v: Option<&str>| {
19842 let mut h = axum::http::HeaderMap::new();
19843 if let Some(v) = v {
19844 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
19845 }
19846 h
19847 };
19848 assert_eq!(
19850 lane_for_tenant(&hdr(None), &interactive).unwrap(),
19851 lanes::Lane::Interactive
19852 );
19853 assert_eq!(
19854 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
19855 lanes::Lane::Judge
19856 );
19857 assert_eq!(
19859 lane_for_tenant(&hdr(None), &batch).unwrap(),
19860 lanes::Lane::Harvest
19861 );
19862 assert_eq!(
19863 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
19864 lanes::Lane::Judge
19865 );
19866 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
19867 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
19868 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
19870 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
19871 }
19872
19873 #[tokio::test]
19874 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
19875 let hdr = |v: &str| {
19880 let mut h = axum::http::HeaderMap::new();
19881 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
19882 h
19883 };
19884 let body = |resp: Response| async move {
19885 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19886 .await
19887 .unwrap();
19888 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
19889 };
19890
19891 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
19892 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
19893 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
19894 let payload = body(resp).await;
19895 assert!(
19896 payload["error"].is_object(),
19897 "bare-string error body: {payload}"
19898 );
19899 assert_eq!(payload["error"]["type"], "invalid_request_error");
19900 assert_eq!(payload["error"]["param"], "x-lane");
19901 assert_eq!(payload["error"]["code"], "invalid_lane");
19902
19903 let batch = auth::TenantCtx {
19904 tenant: "bulk".into(),
19905 lane_class: auth::LaneClass::Batch,
19906 rate_limit: None,
19907 key_prefix: None,
19908 };
19909 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
19910 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
19911 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
19912 let payload = body(resp).await;
19913 assert_eq!(payload["error"]["type"], "authentication_error");
19914 assert_eq!(payload["error"]["param"], "x-lane");
19915 }
19916
19917 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
19920
19921 fn drain_lock() -> std::sync::MutexGuard<'static, ()> {
19934 let guard = DRAIN_LOCK.lock().unwrap_or_else(|poisoned| {
19935 DRAIN_LOCK.clear_poison();
19938 poisoned.into_inner()
19939 });
19940 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
19941 guard
19942 }
19943
19944 struct DrainingRestore;
19950 impl Drop for DrainingRestore {
19951 fn drop(&mut self) {
19952 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
19953 }
19954 }
19955
19956 #[tokio::test]
19957 #[allow(clippy::await_holding_lock)] async fn responses_carry_rate_limit_headers_and_slot_frees() {
19959 let _l = drain_lock();
19960 let st = fake_worker_state();
19961 let resp = chat_completions(
19964 State(st.clone()),
19965 axum::http::HeaderMap::new(),
19966 None,
19967 Json(
19968 serde_json::from_value(serde_json::json!({
19969 "model": "m", "messages": [{"role": "user", "content": "t"}]
19970 }))
19971 .unwrap(),
19972 ),
19973 )
19974 .await;
19975 assert_eq!(resp.status(), StatusCode::OK);
19976 let h = resp.headers();
19977 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
19978 let remaining: usize = h["x-ratelimit-remaining"]
19979 .to_str()
19980 .unwrap()
19981 .parse()
19982 .unwrap();
19983 assert_eq!(remaining, limit - 1);
19984 assert_eq!(h["x-ratelimit-reset"], "0");
19985 assert_eq!(
19986 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
19987 0,
19988 "slot must free at completion"
19989 );
19990 let resp = completions(
19993 State(st.clone()),
19994 axum::http::HeaderMap::new(),
19995 None,
19996 Json(
19997 serde_json::from_value(serde_json::json!({
19998 "model": "m", "prompt": "t", "stream": true
19999 }))
20000 .unwrap(),
20001 ),
20002 )
20003 .await;
20004 assert_eq!(resp.status(), StatusCode::OK);
20005 assert!(resp.headers().contains_key("x-ratelimit-limit"));
20006 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
20007 assert!(resp.headers().contains_key("x-ratelimit-reset"));
20008 assert_eq!(
20009 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
20010 1,
20011 "stream in flight holds the slot"
20012 );
20013 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
20014 .await
20015 .unwrap();
20016 assert_eq!(
20017 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
20018 0,
20019 "slot must free when the stream completes"
20020 );
20021 }
20022
20023 #[tokio::test]
20029 #[allow(clippy::await_holding_lock)] async fn multi_item_capture_requests_open_one_receipt_per_item_under_child_ids() {
20031 let _l = drain_lock();
20032 let mut st = fake_worker_state();
20033 let mock = MockMetering::admit_all();
20034 st.metering = Some(mock.clone());
20035
20036 let resp = embed_api::embeddings_admitted(
20037 State(st.clone()),
20038 HeaderMap::new(),
20039 AdmittedJson(
20040 serde_json::from_value(json!({"model": "m", "input": ["a", "bb", "ccc"]})).unwrap(),
20041 BodyAdmissionLease(None),
20042 ),
20043 )
20044 .await;
20045 assert_eq!(resp.status(), StatusCode::OK);
20046 let parent = resp.headers()["x-request-id"].to_str().unwrap().to_string();
20047 assert!(
20048 !parent.contains('.'),
20049 "the caller sees the parent id: {parent}"
20050 );
20051 let body: serde_json::Value = serde_json::from_slice(
20052 &axum::body::to_bytes(resp.into_body(), usize::MAX)
20053 .await
20054 .unwrap(),
20055 )
20056 .unwrap();
20057 assert_eq!(body["data"].as_array().map(Vec::len), Some(3));
20058 let events = mock.events();
20059 let opened: Vec<(String, &'static str)> = events
20060 .iter()
20061 .filter_map(|e| match e {
20062 MeterEvent::Open {
20063 request_id, route, ..
20064 } => Some((request_id.clone(), *route)),
20065 _ => None,
20066 })
20067 .collect();
20068 assert_eq!(
20069 opened,
20070 vec![
20071 (format!("{parent}.0"), "/v1/embeddings"),
20072 (format!("{parent}.1"), "/v1/embeddings"),
20073 (format!("{parent}.2"), "/v1/embeddings"),
20074 ],
20075 "one receipt per input, each under its own child id: {events:?}"
20076 );
20077 assert_eq!(
20078 events
20079 .iter()
20080 .filter(|e| matches!(e, MeterEvent::Complete { .. }))
20081 .count(),
20082 3,
20083 "every input settles its own receipt: {events:?}"
20084 );
20085
20086 let resp = embed_api::rerank_admitted(
20087 State(st),
20088 HeaderMap::new(),
20089 AdmittedJson(
20090 serde_json::from_value(
20091 json!({"model": "m", "query": "q", "documents": ["d0", "d1"]}),
20092 )
20093 .unwrap(),
20094 BodyAdmissionLease(None),
20095 ),
20096 )
20097 .await;
20098 assert_eq!(resp.status(), StatusCode::OK);
20099 let parent = resp.headers()["x-request-id"].to_str().unwrap().to_string();
20100 let opened: Vec<String> = mock
20101 .events()
20102 .into_iter()
20103 .skip(events.len())
20104 .filter_map(|e| match e {
20105 MeterEvent::Open {
20106 request_id,
20107 route: "/v1/rerank",
20108 ..
20109 } => Some(request_id),
20110 _ => None,
20111 })
20112 .collect();
20113 assert_eq!(opened, vec![format!("{parent}.0"), format!("{parent}.1")]);
20114 }
20115
20116 #[tokio::test]
20117 #[allow(clippy::await_holding_lock)] async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
20119 let _l = drain_lock();
20120 let mut st = fake_worker_state();
20121 let mock = MockMetering::admit_all();
20122 st.metering = Some(mock.clone());
20123
20124 let nonstream = chat_completions(
20125 State(st.clone()),
20126 HeaderMap::new(),
20127 None,
20128 Json(
20129 serde_json::from_value(json!({
20130 "model": "m",
20131 "messages": [{"role": "user", "content": "t"}],
20132 }))
20133 .unwrap(),
20134 ),
20135 )
20136 .await;
20137 assert_eq!(nonstream.status(), StatusCode::OK);
20138 let nonstream_id = nonstream.headers()["x-request-id"]
20139 .to_str()
20140 .unwrap()
20141 .to_string();
20142
20143 let stream = completions(
20144 State(st),
20145 HeaderMap::new(),
20146 None,
20147 Json(
20148 serde_json::from_value(json!({
20149 "model": "m",
20150 "prompt": "t",
20151 "stream": true,
20152 }))
20153 .unwrap(),
20154 ),
20155 )
20156 .await;
20157 assert_eq!(stream.status(), StatusCode::OK);
20158 let stream_id = stream.headers()["x-request-id"]
20159 .to_str()
20160 .unwrap()
20161 .to_string();
20162 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
20163 .await
20164 .unwrap();
20165
20166 let events = mock.events();
20170 let opened: Vec<&str> = events
20171 .iter()
20172 .filter_map(|e| match e {
20173 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
20174 _ => None,
20175 })
20176 .collect();
20177 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
20178 let completes = events
20179 .iter()
20180 .filter(|e| {
20181 matches!(
20182 e,
20183 MeterEvent::Complete {
20184 prompt: 1,
20185 cached: 0,
20186 completion: 1,
20187 }
20188 )
20189 })
20190 .count();
20191 assert_eq!(
20192 completes, 2,
20193 "both surfaces settle complete with worker-truth usage: {events:?}"
20194 );
20195 }
20196
20197 #[tokio::test]
20198 #[allow(clippy::await_holding_lock)] async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
20200 let _l = drain_lock();
20201 let mock = MockMetering::with_limits(vec![
20207 ReserveScript::Insufficient,
20208 ReserveScript::Admit { with_permit: false },
20209 ReserveScript::Blocked,
20210 ReserveScript::Admit { with_permit: true },
20211 ]);
20212 let mut st = fake_worker_state();
20213 st.metering = Some(mock.clone());
20214
20215 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
20217 assert_eq!(metrics.status(), StatusCode::OK);
20218 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
20219 .await
20220 .unwrap();
20221 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
20222 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
20223 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
20224 assert_eq!(metrics_body["budget_source_available"], true);
20225
20226 let request = || {
20227 Json(
20228 serde_json::from_value::<CompletionReq>(json!({
20229 "model": "m",
20230 "prompt_ids": [1],
20231 "max_tokens": 1,
20232 }))
20233 .unwrap(),
20234 )
20235 };
20236
20237 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
20238 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
20239 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
20240 .await
20241 .unwrap();
20242 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
20243 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
20244 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
20245
20246 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
20247 assert_eq!(included.status(), StatusCode::OK);
20248
20249 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
20252 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
20253
20254 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
20255 assert_eq!(admitted.status(), StatusCode::OK);
20256
20257 let events = mock.events();
20258 let terminal: Vec<&MeterEvent> = events
20259 .iter()
20260 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
20261 .collect();
20262 assert_eq!(
20263 terminal.len(),
20264 4,
20265 "four requests, four terminal settles: {events:?}"
20266 );
20267 assert!(matches!(
20268 terminal[0],
20269 MeterEvent::Reject { status: 402, .. }
20270 ));
20271 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
20272 assert!(matches!(
20273 terminal[2],
20274 MeterEvent::Reject { status: 402, .. }
20275 ));
20276 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
20277 let permits: Vec<bool> = events
20279 .iter()
20280 .filter_map(|e| match e {
20281 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
20282 _ => None,
20283 })
20284 .collect();
20285 assert_eq!(
20286 permits,
20287 vec![false, false, false, true],
20288 "the permit rides the receipt exactly when reserve minted one: {events:?}"
20289 );
20290 }
20291
20292 #[tokio::test]
20296 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
20297 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
20298 let mut st = fake_worker_state();
20299 st.metering = Some(mock.clone());
20300 let tenant = auth::TenantCtx {
20301 tenant: "acme".into(),
20302 lane_class: auth::LaneClass::Interactive,
20303 rate_limit: None,
20304 key_prefix: Some("mk-acme-testprefix00".into()),
20305 };
20306 let mut request = gate_request(1, 1);
20307 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
20308 .expect_err("a capped key must be refused at admission");
20309 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
20310 let (response, outcome) = rejection.into_response();
20311 assert_eq!(outcome, "key_spend_cap_reached");
20312 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
20313 let body = body_value(response).await;
20314 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
20315 assert!(
20316 body["error"]["message"].as_str().unwrap().contains("cap"),
20317 "the 402 must point at the KEY's cap, not tenant credit: {body}"
20318 );
20319 let events = mock.events();
20320 assert!(
20321 events.contains(&MeterEvent::Reserve {
20322 tenant: "acme".into(),
20323 principal: Some("mk-acme-testprefix00".into()),
20324 model: "qwen/qwen3.8-27b".into(),
20325 }),
20326 "the key prefix must reach reserve: {events:?}"
20327 );
20328 }
20329
20330 #[tokio::test]
20331 #[allow(clippy::await_holding_lock)] async fn streaming_client_disconnect_records_partial_usage_and_cost() {
20333 let _l = drain_lock();
20334 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
20335 let mock = MockMetering::admit_all();
20336 st.metering = Some(mock.clone());
20337
20338 let response = completions(
20339 State(st),
20340 HeaderMap::new(),
20341 None,
20342 Json(
20343 serde_json::from_value(json!({
20344 "model": "m",
20345 "prompt": "disconnect after one delta",
20346 "stream": true,
20347 }))
20348 .unwrap(),
20349 ),
20350 )
20351 .await;
20352 assert_eq!(response.status(), StatusCode::OK);
20353 let request_id = response.headers()["x-request-id"]
20354 .to_str()
20355 .unwrap()
20356 .to_string();
20357 let mut body = Box::pin(response.into_body().into_data_stream());
20358 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
20359 .await
20360 .expect("stream ended before first delta")
20361 .expect("stream body failed");
20362 assert!(
20363 is_sse_data_frame(&first),
20364 "first frame was not SSE data: {first:?}"
20365 );
20366 drop(body);
20367
20368 let mut dropped = None;
20371 for _ in 0..500 {
20372 if let Some(event) = mock
20373 .events()
20374 .into_iter()
20375 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
20376 {
20377 dropped = Some(event);
20378 break;
20379 }
20380 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
20381 }
20382 let events = mock.events();
20383 assert!(
20384 events
20385 .iter()
20386 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
20387 "the receipt was opened under the caller-visible request id: {events:?}"
20388 );
20389 assert_eq!(
20390 dropped,
20391 Some(MeterEvent::Dropped {
20392 prompt: 1,
20393 cached: 0,
20394 completion: 1,
20395 }),
20396 "a client disconnect must leave the partial counts on the dropped receipt \
20397 (the implementation prices that drop): {events:?}"
20398 );
20399 }
20400
20401 #[tokio::test]
20402 #[allow(clippy::await_holding_lock)] async fn draining_rejects_new_requests_with_503_and_retry_after() {
20404 let _l = drain_lock();
20405 let st = fake_worker_state();
20406 let _down = DrainingRestore;
20409 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
20410 let resp = chat_completions(
20412 State(st.clone()),
20413 axum::http::HeaderMap::new(),
20414 None,
20415 Json(
20416 serde_json::from_value(serde_json::json!({
20417 "model": "m", "messages": [{"role": "user", "content": "t"}]
20418 }))
20419 .unwrap(),
20420 ),
20421 )
20422 .await;
20423 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20424 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
20429 let ra_s: u64 = ra
20430 .parse()
20431 .expect("Retry-After must be integer delay-seconds");
20432 assert!(
20433 ra_s > 0 && ra_s <= 60,
20434 "Retry-After {ra_s}s is outside the honored window"
20435 );
20436 let ra_ms: u64 = resp.headers()["retry-after-ms"]
20437 .to_str()
20438 .unwrap()
20439 .parse()
20440 .unwrap();
20441 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
20442 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20443 .await
20444 .unwrap();
20445 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20446 assert!(
20447 payload["error"]["message"]
20448 .as_str()
20449 .unwrap()
20450 .contains("draining")
20451 );
20452 assert_eq!(payload["error"]["type"], "server_error");
20453 assert_eq!(payload["error"]["code"], "draining");
20454 let resp = completions(
20455 State(st.clone()),
20456 axum::http::HeaderMap::new(),
20457 None,
20458 Json(
20459 serde_json::from_value(serde_json::json!({
20460 "model": "m", "prompt": "t"
20461 }))
20462 .unwrap(),
20463 ),
20464 )
20465 .await;
20466 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20467 assert!(resp.headers().contains_key("retry-after"));
20468 assert_eq!(
20469 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
20470 0,
20471 "rejected requests must not hold slots"
20472 );
20473 let resp = health_live(State(st.clone())).await.into_response();
20476 assert_eq!(
20477 resp.status(),
20478 StatusCode::OK,
20479 "a drain must not look like a liveness fault"
20480 );
20481 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20482 .await
20483 .unwrap();
20484 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20485 assert_eq!(payload["status"], "draining");
20486 let resp = health_ready(State(st.clone())).await.into_response();
20488 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20489 let retry_s = drain_deadline_s().clamp(1, 60);
20490 let retry_s_text = retry_s.to_string();
20491 let retry_ms_text = (retry_s * 1000).to_string();
20492 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
20493 assert_eq!(
20494 resp.headers().get("retry-after-ms").unwrap(),
20495 retry_ms_text.as_str()
20496 );
20497 assert_ne!(
20498 resp.headers()
20499 .get("x-should-retry")
20500 .and_then(|v| v.to_str().ok()),
20501 Some("false")
20502 );
20503 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20504 .await
20505 .unwrap();
20506 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20507 assert_eq!(payload["status"], "not_ready");
20508 assert!(payload["detail"].as_str().unwrap().contains("draining"));
20509 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
20510 let resp = chat_completions(
20512 State(st.clone()),
20513 axum::http::HeaderMap::new(),
20514 None,
20515 Json(
20516 serde_json::from_value(serde_json::json!({
20517 "model": "m", "messages": [{"role": "user", "content": "t"}]
20518 }))
20519 .unwrap(),
20520 ),
20521 )
20522 .await;
20523 assert_eq!(resp.status(), StatusCode::OK);
20524 }
20525
20526 #[tokio::test]
20529 #[allow(clippy::await_holding_lock)] async fn health_is_green_only_while_the_worker_is_alive() {
20531 let _l = drain_lock();
20534 let st = fake_worker_state();
20535 let resp = health_live(State(st.clone())).await.into_response();
20538 assert_eq!(resp.status(), StatusCode::OK);
20539 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20540 .await
20541 .unwrap();
20542 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20543 assert_eq!(payload["status"], "ok");
20544 assert_eq!(payload["worker"]["phase"], "idle");
20545 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
20546 let ready = health_ready(State(st.clone())).await.into_response();
20547 assert_eq!(ready.status(), StatusCode::OK);
20548
20549 st.health.mark_dead("worker thread panicked: test-injected");
20553 let resp = health_live(State(st.clone())).await.into_response();
20554 assert_eq!(
20555 resp.status(),
20556 StatusCode::SERVICE_UNAVAILABLE,
20557 "a dead worker MUST NOT report a healthy liveness"
20558 );
20559 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20560 .await
20561 .unwrap();
20562 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20563 assert_eq!(payload["status"], "unhealthy");
20564 assert!(
20566 payload["detail"]
20567 .as_str()
20568 .unwrap()
20569 .contains("test-injected"),
20570 "cause not surfaced: {payload}"
20571 );
20572 let ready = health_ready(State(st.clone())).await.into_response();
20573 assert_eq!(
20574 ready.status(),
20575 StatusCode::SERVICE_UNAVAILABLE,
20576 "dead is also not ready"
20577 );
20578
20579 st.health.mark_ready();
20582 assert_eq!(
20583 health_live(State(st.clone()))
20584 .await
20585 .into_response()
20586 .status(),
20587 StatusCode::OK,
20588 "mark_ready must clear the latch (a successful respawn)"
20589 );
20590 }
20591
20592 #[tokio::test]
20593 #[allow(clippy::await_holding_lock)] async fn readyz_peer_probe_integrity_is_present_and_advisory() {
20595 let _l = drain_lock();
20596 let st = fake_worker_state();
20597
20598 let ready = health_ready(State(st.clone())).await.into_response();
20599 assert_eq!(ready.status(), StatusCode::OK);
20600 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
20601 .await
20602 .unwrap();
20603 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20604 assert_eq!(payload["peer_probe_integrity"], "ok");
20605
20606 st.health.note_peer_probe_deferral(2, false);
20607 let deferred = health_ready(State(st.clone())).await.into_response();
20608 assert_eq!(deferred.status(), StatusCode::OK);
20609 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
20610 .await
20611 .unwrap();
20612 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20613 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
20614
20615 st.health.note_peer_probe_deferral(4, true);
20616 let degraded = health_ready(State(st.clone())).await.into_response();
20617 assert_eq!(
20618 degraded.status(),
20619 StatusCode::OK,
20620 "peer degradation is advisory while plain serving remains healthy"
20621 );
20622 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
20623 .await
20624 .unwrap();
20625 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20626 assert_eq!(payload["peer_probe_integrity"], "degraded");
20627
20628 st.health.mark_dead("test-injected worker failure");
20629 let unready = health_ready(State(st)).await.into_response();
20630 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
20631 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
20632 .await
20633 .unwrap();
20634 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20635 assert_eq!(
20636 payload["peer_probe_integrity"], "degraded",
20637 "the advisory field must also survive an unrelated readiness failure"
20638 );
20639 }
20640
20641 #[tokio::test]
20642 #[allow(clippy::await_holding_lock)] async fn liveness_failure_obeys_the_retry_contract() {
20644 let _l = drain_lock();
20649 let st = fake_worker_state();
20650 st.health
20651 .mark_dead("worker thread panicked: retry-contract-test");
20652
20653 let resp = health_live(State(st)).await.into_response();
20654 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20655 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
20656 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
20657 assert_ne!(
20658 resp.headers()
20659 .get("x-should-retry")
20660 .and_then(|v| v.to_str().ok()),
20661 Some("false")
20662 );
20663 }
20664
20665 #[tokio::test]
20666 #[allow(clippy::await_holding_lock)] async fn readiness_failure_obeys_the_retry_contract() {
20668 let _l = drain_lock();
20669 let st = fake_worker_state();
20670 st.health
20671 .mark_dead("worker thread panicked: retry-contract-test");
20672
20673 let resp = health_ready(State(st)).await.into_response();
20674 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20675 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
20676 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
20677 assert_ne!(
20678 resp.headers()
20679 .get("x-should-retry")
20680 .and_then(|v| v.to_str().ok()),
20681 Some("false")
20682 );
20683 }
20684
20685 #[tokio::test]
20686 #[allow(clippy::await_holding_lock)] async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
20688 let _l = drain_lock();
20698 let st = fake_worker_state();
20699 assert_eq!(
20700 health_live(State(st.clone()))
20701 .await
20702 .into_response()
20703 .status(),
20704 StatusCode::OK
20705 );
20706 st.health
20707 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
20708 let resp = health_live(State(st.clone())).await.into_response();
20709 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
20710 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
20711 .await
20712 .unwrap();
20713 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20714 assert!(
20715 payload["detail"]
20716 .as_str()
20717 .unwrap()
20718 .contains("probe exceeded")
20719 );
20720 st.health.mark_ready();
20723 assert_eq!(
20724 health_live(State(st.clone()))
20725 .await
20726 .into_response()
20727 .status(),
20728 StatusCode::SERVICE_UNAVAILABLE,
20729 "a GPU fault must not be cleared by an in-process respawn"
20730 );
20731 }
20732
20733 #[test]
20734 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
20735 let caps = ModelCaps {
20737 tools_branch: true,
20738 hy3: false,
20739 qwen_think: true,
20740 think_switch: true,
20741 chat_ok: true,
20742 context_length: 262144,
20743 tokenizer: "qwen2".into(),
20744 instruct_type: Some("chatml".into()),
20745 effort_levels: false,
20746 qwen_effort: false,
20747 gemma_think: false,
20748 dsv4: false,
20749 glm5: false,
20750 chat_temperature_default: None,
20751 chat_top_p_default: None,
20752 n_vocab: 151_936,
20753 think_close: Vec::new(),
20754 };
20755 let e = model_entry_v1("main", Some(&caps), None);
20756 assert_eq!(e["id"], "main");
20757 assert_eq!(e["name"], "main");
20758 assert_eq!(e["object"], "model");
20759 assert_eq!(e["context_length"], 262144);
20760 assert!(e["pricing"]["input"].is_null());
20762 assert!(e["pricing"]["output"].is_null());
20763
20764 let meta = OpenRouterModelMetadata {
20767 pricing: OpenRouterPricing {
20768 prompt: Some("0.00000038".into()),
20769 cached_prompt: Some("0.0000002".into()),
20770 completion: Some("0.0000026".into()),
20771 ..Default::default()
20772 },
20773 input_modalities: vec!["image".into(), "video".into()],
20774 max_output_length: Some(32768),
20775 ..Default::default()
20776 };
20777 let e = model_entry_v1("main", Some(&caps), Some(&meta));
20778 assert_eq!(e["pricing"]["currency"], "USD");
20781 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
20782 assert_eq!(e["pricing"]["input"], "0.38");
20783 assert_eq!(e["pricing"]["output"], "2.60");
20784 assert_eq!(e["pricing"]["cached_input"], "0.20");
20785 assert!(e["pricing"]["cache_write"].is_null());
20786 assert_eq!(e["pricing"]["minimum_request"], "0");
20787 assert_eq!(e["owned_by"], "main");
20788 assert_eq!(e["type"], "chat");
20789 assert_eq!(e["max_output_tokens"], 32768);
20790 assert_eq!(e["endpoints"], json!(["chat/completions"]));
20791 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
20792 assert_eq!(e["output_modalities"], json!(["text"]));
20793 assert_eq!(e["capabilities"]["streaming"], true);
20794 assert_eq!(e["capabilities"]["tools"], true);
20795 assert_eq!(e["lifecycle"]["status"], "active");
20796 assert!(e["lifecycle"]["deprecation_at"].is_null());
20797 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
20798 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
20799 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
20803 keys.sort_unstable();
20804 assert_eq!(
20805 keys,
20806 [
20807 "capabilities",
20808 "context_length",
20809 "endpoints",
20810 "id",
20811 "input_modalities",
20812 "lifecycle",
20813 "max_output_tokens",
20814 "name",
20815 "object",
20816 "output_modalities",
20817 "owned_by",
20818 "pricing",
20819 "reliability",
20820 "type",
20821 ],
20822 "unexpected /v1/models entry keys"
20823 );
20824 let mut price_keys: Vec<&str> = e["pricing"]
20825 .as_object()
20826 .unwrap()
20827 .keys()
20828 .map(String::as_str)
20829 .collect();
20830 price_keys.sort_unstable();
20831 assert_eq!(
20832 price_keys,
20833 [
20834 "cache_write",
20835 "cached_input",
20836 "currency",
20837 "input",
20838 "minimum_request",
20839 "output",
20840 "unit",
20841 ],
20842 "unexpected /v1/models pricing keys"
20843 );
20844
20845 let e = model_entry_v1("m", None, None);
20847 assert!(e["context_length"].is_null());
20848 assert!(e["max_output_tokens"].is_null());
20849 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
20851 assert!(e["context_length"].is_null());
20852 }
20853
20854 #[test]
20860 fn catalog_row_follows_the_declared_surface() {
20861 let caps = ModelCaps {
20862 tools_branch: true,
20863 ..Default::default()
20864 };
20865
20866 let embed = OpenRouterModelMetadata {
20867 surface: Some("embedding".into()),
20868 max_output_length: Some(1),
20869 ..Default::default()
20870 };
20871 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
20872 assert_eq!(e["type"], "embedding");
20873 assert_eq!(e["endpoints"], json!(["embeddings"]));
20874 assert_eq!(e["output_modalities"], json!(["embeddings"]));
20875 assert_eq!(e["capabilities"]["streaming"], false);
20876 assert_eq!(
20877 e["capabilities"]["tools"], false,
20878 "an embedder has no tools"
20879 );
20880 assert_eq!(e["capabilities"]["reasoning"], false);
20881 assert_eq!(e["capabilities"]["structured_output"], false);
20882 assert_eq!(e["capabilities"]["prompt_caching"], false);
20883 assert!(
20884 e["max_output_tokens"].is_null(),
20885 "a surface that emits no completion tokens must not advertise a ceiling"
20886 );
20887
20888 let rerank = OpenRouterModelMetadata {
20889 surface: Some("rerank".into()),
20890 ..Default::default()
20891 };
20892 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
20893 assert_eq!(r["type"], "rerank");
20894 assert_eq!(r["endpoints"], json!(["rerank"]));
20895 assert_eq!(r["output_modalities"], json!(["rerank"]));
20896 assert_eq!(r["capabilities"]["tools"], false);
20897 assert_eq!(r["capabilities"]["reasoning"], false);
20898
20899 let chat = OpenRouterModelMetadata {
20902 max_output_length: Some(32768),
20903 ..Default::default()
20904 };
20905 let c = model_entry_v1("main", Some(&caps), Some(&chat));
20906 assert_eq!(c["type"], "chat");
20907 assert_eq!(c["endpoints"], json!(["chat/completions"]));
20908 assert_eq!(c["output_modalities"], json!(["text"]));
20909 assert_eq!(c["capabilities"]["tools"], true);
20910 assert_eq!(c["max_output_tokens"], 32768);
20911 }
20912
20913 #[test]
20916 fn unknown_surface_is_rejected_at_config_load() {
20917 let bad = OpenRouterModelMetadata {
20918 surface: Some("embeddings".into()), ..Default::default()
20920 };
20921 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
20922 .expect_err("an unknown surface must not load");
20923 assert!(err.contains("surface"), "{err}");
20924
20925 for good in ["chat", "embedding", "rerank"] {
20926 let ok = OpenRouterModelMetadata {
20927 surface: Some(good.into()),
20928 ..Default::default()
20929 };
20930 assert!(
20931 validate_openrouter_metadata("m", &ok).is_ok(),
20932 "{good} must load"
20933 );
20934 }
20935 }
20936
20937 #[test]
20938 fn per_million_price_is_exact_decimal_shift() {
20939 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
20941 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
20942 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
20943 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
20944 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
20945 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
20946 assert_eq!(per_million_price("not-a-price"), None);
20947 assert_eq!(per_million_price(""), None);
20948 }
20949
20950 #[test]
20951 fn metadata_provider_block_parses_and_validates() {
20952 let (_, provider) = OpenRouterMetadataFile::parse(
20953 r#"
20954 [provider]
20955 id = "tiyuvta"
20956 status_url = "https://status.tiyuvta.ai"
20957 support_contact = "mailto:support@tiyuvta.ai"
20958 incident_contact = "mailto:incidents@tiyuvta.ai"
20959 regions = ["eu-central"]
20960 "#,
20961 )
20962 .unwrap();
20963 let provider = provider.unwrap();
20964 assert_eq!(provider.id, "tiyuvta");
20965 assert_eq!(provider.regions, vec!["eu-central"]);
20966 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
20968 assert!(err.contains("provider.id"), "{err}");
20969 let err = OpenRouterMetadataFile::parse(
20971 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
20972 )
20973 .unwrap_err();
20974 assert!(err.contains("must be a URI"), "{err}");
20975 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
20977 assert!(provider.is_none());
20978 }
20979
20980 fn write_metadata_tmp(name: &str, body: &str) -> std::path::PathBuf {
20986 static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20987 let path = std::env::temp_dir().join(format!(
20988 "memra-metadata-reload-test-{}-{}-{name}",
20989 std::process::id(),
20990 COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
20991 ));
20992 std::fs::write(&path, body).unwrap();
20993 path
20994 }
20995
20996 fn reload_test_models() -> Vec<(String, String, Option<String>)> {
20997 vec![("m".to_string(), "x".to_string(), None)]
20998 }
20999
21000 fn reload_test_handle(path: &std::path::Path, set: ModelMetadataSet) -> MetadataReloadHandle {
21001 MetadataReloadHandle {
21002 cell: Arc::new(RwLock::new(Arc::new(set))),
21003 models: Arc::new(reload_test_models()),
21004 path: Some(path.to_path_buf()),
21005 }
21006 }
21007
21008 const RELOAD_V1: &str = r#"
21009[models.m]
21010[models.m.pricing]
21011prompt = "0.00000100"
21012completion = "0.00000300"
21013"#;
21014
21015 const RELOAD_V2: &str = r#"
21016[models.m]
21017[models.m.pricing]
21018prompt = "0.00000050"
21019completion = "0.00000150"
21020"#;
21021
21022 fn load_v1(path: &std::path::Path) -> ModelMetadataSet {
21023 load_model_metadata_file(path, &reload_test_models())
21024 .unwrap()
21025 .0
21026 }
21027
21028 #[test]
21029 fn metadata_reload_swaps_new_readers_and_keeps_inflight_on_old() {
21030 let path = write_metadata_tmp("swap.toml", RELOAD_V1);
21031 let handle = reload_test_handle(&path, load_v1(&path));
21032 let inflight = handle.cell.read().unwrap().clone();
21034 assert_eq!(
21035 inflight.models["m"].pricing.prompt.as_deref(),
21036 Some("0.00000100")
21037 );
21038
21039 std::fs::write(&path, RELOAD_V2).unwrap();
21040 let receipt = handle.reload().unwrap();
21041 assert_eq!(receipt.models, 1);
21042 assert!(!receipt.has_provider);
21043
21044 let live = handle.cell.read().unwrap().clone();
21045 assert_eq!(
21046 live.models["m"].pricing.prompt.as_deref(),
21047 Some("0.00000050"),
21048 "new readers see the reloaded set"
21049 );
21050 assert_eq!(
21051 inflight.models["m"].pricing.prompt.as_deref(),
21052 Some("0.00000100"),
21053 "the pre-swap Arc still serves the old set"
21054 );
21055 assert!(!Arc::ptr_eq(&inflight, &live), "the swap replaced the Arc");
21056 std::fs::remove_file(&path).unwrap();
21057 }
21058
21059 #[test]
21060 fn metadata_reload_rejects_invalid_file_and_keeps_old() {
21061 let path = write_metadata_tmp("bad.toml", RELOAD_V1);
21062 let handle = reload_test_handle(&path, load_v1(&path));
21063
21064 std::fs::write(&path, "[models.m]\n[models.m.pricing]\nprompt = \"free\"\n").unwrap();
21065 let err = handle.reload().unwrap_err();
21066 assert!(err.contains("must be a non-negative"), "{err}");
21067 assert_eq!(
21068 handle.cell.read().unwrap().models["m"]
21069 .pricing
21070 .prompt
21071 .as_deref(),
21072 Some("0.00000100"),
21073 "a failed reload keeps the old set"
21074 );
21075 std::fs::remove_file(&path).unwrap();
21076 }
21077
21078 #[test]
21079 fn metadata_reload_rejects_unknown_alias_like_boot() {
21080 let path = write_metadata_tmp("alias.toml", RELOAD_V1);
21081 let handle = reload_test_handle(&path, load_v1(&path));
21082
21083 std::fs::write(
21084 &path,
21085 "[models.ghost]\n[models.ghost.pricing]\nprompt = \"0.00000100\"\n",
21086 )
21087 .unwrap();
21088 let err = handle.reload().unwrap_err();
21089 assert!(err.contains("not present in MEMRA_MODELS"), "{err}");
21090 assert!(
21091 !handle.cell.read().unwrap().models.contains_key("ghost"),
21092 "a failed reload keeps the old set"
21093 );
21094 std::fs::remove_file(&path).unwrap();
21095 }
21096
21097 #[test]
21098 fn metadata_reload_receipt_hashes_file_bytes() {
21099 let path = write_metadata_tmp("receipt.toml", RELOAD_V1);
21100 let handle = reload_test_handle(&path, load_v1(&path));
21101 let receipt = handle.reload().unwrap();
21102 assert_eq!(receipt.models, 1);
21103 assert!(!receipt.has_provider);
21104 assert!(
21105 receipt.sha256.len() == 64 && receipt.sha256.bytes().all(|b| b.is_ascii_hexdigit()),
21106 "sha256 hex, got {:?}",
21107 receipt.sha256
21108 );
21109 assert_eq!(
21110 receipt.sha256,
21111 sha256_hex_bytes(&std::fs::read(&path).unwrap()),
21112 "the receipt hashes the file bytes"
21113 );
21114
21115 std::fs::write(&path, RELOAD_V2).unwrap();
21116 let receipt2 = handle.reload().unwrap();
21117 assert_ne!(
21118 receipt.sha256, receipt2.sha256,
21119 "equal receipts mean byte-identical files"
21120 );
21121 std::fs::remove_file(&path).unwrap();
21122 }
21123
21124 #[test]
21125 fn metadata_reload_provider_presence_is_reported() {
21126 let path = write_metadata_tmp(
21127 "provider.toml",
21128 "[provider]\nid = \"t\"\nstatus_url = \"https://status.example.invalid\"\n",
21129 );
21130 let set = ModelMetadataSet::default();
21131 let handle = reload_test_handle(&path, set);
21132 let receipt = handle.reload().unwrap();
21133 assert_eq!(receipt.models, 0);
21134 assert!(receipt.has_provider);
21135 assert!(handle.cell.read().unwrap().provider.is_some());
21136 std::fs::remove_file(&path).unwrap();
21137 }
21138
21139 #[test]
21140 fn metadata_reload_without_boot_path_errors() {
21141 let handle = MetadataReloadHandle {
21142 cell: Arc::new(RwLock::new(Arc::new(ModelMetadataSet::default()))),
21143 models: Arc::new(reload_test_models()),
21144 path: None,
21145 };
21146 let err = handle.reload().unwrap_err();
21147 assert!(err.contains("not configured"), "{err}");
21148 }
21149
21150 #[test]
21151 fn models_openai_default_body_stays_byte_identical() {
21152 let body = models_openai_body(&["main".into(), "judge".into()]);
21153 let bytes = serde_json::to_vec(&body).unwrap();
21154 assert_eq!(
21155 bytes,
21156 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
21157 );
21158 }
21159
21160 #[test]
21161 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
21162 let loaded = vec![
21164 "qwen/qwen3.6-27b".to_string(),
21165 "qwen/qwen3.6-35b-a3b".to_string(),
21166 ];
21167 assert_eq!(
21168 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
21169 Some("qwen/qwen3.6-35b-a3b"),
21170 );
21171 assert_eq!(
21172 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
21173 Some("qwen/qwen3.6-27b"),
21174 );
21175 assert_eq!(
21177 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
21178 Some("qwen/qwen3.6-35b-a3b"),
21179 );
21180 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
21182 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
21183 assert_eq!(canonical_model_id(&loaded, ""), None);
21184 }
21185
21186 #[test]
21187 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
21188 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
21191 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
21192 assert_eq!(
21194 canonical_model_id(&loaded, "a/shared-name").as_deref(),
21195 Some("a/shared-name")
21196 );
21197 assert_eq!(
21198 canonical_model_id(&loaded, "b/shared-name").as_deref(),
21199 Some("b/shared-name")
21200 );
21201 let bare = vec!["solo".to_string()];
21203 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
21204 }
21205
21206 #[test]
21207 fn openrouter_models_entry_serializes_complete_metadata() {
21208 let metadata = OpenRouterMetadataFile::from_toml(
21209 r#"
21210[models.main]
21211hugging_face_id = "Qwen/Qwen3.6-27B"
21212created = 1786032000
21213quantization = "nvfp4"
21214description = "Qwen3.6 27B served by memra."
21215max_prompt_length = 245760
21216max_output_length = 16384
21217default_output_length = 4096
21218is_ready = true
21219is_free = false
21220discount_to_user = 0.1
21221openrouter_slug = "qwen/qwen3.6-27b"
21222datacenters = [{ country_code = "US", region = "us-east" }]
21223zdr = true
21224hipaa = false
21225
21226[models.main.pricing]
21227prompt = "0.000000234"
21228cached_prompt = "0.0000000585"
21229cache_write = "0.000000234"
21230completion = "0.000001872"
21231internal_reasoning = "0.000001872"
21232request = "0.01"
21233
21234[models.main.capacity]
21235prompt_tpm = 1000000
21236cached_prompt_tpm = 2000000
21237completion_tpm = 500000
21238request_rpm = 1000
21239concurrency = 64
21240"#,
21241 )
21242 .unwrap();
21243 let caps = ModelCaps {
21244 tools_branch: true,
21245 qwen_think: true,
21246 think_switch: true,
21247 chat_ok: true,
21248 context_length: 262144,
21249 tokenizer: "qwen2".into(),
21250 instruct_type: Some("chatml".into()),
21251 ..Default::default()
21252 };
21253 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
21254
21255 assert_eq!(entry["schema_version"], "2.4");
21256 assert_eq!(entry["id"], "main");
21257 assert_eq!(entry["name"], "main");
21258 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
21259 assert_eq!(entry["created"], 1786032000u64);
21260 assert_eq!(entry["quantization"], "nvfp4");
21261 assert_eq!(entry["tokenizer"], "qwen2");
21262 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
21263 assert!(
21264 entry.get("object").is_none(),
21265 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
21266 );
21267
21268 let input = &entry["input_modalities"][0];
21269 assert_eq!(input["type"], "text");
21270 assert_eq!(
21271 input["supported_inputs"]["max_context_length"]["value"],
21272 262144
21273 );
21274 assert_eq!(
21275 input["supported_inputs"]["max_prompt_length"]["value"],
21276 245760
21277 );
21278 let input_prices = input["pricing"].as_array().unwrap();
21279 let input_price = |kind: &str| {
21280 input_prices
21281 .iter()
21282 .find(|price| price["type"] == kind)
21283 .unwrap()
21284 };
21285 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
21286 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
21287 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
21288 assert_eq!(input["capacity"][0]["value"], 1000000);
21289 assert_eq!(input["capacity"][1]["value"], 2000000);
21290
21291 let output = &entry["output_modalities"][0];
21292 assert_eq!(output["type"], "text");
21293 assert_eq!(output["max_length"]["value"], 16384);
21294 assert_eq!(output["streaming"], true);
21295 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
21296 assert_eq!(
21297 output["supported_parameters"]["structured_outputs"]["type"],
21298 "boolean"
21299 );
21300 assert_eq!(
21301 output["supported_parameters"]["reasoning"]["type"],
21302 "boolean"
21303 );
21304 assert_eq!(output["pricing"][0]["type"], "completion");
21305 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
21306 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
21307 assert_eq!(output["capacity"][0]["value"], 500000);
21308 assert_eq!(output["capacity"][1]["type"], "concurrency");
21309 assert_eq!(output["capacity"][1]["value"], 64);
21310
21311 assert_eq!(entry["pricing"][0]["type"], "request");
21312 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
21313 assert_eq!(entry["capacity"][0]["value"], 1000);
21314 assert_eq!(entry["is_ready"], true);
21315 assert_eq!(entry["is_free"], false);
21316 assert_eq!(entry["discount_to_user"], 0.1);
21317 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
21318 assert_eq!(entry["datacenters"][0]["country_code"], "US");
21319 assert_eq!(entry["compliance"]["zdr"], true);
21320 assert_eq!(entry["compliance"]["hipaa"], false);
21321 }
21322
21323 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
21328[models."qwen/qwen3.6-35b-a3b"]
21329hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
21330created = 1777260255
21331quantization = "int4"
21332description = "Qwen3.6 35B-A3B fixture entry."
21333max_prompt_length = 262144
21334max_output_length = 262144
21335default_output_length = 8192
21336is_ready = true
21337is_free = false
21338discount_to_user = 0.0
21339openrouter_slug = "qwen/qwen3.6-35b-a3b"
21340zdr = false
21341hipaa = false
21342
21343[[models."qwen/qwen3.6-35b-a3b".datacenters]]
21344country_code = "CA"
21345region = "Ontario"
21346
21347[models."qwen/qwen3.6-35b-a3b".pricing]
21348prompt = "0.0000000931"
21349cached_prompt = "0.0000000652"
21350completion = "0.0000009025"
21351
21352[models."qwen/qwen3.6-35b-a3b".capacity]
21353prompt_tpm = 780000
21354cached_prompt_tpm = 310000
21355completion_tpm = 9600
21356request_rpm = 160
21357concurrency = 16
21358
21359[planned_models."qwen/qwen3.8-27b"]
21360description = "Planned fixture entry; must never be emitted."
21361max_prompt_length = 262144
21362max_output_length = 262144
21363default_output_length = 8192
21364is_ready = false
21365is_free = false
21366discount_to_user = 0.0
21367openrouter_slug = "qwen/qwen3.8-27b"
21368zdr = false
21369hipaa = false
21370
21371[planned_models."qwen/qwen3.8-27b".pricing]
21372prompt = "0.0000002745"
21373cached_prompt = "0.0000001922"
21374completion = "0.0000022800"
21375
21376[planned_models."google/gemma-4-26b-a4b-it"]
21377hugging_face_id = "google/gemma-4-26B-A4B-it"
21378created = 1775227989
21379quantization = "int4"
21380description = "Planned fixture entry; must never be emitted."
21381max_prompt_length = 262144
21382max_output_length = 262144
21383default_output_length = 8192
21384is_ready = false
21385is_free = false
21386discount_to_user = 0.0
21387openrouter_slug = "google/gemma-4-26b-a4b-it"
21388zdr = false
21389hipaa = false
21390
21391[planned_models."google/gemma-4-26b-a4b-it".pricing]
21392prompt = "0.0000000665"
21393cached_prompt = "0.0000000466"
21394completion = "0.0000003230"
21395"#;
21396
21397 #[test]
21398 fn gateway_registry_generates_the_staged_active_shape() {
21399 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
21400 let caps = ModelCaps {
21401 tools_branch: true,
21402 qwen_think: true,
21403 think_switch: true,
21404 chat_ok: true,
21405 context_length: 262144,
21406 tokenizer: "qwen2".into(),
21407 instruct_type: Some("chatml".into()),
21408 ..Default::default()
21409 };
21410 let q35_entry = model_entry_openrouter(
21411 "qwen/qwen3.6-35b-a3b",
21412 Some(&caps),
21413 metadata.get("qwen/qwen3.6-35b-a3b"),
21414 );
21415 assert_eq!(q35_entry["created"], 1777260255u64);
21416 assert_eq!(q35_entry["quantization"], "int4");
21417 assert_eq!(q35_entry["is_ready"], true);
21418 assert_eq!(
21419 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
21420 262144
21421 );
21422 assert_eq!(
21423 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
21424 262144
21425 );
21426 assert_eq!(
21427 q35_entry["output_modalities"][0]["max_length"]["value"],
21428 262144
21429 );
21430 let prices = q35_entry["input_modalities"][0]["pricing"]
21431 .as_array()
21432 .unwrap();
21433 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
21434 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
21435 assert_eq!(
21439 q35_entry["input_modalities"][0]["capacity"][0]["value"],
21440 780000
21441 );
21442 assert_eq!(
21443 q35_entry["input_modalities"][0]["capacity"][1]["value"],
21444 310000
21445 );
21446 assert_eq!(
21447 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
21448 262144
21449 );
21450 assert_eq!(
21451 q35_entry["output_modalities"][0]["capacity"][0]["value"],
21452 9600
21453 );
21454 assert_eq!(
21455 q35_entry["output_modalities"][0]["capacity"][1]["value"],
21456 16
21457 );
21458 assert_eq!(
21459 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
21460 "0.0000009025"
21461 );
21462 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
21464
21465 assert_eq!(
21466 metadata.len(),
21467 1,
21468 "planned models must never enter the active map"
21469 );
21470 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
21471 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
21472 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
21473
21474 let openmodels = model_entry_openmodels(
21475 "qwen/qwen3.6-35b-a3b",
21476 Some(&caps),
21477 metadata.get("qwen/qwen3.6-35b-a3b"),
21478 )
21479 .unwrap();
21480 assert_eq!(openmodels["currency"], "USD");
21481 assert_eq!(openmodels["max_output_length"], 262144);
21482 assert_eq!(openmodels["is_ready"], true);
21483 assert_eq!(openmodels["is_free"], false);
21484 assert_eq!(openmodels["discount_to_user"], 0.0);
21485 }
21486
21487 #[test]
21488 fn gateway_registry_limits_are_live_request_limits() {
21489 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
21490 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
21491 let caps = ModelCaps {
21492 context_length: 262_144,
21493 ..Default::default()
21494 };
21495 let build = |value: serde_json::Value| {
21496 let req: CompletionReq = serde_json::from_value(value).unwrap();
21497 let (tx, _rx) = worker::event_channel();
21498 build_request(&req, tx, lanes::Lane::Interactive, None)
21499 };
21500
21501 let mut omitted = build(json!({
21502 "model": "qwen/qwen3.6-35b-a3b",
21503 "prompt_ids": [1, 2, 3]
21504 }));
21505 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
21506 assert_eq!(omitted.params.max_new, 8_192);
21507 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
21508
21509 let mut field_top = build(json!({
21510 "model": "qwen/qwen3.6-35b-a3b",
21511 "prompt_ids": [1],
21512 "max_tokens": 262144
21513 }));
21514 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
21515 assert_eq!(field_top.params.max_new, 262_144);
21516 assert_eq!(
21517 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
21518 262_044,
21519 "the field-top output request is accepted but bounded by remaining trained context",
21520 );
21521
21522 let mut too_much_output = build(json!({
21523 "model": "qwen/qwen3.6-35b-a3b",
21524 "prompt_ids": [1],
21525 "max_tokens": 262145
21526 }));
21527 let (message, param) =
21528 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
21529 .unwrap_err();
21530 assert_eq!(param, "max_tokens");
21531 assert!(message.contains("262145"));
21532
21533 let mut oversized_allocation = build(json!({
21534 "model": "qwen/qwen3.6-35b-a3b",
21535 "prompt_ids": [1],
21536 "max_tokens": 1,
21537 "max_ctx": 262145
21538 }));
21539 let (_, param) =
21540 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
21541 .unwrap_err();
21542 assert_eq!(param, "max_ctx");
21543 }
21544
21545 #[test]
21546 fn planned_registry_entries_are_validated_but_never_activated() {
21547 let parsed = OpenRouterMetadataFile::from_toml(
21548 r#"
21549[planned_models.future]
21550max_output_length = 262144
21551default_output_length = 8192
21552
21553[planned_models.future.pricing]
21554prompt = "0.0000001"
21555"#,
21556 )
21557 .unwrap();
21558 assert!(parsed.is_empty());
21559
21560 let error = OpenRouterMetadataFile::from_toml(
21561 r#"
21562[planned_models.future]
21563default_output_length = 8192
21564"#,
21565 )
21566 .unwrap_err();
21567 assert!(error.contains("requires max_output_length"));
21568 }
21569
21570 #[test]
21575 fn every_catalog_feed_honours_the_declared_surface() {
21576 let metadata = OpenRouterMetadataFile::from_toml(
21577 r#"
21578[models."qwen/qwen3-embedding-8b"]
21579surface = "embedding"
21580created = 1787961600
21581max_output_length = 1
21582is_ready = true
21583is_free = false
21584discount_to_user = 0.0
21585
21586[models."qwen/qwen3-embedding-8b".pricing]
21587prompt = "0.00000001"
21588cached_prompt = "0.0"
21589completion = "0.0"
21590
21591[models."main"]
21592created = 1787443200
21593max_output_length = 32768
21594is_ready = true
21595is_free = false
21596discount_to_user = 0.0
21597
21598[models."main".pricing]
21599prompt = "0.00000025"
21600cached_prompt = "0.00000009"
21601completion = "0.0000012"
21602"#,
21603 )
21604 .unwrap();
21605 let caps = ModelCaps {
21606 tools_branch: true,
21607 qwen_think: true,
21608 think_switch: true,
21613 chat_ok: true,
21614 context_length: 32768,
21615 ..Default::default()
21616 };
21617 let embed = metadata.get("qwen/qwen3-embedding-8b");
21618 let chat = metadata.get("main");
21619
21620 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
21622 let out = &or["output_modalities"][0];
21623 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
21624 assert!(
21625 out.get("streaming").is_none(),
21626 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
21627 );
21628 let params = &out["supported_parameters"];
21633 assert_eq!(
21634 params.as_object().map(|o| o.len()),
21635 Some(0),
21636 "no completion parameter belongs on an embedder row: {params}"
21637 );
21638 for field in [
21639 "tools",
21640 "tool_choice",
21641 "reasoning",
21642 "max_tokens",
21643 "json_mode",
21644 "structured_outputs",
21645 "stop",
21646 "temperature",
21647 "seed",
21648 ] {
21649 assert!(params[field].is_null(), "{field} leaked onto an embedder");
21650 }
21651 assert!(
21652 out["max_length"].is_null(),
21653 "a surface emitting no completion tokens advertises no ceiling: {out}"
21654 );
21655
21656 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
21658 .expect("openmodels entry builds");
21659 assert_eq!(om["output_modalities"], json!(["embeddings"]));
21660 let features = om["supported_features"].as_array().unwrap();
21661 assert!(
21662 !features
21663 .iter()
21664 .any(|f| f == "tool_calling" || f == "reasoning"),
21665 "chat-only features leaked onto an embedder: {features:?}"
21666 );
21667
21668 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
21670 assert_eq!(v1["type"], "embedding");
21671 assert_eq!(v1["capabilities"]["tools"], false);
21672
21673 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
21675 let out_chat = &or_chat["output_modalities"][0];
21676 assert_eq!(out_chat["type"], "text");
21677 assert_eq!(out_chat["streaming"], true);
21678 assert!(!out_chat["supported_parameters"]["tools"].is_null());
21679 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
21680 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
21681 assert_eq!(out_chat["max_length"]["value"], 32768u64);
21682 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
21683 assert_eq!(om_chat["output_modalities"], json!(["text"]));
21684 assert!(
21685 om_chat["supported_features"]
21686 .as_array()
21687 .unwrap()
21688 .iter()
21689 .any(|f| f == "tool_calling")
21690 );
21691 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
21692 }
21693
21694 #[test]
21701 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
21702 let raw = std::fs::read_to_string(concat!(
21703 env!("CARGO_MANIFEST_DIR"),
21704 "/../../research/gateway-20260812/raw/sources/",
21705 "openrouter-provider-schema-v2.4-20260812.json"
21706 ))
21707 .expect("vendored Provider Monitor 2.4 schema is in-tree");
21708 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
21709 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
21710 .as_array()
21711 .expect("OutputModality is a oneOf");
21712
21713 let metadata = OpenRouterMetadataFile::from_toml(
21714 r#"
21715[models."embed"]
21716surface = "embedding"
21717created = 1787961600
21718max_output_length = 1
21719is_ready = true
21720is_free = false
21721discount_to_user = 0.0
21722
21723[models."embed".pricing]
21724prompt = "0.00000001"
21725cached_prompt = "0.0"
21726completion = "0.0"
21727
21728[models."rr"]
21729surface = "rerank"
21730created = 1787961600
21731max_output_length = 1
21732is_ready = true
21733is_free = false
21734discount_to_user = 0.0
21735
21736[models."rr".pricing]
21737prompt = "0.00000003"
21738cached_prompt = "0.0"
21739completion = "0.0"
21740
21741[models."chatty"]
21742created = 1787443200
21743max_output_length = 32768
21744is_ready = true
21745is_free = false
21746discount_to_user = 0.0
21747
21748[models."chatty".pricing]
21749prompt = "0.00000025"
21750cached_prompt = "0.00000009"
21751completion = "0.0000012"
21752"#,
21753 )
21754 .unwrap();
21755 let caps = ModelCaps {
21756 tools_branch: true,
21757 qwen_think: true,
21758 chat_ok: true,
21759 context_length: 32768,
21760 ..Default::default()
21761 };
21762
21763 for (alias, want_type) in [
21764 ("embed", "embeddings"),
21765 ("rr", "rerank"),
21766 ("chatty", "text"),
21767 ] {
21768 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
21769 let modality = &row["output_modalities"][0];
21770 assert_eq!(modality["type"], want_type, "{alias}: {row}");
21771
21772 let branch = branches
21774 .iter()
21775 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
21776 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
21777 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
21778 .as_object()
21779 .expect("branch properties")
21780 .keys()
21781 .map(String::as_str)
21782 .collect();
21783 for key in modality.as_object().expect("modality object").keys() {
21784 assert!(
21785 allowed.contains(key.as_str()),
21786 "{alias}: {key:?} is not a property of the {want_type:?} branch \
21787 (additionalProperties:false); allowed = {allowed:?}"
21788 );
21789 }
21790 for req in branch["required"].as_array().into_iter().flatten() {
21791 let req = req.as_str().expect("required entry is a string");
21792 assert!(
21793 modality.get(req).is_some(),
21794 "{alias}: required property {req:?} missing from the {want_type:?} branch"
21795 );
21796 }
21797 }
21798 }
21799
21800 #[test]
21801 fn openrouter_models_entry_omits_undeclared_optional_fields() {
21802 let entry = model_entry_openrouter("minimal", None, None);
21803 let object = entry.as_object().unwrap();
21804 for field in [
21805 "hugging_face_id",
21806 "created",
21807 "quantization",
21808 "tokenizer",
21809 "description",
21810 "pricing",
21811 "capacity",
21812 "is_ready",
21813 "is_free",
21814 "discount_to_user",
21815 "openrouter",
21816 "datacenters",
21817 "compliance",
21818 ] {
21819 assert!(
21820 !object.contains_key(field),
21821 "optional field {field} must be absent, not null"
21822 );
21823 }
21824 assert_eq!(entry["schema_version"], "2.4");
21825 assert_eq!(entry["input_modalities"][0]["type"], "text");
21826 assert!(
21827 entry["input_modalities"][0]
21828 .get("supported_inputs")
21829 .is_none()
21830 );
21831 assert!(entry["input_modalities"][0].get("pricing").is_none());
21832 assert!(entry["input_modalities"][0].get("capacity").is_none());
21833 assert_eq!(entry["output_modalities"][0]["type"], "text");
21834 assert_eq!(entry["output_modalities"][0]["streaming"], true);
21835 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
21836 assert!(entry["output_modalities"][0].get("max_length").is_none());
21837 assert!(entry["output_modalities"][0].get("pricing").is_none());
21838 assert!(entry["output_modalities"][0].get("capacity").is_none());
21839 }
21840
21841 #[test]
21842 fn openmodels_entry_serializes_standard_provider_shape() {
21843 let metadata = OpenRouterMetadataFile::from_toml(
21844 r#"
21845[models."qwen/qwen3.6-27b"]
21846created = 1786032000
21847max_output_length = 16384
21848is_ready = true
21849is_free = false
21850discount_to_user = 0.05
21851
21852[models."qwen/qwen3.6-27b".pricing]
21853prompt = "0.000000291"
21854cached_prompt = "0.000000291"
21855completion = "0.000002763"
21856request = "0"
21857"#,
21858 )
21859 .unwrap();
21860 let caps = ModelCaps {
21861 tools_branch: true,
21862 qwen_think: true,
21863 chat_ok: true,
21864 context_length: 262144,
21865 ..Default::default()
21866 };
21867 let entry = model_entry_openmodels(
21868 "qwen/qwen3.6-27b",
21869 Some(&caps),
21870 metadata.get("qwen/qwen3.6-27b"),
21871 )
21872 .unwrap();
21873
21874 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
21875 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
21876 assert_eq!(entry["created"], 1786032000u64);
21877 assert_eq!(entry["input_modalities"], json!(["text"]));
21878 assert_eq!(entry["output_modalities"], json!(["text"]));
21879 assert_eq!(entry["context_length"], 262144u64);
21880 assert_eq!(entry["max_output_length"], 16384u64);
21881 assert_eq!(entry["currency"], "USD");
21882 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
21883 assert_eq!(entry["pricing"]["completion"], "0.000002763");
21884 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
21885 assert_eq!(entry["pricing"]["request"], "0");
21886 assert_eq!(
21887 entry["supported_features"],
21888 json!(["tool_calling", "reasoning"])
21889 );
21890 assert_eq!(entry["is_ready"], true);
21891 assert_eq!(entry["is_free"], false);
21892 assert_eq!(entry["discount_to_user"], 0.05);
21893 assert!(entry.get("schema_version").is_none());
21894 assert!(entry.get("quantization").is_none());
21895 }
21896
21897 #[test]
21898 fn openmodels_entry_rejects_missing_operator_metadata() {
21899 let caps = ModelCaps {
21900 context_length: 262144,
21901 ..Default::default()
21902 };
21903 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
21904 assert_eq!(
21905 error,
21906 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
21907 );
21908 }
21909
21910 #[tokio::test]
21911 async fn blocking_response_excludes_stop_text_across_token_events() {
21912 let (tx, rx) = worker::event_channel();
21913 tx.send(Event::Token {
21914 id: 1,
21915 text: "answer\nPro".into(),
21916 })
21917 .unwrap();
21918 tx.send(Event::Token {
21919 id: 2,
21920 text: "blem: leaked prompt".into(),
21921 })
21922 .unwrap();
21923 tx.send(Event::Done {
21924 stop_reason: "Callback".into(),
21925 n_tokens: 2,
21926 n_prompt: 8,
21927 n_cached: 0,
21928 elapsed_s: 0.5,
21929 spec: None,
21930 })
21931 .unwrap();
21932 drop(tx);
21933 let response = blocking_response(
21934 rx,
21935 "plain_quant".into(),
21936 false,
21937 vec!["Problem:".into()],
21938 None,
21939 Envelope::new(false),
21940 )
21941 .await;
21942 assert_eq!(response.status(), StatusCode::OK);
21943 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
21944 .await
21945 .unwrap();
21946 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
21947 assert_eq!(payload["text"], "answer\n");
21948 assert_eq!(payload["stop_reason"], "Callback");
21949 }
21950
21951 #[test]
21955 fn step_walker_expansion_and_separator_law() {
21956 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
21958 let uri = format!("data:image/png;base64,{PNG64}");
21959 let content = serde_json::json!([
21960 {"type": "text", "text": "look at"},
21961 {"type": "text", "text": "this:"},
21962 {"type": "image_url", "image_url": {"url": uri}},
21963 {"type": "text", "text": "what is it?"},
21964 ]);
21965 let mut pending: Vec<PendingStepImage> = Vec::new();
21966 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
21967 let mut expansion = String::from("<im_start>");
21968 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
21969 expansion.push_str("<im_patch>");
21970 }
21971 expansion.push_str("<im_end>");
21972 assert_eq!(out, format!("look at this:{expansion}what is it?"));
21975 assert_eq!(pending.len(), 1);
21976 assert_eq!(pending[0].plan.n_tiles, 0);
21977 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
21978
21979 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
21981 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
21982 let http = serde_json::json!([
21983 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
21984 ]);
21985 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
21986 }
21987}
21988
21989#[cfg(test)]
21995mod build_identity_tests {
21996 use super::{BUILD_GIT_SHA, BUILD_ID_NOTE, BUILD_ID_SRC, SYSTEM_FINGERPRINT, build_id};
21997
21998 #[test]
22000 fn baked_fingerprint_is_real_and_well_formed() {
22001 assert!(!SYSTEM_FINGERPRINT.is_empty());
22002 assert_ne!(SYSTEM_FINGERPRINT, "memra-unknown");
22003 assert!(
22004 !SYSTEM_FINGERPRINT.contains("unknown"),
22005 "fingerprint {SYSTEM_FINGERPRINT:?} still carries the degraded literal"
22006 );
22007 assert!(
22008 build_id::fingerprint_is_well_formed(SYSTEM_FINGERPRINT),
22009 "fingerprint {SYSTEM_FINGERPRINT:?} is not memra-<version>-<12 hex>"
22010 );
22011 assert!(
22014 SYSTEM_FINGERPRINT.starts_with(concat!("memra-", env!("CARGO_PKG_VERSION"), "-")),
22015 "fingerprint {SYSTEM_FINGERPRINT:?} does not name this crate version"
22016 );
22017 }
22018
22019 #[test]
22022 fn the_shape_check_rejects_what_shipped_to_prod() {
22023 assert!(!build_id::fingerprint_is_well_formed("memra-unknown"));
22024 assert!(!build_id::fingerprint_is_well_formed(
22025 "memra-0.123.0-unknown"
22026 ));
22027 let old_form = format!("memra-{}", "0".repeat(12));
22032 assert!(!build_id::fingerprint_is_well_formed(&old_form));
22033 assert!(!build_id::fingerprint_is_well_formed(""));
22034 assert!(!build_id::fingerprint_is_well_formed("memra-"));
22035 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-"));
22036 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-abc"));
22038 assert!(!build_id::fingerprint_is_well_formed(
22039 "memra-0.123.0-ABCDEF012345"
22040 ));
22041 assert!(!build_id::fingerprint_is_well_formed(
22042 "memra-0.123.0-zzzzzzzzzzzz"
22043 ));
22044 assert!(build_id::fingerprint_is_well_formed(
22046 "memra-0.123.0-4b1f9c02d7a3"
22047 ));
22048 }
22049
22050 #[test]
22058 fn build_id_is_rederivable_from_the_source_tree() {
22059 let root = build_id::workspace_root(env!("CARGO_MANIFEST_DIR"));
22060 let scan = root.as_deref().and_then(build_id::content_id);
22061 match scan {
22062 Some(scan) => {
22063 assert_eq!(
22064 BUILD_ID_SRC,
22065 build_id::BUILD_ID_SRC_TREE,
22066 "the source tree is readable, so the baked id must come from it"
22067 );
22068 assert!(BUILD_ID_NOTE.is_empty(), "note set on a non-degraded build");
22069 let expected =
22070 format!(concat!("memra-", env!("CARGO_PKG_VERSION"), "-{}"), scan.id);
22071 assert_eq!(
22072 SYSTEM_FINGERPRINT,
22073 expected,
22074 "baked fingerprint disagrees with a re-derivation over {} files: the id \
22075 is not a pure function of the source tree, or the build script did not \
22076 re-run after an edit",
22077 scan.files.len()
22078 );
22079 assert!(scan.files.len() > 100, "suspiciously small hashed file set");
22080 }
22081 None => {
22082 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
22086 assert!(
22087 !BUILD_ID_NOTE.is_empty(),
22088 "a degraded build must state its reason so the boot WARN can print it"
22089 );
22090 }
22091 }
22092 }
22093
22094 #[test]
22097 fn identity_is_independent_of_git_history() {
22098 let id = SYSTEM_FINGERPRINT.rsplit_once('-').unwrap().1;
22099 assert_ne!(
22100 id, BUILD_GIT_SHA,
22101 "the content id equals the git sha; the identity must not be history, it has to \
22102 survive a rewrite that changes every commit"
22103 );
22104 assert!(
22105 !SYSTEM_FINGERPRINT.contains(BUILD_GIT_SHA),
22106 "the git sha leaked into the customer-visible fingerprint {SYSTEM_FINGERPRINT:?}"
22107 );
22108 assert!(!BUILD_GIT_SHA.is_empty());
22111 }
22112
22113 #[test]
22116 fn content_digest_is_deterministic_and_change_sensitive() {
22117 let a = build_id::degraded_build_id("memra-server", "0.123.0");
22118 let b = build_id::degraded_build_id("memra-server", "0.123.0");
22119 assert_eq!(a, b, "the digest is not deterministic");
22120 assert_eq!(a.len(), build_id::BUILD_ID_HEX);
22121 assert!(
22122 a.chars()
22123 .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
22124 );
22125 assert_ne!(a, build_id::degraded_build_id("memra-server", "0.123.1"));
22126 assert_ne!(a, build_id::degraded_build_id("memra-serve", "r0.123.0"));
22127 assert_eq!(build_id::render_build_id(0).len(), build_id::BUILD_ID_HEX);
22129 assert_eq!(
22130 build_id::render_build_id(0),
22131 "0".repeat(build_id::BUILD_ID_HEX)
22132 );
22133 }
22134
22135 #[test]
22138 fn two_scans_of_one_tree_agree() {
22139 let Some(root) = build_id::workspace_root(env!("CARGO_MANIFEST_DIR")) else {
22140 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
22141 return;
22142 };
22143 let first = build_id::content_id(&root).expect("first scan");
22144 let second = build_id::content_id(&root).expect("second scan");
22145 assert_eq!(first.id, second.id);
22146 assert_eq!(first.files.len(), second.files.len());
22147 }
22148}
22149
22150#[cfg(test)]
22156mod vision_placement_gate_tests {
22157 use super::vision_media_admissible;
22158
22159 #[test]
22160 fn a_media_part_is_admitted_only_when_the_placement_admits() {
22161 assert_eq!(vision_media_admissible(true, "image"), Ok(()));
22162 assert_eq!(vision_media_admissible(true, "video"), Ok(()));
22163 let err = vision_media_admissible(false, "image").unwrap_err();
22164 assert!(
22165 err.starts_with("image input is not enabled on this deployment"),
22166 "same named refusal the armed-off path gives, so clients see one contract: {err}"
22167 );
22168 assert!(
22169 err.contains("placement"),
22170 "the refusal names its cause: {err}"
22171 );
22172 let err = vision_media_admissible(false, "video").unwrap_err();
22173 assert!(
22174 err.starts_with("video input is not enabled on this deployment"),
22175 "{err}"
22176 );
22177 }
22178
22179 fn live_src() -> String {
22180 let src: String = include_str!("lib.rs")
22181 .lines()
22182 .map(|l| match l.find("//") {
22183 Some(i) => &l[..i],
22184 None => l,
22185 })
22186 .collect::<Vec<_>>()
22187 .join("\n");
22188 let end = src
22189 .find("\nmod vision_placement_gate_tests")
22190 .expect("this test module exists");
22191 src[..end].to_string()
22192 }
22193
22194 fn item_body<'a>(live: &'a str, head: &str) -> &'a str {
22196 let start = live
22197 .find(head)
22198 .unwrap_or_else(|| panic!("{head} not found — did it get renamed?"));
22199 let body = &live[start..];
22200 let end = body.find("\n}\n").expect("item body closes");
22201 &body[..end]
22202 }
22203
22204 fn head_of(s: &str, n: usize) -> &str {
22206 match s.char_indices().nth(n) {
22207 Some((i, _)) => &s[..i],
22208 None => s,
22209 }
22210 }
22211
22212 #[test]
22217 fn no_family_switch_reads_the_placement_decision() {
22218 let live = live_src();
22219 for switch in [
22220 "fn vision_enabled()",
22221 "fn gemma_vision_enabled()",
22222 "fn step_vision_enabled()",
22223 ] {
22224 let body = item_body(&live, switch);
22225 assert!(
22226 !body.contains("vision_placement_serving")
22227 && !body.contains("vision_placement_admits"),
22228 "{switch} routes text rendering; it must stay keyed on the operator knobs alone"
22229 );
22230 }
22231 let walker = item_body(&live, "fn content_to_text_vision(");
22232 assert!(
22233 walker.contains(
22234 "if step_vision_enabled() {\n return content_to_text_vision_step(v, step_images);"
22235 ),
22236 "the step walker dispatch is keyed on the armed switch alone"
22237 );
22238 }
22239
22240 #[test]
22243 fn every_media_accepting_arm_passes_the_placement_gate() {
22244 let live = live_src();
22245 let step = item_body(&live, "fn content_to_text_vision_step(");
22246 let arm = step
22247 .split("Some(\"image_url\") => {")
22248 .nth(1)
22249 .expect("the step walker has an image arm");
22250 assert!(
22251 head_of(arm, 120).contains("vision_placement_admits(\"image\")?;"),
22252 "the step image arm must pass the placement gate first: {}",
22253 head_of(arm, 120)
22254 );
22255 let walker = item_body(&live, "fn content_to_text_vision(");
22256 for (head, kind) in [
22257 (
22258 "Some(\"image_url\") if gemma_vision_enabled() => {",
22259 "image",
22260 ),
22261 ("Some(\"image_url\") => {", "image"),
22262 ("Some(\"video_url\") => {", "video"),
22263 ] {
22264 let arm = walker
22265 .split(head)
22266 .nth(1)
22267 .unwrap_or_else(|| panic!("{head} is not an arm of the walker"));
22268 let window = head_of(arm, 400);
22269 assert!(
22270 window.contains(&format!("vision_placement_admits(\"{kind}\")?;")),
22271 "{head} must pass the placement gate before planning anything: {window}"
22272 );
22273 }
22274 assert!(live.contains("GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)"));
22278 let gate = item_body(&live, "fn vision_placement_admits(");
22280 assert!(gate.contains("vision_media_admissible(vision_placement_serving(), kind)"));
22281 }
22282}