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 responses_api;
104mod surfaces;
105mod toolcall;
106mod ttft;
107mod worker;
108
109use std::collections::HashMap;
110use std::net::{SocketAddr, ToSocketAddrs};
111use std::sync::mpsc::Sender;
112use std::sync::{Arc, Mutex};
113
114use axum::{
115 Extension, Json, Router,
116 body::Body,
117 extract::{DefaultBodyLimit, FromRequest, Query, Request as AxumRequest, State},
118 http::{
119 HeaderMap, StatusCode,
120 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
121 },
122 middleware::{self, Next},
123 response::{
124 IntoResponse, Response,
125 sse::{Event as SseEvent, Sse},
126 },
127 routing::{get, post},
128};
129use futures_core::Stream as _;
130use serde::de::DeserializeOwned;
131use serde::{Deserialize, Serialize};
132use serde_json::json;
133use tower::ServiceExt as _;
134
135use memra_engine::decode::GenParams;
136use memra_engine::sampler::SamplerConfig;
137use memra_tokenizer::{
138 Tokenizer,
139 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
140};
141use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
142use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
143
144const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
168const MAX_BODY_ADMISSIONS: usize = 4;
169const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
170#[allow(clippy::identity_op)] const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
175const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
176const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
177const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
178const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
179const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
180const MAX_STOP_SEQUENCES: usize = 16;
181const MAX_STOP_SEQUENCE_BYTES: usize = 1_024;
182const MAX_STOP_SEQUENCES_BYTES: usize = 4 * 1_024;
183const MAX_CLIENT_IDENTIFIER_BYTES: usize = 256;
184const MAX_HTTP_CONNECTIONS: usize = 1_024;
185const MAX_HTTP2_STREAMS_PER_CONNECTION: u32 = 128;
186const HTTP1_HEADER_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
187const HTTP_CONNECTION_MAX_LIFETIME: std::time::Duration = std::time::Duration::from_secs(300);
188
189fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
190 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
191 SEMAPHORE
192 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
193 .clone()
194}
195
196fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
197 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
198 SEMAPHORE
199 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
200 .clone()
201}
202
203#[derive(Clone)]
204pub(crate) struct BodyAdmissionGuard {
205 permit: Arc<Mutex<Option<tokio::sync::OwnedSemaphorePermit>>>,
206}
207
208impl BodyAdmissionGuard {
209 fn new(permit: tokio::sync::OwnedSemaphorePermit) -> Self {
210 Self {
211 permit: Arc::new(Mutex::new(Some(permit))),
212 }
213 }
214
215 pub(crate) fn release(&self) {
216 if let Ok(mut permit) = self.permit.lock() {
217 permit.take();
218 }
219 }
220}
221
222pub(crate) struct BodyAdmissionLease(Option<BodyAdmissionGuard>);
223
224impl BodyAdmissionLease {
225 fn release(&mut self) {
226 if let Some(admission) = self.0.take() {
227 admission.release();
228 }
229 }
230
231 pub(crate) fn guard(&self) -> Option<&BodyAdmissionGuard> {
232 self.0.as_ref()
233 }
234}
235
236impl Drop for BodyAdmissionLease {
237 fn drop(&mut self) {
238 self.release();
239 }
240}
241
242pub(crate) struct AdmittedJson<T>(pub(crate) T, pub(crate) BodyAdmissionLease);
243
244#[axum::async_trait]
245impl<S, T> FromRequest<S> for AdmittedJson<T>
246where
247 S: Send + Sync,
248 T: DeserializeOwned,
249{
250 type Rejection = axum::extract::rejection::JsonRejection;
251
252 async fn from_request(req: AxumRequest, state: &S) -> Result<Self, Self::Rejection> {
253 let admission = req.extensions().get::<BodyAdmissionGuard>().cloned();
254 let parsed = Json::<T>::from_request(req, state).await;
255 parsed.map(|Json(value)| Self(value, BodyAdmissionLease(admission)))
256 }
257}
258
259fn declared_body_length(req: &AxumRequest) -> Option<usize> {
260 req.headers()
261 .get(CONTENT_LENGTH)
262 .and_then(|value| value.to_str().ok())
263 .and_then(|value| value.parse().ok())
264}
265
266fn body_requires_admission(req: &AxumRequest) -> bool {
267 if req.headers().contains_key(TRANSFER_ENCODING) {
271 return true;
272 }
273 declared_body_length(req).is_none_or(|length| length > BODY_ADMISSION_BYPASS_BYTES)
274}
275
276fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
280 let Some(length) = declared_body_length(req) else {
281 return BODY_READ_TIMEOUT;
282 };
283 let bytes = length as u64;
284 let extra_seconds =
285 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
286 let seconds = BODY_READ_TIMEOUT
287 .as_secs()
288 .saturating_add(extra_seconds)
289 .min(BODY_READ_TIMEOUT_MAX.as_secs());
290 std::time::Duration::from_secs(seconds)
291}
292
293async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
298 let resp = next.run(req).await;
299 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
300 return resp;
301 }
302 error_response_coded(
303 StatusCode::PAYLOAD_TOO_LARGE,
304 &format!(
305 "request body exceeds the {} MiB limit",
306 MAX_BODY_BYTES / (1024 * 1024)
307 ),
308 "invalid_request_error",
309 None,
310 Some("request_too_large"),
311 )
312}
313
314fn apply_body_limit(app: Router) -> Router {
317 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
318 .layer(middleware::from_fn(shape_payload_too_large))
319}
320
321fn protected_inference_path(path: &str) -> bool {
322 matches!(
323 path,
324 "/v1/auth/check"
325 | "/v1/completions"
326 | "/v1/chat/completions"
327 | "/v1/messages"
328 | "/v1/responses"
329 | "/v1/embeddings"
330 | "/v1/rerank"
331 )
332}
333
334async fn shape_inference_early_response(path: &str, response: Response) -> Response {
338 let request_id = Envelope::new(path != "/v1/completions");
339 if path == "/v1/messages" {
340 anthropic::with_anthropic_request_id(
341 &request_id.id,
342 anthropic::reshape_error(response, &request_id.id).await,
343 )
344 } else {
345 with_request_id(&request_id.id, response)
346 }
347}
348
349async fn authenticate_inference_before_body(
354 State(st): State<AppState>,
355 mut req: AxumRequest,
356 next: Next,
357) -> Response {
358 if !protected_inference_path(req.uri().path()) {
359 return next.run(req).await;
360 }
361 let path = req.uri().path().to_string();
362 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
366 return shape_inference_early_response(
367 &path,
368 error_response_coded(
369 StatusCode::PAYLOAD_TOO_LARGE,
370 &format!(
371 "request body exceeds the {} MiB limit",
372 MAX_BODY_BYTES / (1024 * 1024)
373 ),
374 "invalid_request_error",
375 None,
376 Some("request_too_large"),
377 ),
378 )
379 .await;
380 }
381 let headers = req.headers();
382 let bearer = bearer_token(headers);
383 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
384 let api_key = headers
385 .get("x-api-key")
386 .and_then(|value| value.to_str().ok());
387 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
388 } else {
389 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
390 };
391 if let Err(why) = auth {
392 return shape_inference_early_response(&path, authentication_error(why)).await;
393 }
394 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
402 let body_admission = if body_requires_admission(&req) {
403 body_admission_semaphore()
404 } else {
405 small_body_admission_semaphore()
406 };
407 let body_permit = match body_admission.try_acquire_owned() {
408 Ok(permit) => permit,
409 Err(tokio::sync::TryAcquireError::Closed) => {
410 let response = retry_contract_response(
411 error_response_coded(
412 StatusCode::SERVICE_UNAVAILABLE,
413 "request body admission is unavailable",
414 "server_error",
415 None,
416 Some("body_admission_unavailable"),
417 ),
418 Some(BODY_ADMISSION_RETRY_AFTER_S),
419 );
420 return shape_inference_early_response(&path, response).await;
421 }
422 Err(tokio::sync::TryAcquireError::NoPermits) => {
423 let response = retry_contract_response(
424 error_response_coded(
425 StatusCode::TOO_MANY_REQUESTS,
426 "request body admission is busy",
427 "rate_limit_error",
428 None,
429 Some("body_admission_busy"),
430 ),
431 Some(BODY_ADMISSION_RETRY_AFTER_S),
432 );
433 return shape_inference_early_response(&path, response).await;
434 }
435 };
436 let body_admission_guard = BodyAdmissionGuard::new(body_permit);
441 req.extensions_mut().insert(body_admission_guard.clone());
442 let body = std::mem::replace(req.body_mut(), Body::empty());
443 let mut body = Box::pin(body.into_data_stream());
444 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
445 let body_timed_out_flag = body_timed_out.clone();
446 let guarded_body = async_stream::stream! {
447 loop {
448 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
449 if remaining.is_zero() {
450 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
451 yield Err(std::io::Error::new(
452 std::io::ErrorKind::TimedOut,
453 "request body read deadline exceeded",
454 ));
455 break;
456 }
457 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
458 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
459 Ok(frame) => frame,
460 Err(_) => {
461 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
462 yield Err(std::io::Error::new(
463 std::io::ErrorKind::TimedOut,
464 "request body idle timeout exceeded",
465 ));
466 break;
467 }
468 };
469 match frame {
470 Some(Ok(bytes)) => yield Ok(bytes),
471 Some(Err(error)) => {
472 yield Err(std::io::Error::other(error.to_string()));
473 break;
474 }
475 None => break,
476 }
477 }
478 };
479 *req.body_mut() = Body::from_stream(guarded_body);
480 let response = next.run(req).await;
481 body_admission_guard.release();
482 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
483 let request_id = Envelope::new(path != "/v1/completions");
484 let timeout = error_response_coded(
485 StatusCode::REQUEST_TIMEOUT,
486 "request body read timed out",
487 "invalid_request_error",
488 None,
489 Some("request_body_timeout"),
490 );
491 return if path == "/v1/messages" {
492 anthropic::with_anthropic_request_id(
493 &request_id.id,
494 anthropic::reshape_error(timeout, &request_id.id).await,
495 )
496 } else {
497 with_request_id(&request_id.id, timeout)
498 };
499 }
500 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
501 let request_id = Envelope::new(true);
502 return anthropic::with_anthropic_request_id(
503 &request_id.id,
504 anthropic::reshape_error(response, &request_id.id).await,
505 );
506 }
507 response
508}
509
510#[cfg(test)]
511mod body_limit_tests {
512 use super::*;
513
514 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
515
516 fn test_app() -> Router {
520 let app = Router::new()
521 .route(
522 "/bytes",
523 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
524 )
525 .route(
526 "/json",
527 post(
528 |AdmittedJson(v, _admission): AdmittedJson<serde_json::Value>| async move {
529 v["pad"].as_str().unwrap_or("").len().to_string()
530 },
531 ),
532 );
533 apply_body_limit(app)
534 }
535
536 fn streamed_body(chunks: usize) -> Body {
537 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
540 Body::from_stream(async_stream::stream! {
541 for _ in 0..chunks {
542 yield Ok::<_, std::io::Error>(chunk.clone());
543 }
544 })
545 }
546
547 #[tokio::test]
548 async fn bodies_past_the_old_2mib_default_are_accepted() {
549 for (path, body) in [
552 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
553 (
554 "/json",
555 Body::from(
556 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
557 ),
558 ),
559 ] {
560 let resp = test_app()
561 .oneshot(
562 axum::http::Request::post(path)
563 .header(CONTENT_TYPE, "application/json")
564 .body(body)
565 .unwrap(),
566 )
567 .await
568 .unwrap();
569 assert_eq!(resp.status(), StatusCode::OK, "{path}");
570 }
571 }
572
573 #[tokio::test]
574 async fn body_at_exactly_the_limit_is_accepted() {
575 let resp = test_app()
576 .oneshot(
577 axum::http::Request::post("/bytes")
578 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
579 .unwrap(),
580 )
581 .await
582 .unwrap();
583 assert_eq!(resp.status(), StatusCode::OK);
584 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
585 .await
586 .unwrap();
587 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
588 }
589
590 #[tokio::test]
591 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
592 for path in ["/bytes", "/json"] {
596 let resp = test_app()
597 .oneshot(
598 axum::http::Request::post(path)
599 .header(CONTENT_TYPE, "application/json")
600 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
601 .unwrap(),
602 )
603 .await
604 .unwrap();
605 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
606 assert_eq!(
607 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
608 Some(b"false".as_ref()),
609 "{path}: retrying identical bytes cannot fix a 413"
610 );
611 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
612 .await
613 .unwrap();
614 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
615 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
616 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
617 assert!(
618 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
619 "{path}: message names the limit"
620 );
621 }
622 }
623
624 #[tokio::test]
625 async fn authenticated_body_admission_is_finite() {
626 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
627 let semaphore = body_admission_semaphore();
628 let mut permits = Vec::new();
629 for _ in 0..MAX_BODY_ADMISSIONS {
630 permits.push(semaphore.clone().acquire_owned().await.unwrap());
631 }
632 assert!(
633 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
634 .await
635 .is_err(),
636 "body parser admission must not be unbounded"
637 );
638 drop(permits);
639 assert!(semaphore.acquire().await.is_ok());
640 }
641
642 #[tokio::test]
643 async fn small_body_admission_is_finite_and_separate() {
644 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
645 let large = body_admission_semaphore();
646 let small = small_body_admission_semaphore();
647 let mut small_permits = Vec::new();
648 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
649 small_permits.push(small.clone().acquire_owned().await.unwrap());
650 }
651 assert!(
652 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
653 .await
654 .is_err(),
655 "small body parser admission must be bounded"
656 );
657 assert!(
658 large.clone().try_acquire().is_ok(),
659 "small uploads must not consume large-upload permits"
660 );
661 drop(small_permits);
662 assert!(small.acquire().await.is_ok());
663 }
664
665 #[test]
666 fn small_declared_bodies_bypass_large_upload_admission() {
667 let request = axum::http::Request::post("/v1/chat/completions")
668 .header(CONTENT_LENGTH, "2048")
669 .body(Body::empty())
670 .unwrap();
671 assert!(!body_requires_admission(&request));
672
673 let request = axum::http::Request::post("/v1/chat/completions")
674 .header(
675 CONTENT_LENGTH,
676 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
677 )
678 .body(Body::empty())
679 .unwrap();
680 assert!(body_requires_admission(&request));
681
682 let request = axum::http::Request::post("/v1/chat/completions")
683 .header(CONTENT_LENGTH, "2048")
684 .header(TRANSFER_ENCODING, "chunked")
685 .body(Body::empty())
686 .unwrap();
687 assert!(body_requires_admission(&request));
688 }
689
690 #[test]
691 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
692 let unknown = axum::http::Request::post("/v1/chat/completions")
693 .body(Body::empty())
694 .unwrap();
695 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
696
697 let large = axum::http::Request::post("/v1/chat/completions")
698 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
699 .body(Body::empty())
700 .unwrap();
701 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
702 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
703
704 let absurd = axum::http::Request::post("/v1/chat/completions")
705 .header(CONTENT_LENGTH, u64::MAX.to_string())
706 .body(Body::empty())
707 .unwrap();
708 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
709 }
710
711 #[tokio::test]
712 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
713 let too_large = shape_inference_early_response(
714 "/v1/messages",
715 error_response_coded(
716 StatusCode::PAYLOAD_TOO_LARGE,
717 "request body exceeds the 192 MiB limit",
718 "invalid_request_error",
719 None,
720 Some("request_too_large"),
721 ),
722 )
723 .await;
724 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
725 let house_id = too_large.headers()["x-request-id"].clone();
726 assert_eq!(too_large.headers()["request-id"], house_id);
727 assert_eq!(too_large.headers()["x-should-retry"], "false");
728 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
729 .await
730 .unwrap();
731 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
732 assert_eq!(payload["type"], "error");
733 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
734
735 let busy = shape_inference_early_response(
736 "/v1/chat/completions",
737 retry_contract_response(
738 error_response_coded(
739 StatusCode::TOO_MANY_REQUESTS,
740 "request body admission is busy",
741 "rate_limit_error",
742 None,
743 Some("body_admission_busy"),
744 ),
745 Some(BODY_ADMISSION_RETRY_AFTER_S),
746 ),
747 )
748 .await;
749 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
750 assert!(!busy.headers()["x-request-id"].is_empty());
751 assert_eq!(busy.headers()["retry-after"], "1");
752 assert_eq!(busy.headers()["retry-after-ms"], "1000");
753 assert!(busy.headers().get("x-should-retry").is_none());
754 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
755 .await
756 .unwrap();
757 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
758 assert_eq!(payload["error"]["code"], "body_admission_busy");
759 }
760
761 #[tokio::test]
762 async fn vision_preprocess_admission_is_fail_fast_and_retryable() {
763 let semaphore = Box::leak(Box::new(tokio::sync::Semaphore::new(1)));
764 let held = semaphore.try_acquire().unwrap();
765 let busy = try_vision_preprocess_with(true, semaphore).unwrap_err();
766 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
767 assert_eq!(busy.headers()["retry-after"], "1");
768 drop(held);
769 assert!(
770 try_vision_preprocess_with(true, semaphore)
771 .unwrap()
772 .is_some()
773 );
774 assert!(
775 try_vision_preprocess_with(false, semaphore)
776 .unwrap()
777 .is_none()
778 );
779 }
780
781 #[tokio::test]
782 async fn typed_json_retains_body_admission_until_handler_validation_releases_it() {
783 #[derive(Clone)]
784 struct Signals {
785 parsed: Arc<tokio::sync::Notify>,
786 finish: Arc<tokio::sync::Notify>,
787 semaphore: Arc<tokio::sync::Semaphore>,
788 }
789
790 let semaphore = Arc::new(tokio::sync::Semaphore::new(1));
791 let guard = BodyAdmissionGuard::new(semaphore.clone().try_acquire_owned().unwrap());
792 let signals = Signals {
793 parsed: Arc::new(tokio::sync::Notify::new()),
794 finish: Arc::new(tokio::sync::Notify::new()),
795 semaphore: semaphore.clone(),
796 };
797 let app = Router::new()
798 .route(
799 "/",
800 post(
801 |Extension(signals): Extension<Signals>,
802 AdmittedJson(_, mut admission): AdmittedJson<serde_json::Value>| async move {
803 assert_eq!(
804 signals.semaphore.available_permits(),
805 0,
806 "typed deserialization alone must not release post-parse admission"
807 );
808 admission.release();
809 assert_eq!(signals.semaphore.available_permits(), 1);
810 signals.parsed.notify_one();
811 signals.finish.notified().await;
812 "ok"
813 },
814 ),
815 )
816 .layer(Extension(signals.clone()))
817 .layer(Extension(guard));
818 let response = tokio::spawn(
819 app.oneshot(
820 axum::http::Request::post("/")
821 .header(CONTENT_TYPE, "application/json")
822 .body(Body::from(r#"{"value":1}"#))
823 .unwrap(),
824 ),
825 );
826 signals.parsed.notified().await;
827 assert_eq!(
828 semaphore.available_permits(),
829 1,
830 "validated work must release admission before generation waits"
831 );
832 signals.finish.notify_one();
833 assert_eq!(response.await.unwrap().unwrap().status(), StatusCode::OK);
834 }
835
836 #[tokio::test]
837 async fn transport_closes_stalled_headers_and_caps_connections() {
838 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
839
840 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
841 let address = listener.local_addr().unwrap();
842 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
843 let server = tokio::spawn(serve_bounded_http_with_limits(
844 listener,
845 Router::new().route("/", get(|| async { "ok" })).route(
846 "/slow",
847 get(|| async {
848 tokio::time::sleep(std::time::Duration::from_millis(140)).await;
849 "slow-ok"
850 }),
851 ),
852 async move {
853 let _ = shutdown_rx.await;
854 },
855 std::time::Duration::from_millis(30),
856 1,
857 std::time::Duration::from_millis(80),
858 ));
859
860 let mut stalled = tokio::net::TcpStream::connect(address).await.unwrap();
861 stalled.write_all(b"GET / HT").await.unwrap();
862 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
863 let mut excess = tokio::net::TcpStream::connect(address).await.unwrap();
864 let mut bytes = Vec::new();
865 tokio::time::timeout(
866 std::time::Duration::from_millis(250),
867 excess.read_to_end(&mut bytes),
868 )
869 .await
870 .expect("connection beyond the cap must be closed promptly")
871 .unwrap();
872
873 bytes.clear();
874 tokio::time::timeout(
875 std::time::Duration::from_millis(500),
876 stalled.read_to_end(&mut bytes),
877 )
878 .await
879 .expect("stalled request headers must hit the configured deadline")
880 .unwrap();
881
882 let mut idle = tokio::net::TcpStream::connect(address).await.unwrap();
883 idle.write_all(b"GET / HTTP/1.1\r\nHost: local\r\n\r\n")
884 .await
885 .unwrap();
886 bytes.clear();
887 tokio::time::timeout(
888 std::time::Duration::from_millis(500),
889 idle.read_to_end(&mut bytes),
890 )
891 .await
892 .expect("an idle keep-alive connection must hit the maximum lifetime")
893 .unwrap();
894 assert!(String::from_utf8_lossy(&bytes).contains("200 OK"));
895
896 let mut active = tokio::net::TcpStream::connect(address).await.unwrap();
897 active
898 .write_all(b"GET /slow HTTP/1.1\r\nHost: local\r\n\r\n")
899 .await
900 .unwrap();
901 bytes.clear();
902 tokio::time::timeout(
903 std::time::Duration::from_millis(500),
904 active.read_to_end(&mut bytes),
905 )
906 .await
907 .expect("an active response must finish across the connection age boundary")
908 .unwrap();
909 let active_response = String::from_utf8_lossy(&bytes);
910 assert!(active_response.contains("200 OK"), "{active_response}");
911 assert!(active_response.contains("slow-ok"), "{active_response}");
912
913 let h2_stream = tokio::net::TcpStream::connect(address).await.unwrap();
917 let (mut h2_client, h2_connection) = h2::client::handshake(h2_stream).await.unwrap();
918 let h2_driver = tokio::spawn(h2_connection);
919 let request = axum::http::Request::builder()
920 .uri(format!("http://{address}/"))
921 .body(())
922 .unwrap();
923 let (response, _) = h2_client.send_request(request, true).unwrap();
924 let response = tokio::time::timeout(std::time::Duration::from_millis(500), response)
925 .await
926 .expect("HTTP/2 handshake and response must complete")
927 .expect("HTTP/2 connection must stay alive through the response");
928 assert_eq!(response.status(), StatusCode::OK);
929 drop(h2_client);
930 h2_driver.abort();
931 let _ = h2_driver.await;
932
933 let _ = shutdown_tx.send(());
934 server.await.unwrap().unwrap();
935 }
936}
937
938#[derive(Clone, Default)]
939struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
940
941fn is_sse_data_frame(bytes: &[u8]) -> bool {
942 bytes
943 .windows(b"data:".len())
944 .any(|window| window == b"data:")
945}
946
947async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
948 let trace = ttft::start(req.uri().path());
949 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
950 let response = next.run(req).await;
951 let Some(trace) = trace else {
952 return response;
953 };
954 let is_sse = response
955 .headers()
956 .get(CONTENT_TYPE)
957 .and_then(|value| value.to_str().ok())
958 .is_some_and(|value| value.starts_with("text/event-stream"));
959 if !is_sse {
960 return response;
961 }
962
963 let (parts, body) = response.into_parts();
966 let mut body = Box::pin(body.into_data_stream());
967 let stream = async_stream::stream! {
968 while let Some(frame) =
969 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
970 {
971 if frame
972 .as_ref()
973 .is_ok_and(|bytes| is_sse_data_frame(bytes))
974 {
975 trace.mark_first_sse_byte();
976 }
977 yield frame;
978 }
979 };
980 Response::from_parts(parts, Body::from_stream(stream))
981}
982
983const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
984const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
985
986#[derive(Debug, Clone, Default, Deserialize)]
987#[serde(deny_unknown_fields)]
988struct OpenRouterMetadataFile {
989 #[serde(default)]
990 models: HashMap<String, OpenRouterModelMetadata>,
991 #[serde(default)]
994 planned_models: HashMap<String, OpenRouterModelMetadata>,
995 #[serde(default)]
998 provider: Option<ProviderMetadata>,
999}
1000
1001#[derive(Debug, Clone, Deserialize)]
1005#[serde(deny_unknown_fields)]
1006struct ProviderMetadata {
1007 id: String,
1008 #[serde(default)]
1009 status_url: Option<String>,
1010 #[serde(default)]
1011 support_contact: Option<String>,
1012 #[serde(default)]
1013 incident_contact: Option<String>,
1014 #[serde(default)]
1015 regions: Vec<String>,
1016}
1017
1018#[derive(Debug, Clone, Default, Deserialize)]
1020#[serde(deny_unknown_fields)]
1021struct LifecycleMetadata {
1022 #[serde(default)]
1023 status: Option<String>,
1024 #[serde(default)]
1025 deprecation_at: Option<String>,
1026 #[serde(default)]
1027 retirement_at: Option<String>,
1028 #[serde(default)]
1029 replacement_model_id: Option<String>,
1030}
1031
1032#[derive(Debug, Clone, Default, Deserialize)]
1034#[serde(deny_unknown_fields)]
1035struct ReliabilityMetadata {
1036 #[serde(default)]
1037 first_token_timeout_seconds: Option<u64>,
1038 #[serde(default)]
1039 completion_timeout_seconds: Option<u64>,
1040 #[serde(default)]
1041 stream_idle_timeout_seconds: Option<u64>,
1042 #[serde(default)]
1043 capacity_scope: Option<String>,
1044}
1045
1046#[derive(Debug, Clone, Default, Deserialize)]
1047#[serde(deny_unknown_fields)]
1048struct OpenRouterModelMetadata {
1049 #[serde(default)]
1051 owned_by: Option<String>,
1052 #[serde(default)]
1053 lifecycle: Option<LifecycleMetadata>,
1054 #[serde(default)]
1055 reliability: Option<ReliabilityMetadata>,
1056 #[serde(default)]
1057 hugging_face_id: Option<String>,
1058 #[serde(default)]
1059 created: Option<u64>,
1060 #[serde(default)]
1061 quantization: Option<String>,
1062 #[serde(default)]
1063 description: Option<String>,
1064 #[serde(default)]
1065 max_prompt_length: Option<u64>,
1066 #[serde(default)]
1067 max_output_length: Option<u64>,
1068 #[serde(default)]
1071 default_output_length: Option<u64>,
1072 #[serde(default)]
1073 pricing: OpenRouterPricing,
1074 #[serde(default)]
1075 capacity: OpenRouterCapacity,
1076 #[serde(default)]
1077 is_ready: Option<bool>,
1078 #[serde(default)]
1079 is_free: Option<bool>,
1080 #[serde(default)]
1081 discount_to_user: Option<f64>,
1082 #[serde(default)]
1083 openrouter_slug: Option<String>,
1084 #[serde(default)]
1085 datacenters: Vec<OpenRouterDatacenter>,
1086 #[serde(default)]
1090 input_modalities: Vec<String>,
1091 #[serde(default)]
1106 surface: Option<String>,
1107 #[serde(default)]
1108 zdr: Option<bool>,
1109 #[serde(default)]
1110 hipaa: Option<bool>,
1111 #[serde(default)]
1121 default_reasoning_effort: Option<String>,
1122 #[serde(default)]
1140 default_temperature: Option<f32>,
1141 #[serde(default)]
1142 default_top_p: Option<f32>,
1143 #[serde(default)]
1145 default_top_k: Option<usize>,
1146 #[serde(default)]
1147 default_min_p: Option<f32>,
1148 #[serde(default)]
1149 default_presence_penalty: Option<f32>,
1150 #[serde(default)]
1151 default_frequency_penalty: Option<f32>,
1152 #[serde(default)]
1154 default_repetition_penalty: Option<f32>,
1155 #[serde(default)]
1176 non_thinking_sampling: Option<SamplingArmMetadata>,
1177}
1178
1179#[derive(Debug, Clone, Default, Deserialize)]
1185#[serde(deny_unknown_fields)]
1186struct SamplingArmMetadata {
1187 #[serde(default)]
1188 temperature: Option<f32>,
1189 #[serde(default)]
1190 top_p: Option<f32>,
1191 #[serde(default)]
1192 top_k: Option<usize>,
1193 #[serde(default)]
1194 min_p: Option<f32>,
1195 #[serde(default)]
1196 presence_penalty: Option<f32>,
1197 #[serde(default)]
1198 frequency_penalty: Option<f32>,
1199 #[serde(default)]
1200 repetition_penalty: Option<f32>,
1201}
1202
1203impl SamplingArmMetadata {
1204 fn is_empty(&self) -> bool {
1205 self.temperature.is_none()
1206 && self.top_p.is_none()
1207 && self.top_k.is_none()
1208 && self.min_p.is_none()
1209 && self.presence_penalty.is_none()
1210 && self.frequency_penalty.is_none()
1211 && self.repetition_penalty.is_none()
1212 }
1213}
1214
1215#[derive(Debug, Clone, Default, Deserialize)]
1216#[serde(deny_unknown_fields)]
1217struct OpenRouterPricing {
1218 #[serde(default)]
1219 prompt: Option<String>,
1220 #[serde(default)]
1221 cached_prompt: Option<String>,
1222 #[serde(default)]
1223 cache_write: Option<String>,
1224 #[serde(default)]
1225 completion: Option<String>,
1226 #[serde(default)]
1227 internal_reasoning: Option<String>,
1228 #[serde(default)]
1229 request: Option<String>,
1230}
1231
1232#[derive(Debug, Clone, Default, Deserialize)]
1233#[serde(deny_unknown_fields)]
1234struct OpenRouterCapacity {
1235 #[serde(default)]
1236 prompt_tpm: Option<u64>,
1237 #[serde(default)]
1238 cached_prompt_tpm: Option<u64>,
1239 #[serde(default)]
1240 completion_tpm: Option<u64>,
1241 #[serde(default)]
1242 request_rpm: Option<u64>,
1243 #[serde(default)]
1244 concurrency: Option<u64>,
1245}
1246
1247#[derive(Debug, Clone, Deserialize, Serialize)]
1248#[serde(deny_unknown_fields)]
1249struct OpenRouterDatacenter {
1250 country_code: String,
1251 #[serde(default, skip_serializing_if = "Option::is_none")]
1252 region: Option<String>,
1253}
1254
1255impl OpenRouterMetadataFile {
1256 fn parse(
1257 text: &str,
1258 ) -> Result<
1259 (
1260 HashMap<String, OpenRouterModelMetadata>,
1261 Option<ProviderMetadata>,
1262 ),
1263 String,
1264 > {
1265 let file: Self =
1266 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
1267 for (alias, metadata) in &file.models {
1268 validate_openrouter_metadata(alias, metadata)?;
1269 }
1270 for (alias, metadata) in &file.planned_models {
1271 validate_openrouter_metadata(alias, metadata)?;
1272 if file.models.contains_key(alias) {
1273 return Err(format!(
1274 "model alias {alias:?} appears in both models and planned_models"
1275 ));
1276 }
1277 }
1278 if let Some(provider) = &file.provider {
1279 if provider.id.is_empty() {
1280 return Err("provider.id must be a non-empty slug".into());
1281 }
1282 for (field, value) in [
1284 ("provider.support_contact", &provider.support_contact),
1285 ("provider.incident_contact", &provider.incident_contact),
1286 ] {
1287 if let Some(value) = value
1288 && !value.contains(':')
1289 {
1290 return Err(format!(
1291 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1292 ));
1293 }
1294 }
1295 }
1296 Ok((file.models, file.provider))
1297 }
1298
1299 #[cfg(test)]
1300 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1301 Self::parse(text).map(|(models, _)| models)
1302 }
1303}
1304
1305fn per_million_price(per_token: &str) -> Option<String> {
1309 if !valid_price_string(per_token) {
1310 return None;
1311 }
1312 let (whole, frac) = match per_token.split_once('.') {
1313 Some((whole, frac)) => (whole, frac),
1314 None => (per_token, ""),
1315 };
1316 let mut digits = format!("{whole}{frac}");
1317 let point = whole.len() + 6;
1318 while digits.len() < point {
1319 digits.push('0');
1320 }
1321 let (int_part, frac_part) = digits.split_at(point);
1322 let int_part = int_part.trim_start_matches('0');
1323 let int_part = if int_part.is_empty() { "0" } else { int_part };
1324 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1325 while frac_out.len() < 2 {
1326 frac_out.push('0');
1327 }
1328 Some(format!("{int_part}.{frac_out}"))
1329}
1330
1331fn valid_price_string(value: &str) -> bool {
1332 let mut parts = value.split('.');
1333 let whole = parts.next().unwrap_or_default();
1334 let fraction = parts.next();
1335 !whole.is_empty()
1336 && whole.bytes().all(|b| b.is_ascii_digit())
1337 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1338 && parts.next().is_none()
1339}
1340
1341fn validate_openrouter_metadata(
1342 alias: &str,
1343 metadata: &OpenRouterModelMetadata,
1344) -> Result<(), String> {
1345 if alias.is_empty() {
1346 return Err("models metadata contains an empty model alias".into());
1347 }
1348 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1351 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1352 {
1353 return Err(format!(
1354 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1355 reasoning_effort level (none|minimal|low|medium|high)"
1356 ));
1357 }
1358 validate_sampling_defaults(alias, metadata)?;
1359 for m in &metadata.input_modalities {
1360 if m != "image" && m != "video" {
1361 return Err(format!(
1362 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1363 ));
1364 }
1365 }
1366 if let Some(sfc) = metadata.surface.as_deref()
1367 && !matches!(sfc, "chat" | "embedding" | "rerank")
1368 {
1369 return Err(format!(
1370 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1371 ));
1372 }
1373 if let Some(q) = metadata.quantization.as_deref()
1374 && !matches!(
1375 q,
1376 "int4"
1377 | "int8"
1378 | "fp4"
1379 | "mxfp4"
1380 | "nvfp4"
1381 | "fp6"
1382 | "fp8"
1383 | "mxfp8"
1384 | "fp16"
1385 | "bf16"
1386 | "fp32"
1387 )
1388 {
1389 return Err(format!(
1390 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1391 ));
1392 }
1393 for (field, value) in [
1394 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1395 (
1396 "pricing.cached_prompt",
1397 metadata.pricing.cached_prompt.as_deref(),
1398 ),
1399 (
1400 "pricing.cache_write",
1401 metadata.pricing.cache_write.as_deref(),
1402 ),
1403 ("pricing.completion", metadata.pricing.completion.as_deref()),
1404 (
1405 "pricing.internal_reasoning",
1406 metadata.pricing.internal_reasoning.as_deref(),
1407 ),
1408 ("pricing.request", metadata.pricing.request.as_deref()),
1409 ] {
1410 if let Some(value) = value
1411 && !valid_price_string(value)
1412 {
1413 return Err(format!(
1414 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1415 ));
1416 }
1417 }
1418 for (field, value) in [
1419 ("created", metadata.created),
1420 ("max_prompt_length", metadata.max_prompt_length),
1421 ("max_output_length", metadata.max_output_length),
1422 ("default_output_length", metadata.default_output_length),
1423 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1424 (
1425 "capacity.cached_prompt_tpm",
1426 metadata.capacity.cached_prompt_tpm,
1427 ),
1428 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1429 ("capacity.request_rpm", metadata.capacity.request_rpm),
1430 ("capacity.concurrency", metadata.capacity.concurrency),
1431 ] {
1432 if let Some(value) = value
1433 && value > JSON_SAFE_INTEGER_MAX
1434 {
1435 return Err(format!(
1436 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1437 ));
1438 }
1439 }
1440 for (field, value) in [
1441 ("max_prompt_length", metadata.max_prompt_length),
1442 ("max_output_length", metadata.max_output_length),
1443 ("default_output_length", metadata.default_output_length),
1444 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1445 (
1446 "capacity.cached_prompt_tpm",
1447 metadata.capacity.cached_prompt_tpm,
1448 ),
1449 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1450 ("capacity.request_rpm", metadata.capacity.request_rpm),
1451 ("capacity.concurrency", metadata.capacity.concurrency),
1452 ] {
1453 if value == Some(0) {
1454 return Err(format!(
1455 "model {alias:?}: {field} must be greater than zero when declared"
1456 ));
1457 }
1458 }
1459 if let (Some(default), Some(maximum)) =
1460 (metadata.default_output_length, metadata.max_output_length)
1461 && default > maximum
1462 {
1463 return Err(format!(
1464 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1465 ));
1466 }
1467 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1468 return Err(format!(
1469 "model {alias:?}: default_output_length requires max_output_length"
1470 ));
1471 }
1472 if let Some(discount) = metadata.discount_to_user
1473 && (!discount.is_finite() || discount >= 1.0)
1474 {
1475 return Err(format!(
1476 "model {alias:?}: discount_to_user must be finite and less than 1"
1477 ));
1478 }
1479 if metadata
1480 .openrouter_slug
1481 .as_deref()
1482 .is_some_and(str::is_empty)
1483 {
1484 return Err(format!(
1485 "model {alias:?}: openrouter_slug must not be empty when declared"
1486 ));
1487 }
1488 for dc in &metadata.datacenters {
1489 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1490 return Err(format!(
1491 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1492 dc.country_code
1493 ));
1494 }
1495 }
1496 Ok(())
1497}
1498
1499fn validate_sampling_defaults(
1514 alias: &str,
1515 metadata: &OpenRouterModelMetadata,
1516) -> Result<(), String> {
1517 validate_sampling_arm(
1518 alias,
1519 &[
1520 "default_temperature",
1521 "default_top_p",
1522 "default_min_p",
1523 "default_presence_penalty",
1524 "default_frequency_penalty",
1525 "default_repetition_penalty",
1526 ],
1527 metadata.default_temperature,
1528 metadata.default_top_p,
1529 metadata.default_min_p,
1530 metadata.default_presence_penalty,
1531 metadata.default_frequency_penalty,
1532 metadata.default_repetition_penalty,
1533 )?;
1534 if let Some(arm) = &metadata.non_thinking_sampling {
1535 if arm.is_empty() {
1539 return Err(format!(
1540 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1541 least one vendor recommendation or delete the table"
1542 ));
1543 }
1544 validate_sampling_arm(
1545 alias,
1546 &[
1547 "non_thinking_sampling.temperature",
1548 "non_thinking_sampling.top_p",
1549 "non_thinking_sampling.min_p",
1550 "non_thinking_sampling.presence_penalty",
1551 "non_thinking_sampling.frequency_penalty",
1552 "non_thinking_sampling.repetition_penalty",
1553 ],
1554 arm.temperature,
1555 arm.top_p,
1556 arm.min_p,
1557 arm.presence_penalty,
1558 arm.frequency_penalty,
1559 arm.repetition_penalty,
1560 )?;
1561 }
1562 Ok(())
1563}
1564
1565#[allow(clippy::too_many_arguments)]
1571fn validate_sampling_arm(
1572 alias: &str,
1573 keys: &[&str; 6],
1574 temperature: Option<f32>,
1575 top_p: Option<f32>,
1576 min_p: Option<f32>,
1577 presence_penalty: Option<f32>,
1578 frequency_penalty: Option<f32>,
1579 repetition_penalty: Option<f32>,
1580) -> Result<(), String> {
1581 if let Some(t) = temperature
1582 && (!t.is_finite() || t <= 0.0 || t > 2.0)
1583 {
1584 return Err(format!(
1585 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1586 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1587 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1588 vendor recommendation, not greedy); clients reach greedy by sending an \
1589 explicit temperature 0.",
1590 keys[0]
1591 ));
1592 }
1593 if let Some(p) = top_p
1594 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1595 {
1596 return Err(format!(
1597 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1598 keys[1]
1599 ));
1600 }
1601 if let Some(m) = min_p
1602 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1603 {
1604 return Err(format!(
1605 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1606 keys[2]
1607 ));
1608 }
1609 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1610 if let Some(v) = value
1611 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1612 {
1613 return Err(format!(
1614 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1615 ));
1616 }
1617 }
1618 if let Some(r) = repetition_penalty
1619 && (!r.is_finite() || r <= 0.0)
1620 {
1621 return Err(format!(
1622 "model {alias:?}: {} {r} must be finite and \
1623 greater than zero (1.0 = off)",
1624 keys[5]
1625 ));
1626 }
1627 Ok(())
1628}
1629
1630fn load_openrouter_metadata(
1631 models: &[(String, String, Option<String>)],
1632) -> Result<
1633 (
1634 HashMap<String, OpenRouterModelMetadata>,
1635 Option<ProviderMetadata>,
1636 ),
1637 String,
1638> {
1639 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1640 Ok(path) => path,
1641 Err(_) => return Ok((HashMap::new(), None)),
1642 };
1643 let p = std::path::Path::new(&path);
1644 if !p.is_file() {
1645 return Err(format!(
1646 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1647 ));
1648 }
1649 let text =
1650 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1651 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1652 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1653 for alias in metadata.keys() {
1654 if !models.iter().any(|(name, _, _)| name == alias) {
1655 return Err(format!(
1656 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1657 ));
1658 }
1659 }
1660 eprintln!(
1661 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1662 metadata.len()
1663 );
1664 Ok((metadata, provider))
1665}
1666
1667#[derive(Clone)]
1668struct AppState {
1669 cmd_tx: Sender<Cmd>,
1670 models: Arc<Vec<String>>,
1671 caps: Arc<HashMap<String, ModelCaps>>,
1672 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1673 provider_metadata: Arc<Option<ProviderMetadata>>,
1675 metering: Option<Arc<dyn metering::Metering>>,
1681 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1684 api_auth: ApiAuth,
1687 metrics_auth: MetricsAuth,
1689 metrics: SharedMetrics,
1690 inflight: InflightCounts,
1694 tenant_inflight: TenantGauge,
1697 health: health::SharedHealth,
1701 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1705}
1706
1707impl AppState {
1708 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1722 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1723 }
1724}
1725
1726#[derive(Clone, Default)]
1727struct ApiAuth {
1728 keyring: Option<&'static auth::KeyStore>,
1729 single_key: Option<Arc<str>>,
1730}
1731
1732impl ApiAuth {
1733 fn from_env() -> Result<ApiAuth, String> {
1734 let single_key = match std::env::var("MEMRA_API_KEY") {
1735 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1736 Ok(key) => Some(Arc::from(key)),
1737 Err(std::env::VarError::NotPresent) => None,
1738 Err(std::env::VarError::NotUnicode(_)) => {
1739 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1740 }
1741 };
1742 Ok(ApiAuth {
1743 keyring: auth::global(),
1744 single_key,
1745 })
1746 }
1747
1748 fn configured(&self) -> bool {
1749 self.keyring.is_some() || self.single_key.is_some()
1750 }
1751}
1752
1753#[derive(Clone, Default)]
1754struct MetricsAuth {
1755 required: bool,
1756 token: Option<Arc<str>>,
1757}
1758
1759impl MetricsAuth {
1760 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1761 let token = token.map(Arc::from);
1762 MetricsAuth {
1763 required: !bind_loopback || api_auth_configured || token.is_some(),
1764 token,
1765 }
1766 }
1767}
1768
1769fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1770 let mut resolved = addr
1771 .to_socket_addrs()
1772 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1773 let first = resolved
1774 .next()
1775 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1776 let mut loopback = first.ip().to_canonical().is_loopback();
1777 for socket in resolved {
1778 loopback &= socket.ip().to_canonical().is_loopback();
1779 }
1780 Ok((first, loopback))
1781}
1782
1783fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1784 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1785}
1786
1787fn validate_bind_security(
1788 addr: &str,
1789 api_auth_configured: bool,
1790 allow_open_bind: bool,
1791) -> Result<bool, String> {
1792 let loopback = bind_is_loopback(addr)?;
1793 if !loopback && !api_auth_configured && !allow_open_bind {
1794 return Err(format!(
1795 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1796 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1797 ));
1798 }
1799 Ok(loopback)
1800}
1801
1802type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1820
1821type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1824
1825struct InflightGuard {
1829 counts: InflightCounts,
1830 idx: usize,
1831 tenants: TenantGauge,
1832 tenant: String,
1833}
1834
1835impl InflightGuard {
1836 fn try_acquire(
1840 counts: InflightCounts,
1841 lane: lanes::Lane,
1842 tenants: TenantGauge,
1843 tenant: &str,
1844 tenant_cap: Option<usize>,
1845 ) -> Result<(Self, usize, usize), usize> {
1846 let idx = lane.idx();
1847 let nt = {
1848 let mut m = tenants.lock().unwrap();
1849 let e = m.entry(tenant.to_string()).or_insert(0);
1850 if tenant_cap.is_some_and(|cap| *e >= cap) {
1851 return Err(*e);
1852 }
1853 *e += 1;
1854 *e
1855 };
1856 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1857 Ok((
1858 InflightGuard {
1859 counts,
1860 idx,
1861 tenants,
1862 tenant: tenant.to_string(),
1863 },
1864 n,
1865 nt,
1866 ))
1867 }
1868}
1869
1870impl Drop for InflightGuard {
1871 fn drop(&mut self) {
1872 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1873 let mut m = self.tenants.lock().unwrap();
1874 if let Some(e) = m.get_mut(&self.tenant) {
1875 *e -= 1;
1876 if *e == 0 {
1877 m.remove(&self.tenant);
1878 }
1879 }
1880 }
1881}
1882
1883fn lane_cap(lane: lanes::Lane) -> usize {
1887 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1888 CAPS.get_or_init(|| {
1889 let batching = std::env::var("MEMRA_SERVE_BATCH")
1890 .map(|v| v != "0")
1891 .unwrap_or(true);
1892 let interactive = if batching {
1893 std::env::var("MEMRA_MAX_SESSIONS")
1894 .ok()
1895 .and_then(|v| v.parse().ok())
1896 .unwrap_or(64)
1897 } else {
1898 worker::MAX_ACTIVE
1899 };
1900 let p = lanes::LanePolicy::from_env();
1901 [interactive, p.max_sessions[1], p.max_sessions[2]]
1902 })[lane.idx()]
1903}
1904
1905fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1908 if m.completed > 0 && m.step_p50_ms > 0.0 {
1909 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1910 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1911 }
1912 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1913 *D.get_or_init(|| {
1914 std::env::var("MEMRA_RL_RESET_S")
1915 .ok()
1916 .and_then(|v| v.parse().ok())
1917 .unwrap_or(2)
1918 })
1919}
1920
1921pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1938pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1939pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1940
1941pub(crate) fn timeout_ms_max() -> u64 {
1951 static V: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1952 *V.get_or_init(|| {
1953 std::env::var("MEMRA_TIMEOUT_MS_MAX")
1954 .ok()
1955 .and_then(|s| s.parse::<u64>().ok())
1956 .filter(|&ms| ms >= TIMEOUT_MS_MIN)
1957 .unwrap_or(TIMEOUT_MS_MAX)
1958 })
1959}
1960
1961pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1965 let max = timeout_ms_max();
1966 let Some(v) = v.filter(|v| !v.is_null()) else {
1967 return Ok(max);
1969 };
1970 let Some(ms) = v.as_u64() else {
1971 return Err(format!(
1972 "timeout_ms must be an integer number of milliseconds in \
1973 {TIMEOUT_MS_MIN}..={max}, got {v}; for work longer than \
1974 {max} ms use \"stream\": true — the deadline then bounds only the \
1975 time to first token and the stream may run as long as it needs"
1976 ));
1977 };
1978 if !(TIMEOUT_MS_MIN..=max).contains(&ms) {
1979 return Err(format!(
1980 "timeout_ms {ms} is outside the accepted range \
1981 {TIMEOUT_MS_MIN}..={max} (milliseconds). {max} is a \
1982 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1983 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1984 lie. For work longer than {max} ms use \"stream\": true — the \
1985 deadline then bounds only the time to first token and the stream may run as \
1986 long as it needs"
1987 ));
1988 }
1989 Ok(ms)
1990}
1991
1992#[derive(Clone, Copy)]
1995pub(crate) struct RequestDeadline {
1996 pub(crate) at: tokio::time::Instant,
1997 pub(crate) ms: u64,
1998}
1999
2000impl RequestDeadline {
2001 pub(crate) fn starting_now(ms: u64) -> Self {
2002 Self {
2003 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
2004 ms,
2005 }
2006 }
2007
2008 pub(crate) fn remaining(&self) -> std::time::Duration {
2009 self.at
2010 .saturating_duration_since(tokio::time::Instant::now())
2011 }
2012}
2013
2014pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
2020 let what = if stream {
2021 "the first token was produced"
2022 } else {
2023 "the response completed"
2024 };
2025 let msg = format!(
2026 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
2027 {what}; generation was cancelled and this request is not billed"
2028 );
2029 error_response_coded(
2030 StatusCode::REQUEST_TIMEOUT,
2031 &msg,
2032 "timeout",
2033 Some("timeout_ms"),
2034 Some("deadline_exceeded"),
2035 )
2036}
2037
2038pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
2078
2079pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
2084
2085pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
2089
2090fn env_flag_on(name: &'static str, default_on: bool) -> bool {
2096 match std::env::var(name) {
2097 Ok(v) => !matches!(
2098 v.trim().to_ascii_lowercase().as_str(),
2099 "0" | "off" | "false"
2100 ),
2101 Err(_) => default_on,
2102 }
2103}
2104
2105pub(crate) fn env_u64(name: &'static str, default: u64) -> u64 {
2108 std::env::var(name)
2109 .ok()
2110 .and_then(|v| v.parse::<u64>().ok())
2111 .filter(|v| *v > 0)
2112 .unwrap_or(default)
2113}
2114
2115const CHARS_PER_TOKEN_FLOOR: usize = 6;
2126
2127pub(crate) fn prompt_tokens_estimate(
2128 request: &worker::Request,
2129 tokenizer: Option<&Tokenizer>,
2130) -> u64 {
2131 if !request.prompt_ids.is_empty() {
2132 return request.prompt_ids.len() as u64;
2133 }
2134 let mut text = String::new();
2135 text.push_str(&request.prompt_text);
2136 for turn in &request.chat_turns {
2137 text.push_str(&turn.content);
2138 }
2139 for tool in &request.tools_json {
2140 text.push_str(tool);
2141 }
2142 if let Some(tokenizer) = tokenizer {
2143 return tokenizer.encode(text.as_str(), false).len() as u64;
2144 }
2145 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
2146}
2147
2148pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
2152 let prefill_ms = prompt_tokens
2153 .saturating_mul(1_000)
2154 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
2155 .unwrap_or(u64::MAX);
2156 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
2157 if decode_ms == 0 {
2158 return None;
2159 }
2160 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
2161}
2162
2163pub(crate) fn nonstream_deadline_gate(
2173 request: &worker::Request,
2174 stream: bool,
2175 deadline: RequestDeadline,
2176 caller_declared_max_tokens: bool,
2177 tokenizer: Option<&Tokenizer>,
2178) -> Result<(), String> {
2179 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
2180 return Ok(());
2181 }
2182 let max_new = request.params.max_new as u64;
2183 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
2190 {
2191 return Ok(());
2192 }
2193 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
2194 let remaining_ms = deadline.remaining().as_millis() as u64;
2195 let prefill_ms = prompt_tokens.saturating_mul(1_000)
2196 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
2197 let decode_ms = max_new.saturating_mul(1_000)
2198 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
2199 let est_ms = prefill_ms.saturating_add(decode_ms);
2200 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
2201 if est_ms <= bound_ms {
2202 return Ok(());
2203 }
2204 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
2205 let advice = match fits {
2206 Some(fits) if fits > 0 => format!(
2207 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
2208 stream's deadline bounds only the time to first token, so it may run as long \
2209 as it needs"
2210 ),
2211 _ => format!(
2212 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
2213 token, so no max_tokens fits: set \"stream\": true"
2214 ),
2215 };
2216 Err(format!(
2217 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
2218 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
2219 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
2220 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
2221 rather than after the deadline: {advice}",
2222 est_ms / 1_000,
2223 ))
2224}
2225
2226fn max_queue_depth(cap: usize) -> usize {
2230 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
2231 D.get_or_init(|| {
2232 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
2233 .ok()
2234 .and_then(|v| v.parse().ok())
2235 })
2236 .unwrap_or(cap.saturating_mul(4))
2237}
2238
2239fn queue_wait_ceiling_s() -> u64 {
2246 static S: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2247 *S.get_or_init(|| {
2248 std::env::var("MEMRA_QUEUE_WAIT_CEILING_S")
2249 .ok()
2250 .and_then(|v| v.parse().ok())
2251 .unwrap_or(0)
2252 })
2253}
2254
2255pub(crate) struct PendingAdmissionGuard {
2276 reserved: bool,
2277 lane: lanes::Lane,
2278}
2279
2280impl PendingAdmissionGuard {
2281 pub(crate) fn commit(mut self) {
2285 self.reserved = false;
2286 std::mem::forget(self);
2287 }
2288}
2289
2290impl Drop for PendingAdmissionGuard {
2291 fn drop(&mut self) {
2292 if self.reserved {
2293 worker::release_pending_admit();
2294 worker::release_admission_reservation(self.lane);
2295 }
2296 }
2297}
2298
2299#[allow(clippy::result_large_err)] pub(crate) fn reserve_pending_admit(
2301 st: &AppState,
2302 lane: lanes::Lane,
2303 rl: &RateLimit,
2304 deadline: RequestDeadline,
2305) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2306 reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())
2307}
2308
2309#[allow(clippy::result_large_err)] fn reserve_pending_admit_with_ceiling(
2314 st: &AppState,
2315 lane: lanes::Lane,
2316 rl: &RateLimit,
2317 deadline: RequestDeadline,
2318 ceiling_s: u64,
2319) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2320 let cap = lane_cap(lane).max(1);
2325 let bound = max_queue_depth(cap);
2326 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2327 loop {
2328 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2329 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2330 let backlog = reservations;
2334 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2335 if backlog >= bound {
2336 let msg = format!(
2337 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2338 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2339 coarse estimate, not a promise)",
2340 lane.as_str()
2341 );
2342 let resp = retry_contract_response(
2343 (
2344 StatusCode::TOO_MANY_REQUESTS,
2345 Json(error_body(
2346 &msg,
2347 "rate_limit_error",
2348 None,
2349 Some("shed_queue"),
2350 )),
2351 )
2352 .into_response(),
2353 Some(est_wait_s),
2354 );
2355 return Err((resp, "shed_queue"));
2356 }
2357 let remaining_ms = deadline.remaining().as_millis() as u64;
2358 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2362 if lane == lanes::Lane::Interactive
2363 && waits_for_capacity
2364 && est_wait_s.saturating_mul(1_000) > remaining_ms
2365 {
2366 let msg = format!(
2367 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2368 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2369 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2370 estimate, not a promise)"
2371 );
2372 let resp = retry_contract_response(
2373 (
2374 StatusCode::TOO_MANY_REQUESTS,
2375 Json(error_body(
2376 &msg,
2377 "rate_limit_error",
2378 None,
2379 Some("shed_deadline"),
2380 )),
2381 )
2382 .into_response(),
2383 Some(est_wait_s),
2384 );
2385 return Err((resp, "shed_deadline"));
2386 }
2387 if lane == lanes::Lane::Interactive
2395 && waits_for_capacity
2396 && ceiling_s > 0
2397 && est_wait_s > ceiling_s
2398 {
2399 let msg = format!(
2400 "estimated queue wait ~{est_wait_s}s exceeds this deployment's queue-wait \
2401 ceiling ({ceiling_s}s); this request was not admitted and is not billed; \
2402 retry after ~{est_wait_s}s (a coarse estimate, not a promise)"
2403 );
2404 let resp = retry_contract_response(
2405 (
2406 StatusCode::TOO_MANY_REQUESTS,
2407 Json(error_body(
2408 &msg,
2409 "rate_limit_error",
2410 None,
2411 Some("shed_queue_wait"),
2412 )),
2413 )
2414 .into_response(),
2415 Some(est_wait_s),
2416 );
2417 return Err((resp, "shed_queue_wait"));
2418 }
2419 if reservations_for_lane
2420 .compare_exchange(
2421 reservations,
2422 reservations.saturating_add(1),
2423 std::sync::atomic::Ordering::AcqRel,
2424 std::sync::atomic::Ordering::Acquire,
2425 )
2426 .is_ok()
2427 {
2428 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2432 return Ok(PendingAdmissionGuard {
2433 reserved: true,
2434 lane,
2435 });
2436 }
2437 }
2438}
2439
2440static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2451
2452fn draining() -> bool {
2453 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2454}
2455
2456fn drain_deadline_s() -> u64 {
2458 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2459 *D.get_or_init(|| {
2460 std::env::var("MEMRA_DRAIN_S")
2461 .ok()
2462 .and_then(|v| v.parse().ok())
2463 .unwrap_or(30)
2464 })
2465}
2466
2467fn drain_response() -> Response {
2477 let resp = (
2478 StatusCode::SERVICE_UNAVAILABLE,
2479 Json(error_body(
2480 "server is draining (shutdown in progress); retry",
2481 "server_error",
2482 None,
2483 Some("draining"),
2484 )),
2485 )
2486 .into_response();
2487 retry_contract_response(resp, Some(drain_deadline_s()))
2488}
2489
2490struct RateLimit {
2492 limit: usize,
2493 remaining: usize,
2494 reset_s: u64,
2495}
2496
2497impl RateLimit {
2498 fn at_admit(
2503 lane: lanes::Lane,
2504 n_inflight: usize,
2505 metrics: &SharedMetrics,
2506 tenant: &auth::TenantCtx,
2507 n_tenant: usize,
2508 ) -> Self {
2509 let global = lane_cap(lane);
2510 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2511 return Self::compute(global, n_inflight, metrics);
2512 };
2513 let headroom = t
2514 .saturating_sub(n_tenant)
2515 .min(global.saturating_sub(n_inflight));
2516 Self::compute(t, t - headroom, metrics)
2518 }
2519
2520 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2521 let remaining = limit.saturating_sub(n_inflight);
2522 let reset_s = if remaining > 0 {
2523 0
2524 } else {
2525 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2526 reset_estimate_s(&m)
2527 };
2528 RateLimit {
2529 limit,
2530 remaining,
2531 reset_s,
2532 }
2533 }
2534
2535 fn attach(&self, mut resp: Response) -> Response {
2537 let h = resp.headers_mut();
2538 for (k, v) in [
2539 ("x-ratelimit-limit", self.limit as u64),
2540 ("x-ratelimit-remaining", self.remaining as u64),
2541 ("x-ratelimit-reset", self.reset_s),
2542 ] {
2543 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2544 h.insert(axum::http::HeaderName::from_static(k), v);
2545 }
2546 }
2547 resp
2548 }
2549}
2550
2551#[allow(clippy::result_large_err)] fn acquire_request_slot(
2556 st: &AppState,
2557 lane: lanes::Lane,
2558 tenant: &auth::TenantCtx,
2559 env: &Envelope,
2560) -> Result<(InflightGuard, RateLimit), Response> {
2561 let global = lane_cap(lane);
2562 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2563 match InflightGuard::try_acquire(
2564 st.inflight.clone(),
2565 lane,
2566 st.tenant_inflight.clone(),
2567 &tenant.tenant,
2568 tenant_cap,
2569 ) {
2570 Ok((guard, n_inflight, n_tenant)) => {
2571 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2572 Ok((guard, rl))
2573 }
2574 Err(n_tenant) => {
2575 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2576 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2577 let error =
2578 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2579 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2580 }
2581 }
2582}
2583
2584#[derive(Deserialize)]
2586struct CompletionReq {
2587 model: String,
2588 #[serde(default)]
2589 prompt: String,
2590 #[serde(default)]
2592 prompt_ids: Vec<u32>,
2593 #[serde(default)]
2596 max_tokens: Option<usize>,
2597 #[serde(default)]
2610 temperature: Option<f32>,
2611 #[serde(default)]
2612 top_p: Option<f32>,
2613 #[serde(default)]
2615 top_k: Option<usize>,
2616 #[serde(default)]
2618 min_p: Option<f32>,
2619 #[serde(default)]
2621 frequency_penalty: Option<f32>,
2622 #[serde(default)]
2623 presence_penalty: Option<f32>,
2624 #[serde(default)]
2626 repetition_penalty: Option<f32>,
2627 #[serde(default)]
2633 seed: Option<u64>,
2634 #[serde(default)]
2635 stop: StopSequences,
2636 #[serde(default)]
2639 logit_bias: Option<serde_json::Value>,
2640 #[serde(default)]
2641 logprobs: Option<serde_json::Value>,
2642 #[serde(default)]
2643 n: Option<usize>,
2644 #[serde(default)]
2645 best_of: Option<usize>,
2646 #[serde(default)]
2648 chat: bool,
2649 #[serde(default)]
2651 stream: bool,
2652 #[serde(default)]
2654 max_ctx: Option<usize>,
2655 #[serde(default)]
2657 trace_id: Option<String>,
2658 #[serde(default)]
2662 cache_salt: Option<String>,
2663 #[serde(default)]
2667 session_id: Option<String>,
2668 #[serde(default)]
2669 user: Option<String>,
2670 #[serde(default)]
2674 timeout_ms: Option<serde_json::Value>,
2675}
2676
2677#[derive(Deserialize)]
2678struct ChatMessage {
2679 role: String,
2680 #[serde(default)]
2682 content: serde_json::Value,
2683 #[serde(default)]
2685 tool_calls: Vec<ReqToolCall>,
2686 #[serde(default)]
2689 tool_call_id: Option<String>,
2690 #[serde(default)]
2693 name: Option<String>,
2694 #[serde(default, alias = "reasoning_content")]
2704 reasoning: Option<String>,
2705}
2706
2707#[derive(Deserialize)]
2708struct ReqToolCall {
2709 #[serde(default)]
2710 #[allow(dead_code)]
2711 id: Option<String>,
2712 function: ReqToolFunction,
2713}
2714
2715#[derive(Deserialize)]
2716struct ReqToolFunction {
2717 name: String,
2718 #[serde(default)]
2720 arguments: serde_json::Value,
2721}
2722
2723#[derive(Clone, Default, Deserialize)]
2724#[serde(untagged)]
2725enum StopSequences {
2726 One(String),
2727 Many(Vec<String>),
2728 #[default]
2729 None,
2730}
2731
2732impl StopSequences {
2733 fn into_vec(self) -> Vec<String> {
2739 let stops = match self {
2740 Self::One(stop) => vec![stop],
2741 Self::Many(stops) => stops,
2742 Self::None => Vec::new(),
2743 };
2744 stops.into_iter().filter(|s| !s.is_empty()).collect()
2745 }
2746
2747 fn validate(&self) -> Result<(), String> {
2748 let stops: &[String] = match self {
2749 Self::One(stop) => std::slice::from_ref(stop),
2750 Self::Many(stops) => stops,
2751 Self::None => &[],
2752 };
2753 if stops.len() > MAX_STOP_SEQUENCES {
2754 return Err(format!(
2755 "stop accepts at most {MAX_STOP_SEQUENCES} sequences"
2756 ));
2757 }
2758 let mut total = 0usize;
2759 for stop in stops {
2760 let bytes = stop.len();
2761 if bytes > MAX_STOP_SEQUENCE_BYTES {
2762 return Err(format!(
2763 "each stop sequence must be at most {MAX_STOP_SEQUENCE_BYTES} UTF-8 bytes"
2764 ));
2765 }
2766 total = total
2767 .checked_add(bytes)
2768 .ok_or_else(|| "stop sequence byte count overflowed".to_string())?;
2769 }
2770 if total > MAX_STOP_SEQUENCES_BYTES {
2771 return Err(format!(
2772 "stop sequences must total at most {MAX_STOP_SEQUENCES_BYTES} UTF-8 bytes"
2773 ));
2774 }
2775 Ok(())
2776 }
2777}
2778
2779#[derive(Deserialize)]
2785struct ChatCompletionReq {
2786 model: String,
2787 messages: Vec<ChatMessage>,
2788 #[serde(default, alias = "max_completion_tokens")]
2791 max_tokens: Option<usize>,
2792 #[serde(default)]
2795 temperature: Option<f32>,
2796 #[serde(default)]
2797 top_p: Option<f32>,
2798 #[serde(default)]
2802 top_k: Option<usize>,
2803 #[serde(default)]
2805 min_p: Option<f32>,
2806 #[serde(default)]
2808 frequency_penalty: Option<f32>,
2809 #[serde(default)]
2810 presence_penalty: Option<f32>,
2811 #[serde(default)]
2813 repetition_penalty: Option<f32>,
2814 #[serde(default)]
2816 seed: Option<u64>,
2817 #[serde(default)]
2818 stop: StopSequences,
2819 #[serde(default)]
2820 stream: bool,
2821 #[serde(default)]
2822 max_ctx: Option<usize>,
2823 #[serde(default)]
2828 response_format: Option<serde_json::Value>,
2829 #[serde(default)]
2830 logit_bias: Option<serde_json::Value>,
2831 #[serde(default)]
2832 logprobs: Option<serde_json::Value>,
2833 #[serde(default)]
2834 top_logprobs: Option<usize>,
2835 #[serde(default)]
2836 n: Option<usize>,
2837 #[serde(default)]
2839 tools: Vec<serde_json::Value>,
2840 #[serde(default)]
2842 tool_choice: Option<serde_json::Value>,
2843 #[serde(default)]
2849 reasoning_effort: Option<String>,
2850 #[serde(default)]
2859 reasoning: Option<serde_json::Value>,
2860 #[serde(default)]
2875 include_reasoning: Option<bool>,
2876 #[serde(default)]
2885 enable_thinking: Option<bool>,
2886 #[serde(default)]
2892 chat_template_kwargs: Option<serde_json::Value>,
2893 #[serde(default)]
2897 cache_salt: Option<String>,
2898 #[serde(default)]
2900 session_id: Option<String>,
2901 #[serde(default)]
2902 user: Option<String>,
2903 #[serde(default)]
2908 timeout_ms: Option<serde_json::Value>,
2909}
2910fn one() -> f32 {
2911 1.0
2912}
2913fn default_temperature() -> f32 {
2919 1.0
2920}
2921
2922#[derive(Debug, Clone, Copy, Default, PartialEq)]
2942struct SamplingDefaults {
2943 temperature: Option<f32>,
2944 top_p: Option<f32>,
2945 top_k: Option<usize>,
2946 min_p: Option<f32>,
2947 frequency_penalty: Option<f32>,
2948 presence_penalty: Option<f32>,
2949 repetition_penalty: Option<f32>,
2950}
2951
2952impl SamplingDefaults {
2953 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2956 SamplingDefaults {
2957 temperature: metadata
2958 .and_then(|m| m.default_temperature)
2959 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2960 top_p: metadata
2961 .and_then(|m| m.default_top_p)
2962 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2963 top_k: metadata.and_then(|m| m.default_top_k),
2964 min_p: metadata.and_then(|m| m.default_min_p),
2965 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2966 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2967 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2968 }
2969 }
2970}
2971
2972#[derive(Debug, Clone, Copy, Default, PartialEq)]
2985struct ModelSamplingDefaults {
2986 thinking: SamplingDefaults,
2987 non_thinking: Option<SamplingDefaults>,
2988}
2989
2990impl ModelSamplingDefaults {
2991 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2992 ModelSamplingDefaults {
2993 thinking: SamplingDefaults::resolve(metadata, caps),
2994 non_thinking: metadata
2999 .and_then(|m| m.non_thinking_sampling.as_ref())
3000 .map(|arm| SamplingDefaults {
3001 temperature: arm.temperature,
3002 top_p: arm.top_p,
3003 top_k: arm.top_k,
3004 min_p: arm.min_p,
3005 frequency_penalty: arm.frequency_penalty,
3006 presence_penalty: arm.presence_penalty,
3007 repetition_penalty: arm.repetition_penalty,
3008 }),
3009 }
3010 }
3011
3012 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
3025 match (think, &self.non_thinking) {
3026 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
3027 _ => &self.thinking,
3028 }
3029 }
3030
3031 #[cfg(test)] fn single(thinking: SamplingDefaults) -> Self {
3035 ModelSamplingDefaults {
3036 thinking,
3037 non_thinking: None,
3038 }
3039 }
3040}
3041
3042#[derive(Debug, Clone, Copy, Default)]
3048struct ClientSampling {
3049 temperature: Option<f32>,
3050 top_p: Option<f32>,
3051 top_k: Option<usize>,
3052 min_p: Option<f32>,
3053 frequency_penalty: Option<f32>,
3054 presence_penalty: Option<f32>,
3055 repetition_penalty: Option<f32>,
3056 seed: Option<u64>,
3057}
3058
3059impl From<&CompletionReq> for ClientSampling {
3060 fn from(r: &CompletionReq) -> Self {
3061 ClientSampling {
3062 temperature: r.temperature,
3063 top_p: r.top_p,
3064 top_k: r.top_k,
3065 min_p: r.min_p,
3066 frequency_penalty: r.frequency_penalty,
3067 presence_penalty: r.presence_penalty,
3068 repetition_penalty: r.repetition_penalty,
3069 seed: r.seed,
3070 }
3071 }
3072}
3073
3074impl From<&ChatCompletionReq> for ClientSampling {
3075 fn from(r: &ChatCompletionReq) -> Self {
3076 ClientSampling {
3077 temperature: r.temperature,
3078 top_p: r.top_p,
3079 top_k: r.top_k,
3080 min_p: r.min_p,
3081 frequency_penalty: r.frequency_penalty,
3082 presence_penalty: r.presence_penalty,
3083 repetition_penalty: r.repetition_penalty,
3084 seed: r.seed,
3085 }
3086 }
3087}
3088
3089fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
3095 sampler_config(
3096 client
3097 .temperature
3098 .or(defaults.temperature)
3099 .unwrap_or_else(default_temperature),
3100 client.top_k.or(defaults.top_k).unwrap_or(0),
3101 client.top_p.or(defaults.top_p).unwrap_or_else(one),
3102 client.min_p.or(defaults.min_p).unwrap_or(0.0),
3103 client
3104 .frequency_penalty
3105 .or(defaults.frequency_penalty)
3106 .unwrap_or(0.0),
3107 client
3108 .presence_penalty
3109 .or(defaults.presence_penalty)
3110 .unwrap_or(0.0),
3111 client
3112 .repetition_penalty
3113 .or(defaults.repetition_penalty)
3114 .unwrap_or_else(one),
3115 client.seed,
3116 )
3117}
3118
3119#[derive(Serialize)]
3120struct CompletionResp {
3121 model: String,
3122 text: String,
3123 tokens: Vec<u32>,
3124 stop_reason: String,
3128 #[serde(default, skip_serializing_if = "Option::is_none")]
3133 error: Option<serde_json::Value>,
3134 n_tokens: usize,
3135 prompt_tokens: usize,
3138 cached_tokens: usize,
3139 elapsed_s: f64,
3140}
3141
3142fn usage_json(
3150 n_prompt: usize,
3151 n_tokens: usize,
3152 n_cached: usize,
3153 elapsed_s: f64,
3154 spec: Option<worker::SpecUsage>,
3155) -> serde_json::Value {
3156 let mut u = json!({
3157 "prompt_tokens": n_prompt,
3158 "completion_tokens": n_tokens,
3159 "total_tokens": n_prompt + n_tokens,
3160 "prompt_tokens_details": { "cached_tokens": n_cached },
3161 "elapsed_s": elapsed_s,
3162 });
3163 if let Some(sp) = spec {
3164 u["spec"] = json!({
3165 "rounds": sp.rounds,
3166 "drafted": sp.drafted,
3167 "accepted": sp.accepted,
3168 "acceptance_rate": if sp.drafted > 0 {
3169 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
3170 });
3171 }
3172 u
3173}
3174
3175pub const SYSTEM_FINGERPRINT: &str = concat!(
3207 "memra-",
3208 env!("CARGO_PKG_VERSION"),
3209 "-",
3210 env!("MEMRA_BUILD_ID")
3211);
3212
3213pub const BUILD_ID_SRC: &str = env!("MEMRA_BUILD_ID_SRC");
3215
3216pub const BUILD_ID_NOTE: &str = env!("MEMRA_BUILD_ID_NOTE");
3218
3219pub const BUILD_GIT_SHA: &str = env!("MEMRA_BUILD_SHA");
3223
3224pub fn build_identity_line() -> String {
3227 format!("[server] build: {SYSTEM_FINGERPRINT} (id: {BUILD_ID_SRC}, git: {BUILD_GIT_SHA})")
3228}
3229
3230fn gen_hex128() -> String {
3233 use std::hash::{BuildHasher, Hasher};
3234 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3235 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3236 let t = std::time::SystemTime::now()
3237 .duration_since(std::time::UNIX_EPOCH)
3238 .map(|d| d.as_nanos() as u64)
3239 .unwrap_or(0);
3240 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
3241 h1.write_u64(n);
3242 h1.write_u64(t);
3243 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
3244 h2.write_u64(t.rotate_left(17));
3245 h2.write_u64(n);
3246 format!("{:016x}{:016x}", h1.finish(), h2.finish())
3247}
3248
3249#[derive(Clone)]
3252struct Envelope {
3253 id: String,
3254 created: u64,
3255}
3256
3257impl Envelope {
3258 fn new(chat: bool) -> Self {
3259 Envelope {
3260 id: format!(
3261 "{}-{}",
3262 if chat { "chatcmpl" } else { "cmpl" },
3263 gen_hex128()
3264 ),
3265 created: std::time::SystemTime::now()
3266 .duration_since(std::time::UNIX_EPOCH)
3267 .map(|d| d.as_secs())
3268 .unwrap_or(0),
3269 }
3270 }
3271
3272 fn capture_child(&self, index: usize) -> Self {
3289 Envelope {
3290 id: format!("{}.{index}", self.id),
3291 created: self.created,
3292 }
3293 }
3294
3295 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
3297 v["id"] = json!(self.id);
3298 v["created"] = json!(self.created);
3299 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
3300 v
3301 }
3302}
3303
3304fn with_request_id(id: &str, mut resp: Response) -> Response {
3306 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
3307 resp.headers_mut()
3308 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
3309 }
3310 resp
3311}
3312
3313fn openai_compat() -> bool {
3321 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3322 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
3323 Ok("openai") => true,
3324 Ok(_) => false,
3325 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
3326 })
3327}
3328
3329fn cache_namespace(cache_salt: &Option<String>) -> String {
3341 cache_salt.clone().unwrap_or_default()
3342}
3343
3344const CACHE_SALT_MAX_BYTES: usize = 64;
3345
3346fn validate_cache_namespace(
3347 cache_salt: &Option<String>,
3348 keyring_configured: bool,
3349) -> Result<String, &'static str> {
3350 let raw = cache_namespace(cache_salt);
3351 if raw.len() > CACHE_SALT_MAX_BYTES {
3352 return Err("cache_salt must be at most 64 bytes");
3353 }
3354 if !keyring_configured && raw.starts_with("t:") {
3355 return Err("cache_salt must not use the reserved t: prefix without a keyring");
3356 }
3357 if !raw
3358 .bytes()
3359 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
3360 {
3361 return Err("cache_salt contains unsupported characters");
3362 }
3363 Ok(raw)
3364}
3365
3366fn affinity_key(
3382 session_id: &Option<String>,
3383 user: &Option<String>,
3384 headers: &axum::http::HeaderMap,
3385) -> Result<Option<String>, String> {
3386 let clean = |s: &str| -> Result<Option<String>, String> {
3387 let t = s.trim();
3388 if t.is_empty() {
3389 Ok(None)
3390 } else if t.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3391 Err(format!(
3392 "session identity must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3393 ))
3394 } else if t.chars().any(char::is_control) {
3395 Err("session identity must not contain control characters".into())
3396 } else {
3397 Ok(Some(t.to_string()))
3398 }
3399 };
3400 if let Some(value) = session_id.as_deref()
3401 && let Some(value) = clean(value)?
3402 {
3403 return Ok(Some(value));
3404 }
3405 if let Some(value) = user.as_deref()
3406 && let Some(value) = clean(value)?
3407 {
3408 return Ok(Some(value));
3409 }
3410 match headers.get("x-session-id") {
3411 Some(value) => clean(
3412 value
3413 .to_str()
3414 .map_err(|_| "x-session-id must contain visible ASCII or UTF-8 text")?,
3415 ),
3416 None => Ok(None),
3417 }
3418}
3419
3420fn validate_client_identifier(value: Option<&str>, name: &str) -> Result<(), String> {
3421 let Some(value) = value else {
3422 return Ok(());
3423 };
3424 if value.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3425 return Err(format!(
3426 "{name} must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3427 ));
3428 }
3429 if value.chars().any(char::is_control) {
3430 return Err(format!("{name} must not contain control characters"));
3431 }
3432 Ok(())
3433}
3434
3435fn error_body(
3440 message: &str,
3441 etype: &str,
3442 param: Option<&str>,
3443 code: Option<&str>,
3444) -> serde_json::Value {
3445 json!({ "error": {
3446 "message": message,
3447 "type": etype,
3448 "param": param,
3449 "code": code,
3450 } })
3451}
3452
3453fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3454 error_response_coded(status, message, etype, param, None)
3455}
3456
3457fn error_response_coded(
3462 status: StatusCode,
3463 message: &str,
3464 etype: &str,
3465 param: Option<&str>,
3466 code: Option<&str>,
3467) -> Response {
3468 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3469 if status.is_client_error()
3470 && status != StatusCode::TOO_MANY_REQUESTS
3471 && status != StatusCode::REQUEST_TIMEOUT
3472 && status != StatusCode::CONFLICT
3473 {
3474 resp.headers_mut().insert(
3475 "x-should-retry",
3476 axum::http::HeaderValue::from_static("false"),
3477 );
3478 }
3479 resp
3480}
3481
3482fn bad_request(message: &str, param: Option<&str>) -> Response {
3483 error_response(
3484 StatusCode::BAD_REQUEST,
3485 message,
3486 "invalid_request_error",
3487 param,
3488 )
3489}
3490
3491const 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>) {
3520 use worker::ErrClass as C;
3521 match class {
3522 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3523 C::ContextLength => (
3524 StatusCode::BAD_REQUEST,
3525 "invalid_request_error",
3526 Some("context_length_exceeded"),
3527 ),
3528 C::ModelNotFound => (
3529 StatusCode::BAD_REQUEST,
3530 "invalid_request_error",
3531 Some("model_not_found"),
3532 ),
3533 C::RateLimit => (
3534 StatusCode::TOO_MANY_REQUESTS,
3535 "rate_limit_error",
3536 Some("rate_limit_exceeded"),
3537 ),
3538 C::Overloaded => (
3539 StatusCode::SERVICE_UNAVAILABLE,
3540 "server_error",
3541 Some("overloaded"),
3542 ),
3543 C::Engine => (
3544 StatusCode::INTERNAL_SERVER_ERROR,
3545 "server_error",
3546 Some("engine_error"),
3547 ),
3548 }
3549}
3550
3551fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3553 use worker::ErrClass as C;
3554 match class {
3555 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3556 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3557 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3561 }
3562}
3563
3564fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3567 let (_, etype, code) = class_http(e.class);
3568 error_body(&e.message, etype, e.param, code)
3569}
3570
3571fn engine_error_response(e: &worker::EngineError) -> Response {
3577 engine_error_response_with_retry_after(
3578 e,
3579 e.retry_after_s.or_else(|| class_retry_after_s(e.class)),
3580 )
3581}
3582
3583fn engine_error_response_with_retry_after(
3584 e: &worker::EngineError,
3585 retry_after_s: Option<u64>,
3586) -> Response {
3587 let (status, _, _) = class_http(e.class);
3588 let resp = (status, Json(engine_error_body(e))).into_response();
3589 retry_contract_response(resp, retry_after_s)
3590}
3591
3592fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3594 let status = resp.status();
3595 let h = resp.headers_mut();
3596 match retry_after_s {
3597 Some(secs) => {
3598 let secs = secs.clamp(1, 60);
3600 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3601 h.insert(axum::http::header::RETRY_AFTER, v);
3602 }
3603 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3604 h.insert("retry-after-ms", v);
3605 }
3606 }
3607 None if status.is_client_error() => {
3608 h.insert(
3611 "x-should-retry",
3612 axum::http::HeaderValue::from_static("false"),
3613 );
3614 }
3615 None => {}
3616 }
3617 resp
3618}
3619
3620fn worker_unavailable_response() -> Response {
3621 engine_error_response_with_retry_after(
3622 &worker::EngineError::overloaded("worker unavailable"),
3623 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3624 )
3625}
3626
3627fn stop_reason_to_finish(r: &str) -> &'static str {
3628 match r {
3629 "Eos" | "Callback" => "stop",
3630 "MaxNew" | "ContextFull" => "length",
3631 _ => "stop",
3632 }
3633}
3634
3635fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3639 match v {
3640 serde_json::Value::Null => Ok(String::new()),
3641 serde_json::Value::String(s) => Ok(s.clone()),
3642 serde_json::Value::Array(parts) => {
3643 let mut out = String::new();
3644 for p in parts {
3645 match p.get("type").and_then(|t| t.as_str()) {
3646 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3647 Some(t) => out.push_str(t),
3648 None => return Err("content part has no text field".into()),
3649 },
3650 Some(other) => {
3651 return Err(format!(
3652 "unsupported content part type {other:?} (text only)"
3653 ));
3654 }
3655 }
3656 }
3657 Ok(out)
3658 }
3659 _ => Err("content must be a string, null, or an array of text parts".into()),
3660 }
3661}
3662
3663pub(crate) static VISION_PLACEMENT_SERVING: std::sync::atomic::AtomicBool =
3681 std::sync::atomic::AtomicBool::new(true);
3682
3683fn vision_placement_serving() -> bool {
3684 VISION_PLACEMENT_SERVING.load(std::sync::atomic::Ordering::Acquire)
3685}
3686
3687fn vision_media_admissible(placement: bool, kind: &str) -> Result<(), String> {
3693 if placement {
3694 Ok(())
3695 } else {
3696 Err(format!(
3697 "{kind} input is not enabled on this deployment (vision overlay placement \
3698 inadmissible at boot: see the worker's IMAGE INPUT DISABLED line)"
3699 ))
3700 }
3701}
3702
3703fn vision_placement_admits(kind: &str) -> Result<(), String> {
3704 vision_media_admissible(vision_placement_serving(), kind)
3705}
3706
3707fn vision_enabled() -> bool {
3711 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3712 *ON.get_or_init(|| {
3713 std::env::var("MEMRA_VISION_DIR").is_ok()
3714 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3715 })
3716}
3717
3718fn gemma_vision_enabled() -> bool {
3723 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3724 *ON.get_or_init(|| {
3725 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3726 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3727 })
3728}
3729
3730pub(crate) static GLM5_VISION_SERVING: std::sync::atomic::AtomicBool =
3740 std::sync::atomic::AtomicBool::new(false);
3741
3742fn glm5_vision_enabled() -> bool {
3745 GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)
3746}
3747
3748fn step_vision_enabled() -> bool {
3756 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3757 *ON.get_or_init(|| {
3758 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3759 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3760 })
3761}
3762
3763const VISION_MAX_IMAGES: usize = 8;
3765
3766pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3772 std::sync::atomic::AtomicUsize::new(0);
3773pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3777 tokio::sync::Semaphore::const_new(1);
3778
3779#[allow(clippy::result_large_err)]
3782pub(crate) fn try_vision_preprocess(
3783 required: bool,
3784) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3785 try_vision_preprocess_with(required, &VISION_PREPROCESS_SEMAPHORE)
3786}
3787
3788#[allow(clippy::result_large_err)]
3789fn try_vision_preprocess_with(
3790 required: bool,
3791 semaphore: &'static tokio::sync::Semaphore,
3792) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3793 if !required {
3794 return Ok(None);
3795 }
3796 match semaphore.try_acquire() {
3797 Ok(permit) => Ok(Some(permit)),
3798 Err(tokio::sync::TryAcquireError::NoPermits) => Err(retry_contract_response(
3799 error_response_coded(
3800 StatusCode::TOO_MANY_REQUESTS,
3801 "vision preprocessing is busy",
3802 "rate_limit_error",
3803 Some("messages"),
3804 Some("vision_preprocess_busy"),
3805 ),
3806 Some(BODY_ADMISSION_RETRY_AFTER_S),
3807 )),
3808 Err(tokio::sync::TryAcquireError::Closed) => Err(error_response_coded(
3809 StatusCode::SERVICE_UNAVAILABLE,
3810 "vision preprocessing is unavailable",
3811 "server_error",
3812 Some("messages"),
3813 Some("vision_preprocess_unavailable"),
3814 )),
3815 }
3816}
3817
3818pub(crate) struct VisionMemoryPermit {
3819 bytes: usize,
3820}
3821
3822#[derive(Debug)]
3823pub(crate) enum VisionMemoryError {
3824 Request(String),
3825 Capacity(String),
3826}
3827
3828impl Drop for VisionMemoryPermit {
3829 fn drop(&mut self) {
3830 if self.bytes != 0 {
3831 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3832 }
3833 }
3834}
3835
3836fn try_reserve_vision_memory(
3837 bytes: usize,
3838) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3839 if bytes == 0 {
3840 return Ok(None);
3841 }
3842 if bytes > MAX_VISION_PATCH_BYTES {
3843 return Err(VisionMemoryError::Request(format!(
3844 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3845 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3846 )));
3847 }
3848 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3849 loop {
3850 let Some(next) = in_use.checked_add(bytes) else {
3851 return Err(VisionMemoryError::Capacity(
3852 "vision patch memory reservation overflowed".into(),
3853 ));
3854 };
3855 if next > MAX_VISION_PATCH_BYTES {
3856 return Err(VisionMemoryError::Capacity(format!(
3857 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3858 in_use / (1024 * 1024),
3859 bytes / (1024 * 1024)
3860 )));
3861 }
3862 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3863 in_use,
3864 next,
3865 std::sync::atomic::Ordering::AcqRel,
3866 std::sync::atomic::Ordering::Acquire,
3867 ) {
3868 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3869 Err(actual) => in_use = actual,
3870 }
3871 }
3872}
3873
3874pub(crate) fn vision_memory_error_response(
3875 error: VisionMemoryError,
3876 param: Option<&str>,
3877) -> Response {
3878 match error {
3879 VisionMemoryError::Request(message) => bad_request(&message, param),
3880 VisionMemoryError::Capacity(message) => retry_contract_response(
3881 error_response_coded(
3882 StatusCode::SERVICE_UNAVAILABLE,
3883 &message,
3884 "server_error",
3885 None,
3886 Some("vision_memory_busy"),
3887 ),
3888 Some(RETRY_AFTER_S_OVERLOADED),
3889 ),
3890 }
3891}
3892
3893enum PendingVisionUnit {
3899 Still {
3900 bytes: Vec<u8>,
3901 gh: usize,
3902 gw: usize,
3903 },
3904 Video {
3905 bytes: Vec<u8>,
3906 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3907 video: usize,
3908 },
3909}
3910
3911struct PendingGemmaImage {
3913 bytes: Vec<u8>,
3914 gw: usize,
3915 gh: usize,
3916}
3917
3918struct PendingGlm5Image {
3921 bytes: Vec<u8>,
3922 gh: usize,
3923 gw: usize,
3924}
3925
3926struct PendingStepImage {
3929 bytes: Vec<u8>,
3930 plan: memra_engine::vision_step::StepImagePlan,
3931}
3932
3933fn content_to_text_vision_step(
3943 v: &serde_json::Value,
3944 step_images: &mut Vec<PendingStepImage>,
3945) -> Result<String, String> {
3946 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
3947 let parts = match v {
3948 serde_json::Value::Array(parts) => parts,
3949 _ => return content_to_text(v),
3950 };
3951 let mut out = String::new();
3952 let mut needs_sep = false;
3953 for p in parts {
3954 match p.get("type").and_then(|t| t.as_str()) {
3955 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3956 Some(t) => {
3957 if needs_sep {
3958 out.push(' ');
3959 }
3960 out.push_str(t);
3961 needs_sep = true;
3962 }
3963 None => return Err("content part has no text field".into()),
3964 },
3965 Some("image_url") => {
3966 vision_placement_admits("image")?;
3967 let url = p
3968 .get("image_url")
3969 .and_then(|u| {
3970 if u.is_string() {
3971 u.as_str()
3972 } else {
3973 u.get("url").and_then(|x| x.as_str())
3974 }
3975 })
3976 .ok_or("image_url part has no url")?;
3977 if !url.starts_with("data:") {
3978 return Err(
3979 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3980 );
3981 }
3982 if step_images.len() >= VISION_MAX_IMAGES {
3983 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3984 }
3985 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3989 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3990 let plan = memra_engine::vision_step::step_plan_image(&bytes)
3991 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3992 for i in 0..plan.n_tiles {
3993 out.push_str("<patch_start>");
3994 for _ in 0..SV_TILE_ROWS {
3995 out.push_str("<im_patch>");
3996 }
3997 out.push_str("<patch_end>");
3998 if plan.newline_mask[i] {
3999 out.push_str("<patch_newline>");
4000 }
4001 }
4002 out.push_str("<im_start>");
4003 for _ in 0..SV_MAIN_ROWS {
4004 out.push_str("<im_patch>");
4005 }
4006 out.push_str("<im_end>");
4007 step_images.push(PendingStepImage { bytes, plan });
4008 needs_sep = false;
4009 }
4010 Some("video_url") => {
4011 return Err("step37 has no video input (image-only processor)".into());
4012 }
4013 Some(other) => {
4014 return Err(format!("unsupported content part type {other:?}"));
4015 }
4016 }
4017 }
4018 Ok(out)
4019}
4020
4021fn content_to_text_vision(
4030 v: &serde_json::Value,
4031 images: &mut Vec<PendingVisionUnit>,
4032 gemma_images: &mut Vec<PendingGemmaImage>,
4033 glm5_images: &mut Vec<PendingGlm5Image>,
4034 step_images: &mut Vec<PendingStepImage>,
4035 next_video: &mut usize,
4036) -> Result<String, String> {
4037 if step_vision_enabled() {
4041 return content_to_text_vision_step(v, step_images);
4042 }
4043 let parts = match v {
4044 serde_json::Value::Array(parts) => parts,
4045 _ => return content_to_text(v),
4046 };
4047 let mut out = String::new();
4048 for p in parts {
4049 match p.get("type").and_then(|t| t.as_str()) {
4050 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
4051 Some(t) => out.push_str(t),
4052 None => return Err("content part has no text field".into()),
4053 },
4054 Some("image_url") if glm5_vision_enabled() => {
4055 let url = p
4056 .get("image_url")
4057 .and_then(|u| {
4058 if u.is_string() {
4059 u.as_str()
4060 } else {
4061 u.get("url").and_then(|x| x.as_str())
4062 }
4063 })
4064 .ok_or("image_url part has no url")?;
4065 if !url.starts_with("data:") {
4066 return Err(
4067 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4068 );
4069 }
4070 if glm5_images.len() >= VISION_MAX_IMAGES {
4071 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4072 }
4073 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4077 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4078 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes)
4079 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4080 out.push_str("<|begin_of_image|>");
4084 for _ in 0..memra_engine::vision_glm5::n_merged_for_grid(gh, gw) {
4085 out.push_str("<|image|>");
4086 }
4087 out.push_str("<|end_of_image|>");
4088 glm5_images.push(PendingGlm5Image { bytes, gh, gw });
4089 }
4090 Some("video_url") if glm5_vision_enabled() => {
4091 return Err(
4092 "glm5 video input is not served (tensor census only; image input is the \
4093 supported surface)"
4094 .into(),
4095 );
4096 }
4097 Some("image_url") if gemma_vision_enabled() => {
4098 vision_placement_admits("image")?;
4099 let url = p
4100 .get("image_url")
4101 .and_then(|u| {
4102 if u.is_string() {
4103 u.as_str()
4104 } else {
4105 u.get("url").and_then(|x| x.as_str())
4106 }
4107 })
4108 .ok_or("image_url part has no url")?;
4109 if !url.starts_with("data:") {
4110 return Err(
4111 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4112 );
4113 }
4114 if gemma_images.len() >= VISION_MAX_IMAGES {
4115 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4116 }
4117 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
4121 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4122 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
4123 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4124 out.push_str("<|image>");
4126 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
4127 out.push_str("<|image|>");
4128 }
4129 out.push_str("<image|>");
4130 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
4131 }
4132 Some("image_url") => {
4133 if !vision_enabled() {
4134 return Err("image input is not enabled on this deployment".into());
4135 }
4136 vision_placement_admits("image")?;
4137 let url = p
4138 .get("image_url")
4139 .and_then(|u| {
4140 if u.is_string() {
4141 u.as_str()
4142 } else {
4143 u.get("url").and_then(|x| x.as_str())
4144 }
4145 })
4146 .ok_or("image_url part has no url")?;
4147 if !url.starts_with("data:") {
4148 return Err(
4149 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4150 );
4151 }
4152 if images
4153 .iter()
4154 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
4155 .count()
4156 >= VISION_MAX_IMAGES
4157 {
4158 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4159 }
4160 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4165 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4166 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
4167 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4168 out.push_str("<|vision_start|>");
4169 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
4170 out.push_str("<|image_pad|>");
4171 }
4172 out.push_str("<|vision_end|>");
4173 images.push(PendingVisionUnit::Still { bytes, gh, gw });
4174 }
4175 Some("video_url") if gemma_vision_enabled() => {
4176 return Err("gemma-4 has no video input (image-only projector)".into());
4177 }
4178 Some("video_url") => {
4179 if !vision_enabled() {
4180 return Err("video input is not enabled on this deployment".into());
4181 }
4182 vision_placement_admits("video")?;
4183 let url = p
4184 .get("video_url")
4185 .and_then(|u| {
4186 if u.is_string() {
4187 u.as_str()
4188 } else {
4189 u.get("url").and_then(|x| x.as_str())
4190 }
4191 })
4192 .ok_or("video_url part has no url")?;
4193 if !url.starts_with("data:") {
4194 return Err(
4195 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4196 );
4197 }
4198 if *next_video >= 2 {
4199 return Err("too many videos (max 2)".into());
4200 }
4201 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
4204 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
4205 .map_err(|e| format!("video: {e}"))?;
4206 let vidx = *next_video;
4207 *next_video += 1;
4208 for group in &vid.groups {
4210 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
4211 out.push_str("<|vision_start|>");
4212 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
4213 out.push_str("<|video_pad|>");
4214 }
4215 out.push_str("<|vision_end|>");
4216 }
4217 images.push(PendingVisionUnit::Video {
4220 bytes,
4221 groups: vid.groups,
4222 video: vidx,
4223 });
4224 }
4225 Some(other) => {
4226 return Err(format!("unsupported content part type {other:?}"));
4227 }
4228 }
4229 }
4230 Ok(out)
4231}
4232
4233fn pyjson(v: &serde_json::Value, out: &mut String) {
4237 match v {
4238 serde_json::Value::Object(m) => {
4239 out.push('{');
4240 for (i, (k, val)) in m.iter().enumerate() {
4241 if i > 0 {
4242 out.push_str(", ");
4243 }
4244 out.push_str(&serde_json::Value::String(k.clone()).to_string());
4245 out.push_str(": ");
4246 pyjson(val, out);
4247 }
4248 out.push('}');
4249 }
4250 serde_json::Value::Array(a) => {
4251 out.push('[');
4252 for (i, val) in a.iter().enumerate() {
4253 if i > 0 {
4254 out.push_str(", ");
4255 }
4256 pyjson(val, out);
4257 }
4258 out.push(']');
4259 }
4260 scalar => out.push_str(&scalar.to_string()),
4261 }
4262}
4263
4264fn pyjson_str(v: &serde_json::Value) -> String {
4265 let mut s = String::new();
4266 pyjson(v, &mut s);
4267 s
4268}
4269
4270#[allow(clippy::too_many_arguments)] fn sampler_config(
4278 temperature: f32,
4279 top_k: usize,
4280 top_p: f32,
4281 min_p: f32,
4282 frequency_penalty: f32,
4283 presence_penalty: f32,
4284 repetition_penalty: f32,
4285 seed: Option<u64>,
4286) -> SamplerConfig {
4287 let penalties_on =
4288 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
4289 SamplerConfig {
4290 temperature,
4291 top_k,
4292 top_p,
4293 min_p,
4294 penalty_last_n: if penalties_on {
4295 memra_engine::spec::PEN_WINDOW_MAX
4296 } else {
4297 0
4298 },
4299 penalty_repeat: repetition_penalty,
4300 penalty_freq: frequency_penalty,
4301 penalty_present: presence_penalty,
4302 seed: seed.unwrap_or_else(fresh_seed),
4305 }
4306}
4307
4308fn fresh_seed() -> u64 {
4313 use std::sync::atomic::{AtomicU64, Ordering};
4314 static COUNTER: AtomicU64 = AtomicU64::new(0);
4315 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
4316 let nanos = std::time::SystemTime::now()
4317 .duration_since(std::time::UNIX_EPOCH)
4318 .map(|d| d.as_nanos() as u64)
4319 .unwrap_or(0);
4320 let mut z = nanos
4321 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
4322 .wrapping_add(0x9E3779B97F4A7C15);
4323 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
4324 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
4325 z ^= z >> 31;
4326 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
4329}
4330
4331fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
4336 for (param, present, why) in fields {
4337 if *present {
4338 return Err((format!("{param} is not supported{why}"), param.to_string()));
4339 }
4340 }
4341 Ok(())
4342}
4343
4344#[derive(PartialEq)]
4345enum ToolChoice {
4346 Auto,
4347 None,
4348}
4349
4350fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
4351 match v {
4352 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
4353 Some(serde_json::Value::String(s)) => match s.as_str() {
4354 "auto" => Ok(ToolChoice::Auto),
4355 "none" => Ok(ToolChoice::None),
4356 "required" => Err("tool_choice \"required\" is not supported (no constrained \
4357 decoding); use \"auto\""
4358 .into()),
4359 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
4360 },
4361 Some(serde_json::Value::Object(_)) => {
4362 Err("named-function tool_choice is not supported; use \"auto\"".into())
4363 }
4364 Some(other) => Err(format!("bad tool_choice: {other}")),
4365 }
4366}
4367
4368fn parse_think(
4418 reasoning_effort: &Option<String>,
4419 reasoning: &Option<serde_json::Value>,
4420 vllm_switch: Option<bool>,
4421 suppress_switch: Option<bool>,
4422 default_effort: Option<&str>,
4423 max_tier: bool,
4424) -> Result<(ThinkMode, Option<String>, bool), String> {
4425 let mut effort = reasoning_effort.clone();
4426 let ReasoningObject {
4427 mut enabled,
4428 effort: object_effort,
4429 exclude,
4430 } = parse_reasoning_object(reasoning)?;
4431 if let Some(e) = object_effort {
4432 effort = Some(e);
4433 }
4434 match (enabled, vllm_switch) {
4439 (Some(a), Some(b)) if a != b => {
4440 return Err(format!(
4441 "contradictory reasoning switches: reasoning.enabled={a} and \
4442 enable_thinking={b} — send one"
4443 ));
4444 }
4445 (None, Some(b)) => enabled = Some(b),
4446 _ => {}
4447 }
4448 let suppress = match (exclude, suppress_switch) {
4460 (Some(true), _) | (_, Some(false)) => Some(false),
4461 _ => None,
4462 };
4463 match (enabled, suppress) {
4464 (Some(true), Some(false)) => {
4465 return Err(
4466 "contradictory reasoning switches: reasoning is enabled but \
4467 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
4468 on this server not delivering reasoning means not generating it, so send one"
4469 .into(),
4470 );
4471 }
4472 (None, Some(b)) => enabled = Some(b),
4473 _ => {}
4474 }
4475 let client_explicit = effort.is_some() || enabled.is_some();
4479 if effort.is_none() && enabled.is_none() {
4484 effort = default_effort.map(str::to_string);
4485 }
4486 let effort_arm = match effort.as_deref() {
4491 None => None,
4492 Some(raw) => {
4493 let level = canonical_effort_for(raw, max_tier).ok_or_else(|| {
4494 format!(
4495 "bad reasoning_effort {raw:?} \
4496 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
4497 highest level this model's template distinguishes)"
4498 )
4499 })?;
4500 Some(match level {
4501 "none" | "minimal" => (ThinkMode::NoThink, "low"),
4502 "low" => (ThinkMode::Think, "low"),
4503 "medium" => (ThinkMode::Think, "medium"),
4504 "max" => (ThinkMode::Think, "max"),
4505 _ => (ThinkMode::Think, "high"),
4506 })
4507 }
4508 };
4509 let (think, level) = match (enabled, effort_arm) {
4510 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
4513 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
4514 (None, Some((think, level))) => (think, Some(level.to_string())),
4515 (None, None) => (ThinkMode::Default, None),
4516 };
4517 Ok((think, level, client_explicit))
4518}
4519
4520struct ReasoningObject {
4522 enabled: Option<bool>,
4523 effort: Option<String>,
4524 exclude: Option<bool>,
4525}
4526
4527fn parse_reasoning_object(
4546 reasoning: &Option<serde_json::Value>,
4547) -> Result<ReasoningObject, String> {
4548 let mut out = ReasoningObject {
4549 enabled: None,
4550 effort: None,
4551 exclude: None,
4552 };
4553 let Some(v) = reasoning else { return Ok(out) };
4554 let obj = match v {
4555 serde_json::Value::Null => return Ok(out),
4556 serde_json::Value::Object(obj) => obj,
4557 _ => return Err("reasoning must be an object".into()),
4558 };
4559 for (key, value) in obj {
4560 match key.as_str() {
4568 "enabled" => {
4569 if !value.is_null() {
4570 out.enabled = Some(
4571 value
4572 .as_bool()
4573 .ok_or("reasoning.enabled must be true or false")?,
4574 );
4575 }
4576 }
4577 "exclude" => {
4578 if !value.is_null() {
4579 out.exclude = Some(
4580 value
4581 .as_bool()
4582 .ok_or("reasoning.exclude must be true or false")?,
4583 );
4584 }
4585 }
4586 "effort" => {
4587 if !value.is_null() {
4588 out.effort = Some(
4589 value
4590 .as_str()
4591 .ok_or("reasoning.effort must be a string")?
4592 .to_string(),
4593 );
4594 }
4595 }
4596 "max_tokens" => {
4597 return Err(
4598 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4599 are output tokens here, and max_tokens is the ONE output budget covering \
4600 reasoning and content together — there is no separate reasoning budget to \
4601 spend against, so honouring this field is impossible rather than merely \
4602 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4603 reasoning.enabled:false) to spend less of it on reasoning"
4604 .into(),
4605 );
4606 }
4607 other => {
4608 return Err(format!(
4609 "reasoning.{other} is not a field this server implements (it would change \
4610 nothing about the request); the supported keys are enabled, effort and \
4611 exclude"
4612 ));
4613 }
4614 }
4615 }
4616 Ok(out)
4617}
4618
4619fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4640 let Some(v) = kwargs else { return Ok(None) };
4641 let obj = match v {
4642 serde_json::Value::Null => return Ok(None),
4643 serde_json::Value::Object(obj) => obj,
4644 _ => return Err("chat_template_kwargs must be an object".into()),
4645 };
4646 let mut switch = None;
4647 for (key, value) in obj {
4648 match key.as_str() {
4649 "enable_thinking" => {
4650 switch = Some(
4651 value
4652 .as_bool()
4653 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4654 );
4655 }
4656 "preserve_thinking" => {
4657 let preserve = value
4658 .as_bool()
4659 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4660 if !preserve {
4661 return Err(
4662 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4663 server: the renderer implements the vendor DEFAULT (replay every prior \
4664 assistant turn's <think> block, empty when no reasoning was sent) but \
4665 not the strip arm — serving replay bytes under a strip request would \
4666 misdescribe the prompt. Omit the flag or send true"
4667 .into(),
4668 );
4669 }
4670 }
4672 other => {
4673 return Err(format!(
4674 "chat_template_kwargs.{other} is not supported by this server's \
4675 template renderer (it would change nothing about the prompt); the only \
4676 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4677 refuses in both directions — see its own message)"
4678 ));
4679 }
4680 }
4681 }
4682 Ok(switch)
4683}
4684
4685fn resolve_vllm_think_switch(
4689 enable_thinking: Option<bool>,
4690 kwargs: &Option<serde_json::Value>,
4691) -> Result<Option<bool>, String> {
4692 let from_kwargs = parse_template_kwargs(kwargs)?;
4693 match (enable_thinking, from_kwargs) {
4694 (Some(a), Some(b)) if a != b => Err(format!(
4695 "contradictory reasoning switches: enable_thinking={a} and \
4696 chat_template_kwargs.enable_thinking={b} — send one"
4697 )),
4698 (Some(a), _) => Ok(Some(a)),
4699 (None, b) => Ok(b),
4700 }
4701}
4702
4703pub(crate) fn canonical_effort_for(value: &str, max_tier: bool) -> Option<&'static str> {
4723 match value {
4724 "none" => Some("none"),
4725 "minimal" => Some("minimal"),
4726 "low" => Some("low"),
4727 "medium" => Some("medium"),
4728 "high" => Some("high"),
4729 "xhigh" | "max" | "ultra" => Some(if max_tier { "max" } else { "high" }),
4735 _ => None,
4736 }
4737}
4738
4739pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4742 canonical_effort_for(value, false)
4743}
4744
4745fn json_to_val(v: &serde_json::Value) -> chat::Val {
4748 match v {
4749 serde_json::Value::Null => chat::Val::Null,
4750 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4751 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4752 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4753 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4754 serde_json::Value::Object(o) => chat::Val::Obj(
4757 o.iter()
4758 .map(|(k, val)| (k.clone(), json_to_val(val)))
4759 .collect(),
4760 ),
4761 }
4762}
4763
4764#[allow(clippy::type_complexity)]
4768fn prepare_tools(
4769 tools: &[serde_json::Value],
4770) -> Result<
4771 (
4772 Vec<String>,
4773 Vec<chat::Val>,
4774 HashMap<String, HashMap<String, String>>,
4775 ),
4776 String,
4777> {
4778 let mut tools_json = Vec::with_capacity(tools.len());
4779 let mut tools_struct = Vec::with_capacity(tools.len());
4780 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4781 for t in tools {
4782 let f = t
4783 .get("function")
4784 .ok_or("each tool needs a function object")?;
4785 let name = f
4786 .get("name")
4787 .and_then(|n| n.as_str())
4788 .ok_or("each tool needs function.name")?;
4789 let mut params: HashMap<String, String> = HashMap::new();
4790 if let Some(props) = f
4791 .get("parameters")
4792 .and_then(|p| p.get("properties"))
4793 .and_then(|p| p.as_object())
4794 {
4795 for (p, def) in props {
4796 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4797 params.insert(p.clone(), ty.to_string());
4798 }
4799 }
4800 }
4801 schemas.insert(name.to_string(), params);
4802 tools_json.push(pyjson_str(t));
4803 tools_struct.push(json_to_val(f));
4805 }
4806 Ok((tools_json, tools_struct, schemas))
4807}
4808
4809fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4814 let parsed: serde_json::Value = match &tc.function.arguments {
4815 serde_json::Value::Null => json!({}),
4816 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4817 serde_json::Value::String(s) => serde_json::from_str(s)
4818 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4819 v @ serde_json::Value::Object(_) => v.clone(),
4820 _ => return Err("tool_calls arguments must be a JSON object".into()),
4821 };
4822 let obj = parsed
4823 .as_object()
4824 .ok_or("tool_calls arguments must decode to a JSON object")?;
4825 let params = obj
4826 .iter()
4827 .map(|(k, v)| {
4828 let rendered = match v {
4829 serde_json::Value::String(s) => s.clone(),
4830 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4831 scalar => scalar.to_string(),
4832 };
4833 (k.clone(), rendered)
4834 })
4835 .collect();
4836 let args = obj
4839 .iter()
4840 .map(|(k, v)| (k.clone(), json_to_val(v)))
4841 .collect();
4842 Ok(TmplToolCall {
4843 name: tc.function.name.clone(),
4844 params,
4845 args,
4846 id: tc.id.clone(),
4847 })
4848}
4849
4850fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4852 json!({ "id": c.id, "type": "function",
4853 "function": { "name": c.name, "arguments": c.arguments } })
4854}
4855
4856async fn serve_bounded_http_with_limits<F>(
4860 listener: tokio::net::TcpListener,
4861 app: Router,
4862 shutdown: F,
4863 header_read_timeout: std::time::Duration,
4864 max_connections: usize,
4865 connection_max_lifetime: std::time::Duration,
4866) -> std::io::Result<()>
4867where
4868 F: std::future::Future<Output = ()> + Send,
4869{
4870 let connections = Arc::new(tokio::sync::Semaphore::new(max_connections));
4871 let (connection_shutdown, _) = tokio::sync::watch::channel(false);
4872 let mut connection_tasks = tokio::task::JoinSet::new();
4873 let mut shutdown = Box::pin(shutdown);
4874
4875 loop {
4876 tokio::select! {
4877 _ = &mut shutdown => break,
4878 joined = connection_tasks.join_next(), if !connection_tasks.is_empty() => {
4879 if let Some(Err(error)) = joined {
4880 eprintln!("[server] connection task failed: {error}");
4881 }
4882 }
4883 accepted = listener.accept() => {
4884 let (stream, _) = match accepted {
4885 Ok(connection) => connection,
4886 Err(error) => {
4887 eprintln!("[server] accept failed: {error}");
4888 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4889 continue;
4890 }
4891 };
4892 let permit = match connections.clone().try_acquire_owned() {
4893 Ok(permit) => permit,
4894 Err(_) => {
4895 drop(stream);
4896 continue;
4897 }
4898 };
4899 let service = app.clone().map_request(
4900 |request: hyper::Request<hyper::body::Incoming>| request.map(Body::new),
4901 );
4902 let service = hyper_util::service::TowerToHyperService::new(service);
4903 let io = hyper_util::rt::TokioIo::new(stream);
4904 let mut builder = hyper_util::server::conn::auto::Builder::new(
4905 hyper_util::rt::TokioExecutor::new(),
4906 );
4907 builder
4908 .http1()
4909 .timer(hyper_util::rt::TokioTimer::new())
4910 .header_read_timeout(header_read_timeout)
4911 .max_headers(64);
4912 builder
4913 .http2()
4914 .timer(hyper_util::rt::TokioTimer::new())
4915 .max_concurrent_streams(MAX_HTTP2_STREAMS_PER_CONNECTION)
4916 .keep_alive_interval(Some(std::time::Duration::from_secs(30)))
4917 .keep_alive_timeout(std::time::Duration::from_secs(10));
4918 let mut connection = Box::pin(builder
4919 .serve_connection_with_upgrades(io, service)
4920 .into_owned());
4921 let mut shutdown_rx = connection_shutdown.subscribe();
4922 connection_tasks.spawn(async move {
4923 let _permit = permit;
4924 tokio::select! {
4925 result = connection.as_mut() => {
4926 let _ = result;
4927 }
4928 _ = tokio::time::sleep(connection_max_lifetime) => {
4929 connection.as_mut().graceful_shutdown();
4934 let _ = connection.await;
4935 }
4936 _ = shutdown_rx.changed() => {
4937 connection.as_mut().graceful_shutdown();
4938 let _ = connection.await;
4939 }
4940 }
4941 });
4942 }
4943 }
4944 }
4945 drop(listener);
4946 let _ = connection_shutdown.send(true);
4947 let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
4948 while connection_tasks.join_next().await.is_some() {}
4949 })
4950 .await;
4951 if drained.is_err() {
4952 connection_tasks.abort_all();
4953 eprintln!("[server] WARN: HTTP connections exceeded the 5s graceful close deadline");
4954 }
4955 Ok(())
4956}
4957
4958async fn serve_bounded_http<F>(
4959 listener: tokio::net::TcpListener,
4960 app: Router,
4961 shutdown: F,
4962) -> std::io::Result<()>
4963where
4964 F: std::future::Future<Output = ()> + Send,
4965{
4966 serve_bounded_http_with_limits(
4967 listener,
4968 app,
4969 shutdown,
4970 HTTP1_HEADER_READ_TIMEOUT,
4971 MAX_HTTP_CONNECTIONS,
4972 HTTP_CONNECTION_MAX_LIFETIME,
4973 )
4974 .await
4975}
4976
4977#[tokio::main]
4978pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4979 serve_with(ServerWiring::stock()).await
4980}
4981
4982enum MeteringWiring {
4984 Stock,
4988 Custom(metering::MeteringFactory),
4993}
4994
4995pub struct ServerWiring {
4999 metering: MeteringWiring,
5000 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
5003 claimed_env: Vec<&'static str>,
5007}
5008
5009impl ServerWiring {
5010 pub fn stock() -> Self {
5012 ServerWiring {
5013 metering: MeteringWiring::Stock,
5014 on_ready: None,
5015 claimed_env: Vec::new(),
5016 }
5017 }
5018
5019 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
5022 ServerWiring {
5023 metering: MeteringWiring::Custom(factory),
5024 on_ready: None,
5025 claimed_env: Vec::new(),
5026 }
5027 }
5028
5029 pub fn claiming(mut self, var: &'static str) -> Self {
5033 self.claimed_env.push(var);
5034 self
5035 }
5036
5037 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
5038 self.on_ready = Some(Box::new(hook));
5039 self
5040 }
5041}
5042
5043pub struct RuntimeHandles {
5046 pub trim: TrimHandle,
5047 pub purge: PurgeHandle,
5050 pub kv_handoff: HostHandoffHandle,
5056 pub shutdown: tokio::sync::watch::Receiver<bool>,
5061}
5062
5063#[derive(Clone)]
5066pub struct TrimHandle {
5067 cmd_tx: Sender<Cmd>,
5068}
5069
5070impl TrimHandle {
5071 pub async fn trim(&self) -> Result<serde_json::Value, String> {
5073 let (tx, rx) = tokio::sync::oneshot::channel();
5074 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
5075 return Err("worker is down".into());
5076 }
5077 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5078 Ok(Ok(report)) => Ok(json!(report)),
5079 _ => Err("worker did not answer the trim within 30s".into()),
5080 }
5081 }
5082}
5083
5084#[derive(Clone)]
5093pub struct PurgeHandle {
5094 cmd_tx: Sender<Cmd>,
5095}
5096
5097impl PurgeHandle {
5098 pub async fn purge_tenant(&self, tenant: &str) -> Result<serde_json::Value, String> {
5100 let (tx, rx) = tokio::sync::oneshot::channel();
5101 let cmd = Cmd::PurgeTenantHost {
5102 tenant: tenant.to_string(),
5103 tx,
5104 };
5105 if self.cmd_tx.send(cmd).is_err() {
5106 return Err("worker is down".into());
5107 }
5108 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5109 Ok(Ok(report)) => Ok(json!(report)),
5110 _ => Err("worker did not answer the purge within 30s".into()),
5111 }
5112 }
5113}
5114
5115#[derive(Clone)]
5124pub struct HostHandoffHandle {
5125 cmd_tx: Sender<Cmd>,
5126}
5127
5128impl HostHandoffHandle {
5129 pub async fn export(&self, force: bool) -> Result<serde_json::Value, String> {
5132 let (tx, rx) = tokio::sync::oneshot::channel();
5133 if self
5134 .cmd_tx
5135 .send(Cmd::ExportHostHandoff { force, tx })
5136 .is_err()
5137 {
5138 return Err("worker is down".into());
5139 }
5140 match tokio::time::timeout(std::time::Duration::from_secs(900), rx).await {
5141 Ok(Ok(Ok(report))) => Ok(json!(report)),
5142 Ok(Ok(Err(refused))) => Err(refused),
5143 _ => Err("worker did not answer the export within 900s".into()),
5144 }
5145 }
5146
5147 pub async fn import(&self) -> Result<serde_json::Value, String> {
5150 let (tx, rx) = tokio::sync::oneshot::channel();
5151 if self.cmd_tx.send(Cmd::ImportHostHandoff { tx }).is_err() {
5152 return Err("worker is down".into());
5153 }
5154 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5155 Ok(Ok(Ok(start))) => Ok(json!(start)),
5156 Ok(Ok(Err(refused))) => Err(refused),
5157 _ => Err("worker did not answer the import within 30s".into()),
5158 }
5159 }
5160}
5161
5162pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
5163 let args: Vec<String> = std::env::args().skip(1).collect();
5166 if args.iter().any(|a| a == "--version" || a == "-V") {
5171 println!("memra-server {}", env!("CARGO_PKG_VERSION"));
5172 println!("system_fingerprint {SYSTEM_FINGERPRINT}");
5173 println!("build_id_src {BUILD_ID_SRC}");
5174 println!("git_sha {BUILD_GIT_SHA}");
5175 if !BUILD_ID_NOTE.is_empty() {
5176 println!("degraded {BUILD_ID_NOTE}");
5177 }
5178 return Ok(());
5179 }
5180 if let Some(code) = auth::run_cli(&args) {
5181 std::process::exit(code);
5182 }
5183 eprintln!("{}", build_identity_line());
5187 if BUILD_ID_SRC != build_id::BUILD_ID_SRC_TREE {
5188 eprintln!(
5189 "[server] WARNING: build identity is DEGRADED: {BUILD_ID_NOTE}. \
5190 system_fingerprint {SYSTEM_FINGERPRINT} carries a version-only id, so it does \
5191 NOT identify the source this binary was compiled from and published \
5192 performance pins cannot be verified against it (darklanes \
5193 tools/check-claim-builds.mjs --live). Rebuild where the workspace source tree \
5194 is readable."
5195 );
5196 }
5197 auth::init_from_env();
5200 let api_auth = match ApiAuth::from_env() {
5201 Ok(auth) => auth,
5202 Err(err) => {
5203 eprintln!("[server] FATAL: {err}");
5204 std::process::exit(1);
5205 }
5206 };
5207 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
5208 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
5209 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
5210 Ok(resolved) => resolved,
5211 Err(err) => {
5212 eprintln!("[server] FATAL: {err}");
5213 std::process::exit(1);
5214 }
5215 };
5216 if let Err(message) = validate_bind_security(&addr, api_auth.configured(), allow_open_bind) {
5221 eprintln!("[server] FATAL: {message}");
5222 std::process::exit(1);
5223 }
5224 if !bind_loopback && !api_auth.configured() {
5225 eprintln!(
5226 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
5227 metrics remain bearer-protected"
5228 );
5229 }
5230 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
5231 Ok(token) if token.is_empty() => {
5232 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
5233 std::process::exit(1);
5234 }
5235 Ok(token) => Some(token),
5236 Err(std::env::VarError::NotPresent) => None,
5237 Err(std::env::VarError::NotUnicode(_)) => {
5238 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
5239 std::process::exit(1);
5240 }
5241 };
5242 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
5243
5244 let models = parse_models_config();
5245 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
5246 Ok(loaded) => loaded,
5247 Err(err) => {
5248 eprintln!("[server] FATAL: {err}");
5249 std::process::exit(1);
5250 }
5251 };
5252 let metering_obj: Option<Arc<dyn metering::Metering>> = {
5258 let factory = match wiring.metering {
5259 MeteringWiring::Stock => None,
5260 MeteringWiring::Custom(factory) => Some(factory),
5261 };
5262 for deployment_only in [
5263 "MEMRA_REQUEST_LEDGER",
5264 "MEMRA_TENANT_BUDGETS",
5265 "MEMRA_ADMIN_ADDR",
5266 "MEMRA_ADMIN_TOKEN_FILE",
5267 "MEMRA_CAPTURE_DIR",
5268 ] {
5269 if std::env::var_os(deployment_only).is_some()
5270 && !wiring.claimed_env.contains(&deployment_only)
5271 {
5272 eprintln!(
5273 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
5274 build ships no accounting/admin/capture. Wire a Metering implementation \
5275 through ServerWiring and claim the vars it consumes."
5276 );
5277 std::process::exit(1);
5278 }
5279 }
5280 match factory {
5281 None => None,
5282 Some(factory) => {
5283 let model_ids: Vec<String> =
5284 models.iter().map(|(name, _, _)| name.clone()).collect();
5285 match factory(&metering::MeteringInit { models: &model_ids }) {
5286 Ok(metering_obj) => metering_obj,
5287 Err(err) => {
5288 eprintln!("[server] FATAL: metering wiring: {err}");
5289 std::process::exit(1);
5290 }
5291 }
5292 }
5293 }
5294 };
5295 let budget_tokenizers = if metering_obj
5296 .as_ref()
5297 .is_some_and(|manager| manager.enforces_limits())
5298 {
5299 match load_budget_tokenizers(&models) {
5300 Ok(tokenizers) => Some(tokenizers),
5301 Err(err) => {
5302 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
5303 std::process::exit(1);
5304 }
5305 }
5306 } else {
5307 None
5308 };
5309 eprintln!("[server] starting; models config = {models:?}");
5310
5311 let health_state = health::WorkerHealth::new();
5316 health::spawn_gpu_watch(health_state.clone());
5320 health::spawn_sd_watchdog(health_state.clone());
5321
5322 let (cmd_tx, model_names, caps, metrics, worker_thread) =
5324 match worker::spawn(models, health_state.clone()) {
5325 Ok(v) => v,
5326 Err(err) => {
5327 eprintln!("[server] FATAL: worker init failed: {err}");
5328 health_state.mark_dead(format!("worker init failed: {err}"));
5329 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
5330 std::process::exit(1);
5331 }
5332 };
5333 eprintln!("[server] worker ready; serving models: {model_names:?}");
5334
5335 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
5342 if let Some(on_ready) = wiring.on_ready {
5343 on_ready(RuntimeHandles {
5344 trim: TrimHandle {
5345 cmd_tx: cmd_tx.clone(),
5346 },
5347 purge: PurgeHandle {
5348 cmd_tx: cmd_tx.clone(),
5349 },
5350 kv_handoff: HostHandoffHandle {
5351 cmd_tx: cmd_tx.clone(),
5352 },
5353 shutdown: drain_shutdown_rx.clone(),
5354 });
5355 }
5356
5357 let bg_handle = darklane::spawn_from_env(health_state.clone());
5360 let bg_state = bg_handle.as_ref().map(|h| {
5361 let mode = darklane::BgConfig::from_env()
5362 .map(|c| c.yield_mode.as_str())
5363 .unwrap_or("stop");
5364 (h.state.clone(), mode)
5365 });
5366
5367 let state = AppState {
5368 cmd_tx,
5369 models: model_names,
5370 caps,
5371 openrouter_metadata: Arc::new(openrouter_metadata),
5372 provider_metadata: Arc::new(provider_metadata),
5373 metering: metering_obj,
5374 budget_tokenizers,
5375 api_auth,
5376 metrics_auth,
5377 metrics,
5378 inflight: Arc::new(Default::default()),
5379 tenant_inflight: Arc::new(Default::default()),
5380 health: health_state.clone(),
5381 bg: bg_state,
5382 };
5383 let inflight_handle = state.inflight.clone();
5384 let drain_metering = state.metering.clone();
5387 worker::register_http_inflight(state.inflight.clone());
5392 let app = Router::new()
5393 .route("/health", get(health_live))
5398 .route("/livez", get(health_live))
5399 .route("/readyz", get(health_ready))
5400 .route("/models", get(list_models))
5401 .route("/v1/models", get(list_models_v1))
5402 .route("/v1/auth/check", get(auth_check))
5403 .route("/v1/completions", post(completions_admitted))
5404 .route("/v1/embeddings", post(embed_api::embeddings_admitted))
5405 .route("/v1/rerank", post(embed_api::rerank_admitted))
5406 .route("/v1/chat/completions", post(chat_completions_admitted))
5407 .route("/v1/messages", post(anthropic::messages_admitted))
5411 .route("/v1/responses", post(responses_api::responses_admitted))
5412 .route("/metrics", get(get_metrics))
5413 .route("/yield/metrics", get(yield_metrics))
5414 .with_state(state.clone());
5415 let app = apply_body_limit(app);
5418 let app = app.layer(middleware::from_fn_with_state(
5421 state,
5422 authenticate_inference_before_body,
5423 ));
5424 let app = if ttft::enabled() {
5425 app.layer(middleware::from_fn(ttft_request_start))
5426 } else {
5427 app
5428 };
5429
5430 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
5431 eprintln!("[server] listening on http://{bind_addr}");
5432 drop(drain_shutdown_rx);
5433 health::sd_notify("READY=1\nSTATUS=serving");
5437 let inflight = inflight_handle;
5444 let signal_admin_shutdown = drain_shutdown_tx.clone();
5445 let serve_result = serve_bounded_http(listener, app, async move {
5446 let mut sigterm =
5447 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
5448 Ok(s) => s,
5449 Err(err) => {
5450 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
5451 std::future::pending::<()>().await;
5452 unreachable!()
5453 }
5454 };
5455 sigterm.recv().await;
5456 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
5457 let _ = signal_admin_shutdown.send(true);
5458 health::sd_notify(&format!(
5462 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
5463 (drain_deadline_s() + 5) * 1_000_000
5464 ));
5465 let n: usize = inflight
5466 .iter()
5467 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5468 .sum();
5469 eprintln!(
5470 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
5471 drain_deadline_s()
5472 );
5473 let deadline = std::time::Duration::from_secs(drain_deadline_s());
5474 let t0 = std::time::Instant::now();
5475 loop {
5476 let n: usize = inflight
5477 .iter()
5478 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5479 .sum();
5480 if n == 0 {
5481 eprintln!(
5482 "[server] drain complete in {:.1}s; exiting",
5483 t0.elapsed().as_secs_f64()
5484 );
5485 break;
5486 }
5487 if t0.elapsed() >= deadline {
5488 eprintln!(
5489 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
5490 drain_deadline_s()
5491 );
5492 if let Some(metering) = drain_metering.as_ref() {
5499 metering.drain_kill();
5500 }
5501 break;
5502 }
5503 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
5504 }
5505 })
5506 .await;
5507 let _ = drain_shutdown_tx.send(true);
5510 serve_result?;
5511 if let Some(h) = bg_handle {
5515 h.shutdown();
5516 }
5517 worker_thread.join().map_err(|_| {
5526 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
5527 })?;
5528 eprintln!("[server] GPU worker shutdown complete");
5529 Ok(())
5530}
5531
5532fn validate_model_path(path: &str) -> Result<(), String> {
5538 let p = std::path::Path::new(path);
5539 if !p.exists() {
5540 return Err(format!("model path {path:?} does not exist"));
5541 }
5542 if p.is_file() {
5543 return Ok(()); }
5545 if p.join("manifest.json").exists() {
5546 return Ok(()); }
5548 let has_st =
5549 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
5550 if !has_st {
5551 return Err(format!(
5552 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
5553 model.safetensors.index.json + config.json (HF safetensors dir), or \
5554 manifest.json (memra repack dir)"
5555 ));
5556 }
5557 if !p.join("config.json").exists() {
5558 return Err(format!(
5559 "model dir {path:?} has safetensors weights but no config.json"
5560 ));
5561 }
5562 Ok(())
5563}
5564
5565fn parse_models_config() -> Vec<(String, String, Option<String>)> {
5573 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
5574 let mut out = Vec::new();
5575 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
5576 if let Some((name, path)) = entry.split_once('=') {
5577 let (mpath, dpath) = match path.trim().split_once('+') {
5580 Some((m, d)) => (m.trim(), Some(d.trim())),
5581 None => (path.trim(), None),
5582 };
5583 let resolve = |p: &str| {
5584 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
5585 eprintln!("[server] FATAL: model {name:?}: {err}");
5586 std::process::exit(1);
5587 })
5588 };
5589 let mpath = resolve(mpath);
5590 if let Err(err) = validate_model_path(&mpath) {
5591 eprintln!("[server] FATAL: model {name:?}: {err}");
5592 std::process::exit(1);
5593 }
5594 let dpath = dpath.map(|d| {
5604 let d = resolve(d);
5605 let p = std::path::Path::new(&d);
5606 if !p.exists() {
5607 eprintln!(
5608 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
5609 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
5610 rather than serving plain decode under a config that asked \
5611 for speculative decoding."
5612 );
5613 std::process::exit(1);
5614 }
5615 if !p.is_file() {
5616 eprintln!(
5617 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
5618 file — a '+draft' attach must be a NextN/MTP GGUF file."
5619 );
5620 std::process::exit(1);
5621 }
5622 d
5623 });
5624 out.push((name.trim().to_string(), mpath, dpath));
5625 } else {
5626 eprintln!(
5627 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
5628 );
5629 }
5630 }
5631 if !out.is_empty() {
5632 return out;
5633 }
5634 }
5635 vec![
5637 (
5638 "main".into(),
5639 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
5640 None,
5641 ),
5642 (
5643 "judge".into(),
5644 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
5645 None,
5646 ),
5647 ]
5648}
5649
5650fn load_budget_tokenizers(
5651 models: &[(String, String, Option<String>)],
5652) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
5653 let mut tokenizers = HashMap::new();
5654 for (alias, path, _) in models {
5655 let path = std::path::Path::new(path);
5656 let tokenizer = if path.is_dir() {
5657 let tokenizer_dir = if path.join("manifest.json").exists() {
5658 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
5659 format!("model {alias:?}: open repack tokenizer source: {err}")
5660 })?;
5661 repack
5662 .source_dir()
5663 .filter(|source| source.join("tokenizer.json").exists())
5664 .unwrap_or(path)
5665 .to_path_buf()
5666 } else {
5667 path.to_path_buf()
5668 };
5669 Tokenizer::from_hf_dir(&tokenizer_dir)
5670 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5671 } else {
5672 let gguf = memra_gguf::GgufFile::open(path)
5673 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
5674 Tokenizer::from_gguf(&gguf)
5675 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5676 };
5677 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
5678 }
5679 Ok(Arc::new(tokenizers))
5680}
5681
5682fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5684 let s = st.health.snapshot();
5685 let mut v = json!({
5686 "status": status,
5687 "models": *st.models,
5688 "worker": {
5689 "phase": health::phase_name(s.phase),
5690 "beat_age_ms": s.beat_age_ms,
5691 "tick_max_ms": s.tick_max_ms,
5692 "stall_threshold_ms": s.stall_threshold_ms,
5693 "forward_progress_age_ms": s.forward_progress_age_ms,
5697 "prime_progress": s.progress.map(|p| json!({
5698 "rows": p.rows,
5699 "chunks": p.events,
5700 "age_ms": p.age_ms,
5701 })),
5702 "generation": s.generation,
5703 "xid_warnings": s.xid_warns,
5704 },
5705 });
5706 if let Some(d) = detail {
5707 v["detail"] = json!(d);
5708 }
5709 v
5710}
5711
5712fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5716 let mut v = health_payload(st, status, detail);
5717 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
5718 v
5719}
5720
5721async fn auth_check() -> impl IntoResponse {
5725 StatusCode::NO_CONTENT
5726}
5727
5728async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
5744 if draining() {
5745 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
5748 }
5749 match st.health.live() {
5750 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
5751 Err(why) => retry_contract_response(
5752 (
5753 StatusCode::SERVICE_UNAVAILABLE,
5754 Json(health_payload(&st, "unhealthy", Some(&why))),
5755 )
5756 .into_response(),
5757 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
5758 ),
5759 }
5760}
5761
5762async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
5773 let is_draining = draining();
5774 match st.health.ready(is_draining) {
5775 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
5776 Err(why) => retry_contract_response(
5777 (
5778 StatusCode::SERVICE_UNAVAILABLE,
5779 Json(readiness_payload(&st, "not_ready", Some(&why))),
5780 )
5781 .into_response(),
5782 Some(if is_draining {
5783 drain_deadline_s()
5784 } else {
5785 worker::WORKER_RESPAWN_BACKOFF_BASE_S
5786 }),
5787 ),
5788 }
5789}
5790
5791#[derive(Clone, Copy)]
5792struct DualPpMetricsSnapshot {
5793 stage_ns: [u64; 4],
5794 stage_samples: [usize; 4],
5795 dropped_timing_samples: usize,
5796 overlaps: usize,
5797 slot_pairs: usize,
5798 slot_uses: [usize; 2],
5799 slot_collisions: usize,
5800}
5801
5802impl DualPpMetricsSnapshot {
5803 fn current() -> Self {
5804 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
5805 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
5806 Self {
5807 stage_ns,
5808 stage_samples,
5809 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
5810 overlaps: memra_engine::pp::dual_pp_overlaps(),
5811 slot_pairs,
5812 slot_uses,
5813 slot_collisions,
5814 }
5815 }
5816
5817 fn populated(self) -> bool {
5818 self.stage_samples.iter().any(|&n| n > 0)
5819 || self.dropped_timing_samples > 0
5820 || self.slot_pairs > 0
5821 || self.slot_collisions > 0
5822 }
5823}
5824
5825fn insert_dual_pp_metrics(
5826 body: &mut serde_json::Value,
5827 metrics_scope: &MetricsScope,
5828 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
5829) {
5830 if !metrics_scope.operator() {
5833 return;
5834 }
5835 let snapshot = snapshot();
5836 if !snapshot.populated() {
5837 return;
5838 }
5839 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
5840 .iter()
5841 .enumerate()
5842 .map(|(i, name)| {
5843 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
5844 (
5845 name.to_string(),
5846 json!({
5847 "samples": snapshot.stage_samples[i],
5848 "total_ms": total_ms,
5849 "mean_ms": if snapshot.stage_samples[i] > 0 {
5850 total_ms / snapshot.stage_samples[i] as f64
5851 } else { 0.0 },
5852 }),
5853 )
5854 })
5855 .collect();
5856 body["dual_pp"] = json!({
5857 "overlaps": snapshot.overlaps,
5858 "slot_pairs": snapshot.slot_pairs,
5859 "slot_uses": snapshot.slot_uses,
5860 "slot_collisions": snapshot.slot_collisions,
5861 "cuda_event_spans": timings,
5862 "dropped_timing_samples": snapshot.dropped_timing_samples,
5863 });
5864}
5865
5866#[derive(Clone, Copy)]
5867struct PpWaveMetricsSnapshot {
5868 ticks: usize,
5869 cells: usize,
5870 overlaps: usize,
5871}
5872
5873impl PpWaveMetricsSnapshot {
5874 fn current() -> Self {
5875 let (ticks, cells, overlaps) = memra_engine::pp::pp_wave_snapshot();
5876 Self {
5877 ticks,
5878 cells,
5879 overlaps,
5880 }
5881 }
5882}
5883
5884fn insert_pp_wave_metrics(
5885 body: &mut serde_json::Value,
5886 metrics_scope: &MetricsScope,
5887 snapshot: impl FnOnce() -> PpWaveMetricsSnapshot,
5888) {
5889 if !metrics_scope.operator() {
5890 return;
5891 }
5892 let snapshot = snapshot();
5893 if snapshot.ticks == 0 && snapshot.cells == 0 {
5894 return;
5895 }
5896 body["pp_wave"] = json!({
5897 "ticks": snapshot.ticks,
5898 "cells": snapshot.cells,
5899 "overlaps": snapshot.overlaps,
5900 });
5901}
5902
5903fn insert_spec_acceptance_metrics(
5904 body: &mut serde_json::Value,
5905 metrics_scope: &MetricsScope,
5906 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
5907) {
5908 if !metrics_scope.operator() {
5911 return;
5912 }
5913 let snapshot = snapshot();
5914 if snapshot.is_empty() {
5915 return;
5916 }
5917
5918 let mut tau = serde_json::Map::new();
5919 let mut by_position = serde_json::Map::new();
5920 for (model, telemetry) in snapshot {
5921 if telemetry.rounds == 0 {
5922 continue;
5923 }
5924 let n_pos = telemetry
5925 .pos_drafted
5926 .iter()
5927 .rposition(|&n| n > 0)
5928 .map_or(0, |position| position + 1);
5929 tau.insert(model.clone(), json!(telemetry.tau()));
5930 by_position.insert(
5931 model,
5932 json!({
5933 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
5934 "rounds": telemetry.rounds,
5935 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
5936 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
5937 "accept_rate": (0..n_pos).map(|position| {
5938 let offered = telemetry.pos_drafted[position];
5939 if offered > 0 {
5940 telemetry.pos_accepted[position] as f64 / offered as f64
5941 } else {
5942 0.0
5943 }
5944 }).collect::<Vec<f64>>(),
5945 }),
5946 );
5947 }
5948 if !tau.is_empty() {
5949 body["spec_tau"] = serde_json::Value::Object(tau);
5950 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
5951 }
5952}
5953
5954fn insert_peer_probe_metrics(
5955 body: &mut serde_json::Value,
5956 metrics_scope: &MetricsScope,
5957 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
5958) {
5959 if !metrics_scope.operator() {
5962 return;
5963 }
5964 let snapshot = snapshot();
5965 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
5966 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
5967 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
5968 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
5969 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
5970 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
5971 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
5972}
5973
5974async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5976 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5977 Ok(scope) => scope,
5978 Err(response) => return response,
5979 };
5980 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5981 let mut body = if metrics_scope.process_wide() {
5985 json!({
5986 "admitted": m.admitted,
5987 "completed": m.completed,
5988 "tokens_out": m.tokens_out,
5989 "step_p50_ms": m.step_p50_ms,
5990 "step_p99_ms": m.step_p99_ms,
5991 "prompt_tokens_in": m.prompt_tokens_in,
5993 "cached_tokens_in": m.cached_tokens_in,
5994 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5997 "admission_session_defers": m.admission_session_defers,
6000 "admission_vram_defers": m.admission_vram_defers,
6001 "step_oom_parks": m.step_oom_parks,
6002 "continuation_pool_hits": m.continuation_pool_hits,
6003 "continuation_pool_evictions": m.continuation_pool_evictions,
6004 "plain_affinity_rewinds": m.plain_affinity_rewinds,
6005 "served_dspark": m.served_dspark,
6006 "served_spec": m.served_spec,
6007 "served_plain": m.served_plain,
6008 "spec_pool_hits": m.spec_pool_hits,
6009 "spec_pool_misses": m.spec_pool_misses,
6010 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
6011 "spec_pool_evictions": m.spec_pool_evictions,
6012 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
6015 })
6016 } else {
6017 json!({})
6018 };
6019 if metrics_scope.operator() {
6023 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
6024 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
6025 body["budget_source_reload_consecutive"] =
6026 json!(budget_health.source_reload_consecutive);
6027 body["budget_source_available"] = json!(budget_health.source_available);
6028 }
6029 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
6031 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
6032 } else {
6033 0.0
6034 });
6035 body["prefix_cache_hits"] = json!(m.prefix_hits);
6036 body["prefix_cache_misses"] = json!(m.prefix_misses);
6037 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
6038 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
6039 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
6040 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
6041 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
6042 body["prefix_host_entries"] = json!(m.prefix_host_entries);
6046 body["prefix_host_bytes"] = json!(m.prefix_host_bytes);
6047 body["prefix_host_demotions"] = json!(m.prefix_host_demotions);
6048 body["prefix_host_promotions"] = json!(m.prefix_host_promotions);
6049 body["prefix_host_demote_ms"] = json!(m.prefix_host_demote_ms);
6050 body["prefix_host_promote_ms"] = json!(m.prefix_host_promote_ms);
6051 body["prefix_host_rejected_allocs"] = json!(m.prefix_host_rejected_allocs);
6052 body["prefix_host_purges"] = json!(m.prefix_host_purges);
6053 body["prefix_host_purged_entries"] = json!(m.prefix_host_purged_entries);
6054 body["prefix_host_purged_bytes"] = json!(m.prefix_host_purged_bytes);
6055 body["prefix_host_tenant_rejects"] = json!(m.prefix_host_tenant_rejects);
6056 body["prefix_host_pause_demotes"] = json!(m.prefix_host_pause_demotes);
6060 body["prefix_host_pause_cancels"] = json!(m.prefix_host_pause_cancels);
6061 body["prefix_host_handoff_exports"] = json!(m.prefix_host_handoff_exports);
6062 body["prefix_host_handoff_imported_entries"] =
6063 json!(m.prefix_host_handoff_imported_entries);
6064 body["prefix_host_handoff_imported_bytes"] = json!(m.prefix_host_handoff_imported_bytes);
6065 body["prefix_host_handoff_skips"] = json!(m.prefix_host_handoff_skips);
6066 body["kv_flex_borrowed_bytes"] = json!(m.kv_flex_borrowed_bytes);
6071 body["kv_flex_sheds"] = json!(m.kv_flex_sheds);
6072 body["kv_flex_shed_ms"] = json!(m.kv_flex_shed_ms);
6073 body["lcp_histogram"] = json!({
6076 "edges": worker::LCP_HIST_EDGES.to_vec(),
6077 "counts": m.lcp_hist.to_vec(),
6078 });
6079 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
6083 body["prefix_cache_entries"] = json!(m.prefix_entries);
6084 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
6085 body["active_sessions"] = json!(m.active_sessions);
6086 body["queued_requests"] = json!(m.queued_requests);
6087 body["admission_inflight"] = json!(m.admission_inflight);
6091 body["admission_booked_bytes"] = json!(m.admission_booked_bytes);
6092 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
6093 body["spec_pool_entries"] = json!(m.spec_pool_entries);
6094 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
6095 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
6096 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
6097 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
6098 if !m.constraint_compiler_fail_closed.is_empty() {
6099 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
6100 m.constraint_compiler_fail_closed
6101 .iter()
6102 .map(|(model, gauge)| {
6103 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
6104 (model.clone(), json!(value))
6105 })
6106 .collect(),
6107 );
6108 }
6109 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
6110 }
6111 if !m.ns_tokens.is_empty() {
6116 let tenants: serde_json::Map<String, serde_json::Value> = m
6117 .ns_tokens
6118 .iter()
6119 .filter(|(ns, _)| metrics_scope.includes(ns))
6120 .map(|(ns, [p, c])| {
6121 (
6122 ns.clone(),
6123 json!({
6124 "prompt_tokens_in": p,
6125 "cached_tokens_in": c,
6126 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
6127 }),
6128 )
6129 })
6130 .collect();
6131 if !tenants.is_empty() {
6132 body["tenants"] = serde_json::Value::Object(tenants);
6133 }
6134 }
6135 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
6136 .adsd_suspect_total
6137 .iter()
6138 .filter(|(tenant, _)| metrics_scope.includes(tenant))
6139 .map(|(tenant, total)| (tenant.clone(), json!(total)))
6140 .collect();
6141 if !adsd_suspect_total.is_empty() {
6142 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
6143 }
6144 if metrics_scope.operator()
6146 && let Some((bg, mode)) = &st.bg
6147 {
6148 body["bg"] = bg.to_json(mode);
6149 }
6150 if metrics_scope.operator() {
6157 let spec: serde_json::Map<String, serde_json::Value> = m
6158 .spec
6159 .iter()
6160 .map(|(model, t)| {
6161 let n_pos = t
6162 .pos_drafted
6163 .iter()
6164 .rposition(|&d| d > 0)
6165 .map_or(0, |p| p + 1);
6166 (
6167 model.clone(),
6168 json!({
6169 "rounds": t.rounds,
6170 "drafted": t.drafted,
6171 "accepted": t.accepted,
6172 "acceptance_rate": if t.drafted > 0 {
6173 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
6174 "tokens_per_round": if t.rounds > 0 {
6175 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
6176 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
6177 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
6178 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
6179 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
6180 .collect::<Vec<f64>>(),
6181 }),
6182 )
6183 })
6184 .collect();
6185 if !spec.is_empty() {
6186 body["spec"] = serde_json::Value::Object(spec);
6187 }
6188 }
6189 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
6190 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
6191 insert_pp_wave_metrics(&mut body, &metrics_scope, PpWaveMetricsSnapshot::current);
6192 insert_peer_probe_metrics(
6193 &mut body,
6194 &metrics_scope,
6195 memra_engine::pp::peer_probe_metrics,
6196 );
6197 Json(body).into_response()
6198}
6199
6200#[derive(Debug, Default, Deserialize)]
6201struct ModelsQuery {
6202 #[serde(default)]
6203 schema: Option<String>,
6204}
6205
6206fn models_openai_body(models: &[String]) -> serde_json::Value {
6207 let data: Vec<_> = models
6208 .iter()
6209 .map(|m| json!({ "id": m, "object": "model" }))
6210 .collect();
6211 json!({ "object": "list", "data": data })
6212}
6213
6214fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
6219 match metadata.and_then(|m| m.surface.as_deref()) {
6220 Some("embedding") => "embedding",
6221 Some("rerank") => "rerank",
6222 _ => "chat",
6223 }
6224}
6225
6226fn openrouter_supported_parameters(
6227 caps: Option<&ModelCaps>,
6228 max_output_length: Option<u64>,
6229 is_chat: bool,
6230) -> serde_json::Value {
6231 let mut parameters = serde_json::Map::new();
6232 if !is_chat {
6240 return serde_json::Value::Object(parameters);
6241 }
6242 for name in [
6243 "temperature",
6244 "top_p",
6245 "min_p",
6246 "frequency_penalty",
6247 "presence_penalty",
6248 "repetition_penalty",
6249 "stop",
6250 ] {
6251 parameters.insert(name.into(), json!({ "type": "unknown" }));
6252 }
6253 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
6254 parameters.insert(
6255 "seed".into(),
6256 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
6257 );
6258 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
6259 if let Some(max) = max_output_length {
6260 max_tokens["max"] = json!(max);
6261 }
6262 parameters.insert("max_tokens".into(), max_tokens);
6263 if is_chat
6276 && caps.is_some_and(|c| {
6277 !c.dsv4 && !(c.qwen_think && !c.think_switch && c.think_close.is_empty())
6278 })
6279 {
6280 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
6281 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
6282 }
6283 if is_chat && caps.is_some_and(|c| c.tools_branch) {
6284 parameters.insert("tools".into(), json!({ "type": "boolean" }));
6285 parameters.insert(
6286 "tool_choice".into(),
6287 json!({ "type": "enum", "values": ["auto", "none"] }),
6288 );
6289 }
6290 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
6291 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
6292 }
6293 if is_chat && caps.is_some_and(|c| c.glm5) {
6302 parameters.insert(
6303 "reasoning_effort".into(),
6304 json!({ "type": "enum", "values": ["low", "high", "max"] }),
6305 );
6306 }
6307 serde_json::Value::Object(parameters)
6308}
6309
6310fn published_context_length(
6326 caps: Option<&ModelCaps>,
6327 metadata: Option<&OpenRouterModelMetadata>,
6328) -> Option<u64> {
6329 let trained = caps
6330 .map(|c| c.context_length as u64)
6331 .filter(|&value| value > 0)?;
6332 let envelope = metadata.and_then(|m| {
6333 let prompt = m.max_prompt_length?;
6334 let output = m.max_output_length?;
6335 prompt.checked_add(output)
6336 });
6337 Some(envelope.map_or(trained, |envelope| trained.min(envelope)))
6338}
6339
6340fn model_entry_openrouter(
6341 name: &str,
6342 caps: Option<&ModelCaps>,
6343 metadata: Option<&OpenRouterModelMetadata>,
6344) -> serde_json::Value {
6345 let empty = OpenRouterModelMetadata::default();
6346 let metadata = metadata.unwrap_or(&empty);
6347 let context_length =
6348 published_context_length(caps, Some(metadata)).filter(|&v| v <= JSON_SAFE_INTEGER_MAX);
6349 let tokenizer = caps
6350 .map(|c| c.tokenizer.as_str())
6351 .filter(|tokenizer| !tokenizer.is_empty());
6352
6353 let mut input = serde_json::Map::new();
6354 input.insert("type".into(), json!("text"));
6355 let mut supported_inputs = serde_json::Map::new();
6356 if let Some(value) = context_length {
6357 supported_inputs.insert(
6358 "max_context_length".into(),
6359 json!({ "value": value, "unit": "token" }),
6360 );
6361 }
6362 if let Some(value) = metadata.max_prompt_length {
6363 supported_inputs.insert(
6364 "max_prompt_length".into(),
6365 json!({ "value": value, "unit": "token" }),
6366 );
6367 }
6368 if !supported_inputs.is_empty() {
6369 input.insert(
6370 "supported_inputs".into(),
6371 serde_json::Value::Object(supported_inputs),
6372 );
6373 }
6374 let mut input_pricing = Vec::new();
6375 for (kind, cost) in [
6376 ("prompt", metadata.pricing.prompt.as_deref()),
6377 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
6378 ("cache_write", metadata.pricing.cache_write.as_deref()),
6379 ] {
6380 if let Some(cost) = cost {
6381 input_pricing.push(json!({
6382 "type": kind,
6383 "unit": "token",
6384 "cost_usd": cost,
6385 }));
6386 }
6387 }
6388 if !input_pricing.is_empty() {
6389 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
6390 }
6391 let mut input_capacity = Vec::new();
6392 for (kind, value) in [
6393 ("prompt", metadata.capacity.prompt_tpm),
6394 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
6395 ] {
6396 if let Some(value) = value {
6397 input_capacity.push(json!({
6398 "type": kind,
6399 "unit": "token",
6400 "per": "minute",
6401 "value": value,
6402 }));
6403 }
6404 }
6405 if !input_capacity.is_empty() {
6406 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
6407 }
6408
6409 let or_surface = declared_surface(Some(metadata));
6410 let or_is_chat = or_surface == "chat";
6411 let mut output = serde_json::Map::new();
6412 output.insert(
6420 "type".into(),
6421 json!(match or_surface {
6422 "embedding" => "embeddings",
6423 "rerank" => "rerank",
6424 _ => "text",
6425 }),
6426 );
6427 output.insert(
6428 "supported_parameters".into(),
6429 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
6430 );
6431 if or_is_chat {
6435 output.insert("streaming".into(), json!(true));
6436 }
6437 if let Some(value) = metadata.max_output_length
6440 && or_is_chat
6441 {
6442 output.insert(
6443 "max_length".into(),
6444 json!({ "value": value, "unit": "token" }),
6445 );
6446 }
6447 let mut output_pricing = Vec::new();
6448 for (kind, cost) in [
6449 ("completion", metadata.pricing.completion.as_deref()),
6450 (
6451 "internal_reasoning",
6452 metadata.pricing.internal_reasoning.as_deref(),
6453 ),
6454 ] {
6455 if let Some(cost) = cost {
6456 output_pricing.push(json!({
6457 "type": kind,
6458 "unit": "token",
6459 "cost_usd": cost,
6460 }));
6461 }
6462 }
6463 if !output_pricing.is_empty() {
6464 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
6465 }
6466 let mut output_capacity = Vec::new();
6467 if let Some(value) = metadata.capacity.completion_tpm {
6468 output_capacity.push(json!({
6469 "type": "completion",
6470 "unit": "token",
6471 "per": "minute",
6472 "value": value,
6473 }));
6474 }
6475 if let Some(value) = metadata.capacity.concurrency {
6476 output_capacity.push(json!({
6477 "type": "concurrency",
6478 "unit": "request",
6479 "value": value,
6480 }));
6481 }
6482 if !output_capacity.is_empty() {
6483 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
6484 }
6485
6486 let mut entry = serde_json::Map::new();
6487 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
6488 entry.insert("id".into(), json!(name));
6489 entry.insert("name".into(), json!(name));
6490 if let Some(value) = metadata.hugging_face_id.as_deref() {
6491 entry.insert("hugging_face_id".into(), json!(value));
6492 }
6493 if let Some(value) = metadata.created {
6494 entry.insert("created".into(), json!(value));
6495 }
6496 if let Some(value) = metadata.quantization.as_deref() {
6497 entry.insert("quantization".into(), json!(value));
6498 }
6499 if let Some(value) = tokenizer {
6500 entry.insert("tokenizer".into(), json!(value));
6501 }
6502 if let Some(value) = metadata.description.as_deref() {
6503 entry.insert("description".into(), json!(value));
6504 }
6505 let mut input_modalities = vec![serde_json::Value::Object(input)];
6506 for m in &metadata.input_modalities {
6507 let mut extra = serde_json::Map::new();
6508 extra.insert("type".into(), json!(m));
6509 if let Some(cost) = metadata.pricing.prompt.as_deref() {
6510 extra.insert(
6512 "pricing".into(),
6513 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
6514 );
6515 }
6516 input_modalities.push(serde_json::Value::Object(extra));
6517 }
6518 entry.insert(
6519 "input_modalities".into(),
6520 serde_json::Value::Array(input_modalities),
6521 );
6522 entry.insert(
6523 "output_modalities".into(),
6524 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
6525 );
6526 if let Some(cost) = metadata.pricing.request.as_deref() {
6527 entry.insert(
6528 "pricing".into(),
6529 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
6530 );
6531 }
6532 if let Some(value) = metadata.capacity.request_rpm {
6533 entry.insert(
6534 "capacity".into(),
6535 json!([{
6536 "type": "request",
6537 "unit": "request",
6538 "per": "minute",
6539 "value": value,
6540 }]),
6541 );
6542 }
6543 if let Some(value) = metadata.is_ready {
6544 entry.insert("is_ready".into(), json!(value));
6545 }
6546 if let Some(value) = metadata.is_free {
6547 entry.insert("is_free".into(), json!(value));
6548 }
6549 if let Some(value) = metadata.discount_to_user {
6550 entry.insert("discount_to_user".into(), json!(value));
6551 }
6552 if let Some(value) = metadata.openrouter_slug.as_deref() {
6553 entry.insert("openrouter".into(), json!({ "slug": value }));
6554 }
6555 if !metadata.datacenters.is_empty() {
6556 entry.insert("datacenters".into(), json!(metadata.datacenters));
6557 }
6558 let mut compliance = serde_json::Map::new();
6559 if let Some(value) = metadata.zdr {
6560 compliance.insert("zdr".into(), json!(value));
6561 }
6562 if let Some(value) = metadata.hipaa {
6563 compliance.insert("hipaa".into(), json!(value));
6564 }
6565 if !compliance.is_empty() {
6566 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
6567 }
6568 serde_json::Value::Object(entry)
6569}
6570
6571fn models_openrouter_body(st: &AppState) -> serde_json::Value {
6572 let data: Vec<_> = st
6573 .models
6574 .iter()
6575 .map(|model| {
6576 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
6577 })
6578 .collect();
6579 json!({ "data": data })
6580}
6581
6582fn model_entry_openmodels(
6583 name: &str,
6584 caps: Option<&ModelCaps>,
6585 metadata: Option<&OpenRouterModelMetadata>,
6586) -> Result<serde_json::Value, String> {
6587 let metadata = metadata.ok_or_else(|| {
6588 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
6589 })?;
6590 let context_length = published_context_length(caps, Some(metadata))
6591 .filter(|&value| value <= JSON_SAFE_INTEGER_MAX)
6592 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
6593 let created = metadata
6594 .created
6595 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
6596 let max_output_length = metadata
6597 .max_output_length
6598 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
6599 let prompt = metadata
6600 .pricing
6601 .prompt
6602 .as_deref()
6603 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
6604 let completion =
6605 metadata.pricing.completion.as_deref().ok_or_else(|| {
6606 format!("OpenModels feed requires pricing.completion for model {name:?}")
6607 })?;
6608 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
6609 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
6610 })?;
6611 let is_ready = metadata
6612 .is_ready
6613 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
6614 let is_free = metadata
6615 .is_free
6616 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
6617 let discount_to_user = metadata
6618 .discount_to_user
6619 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
6620
6621 let mut pricing = serde_json::Map::new();
6622 pricing.insert("prompt".into(), json!(prompt));
6623 pricing.insert("completion".into(), json!(completion));
6624 pricing.insert("input_cache_read".into(), json!(input_cache_read));
6625 if let Some(value) = metadata.pricing.request.as_deref() {
6626 pricing.insert("request".into(), json!(value));
6627 }
6628
6629 let om_surface = declared_surface(Some(metadata));
6630 let om_is_chat = om_surface == "chat";
6631 let mut supported_features = Vec::new();
6632 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
6633 supported_features.push("tool_calling");
6634 }
6635 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
6636 supported_features.push("reasoning");
6637 }
6638
6639 let mut entry = serde_json::Map::new();
6640 entry.insert("id".into(), json!(name));
6641 entry.insert("name".into(), json!(name));
6642 entry.insert("created".into(), json!(created));
6643 entry.insert("input_modalities".into(), json!(["text"]));
6644 entry.insert(
6645 "output_modalities".into(),
6646 json!(match om_surface {
6647 "embedding" => ["embeddings"],
6648 "rerank" => ["rerank"],
6649 _ => ["text"],
6650 }),
6651 );
6652 entry.insert("context_length".into(), json!(context_length));
6653 entry.insert("max_output_length".into(), json!(max_output_length));
6654 entry.insert("currency".into(), json!("USD"));
6657 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
6658 entry.insert("supported_features".into(), json!(supported_features));
6659 entry.insert("is_ready".into(), json!(is_ready));
6660 entry.insert("is_free".into(), json!(is_free));
6661 entry.insert("discount_to_user".into(), json!(discount_to_user));
6662 Ok(serde_json::Value::Object(entry))
6663}
6664
6665fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
6666 let data: Result<Vec<_>, _> = st
6667 .models
6668 .iter()
6669 .map(|model| {
6670 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
6671 })
6672 .collect();
6673 Ok(json!({ "data": data? }))
6674}
6675
6676async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
6677 match query.schema.as_deref() {
6678 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
6679 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
6680 Some("openmodels") => match models_openmodels_body(&st) {
6681 Ok(body) => Json(body).into_response(),
6682 Err(error) => bad_request(&error, Some("schema")),
6683 },
6684 Some(schema) => bad_request(
6685 &format!(
6686 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
6687 ),
6688 Some("schema"),
6689 ),
6690 }
6691}
6692
6693fn model_entry_v1(
6701 name: &str,
6702 caps: Option<&ModelCaps>,
6703 metadata: Option<&OpenRouterModelMetadata>,
6704) -> serde_json::Value {
6705 let ctx = published_context_length(caps, metadata);
6706 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
6710 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
6713 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
6714 Some(p) => json!(p),
6715 None => serde_json::Value::Null,
6716 };
6717 let owned_by = metadata
6718 .and_then(|m| m.owned_by.as_deref())
6719 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
6720 let mut input_modalities = vec!["text"];
6721 if let Some(meta) = metadata {
6722 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
6723 }
6724 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
6725 let reliability = metadata.and_then(|m| m.reliability.as_ref());
6726 let surface = declared_surface(metadata);
6732 let (model_type, endpoints, output_modalities) = match surface {
6733 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
6737 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
6738 _ => ("chat", vec!["chat/completions"], vec!["text"]),
6739 };
6740 let is_chat = surface == "chat";
6741 json!({
6742 "id": name,
6743 "name": name,
6744 "object": "model",
6745 "owned_by": owned_by,
6746 "type": model_type,
6747 "context_length": ctx,
6748 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
6751 "endpoints": endpoints,
6752 "input_modalities": input_modalities,
6753 "output_modalities": output_modalities,
6754 "capabilities": {
6755 "streaming": is_chat,
6758 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
6759 "structured_output": is_chat
6767 && !is_dsv4
6768 && !caps.is_some_and(|c| c.qwen_think && !c.think_switch && c.think_close.is_empty()),
6769 "reasoning": is_chat && thinking,
6770 "prompt_caching": is_chat && !is_dsv4,
6771 },
6772 "pricing": {
6773 "currency": "USD",
6774 "unit": "per_1m_tokens",
6775 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
6776 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
6777 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
6778 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
6779 "minimum_request": metadata
6781 .and_then(|m| m.pricing.request.as_deref())
6782 .unwrap_or("0"),
6783 },
6784 "lifecycle": {
6785 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
6786 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
6787 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
6788 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
6789 },
6790 "reliability": {
6791 "first_token_timeout_seconds":
6792 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
6793 "completion_timeout_seconds":
6794 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
6795 "stream_idle_timeout_seconds":
6796 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
6797 "capacity_scope":
6798 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
6799 },
6800 })
6801}
6802
6803async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
6806 let data: Vec<_> = st
6807 .models
6808 .iter()
6809 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
6810 .collect();
6811 let mut body = json!({
6812 "object": "list",
6813 "contract_version": "2.0",
6814 "data": data,
6815 });
6816 if let Some(provider) = st.provider_metadata.as_ref() {
6821 body["provider"] = json!({
6822 "id": provider.id,
6823 "status_url": provider.status_url,
6824 "support_contact": provider.support_contact,
6825 "incident_contact": provider.incident_contact,
6826 "regions": provider.regions,
6827 "request_id_header": "x-request-id",
6828 "error_contract": {
6829 "rate_limit_status": 429,
6830 "overload_status": 503,
6831 "retry_after_header": "Retry-After",
6832 "account_quota_error_codes": ["insufficient_balance"],
6833 },
6834 });
6835 }
6836 Json(body)
6837}
6838
6839async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
6842 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
6843 Ok(scope) => scope,
6844 Err(response) => return response,
6845 };
6846 if !metrics_scope.process_wide() {
6847 return error_response(
6848 StatusCode::FORBIDDEN,
6849 "completion api keys do not authorize process-wide yield metrics; configure \
6850 MEMRA_METRICS_TOKEN",
6851 "authentication_error",
6852 None,
6853 );
6854 }
6855 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
6856 let lane = |i: usize| {
6857 json!({
6858 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
6859 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
6860 })
6861 };
6862 let mut body = json!({
6863 "lanes": {
6864 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
6865 },
6866 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
6867 });
6868 if metrics_scope.operator() {
6869 body["batch_size_last"] = json!(m.batch_size_last);
6870 }
6871 Json(body).into_response()
6872}
6873
6874async fn peek_admission(
6887 mut rx: worker::EventReceiver,
6888) -> Result<worker::EventReceiver, (Response, &'static str)> {
6889 match rx.recv().await {
6890 Some(Event::Error(e)) => {
6895 let error_code = engine_error_code(e.class);
6896 Err((engine_error_response(&e), error_code))
6897 }
6898 first => {
6899 let (tx2, rx2) = worker::event_channel();
6900 if let Some(ev) = first {
6901 let _ = tx2.send(ev);
6902 }
6903 tokio::spawn(forward_events(rx, tx2));
6904 Ok(rx2)
6905 }
6906 }
6907}
6908
6909async fn forward_events(mut rx: worker::EventReceiver, tx2: worker::EventSender) {
6920 loop {
6921 tokio::select! {
6922 biased;
6923 () = tx2.closed() => break,
6924 ev = rx.recv() => match ev {
6925 Some(ev) => {
6926 if tx2.send(ev).is_err() {
6927 break;
6928 }
6929 }
6930 None => break,
6931 },
6932 }
6933 }
6934}
6935
6936async fn peek_first_token(
6951 mut rx: worker::EventReceiver,
6952 deadline: RequestDeadline,
6953) -> Result<worker::EventReceiver, ()> {
6954 let mut buffered: Vec<Event> = Vec::new();
6955 loop {
6956 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
6957 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
6960 let first_delivery = matches!(
6961 ev,
6962 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
6963 );
6964 buffered.push(ev);
6965 if first_delivery {
6966 break;
6967 }
6968 }
6969 }
6970 }
6971 let (tx2, rx2) = worker::event_channel();
6972 for ev in buffered {
6973 let _ = tx2.send(ev);
6974 }
6975 tokio::spawn(forward_events(rx, tx2));
6976 Ok(rx2)
6977}
6978
6979#[cfg(test)]
6981fn build_request(
6985 req: &CompletionReq,
6986 tx: worker::EventSender,
6987 lane: lanes::Lane,
6988 affinity: Option<String>,
6989) -> Request {
6990 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
6991}
6992
6993fn build_request_with_trace(
6994 req: &CompletionReq,
6995 tx: worker::EventSender,
6996 lane: lanes::Lane,
6997 affinity: Option<String>,
6998 ttft: Option<Arc<ttft::Trace>>,
6999 sampling_defaults: &SamplingDefaults,
7000) -> Request {
7001 let params = GenParams {
7002 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
7003 max_ctx: req.max_ctx,
7004 eos: Vec::new(), };
7006 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
7011 Request {
7012 model: req.model.clone(),
7013 prompt_ids: req.prompt_ids.clone(),
7014 prompt_text: req.prompt.clone(),
7015 chat: req.chat,
7016 chat_turns: Vec::new(),
7017 tools_json: Vec::new(),
7018 tools_struct: Vec::new(),
7019 think: ThinkMode::Default,
7020 reasoning_effort: None, params,
7022 sampler_cfg,
7023 stop_strings: req.stop.clone().into_vec(),
7024 trace_id: req.trace_id.clone(),
7025 request_id: String::new(),
7028 admit_predict_logged: false,
7029 max_prompt_tokens: None,
7030 cache_ns: cache_namespace(&req.cache_salt),
7031 affinity,
7032 lane,
7033 grammar: None, prepared_constraint: None,
7035 constraint_ready: None,
7036 oom_retries: 0, spec_k_replay: None,
7038 prepared_prompt: None,
7039 capture: None, images: Vec::new(), gemma_images: Vec::new(),
7042 glm5_images: Vec::new(),
7043 step_images: Vec::new(),
7044 vision_memory: None,
7045 wire_deadline: None, ttft,
7047 tx,
7048 }
7049}
7050
7051struct ChatPlan {
7054 request: Request,
7055 parser: Option<ToolStreamParser>,
7058 pending_images: Vec<PendingVisionUnit>,
7061 pending_gemma: Vec<PendingGemmaImage>,
7062 pending_glm5: Vec<PendingGlm5Image>,
7063 pending_step: Vec<PendingStepImage>,
7064 vision_memory: Option<VisionMemoryPermit>,
7068}
7069
7070pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
7071 req.messages.iter().any(|message| {
7072 message.content.as_array().is_some_and(|parts| {
7073 parts.iter().any(|part| {
7074 matches!(
7075 part.get("type").and_then(serde_json::Value::as_str),
7076 Some("image_url" | "video_url")
7077 )
7078 })
7079 })
7080 })
7081}
7082
7083fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
7084 let mut total = 0usize;
7085 let mut add = |bytes: usize| {
7086 total = total.checked_add(bytes).ok_or_else(|| {
7087 "vision patch memory reservation overflowed while planning".to_string()
7088 })?;
7089 Ok::<(), String>(())
7090 };
7091 for unit in &plan.pending_images {
7092 let bytes = match unit {
7093 PendingVisionUnit::Still { gh, gw, .. } => gh
7094 .checked_mul(*gw)
7095 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7096 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7097 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
7098 PendingVisionUnit::Video { groups, .. } => {
7099 groups.iter().try_fold(0usize, |total, group| {
7100 let bytes = group
7101 .gh
7102 .checked_mul(group.gw)
7103 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7104 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7105 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7106 total.checked_add(bytes).ok_or_else(|| {
7107 "vision patch memory reservation overflowed while planning".to_string()
7108 })
7109 })?
7110 }
7111 };
7112 add(bytes)?;
7113 }
7114 for unit in &plan.pending_gemma {
7115 let bytes = unit
7116 .gw
7117 .checked_mul(unit.gh)
7118 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
7119 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7120 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7121 add(bytes)?;
7122 }
7123 for unit in &plan.pending_glm5 {
7124 let bytes = unit
7125 .gh
7126 .checked_mul(unit.gw)
7127 .and_then(|n| n.checked_mul(memra_engine::vision_glm5::G5V_PATCH_IN))
7128 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7129 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7130 add(bytes)?;
7131 }
7132 for unit in &plan.pending_step {
7133 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
7134 let patches = unit
7136 .plan
7137 .n_tiles
7138 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
7139 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
7140 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7141 let bytes = patches
7142 .checked_mul(SV_PATCH_IN)
7143 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7144 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7145 add(bytes)?;
7146 }
7147 Ok(total)
7148}
7149
7150pub(crate) fn reserve_vision_memory(
7151 plan: &ChatPlan,
7152) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
7153 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
7154 try_reserve_vision_memory(bytes)
7155}
7156
7157#[cfg(test)]
7158fn build_chat_request(
7159 req: ChatCompletionReq,
7160 caps: Option<&ModelCaps>,
7161 tx: worker::EventSender,
7162 lane: lanes::Lane,
7163 affinity: Option<String>,
7164) -> Result<ChatPlan, String> {
7165 let defaults = ModelSamplingDefaults::resolve(None, caps);
7168 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
7169}
7170
7171#[allow(clippy::too_many_arguments)]
7185fn build_chat_request_with_trace(
7186 req: ChatCompletionReq,
7187 caps: Option<&ModelCaps>,
7188 tx: worker::EventSender,
7189 lane: lanes::Lane,
7190 affinity: Option<String>,
7191 ttft: Option<Arc<ttft::Trace>>,
7192 default_effort: Option<&str>,
7193 sampling_defaults: &ModelSamplingDefaults,
7194) -> Result<ChatPlan, String> {
7195 req.stop.validate()?;
7196 let client_sampling: ClientSampling = (&req).into();
7199 let tool_choice = parse_tool_choice(&req.tool_choice)?;
7200 if let Some(c) = caps
7205 && !c.chat_ok
7206 {
7207 return Err(format!(
7208 "model {:?} has no chat template (checkpoint carries neither \
7209 tokenizer_config.json chat_template nor chat_template.jinja) — \
7210 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
7211 req.model
7212 ));
7213 }
7214 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
7215 let (mut think, effort_level, think_client_explicit) = parse_think(
7216 &req.reasoning_effort,
7217 &req.reasoning,
7218 vllm_switch,
7219 req.include_reasoning,
7220 default_effort,
7221 caps.is_some_and(|c| c.dsv4 || c.glm5),
7225 )?;
7226 let level_template = caps
7230 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort || c.glm5)
7231 .unwrap_or(false);
7232 if think_client_explicit
7253 && think == ThinkMode::NoThink
7254 && let Some(c) = caps
7255 && c.qwen_think
7256 && !c.think_switch
7257 && !c.dsv4
7258 {
7259 return Err(format!(
7260 "model {:?} cannot disable reasoning: its chat template opens a think \
7261 tail unconditionally and carries no enable_thinking switch, so \
7262 reasoning_effort/enable_thinking cannot turn it off on this model",
7263 req.model
7264 ));
7265 }
7266 let reasoning_effort = if level_template { effort_level } else { None };
7299 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
7302 if grammar.is_some()
7316 && let Some(c) = caps
7317 && c.qwen_think
7318 && think != ThinkMode::NoThink
7319 {
7320 if c.think_switch {
7321 think = ThinkMode::NoThink;
7322 } else if c.think_close.is_empty() {
7323 return Err(
7324 "response_format requires the model's think channel to close \
7325 before the grammar can engage, but this chat template has \
7326 neither an enable_thinking switch nor a recognizable \
7327 think-close token sequence"
7328 .into(),
7329 );
7330 }
7331 }
7335
7336 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
7343
7344 let (tools_json, tools_struct, schemas) =
7347 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
7348 prepare_tools(&req.tools)?
7349 } else {
7350 (Vec::new(), Vec::new(), HashMap::new())
7351 };
7352
7353 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
7354 let mut images: Vec<PendingVisionUnit> = Vec::new();
7355 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
7356 let mut glm5_images: Vec<PendingGlm5Image> = Vec::new();
7357 let mut step_images: Vec<PendingStepImage> = Vec::new();
7358 let mut next_video = 0usize;
7359 for msg in &req.messages {
7360 let content = content_to_text_vision(
7361 &msg.content,
7362 &mut images,
7363 &mut gemma_images,
7364 &mut glm5_images,
7365 &mut step_images,
7366 &mut next_video,
7367 )
7368 .map_err(|e| format!("{} message: {e}", msg.role))?;
7369 let tool_calls = msg
7370 .tool_calls
7371 .iter()
7372 .map(render_req_tool_call)
7373 .collect::<Result<Vec<_>, _>>()?;
7374 if !tool_calls.is_empty() && msg.role != "assistant" {
7375 return Err("tool_calls are only valid on assistant messages".into());
7376 }
7377 let role = if msg.role == "developer" {
7380 "system".to_string()
7381 } else {
7382 msg.role.clone()
7383 };
7384 turns.push(TmplTurn {
7385 role,
7386 content,
7387 tool_calls,
7388 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
7390 tool_call_id: msg.tool_call_id.clone(),
7391 tool_name: msg.name.clone(),
7392 tool_responses: Vec::new(),
7393 task: None,
7397 tools: Vec::new(),
7398 });
7399 }
7400
7401 let has_tool_features = !tools_json.is_empty()
7404 || turns
7405 .iter()
7406 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
7407 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
7408 return Err(format!(
7409 "model {:?} chat template has no tools branch",
7410 req.model
7411 ));
7412 }
7413
7414 let think_open = caps
7417 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
7418 .unwrap_or(false);
7419 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
7437 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
7443 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
7444 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
7445 let glm5 = caps.map(|c| c.glm5).unwrap_or(false);
7451 let is_hy3 = caps.map(|c| c.hy3).unwrap_or(false);
7454 let hy3_think_open = is_hy3 && think == ThinkMode::Think;
7455 let hy3_tools = is_hy3 && !tools_json.is_empty();
7456 let parser = if glm5 {
7457 Some(ToolStreamParser::glm5(think_open, schemas))
7458 } else if is_hy3 && (hy3_tools || hy3_think_open) {
7459 Some(ToolStreamParser::hy3(schemas, hy3_think_open))
7460 } else if is_dsv4 && (dsv4_tools || dsv4_think_open) {
7461 Some(ToolStreamParser::dsv4(dsv4_think_open))
7462 } else if gemma_tools {
7463 Some(ToolStreamParser::gemma_tools())
7464 } else if !tools_json.is_empty() {
7465 Some(ToolStreamParser::new(schemas, think_open))
7466 } else if think_open {
7467 Some(ToolStreamParser::reasoning_only())
7468 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
7469 Some(ToolStreamParser::gemma_thought())
7477 } else {
7478 None
7479 };
7480
7481 Ok(ChatPlan {
7482 request: Request {
7483 model: req.model,
7484 prompt_ids: Vec::new(),
7485 prompt_text: String::new(),
7486 chat: false,
7487 chat_turns: turns,
7488 tools_json,
7489 tools_struct,
7490 think,
7491 reasoning_effort,
7492 params: GenParams {
7493 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
7494 max_ctx: req.max_ctx,
7495 eos: Vec::new(),
7496 },
7497 sampler_cfg,
7498 stop_strings: {
7499 let mut stops = req.stop.into_vec();
7504 if gemma_tools {
7505 stops.push("<tool_call|>".to_string());
7506 }
7507 if dsv4_tools {
7512 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
7513 }
7514 if hy3_tools {
7517 stops.push("</tool_calls:opensource>".to_string());
7518 }
7519 stops
7520 },
7521 trace_id: None,
7522 request_id: String::new(),
7525 admit_predict_logged: false,
7526 max_prompt_tokens: None,
7527 cache_ns: cache_namespace(&req.cache_salt),
7528 affinity,
7529 lane,
7530 grammar,
7531 prepared_constraint: None,
7532 constraint_ready: None,
7533 oom_retries: 0, spec_k_replay: None,
7535 prepared_prompt: None,
7536 images: Vec::new(),
7541 gemma_images: Vec::new(),
7542 glm5_images: Vec::new(),
7543 step_images: Vec::new(),
7544 capture: None, vision_memory: None,
7546 wire_deadline: None, ttft,
7548 tx,
7549 },
7550 parser,
7551 pending_images: images,
7552 pending_gemma: gemma_images,
7553 pending_glm5: glm5_images,
7554 pending_step: step_images,
7555 vision_memory: None,
7556 })
7557}
7558
7559fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
7565 for (i, unit) in plan.pending_images.drain(..).enumerate() {
7566 match unit {
7567 PendingVisionUnit::Still { bytes, gh, gw } => {
7568 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
7569 .map_err(|e| format!("image {}: {e}", i + 1))?;
7570 if (prep.gh, prep.gw) != (gh, gw) {
7571 return Err(format!(
7572 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
7573 i + 1,
7574 prep.gh,
7575 prep.gw
7576 ));
7577 }
7578 plan.request
7579 .images
7580 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
7581 }
7582 PendingVisionUnit::Video {
7583 bytes,
7584 groups,
7585 video,
7586 } => {
7587 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
7588 .map_err(|e| format!("video {}: {e}", i + 1))?;
7589 if prepared.groups.len() != groups.len() {
7590 return Err(format!(
7591 "video {}: decoded {} groups differ from its header-planned {} groups",
7592 i + 1,
7593 prepared.groups.len(),
7594 groups.len()
7595 ));
7596 }
7597 for ((group, prep), timestamp) in
7598 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
7599 {
7600 if (prep.gh, prep.gw) != (group.gh, group.gw) {
7601 return Err(format!(
7602 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
7603 i + 1,
7604 prep.gh,
7605 prep.gw,
7606 group.gh,
7607 group.gw
7608 ));
7609 }
7610 if (timestamp - group.timestamp).abs() > 0.001 {
7611 return Err(format!(
7612 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
7613 i + 1,
7614 group.timestamp
7615 ));
7616 }
7617 plan.request
7618 .images
7619 .push(memra_engine::vision_pre::VisionUnit {
7620 prep,
7621 video: Some(video),
7622 });
7623 }
7624 }
7625 }
7626 }
7627 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
7628 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
7629 .map_err(|e| format!("image {}: {e}", i + 1))?;
7630 if (gw, gh) != (unit.gw, unit.gh) {
7631 return Err(format!(
7632 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
7633 i + 1,
7634 unit.gw,
7635 unit.gh
7636 ));
7637 }
7638 plan.request
7639 .gemma_images
7640 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
7641 }
7642 for (i, unit) in plan.pending_glm5.drain(..).enumerate() {
7643 let (patches, gh, gw) = memra_engine::vision_glm5::glm5_prep_image(&unit.bytes)
7644 .map_err(|e| format!("image {}: {e}", i + 1))?;
7645 if (gh, gw) != (unit.gh, unit.gw) {
7646 return Err(format!(
7647 "image {}: decoded grid {gh}x{gw} differs from its header-planned grid {}x{} — refusing (placeholder runs already rendered)",
7648 i + 1,
7649 unit.gh,
7650 unit.gw
7651 ));
7652 }
7653 plan.request
7654 .glm5_images
7655 .push(memra_engine::vision_glm5::Glm5VisionUnit { patches, gh, gw });
7656 }
7657 for (i, unit) in plan.pending_step.drain(..).enumerate() {
7658 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
7659 .map_err(|e| format!("image {}: {e}", i + 1))?;
7660 if prepped.tiles.len() != unit.plan.n_tiles
7661 || prepped.newline_mask != unit.plan.newline_mask
7662 {
7663 return Err(format!(
7664 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
7665 ({} tiles) — refusing (pad runs already rendered)",
7666 i + 1,
7667 prepped.tiles.len(),
7668 unit.plan.n_tiles
7669 ));
7670 }
7671 plan.request.step_images.push(prepped);
7672 }
7673 Ok(())
7674}
7675
7676fn bearer_token(headers: &HeaderMap) -> Option<&str> {
7684 headers
7685 .get("authorization")
7686 .and_then(|value| value.to_str().ok())
7687 .and_then(|value| value.strip_prefix("Bearer "))
7688}
7689
7690fn authentication_error(why: auth::AuthDenied) -> Response {
7691 match why {
7692 auth::AuthDenied::Unknown => error_response(
7693 StatusCode::UNAUTHORIZED,
7694 "invalid api key",
7695 "authentication_error",
7696 None,
7697 ),
7698 auth::AuthDenied::Disabled => error_response(
7699 StatusCode::FORBIDDEN,
7700 "api key is disabled",
7701 "authentication_error",
7702 None,
7703 ),
7704 }
7705}
7706
7707#[allow(clippy::result_large_err)] fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
7709 auth::authenticate_with(
7710 api_auth.keyring,
7711 api_auth.single_key.as_deref(),
7712 bearer_token(headers),
7713 )
7714 .map_err(authentication_error)
7715}
7716
7717#[derive(Debug, Clone, PartialEq, Eq)]
7718enum MetricsScope {
7719 All,
7720 CompletionDomain,
7721 Tenant(String),
7722}
7723
7724impl MetricsScope {
7725 fn operator(&self) -> bool {
7726 matches!(self, MetricsScope::All)
7727 }
7728
7729 fn process_wide(&self) -> bool {
7730 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
7731 }
7732
7733 fn includes(&self, tenant_row: &str) -> bool {
7734 match self {
7735 MetricsScope::All | MetricsScope::CompletionDomain => true,
7736 MetricsScope::Tenant(tenant) => tenant == tenant_row,
7737 }
7738 }
7739}
7740
7741#[allow(clippy::result_large_err)] fn authorize_metrics(
7743 api_auth: &ApiAuth,
7744 metrics_auth: &MetricsAuth,
7745 headers: &HeaderMap,
7746) -> Result<MetricsScope, Response> {
7747 if !metrics_auth.required {
7748 return Ok(MetricsScope::All);
7749 }
7750 let Some(candidate) = bearer_token(headers) else {
7751 return Err(authentication_error(auth::AuthDenied::Unknown));
7752 };
7753 if let Some(token) = metrics_auth.token.as_deref() {
7754 if auth::constant_time_secret_eq(token, candidate) {
7755 return Ok(MetricsScope::All);
7756 }
7757 if api_auth.configured() {
7758 return match auth::authenticate_with(
7759 api_auth.keyring,
7760 api_auth.single_key.as_deref(),
7761 Some(candidate),
7762 ) {
7763 Ok(_) => Err(error_response(
7764 StatusCode::FORBIDDEN,
7765 "completion api keys do not authorize metrics while \
7766 MEMRA_METRICS_TOKEN is configured",
7767 "authentication_error",
7768 None,
7769 )),
7770 Err(why) => Err(authentication_error(why)),
7771 };
7772 }
7773 return Err(authentication_error(auth::AuthDenied::Unknown));
7774 }
7775 if api_auth.configured() {
7776 let tenant = authenticate(api_auth, headers)?;
7777 return Ok(if api_auth.keyring.is_some() {
7778 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
7779 } else {
7780 MetricsScope::CompletionDomain
7784 });
7785 }
7786 Err(authentication_error(auth::AuthDenied::Unknown))
7787}
7788
7789#[allow(clippy::result_large_err)] fn lane_for_tenant(
7796 headers: &axum::http::HeaderMap,
7797 tenant: &auth::TenantCtx,
7798) -> Result<lanes::Lane, Response> {
7799 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
7800 None => None,
7801 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
7806 error_response_coded(
7807 StatusCode::BAD_REQUEST,
7808 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
7809 "invalid_request_error",
7810 Some("x-lane"),
7811 Some("invalid_lane"),
7812 )
7813 })?),
7814 };
7815 match tenant.lane_class {
7816 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
7817 auth::LaneClass::Batch => match requested {
7818 None => Ok(lanes::Lane::Harvest),
7819 Some(lanes::Lane::Interactive) => Err(error_response(
7820 StatusCode::FORBIDDEN,
7821 "this api key is batch-class: x-lane interactive is not permitted \
7822 (use judge or harvest)",
7823 "authentication_error",
7824 Some("x-lane"),
7825 )),
7826 Some(l) => Ok(l),
7827 },
7828 }
7829}
7830
7831fn tenant_namespace(
7835 tenant: &auth::TenantCtx,
7836 cache_salt: &Option<String>,
7837) -> Result<String, &'static str> {
7838 let keyring_configured = auth::global().is_some();
7839 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
7840 if keyring_configured {
7841 Ok(auth::scope_namespace(&tenant.tenant, &raw))
7842 } else {
7843 Ok(raw)
7844 }
7845}
7846
7847fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
7852 eprintln!(
7853 "[meter] admit id={} tenant={} lane={} model={:?}",
7854 env.id,
7855 tenant.tenant,
7856 lane.as_str(),
7857 model
7858 );
7859}
7860
7861fn apply_model_request_limits(
7862 request: &mut Request,
7863 metadata: Option<&OpenRouterModelMetadata>,
7864 caps: Option<&ModelCaps>,
7865) -> Result<(), (String, &'static str)> {
7866 let Some(metadata) = metadata else {
7867 return Ok(());
7868 };
7869 let max_prompt = metadata
7870 .max_prompt_length
7871 .map(usize::try_from)
7872 .transpose()
7873 .map_err(|_| {
7874 (
7875 "configured model prompt limit does not fit this platform".into(),
7876 "model",
7877 )
7878 })?;
7879 let max_output = metadata
7880 .max_output_length
7881 .map(usize::try_from)
7882 .transpose()
7883 .map_err(|_| {
7884 (
7885 "configured model output limit does not fit this platform".into(),
7886 "model",
7887 )
7888 })?;
7889
7890 request.max_prompt_tokens = max_prompt;
7891 if let Some(max_output) = max_output {
7892 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
7893 request.params.max_new = metadata
7894 .default_output_length
7895 .map(usize::try_from)
7896 .transpose()
7897 .map_err(|_| {
7898 (
7899 "configured default output length does not fit this platform".into(),
7900 "model",
7901 )
7902 })?
7903 .unwrap_or(max_output);
7904 } else if request.params.max_new > max_output {
7905 return Err((
7906 format!(
7907 "max_tokens {} exceeds configured model maximum {max_output}",
7908 request.params.max_new
7909 ),
7910 "max_tokens",
7911 ));
7912 }
7913 }
7914
7915 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
7919 (max_prompt, max_output, request.params.max_ctx)
7920 {
7921 let operational_ctx = max_prompt
7922 .checked_add(max_output)
7923 .and_then(|value| value.checked_add(8))
7924 .ok_or_else(|| {
7925 (
7926 "configured model context envelope overflowed".into(),
7927 "model",
7928 )
7929 })?;
7930 let operational_ctx = caps
7931 .map(|caps| caps.context_length)
7932 .filter(|&context| context > 0)
7933 .map_or(operational_ctx, |context| operational_ctx.min(context));
7934 if requested_ctx > operational_ctx {
7935 return Err((
7936 format!(
7937 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
7938 ),
7939 "max_ctx",
7940 ));
7941 }
7942 }
7943 Ok(())
7944}
7945
7946fn effective_max_tokens(request: &worker::Request) -> Option<u64> {
7950 (request.params.max_new != worker::MAX_NEW_CTX_BOUNDED).then_some(request.params.max_new as u64)
7951}
7952
7953#[allow(clippy::too_many_arguments)]
7954fn start_request_receipt(
7955 st: &AppState,
7956 env: &Envelope,
7957 tenant: &auth::TenantCtx,
7958 model: &str,
7959 route: &'static str,
7960 lane: lanes::Lane,
7961 stream: bool,
7962 max_tokens: Option<u64>,
7963 reserved_ctx: Option<u64>,
7964 budget_permit: Option<metering::Permit>,
7965) -> Option<Box<dyn metering::Receipt>> {
7966 st.metering.as_ref().map(|accounting| {
7967 accounting.open(
7968 &metering::RequestMeta {
7969 request_id: &env.id,
7970 tenant: &tenant.tenant,
7971 principal: tenant.key_prefix.as_deref(),
7972 model,
7973 route,
7974 lane: lane.as_str(),
7975 stream,
7976 max_tokens,
7977 reserved_ctx,
7978 },
7979 budget_permit,
7980 )
7981 })
7982}
7983
7984fn arm_capture(
7990 mut receipt: Option<Box<dyn metering::Receipt>>,
7991 prompt: impl FnOnce() -> serde_json::Value,
7992) -> Option<Box<dyn metering::Receipt>> {
7993 if let Some(receipt) = receipt.as_mut()
7994 && receipt.wants_capture()
7995 {
7996 receipt.arm_capture(prompt());
7997 }
7998 receipt
7999}
8000
8001fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
8005 serde_json::Value::Array(
8006 messages
8007 .iter()
8008 .map(|message| {
8009 let mut row = json!({ "role": message.role, "content": message.content });
8010 if !message.tool_calls.is_empty() {
8011 row["tool_calls"] = serde_json::Value::Array(
8012 message
8013 .tool_calls
8014 .iter()
8015 .map(|call| {
8016 json!({
8017 "id": call.id,
8018 "function": {
8019 "name": call.function.name,
8020 "arguments": call.function.arguments,
8021 },
8022 })
8023 })
8024 .collect(),
8025 );
8026 }
8027 row
8028 })
8029 .collect(),
8030 )
8031}
8032
8033enum BudgetRejection {
8034 Invalid(String),
8035 Insufficient,
8036 Unenrolled,
8037 PrincipalCapped,
8040 Unavailable(String),
8041}
8042
8043impl BudgetRejection {
8044 fn into_response(self) -> (Response, &'static str) {
8045 match self {
8046 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
8047 Self::Insufficient => (
8048 error_response_coded(
8049 StatusCode::PAYMENT_REQUIRED,
8050 "tenant prepaid balance is insufficient for this request",
8051 "insufficient_balance",
8052 None,
8053 Some("insufficient_balance"),
8054 ),
8055 "insufficient_balance",
8056 ),
8057 Self::Unenrolled => (
8058 error_response_coded(
8059 StatusCode::PAYMENT_REQUIRED,
8060 "tenant is not enrolled for prepaid billing",
8061 "tenant_not_enrolled",
8062 None,
8063 Some("tenant_not_enrolled"),
8064 ),
8065 "tenant_not_enrolled",
8066 ),
8067 Self::PrincipalCapped => (
8068 error_response_coded(
8069 StatusCode::PAYMENT_REQUIRED,
8070 "this API key's spend cap is reached; raise or clear the key's cap to continue",
8071 "key_spend_cap_reached",
8072 None,
8073 Some("key_spend_cap_reached"),
8074 ),
8075 "key_spend_cap_reached",
8076 ),
8077 Self::Unavailable(err) => {
8078 eprintln!("[budget] ERROR: admission unavailable: {err}");
8079 (
8080 error_response_coded(
8081 StatusCode::SERVICE_UNAVAILABLE,
8082 "tenant budget accounting is unavailable",
8083 "server_error",
8084 None,
8085 Some("tenant_budget_unavailable"),
8086 ),
8087 "tenant_budget_unavailable",
8088 )
8089 }
8090 }
8091 }
8092}
8093
8094fn prepare_budget_prompt(
8095 request: &mut Request,
8096 tokenizer: Option<&Tokenizer>,
8097) -> Result<usize, String> {
8098 if let Some(error) = worker::prompt_source_limit_error(request) {
8099 return Err(error);
8100 }
8101 if request.prepared_prompt.is_none() {
8102 if let Some(trace) = request.ttft.as_ref() {
8103 trace.mark_tokenize_start();
8104 }
8105 let prompt = if !request.prompt_ids.is_empty() {
8106 request.prompt_ids.clone()
8107 } else if !request.chat_turns.is_empty() {
8108 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8109 let plain = worker::plain_chat_render_path(
8116 &request.tools_json,
8117 &request.think,
8118 request.reasoning_effort.as_deref(),
8119 &request.chat_turns,
8120 tokenizer.has_qwen_effort_ladder(),
8121 );
8122 let rendered = if plain {
8123 let messages: Vec<_> = request
8124 .chat_turns
8125 .iter()
8126 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
8127 .collect();
8128 tokenizer.apply_chat_template(&messages, true)
8129 } else {
8130 tokenizer
8131 .apply_chat_template_tools_ex(
8132 &request.chat_turns,
8133 true,
8134 &request.tools_json,
8135 &request.tools_struct,
8136 request.think,
8137 request.reasoning_effort.as_deref(),
8138 )
8139 .map_err(|err| format!("chat template: {err}"))?
8140 };
8141 tokenizer.encode(&rendered, true)
8142 } else if request.chat {
8143 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8144 let rendered =
8145 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
8146 tokenizer.encode(&rendered, true)
8147 } else {
8148 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8149 tokenizer.encode(&request.prompt_text, true)
8150 };
8151 if prompt.is_empty() {
8152 return Err("empty prompt after tokenization".into());
8153 }
8154 if let Some(trace) = request.ttft.as_ref() {
8155 trace.mark_tokenize_end(prompt.len());
8156 }
8157 request.prepared_prompt = Some(prompt);
8158 }
8159 let prompt_tokens = request
8160 .prepared_prompt
8161 .as_ref()
8162 .expect("budget prompt was prepared")
8163 .len();
8164 if let Some(limit) = request.max_prompt_tokens
8165 && prompt_tokens > limit
8166 {
8167 return Err(format!(
8168 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
8169 ));
8170 }
8171 Ok(prompt_tokens)
8172}
8173
8174fn budget_completion_bound(
8175 request: &Request,
8176 prompt_tokens: usize,
8177 caps: Option<&ModelCaps>,
8178) -> Result<usize, String> {
8179 let max_new = request.params.max_new;
8180 let requested_ctx = match (request.params.max_ctx, max_new) {
8181 (Some(cap), _) => cap,
8182 (None, worker::MAX_NEW_CTX_BOUNDED) => {
8183 let server_ctx = std::env::var("MEMRA_CTX")
8184 .ok()
8185 .and_then(|value| value.parse().ok())
8186 .unwrap_or(8192usize);
8187 let mut cap = server_ctx;
8188 if prompt_tokens.saturating_add(16) > cap {
8189 cap = prompt_tokens.saturating_add(server_ctx);
8190 }
8191 cap
8192 }
8193 (None, max_new) => prompt_tokens
8194 .checked_add(max_new)
8195 .and_then(|value| value.checked_add(8))
8196 .ok_or_else(|| "request context bound overflowed".to_string())?,
8197 };
8198 let ctx_cap = caps
8199 .map(|caps| caps.context_length)
8200 .filter(|&context| context > 0)
8201 .map_or(requested_ctx, |context| requested_ctx.min(context));
8202 if prompt_tokens >= ctx_cap {
8203 return Err(format!(
8204 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
8205 ));
8206 }
8207 Ok(max_new.min(ctx_cap - prompt_tokens))
8208}
8209
8210struct BudgetAdmission {
8215 permit: Option<metering::Permit>,
8216 reserved_ctx: Option<u64>,
8217}
8218
8219impl std::fmt::Debug for BudgetAdmission {
8221 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
8222 f.debug_struct("BudgetAdmission")
8223 .field("permit", &self.permit.is_some())
8224 .field("reserved_ctx", &self.reserved_ctx)
8225 .finish()
8226 }
8227}
8228
8229fn admit_tenant_budget(
8230 st: &AppState,
8231 tenant: &auth::TenantCtx,
8232 request: &mut Request,
8233) -> Result<BudgetAdmission, BudgetRejection> {
8234 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
8235 return Ok(BudgetAdmission {
8236 permit: None,
8237 reserved_ctx: None,
8238 });
8239 };
8240 match accounting.is_limited(&tenant.tenant) {
8241 Ok(false) => return Err(BudgetRejection::Unenrolled),
8242 Ok(true) => {}
8243 Err(metering::AdmitError::Unavailable(err)) => {
8244 return Err(BudgetRejection::Unavailable(err));
8245 }
8246 Err(other) => {
8247 return Err(BudgetRejection::Unavailable(format!(
8248 "unexpected budget enrollment result: {other:?}"
8249 )));
8250 }
8251 }
8252 let tokenizer = st
8253 .budget_tokenizers
8254 .as_ref()
8255 .and_then(|tokenizers| tokenizers.get(&request.model))
8256 .map(Arc::as_ref);
8257 if request.prompt_ids.is_empty() && tokenizer.is_none() {
8258 return Err(BudgetRejection::Unavailable(format!(
8259 "no reservation tokenizer for model {:?}",
8260 request.model
8261 )));
8262 }
8263 let prompt_tokens =
8264 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
8265 let completion_tokens =
8266 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
8267 .map_err(BudgetRejection::Invalid)?;
8268 let prompt_tokens = u64::try_from(prompt_tokens)
8269 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
8270 let completion_tokens = u64::try_from(completion_tokens)
8271 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
8272 match accounting.reserve(
8273 &tenant.tenant,
8274 tenant.key_prefix.as_deref(),
8275 &request.model,
8276 prompt_tokens,
8277 completion_tokens,
8278 ) {
8279 Ok(permit) => Ok(BudgetAdmission {
8280 permit,
8281 reserved_ctx: Some(prompt_tokens.saturating_add(completion_tokens)),
8282 }),
8283 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
8284 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
8285 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
8289 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
8290 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
8291 }
8292}
8293
8294fn request_ledger_error_response() -> Response {
8295 error_response_coded(
8296 StatusCode::INTERNAL_SERVER_ERROR,
8297 "request completion could not be committed to the billing ledger",
8298 "server_error",
8299 None,
8300 Some("request_ledger_unavailable"),
8301 )
8302}
8303
8304fn request_ledger_error_body() -> serde_json::Value {
8305 error_body(
8306 "request completion could not be committed to the billing ledger",
8307 "server_error",
8308 None,
8309 Some("request_ledger_unavailable"),
8310 )
8311}
8312
8313fn ledger_rejected(
8314 mut receipt: Option<Box<dyn metering::Receipt>>,
8315 response: Response,
8316 error_code: &str,
8317 request_id: &str,
8318) -> Response {
8319 let status = response.status().as_u16();
8320 if let Some(receipt) = receipt.as_mut()
8321 && let Err(err) = receipt.reject(status, error_code)
8322 {
8323 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
8324 return with_request_id(request_id, request_ledger_error_response());
8325 }
8326 with_request_id(request_id, response)
8327}
8328
8329fn ledger_unbilled(
8334 mut receipt: Option<Box<dyn metering::Receipt>>,
8335 response: Response,
8336 outcome: &'static str,
8337 error_code: &str,
8338 request_id: &str,
8339) -> Response {
8340 let status = response.status().as_u16();
8341 if let Some(receipt) = receipt.as_mut()
8342 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
8343 {
8344 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
8345 return with_request_id(request_id, request_ledger_error_response());
8346 }
8347 with_request_id(request_id, response)
8348}
8349
8350fn engine_error_code(class: worker::ErrClass) -> &'static str {
8351 use worker::ErrClass as C;
8352 match class {
8353 C::InvalidRequest => "invalid_request",
8354 C::ContextLength => "context_length_exceeded",
8355 C::ModelNotFound => "model_not_found",
8356 C::RateLimit => "rate_limit_exceeded",
8357 C::Overloaded => "overloaded",
8358 C::Engine => "engine_error",
8359 }
8360}
8361
8362fn model_not_found_response(models: &[String], requested: &str) -> Response {
8382 error_response_coded(
8383 StatusCode::BAD_REQUEST,
8384 &format!("unknown model {requested:?}; loaded: {models:?}"),
8385 "invalid_request_error",
8386 Some("model"),
8387 Some("model_not_found"),
8388 )
8389}
8390
8391fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
8400 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
8401 return Ok(());
8402 };
8403 if let Some((pos, &id)) = ids
8404 .iter()
8405 .enumerate()
8406 .find(|&(_, &id)| id as usize >= n_vocab)
8407 {
8408 return Err(format!(
8409 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
8410 ));
8411 }
8412 Ok(())
8413}
8414
8415#[cfg(test)]
8416mod prompt_ids_tests {
8417 use super::*;
8418
8419 #[test]
8420 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
8421 let caps = ModelCaps {
8422 n_vocab: 8,
8423 ..Default::default()
8424 };
8425 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
8427 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
8428 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
8430 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
8431 assert!(err.contains("vocab size 8"), "{err}");
8432 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
8433 assert!(err.contains("4294967295"), "{err}");
8434 let unknown = ModelCaps::default();
8436 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
8437 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
8438 }
8439}
8440
8441fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
8442 if models.iter().any(|m| m == requested) {
8443 return Some(requested.to_string());
8444 }
8445 if requested.is_empty() || requested.contains('/') {
8446 return None;
8447 }
8448 let mut matches = models.iter().filter(|m| {
8449 m.rsplit('/')
8450 .next()
8451 .is_some_and(|suffix| suffix == requested)
8452 });
8453 match (matches.next(), matches.next()) {
8454 (Some(only), None) => Some(only.clone()),
8455 _ => None,
8456 }
8457}
8458
8459async fn completions_admitted(
8460 state: State<AppState>,
8461 headers: axum::http::HeaderMap,
8462 trace: Option<Extension<TtftRequestTrace>>,
8463 AdmittedJson(req, admission): AdmittedJson<CompletionReq>,
8464) -> Response {
8465 completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
8466}
8467
8468#[cfg(test)]
8469async fn completions(
8470 State(st): State<AppState>,
8471 headers: axum::http::HeaderMap,
8472 trace: Option<Extension<TtftRequestTrace>>,
8473 request: Json<CompletionReq>,
8474) -> Response {
8475 completions_with_admission(State(st), headers, trace, request, None).await
8476}
8477
8478async fn completions_with_admission(
8479 State(st): State<AppState>,
8480 headers: axum::http::HeaderMap,
8481 trace: Option<Extension<TtftRequestTrace>>,
8482 Json(mut req): Json<CompletionReq>,
8483 mut body_admission: Option<BodyAdmissionLease>,
8484) -> Response {
8485 let env = Envelope::new(false);
8486 if let Err(msg) = req.stop.validate() {
8487 return with_request_id(&env.id, bad_request(&msg, Some("stop")));
8488 }
8489 if let Err(msg) = validate_client_identifier(req.trace_id.as_deref(), "trace_id") {
8490 return with_request_id(&env.id, bad_request(&msg, Some("trace_id")));
8491 }
8492 match canonical_model_id(&st.models, &req.model) {
8493 Some(canonical) => req.model = canonical,
8494 None => {
8495 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
8496 }
8497 }
8498 let ttft = trace.and_then(|Extension(trace)| trace.0);
8501 if let Some(trace) = ttft.as_ref() {
8502 trace.mark_parsed();
8503 trace.bind_request(&env.id, &req.model);
8504 }
8505 let tenant = match authenticate(&st.api_auth, &headers) {
8506 Ok(t) => t,
8507 Err(resp) => return with_request_id(&env.id, resp),
8508 };
8509 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
8510 Ok(ns) => ns,
8511 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
8512 };
8513 if let Err((msg, param)) = reject_unsupported(&[
8515 (
8516 "logit_bias",
8517 req.logit_bias.is_some(),
8518 " (device-side sampling has no bias hook yet)",
8519 ),
8520 ("logprobs", req.logprobs.is_some(), ""),
8521 (
8522 "n",
8523 req.n.is_some_and(|n| n != 1),
8524 " for n != 1 (single choice only)",
8525 ),
8526 (
8527 "best_of",
8528 req.best_of.is_some_and(|n| n != 1),
8529 " (single choice only)",
8530 ),
8531 ]) {
8532 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8533 }
8534 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
8537 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
8538 }
8539 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
8542 Ok(ms) => RequestDeadline::starting_now(ms),
8543 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8544 };
8545 let lane = match lane_for_tenant(&headers, &tenant) {
8546 Ok(l) => l,
8547 Err(resp) => return resp,
8548 };
8549 let (tx, rx) = worker::event_channel();
8550 let model = req.model.clone();
8551 let stream = req.stream;
8552 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
8553 Ok(affinity) => affinity,
8554 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
8555 };
8556 let mut request = build_request_with_trace(
8557 &req,
8558 tx,
8559 lane,
8560 affinity,
8561 ttft.clone(),
8562 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
8566 );
8567 request.cache_ns = cache_ns;
8568 request.request_id = env.id.clone();
8569 request.wire_deadline = Some(deadline.at.into_std());
8572 if let Err((message, param)) = apply_model_request_limits(
8573 &mut request,
8574 st.openrouter_metadata.get(&model),
8575 st.caps.get(&model),
8576 ) {
8577 return with_request_id(&env.id, bad_request(&message, Some(param)));
8578 }
8579 if let Err(msg) = nonstream_deadline_gate(
8584 &request,
8585 req.stream,
8586 deadline,
8587 req.max_tokens.is_some(),
8588 st.budget_tokenizers
8589 .as_ref()
8590 .and_then(|t| t.get(&req.model))
8591 .map(Arc::as_ref),
8592 ) {
8593 return with_request_id(
8594 &env.id,
8595 error_response_coded(
8596 StatusCode::BAD_REQUEST,
8597 &msg,
8598 "invalid_request_error",
8599 Some("max_tokens"),
8600 Some("nonstream_deadline_infeasible"),
8601 ),
8602 );
8603 }
8604 if draining() {
8607 let receipt = start_request_receipt(
8608 &st,
8609 &env,
8610 &tenant,
8611 &req.model,
8612 "/v1/completions",
8613 lane,
8614 req.stream,
8615 effective_max_tokens(&request),
8616 None,
8617 None,
8618 );
8619 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
8620 }
8621 let budget = match admit_tenant_budget(&st, &tenant, &mut request) {
8622 Ok(budget) => budget,
8623 Err(rejection) => {
8624 let (response, error_code) = rejection.into_response();
8625 let receipt = start_request_receipt(
8626 &st,
8627 &env,
8628 &tenant,
8629 &req.model,
8630 "/v1/completions",
8631 lane,
8632 req.stream,
8633 effective_max_tokens(&request),
8634 None,
8635 None,
8636 );
8637 return ledger_rejected(receipt, response, error_code, &env.id);
8638 }
8639 };
8640 let receipt = start_request_receipt(
8641 &st,
8642 &env,
8643 &tenant,
8644 &req.model,
8645 "/v1/completions",
8646 lane,
8647 req.stream,
8648 effective_max_tokens(&request),
8649 budget.reserved_ctx,
8650 budget.permit,
8651 );
8652 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
8653 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
8656 Ok(slot) => slot,
8657 Err(resp) => {
8658 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
8659 }
8660 };
8661 if let Some(admission) = body_admission.as_mut() {
8662 admission.release();
8663 }
8664 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
8667 Ok(guard) => guard,
8668 Err((resp, outcome)) => {
8669 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
8670 }
8671 };
8672 meter_admit(&env, &tenant, &model, lane);
8673 let stop_strings = request.stop_strings.clone();
8674
8675 if let Some(trace) = ttft.as_ref() {
8680 trace.mark_submitted();
8681 }
8682 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
8683 drop(pending_admit);
8684 return ledger_rejected(
8685 receipt,
8686 rl.attach(worker_unavailable_response()),
8687 "worker_unavailable",
8688 &env.id,
8689 );
8690 }
8691 pending_admit.commit();
8692 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
8696 Ok(Ok(rx)) => rx,
8697 Ok(Err((resp, error_code))) => {
8698 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
8699 }
8700 Err(_) => {
8701 return ledger_unbilled(
8702 receipt,
8703 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8704 "deadline_exceeded",
8705 "deadline_exceeded",
8706 &env.id,
8707 );
8708 }
8709 };
8710
8711 let resp = if stream {
8712 let rx = match peek_first_token(rx, deadline).await {
8716 Ok(rx) => rx,
8717 Err(()) => {
8718 return ledger_unbilled(
8719 receipt,
8720 rl.attach(deadline_exceeded_response(deadline.ms, true)),
8721 "deadline_exceeded",
8722 "deadline_exceeded",
8723 &env.id,
8724 );
8725 }
8726 };
8727 sse_response_with_receipt(
8728 rx,
8729 model,
8730 false,
8731 None,
8732 env.clone(),
8733 stop_strings,
8734 Some(guard),
8735 receipt,
8736 )
8737 .into_response()
8738 } else {
8739 let mut receipt = receipt;
8745 let resp = blocking_response_with_receipt(
8746 rx,
8747 model,
8748 false,
8749 stop_strings,
8750 None,
8751 env.clone(),
8752 &mut receipt,
8753 Some(deadline),
8754 )
8755 .await;
8756 drop(guard); resp.into_response()
8758 };
8759 rl.attach(with_request_id(&env.id, resp))
8760}
8761
8762async fn chat_completions_admitted(
8763 state: State<AppState>,
8764 headers: axum::http::HeaderMap,
8765 trace: Option<Extension<TtftRequestTrace>>,
8766 AdmittedJson(req, admission): AdmittedJson<ChatCompletionReq>,
8767) -> Response {
8768 chat_completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
8769}
8770
8771#[cfg(test)]
8772async fn chat_completions(
8773 State(st): State<AppState>,
8774 headers: axum::http::HeaderMap,
8775 trace: Option<Extension<TtftRequestTrace>>,
8776 request: Json<ChatCompletionReq>,
8777) -> Response {
8778 chat_completions_with_admission(State(st), headers, trace, request, None).await
8779}
8780
8781async fn chat_completions_with_admission(
8782 State(st): State<AppState>,
8783 headers: axum::http::HeaderMap,
8784 trace: Option<Extension<TtftRequestTrace>>,
8785 Json(mut req): Json<ChatCompletionReq>,
8786 mut body_admission: Option<BodyAdmissionLease>,
8787) -> Response {
8788 let env = Envelope::new(true);
8789 match canonical_model_id(&st.models, &req.model) {
8794 Some(canonical) => req.model = canonical,
8795 None => {
8796 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
8797 }
8798 }
8799 let ttft = trace.and_then(|Extension(trace)| trace.0);
8800 if let Some(trace) = ttft.as_ref() {
8801 trace.mark_parsed();
8802 trace.bind_request(&env.id, &req.model);
8803 }
8804 let tenant = match authenticate(&st.api_auth, &headers) {
8805 Ok(t) => t,
8806 Err(resp) => return with_request_id(&env.id, resp),
8807 };
8808 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
8809 Ok(ns) => ns,
8810 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
8811 };
8812 if req.messages.is_empty()
8813 || req.messages.iter().any(|message| {
8814 !matches!(
8815 message.role.as_str(),
8816 "system" | "developer" | "user" | "assistant" | "tool"
8817 )
8818 })
8819 {
8820 return with_request_id(
8821 &env.id,
8822 bad_request(
8823 "messages must use system/developer/user/assistant/tool roles",
8824 Some("messages"),
8825 ),
8826 );
8827 }
8828 if let Err((msg, param)) = reject_unsupported(&[
8833 (
8834 "logit_bias",
8835 req.logit_bias.is_some(),
8836 " (device-side sampling has no bias hook yet)",
8837 ),
8838 (
8839 "logprobs",
8840 req.logprobs
8841 .as_ref()
8842 .is_some_and(|v| v.as_bool() != Some(false)),
8843 "",
8844 ),
8845 ("top_logprobs", req.top_logprobs.is_some(), ""),
8846 (
8847 "n",
8848 req.n.is_some_and(|n| n != 1),
8849 " for n != 1 (single choice only)",
8850 ),
8851 ]) {
8852 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8853 }
8854 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
8857 Ok(ms) => RequestDeadline::starting_now(ms),
8858 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8859 };
8860 let lane = match lane_for_tenant(&headers, &tenant) {
8861 Ok(l) => l,
8862 Err(resp) => return resp,
8863 };
8864 let model = req.model.clone();
8865 let stream = req.stream;
8866 let capture_prompt = st
8869 .metering
8870 .as_ref()
8871 .filter(|m| m.captures(&tenant.tenant))
8872 .map(|_| capture_chat_messages(&req.messages));
8873 let declared_max_tokens = req.max_tokens.is_some();
8877 let vision_preprocess_permit = match try_vision_preprocess(request_has_vision(&req)) {
8881 Ok(permit) => permit,
8882 Err(response) => return with_request_id(&env.id, response),
8883 };
8884 let (tx, rx) = worker::event_channel();
8885 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
8886 Ok(affinity) => affinity,
8887 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
8888 };
8889 let mut plan = match build_chat_request_with_trace(
8890 req,
8891 st.caps.get(&model),
8892 tx,
8893 lane,
8894 affinity,
8895 ttft.clone(),
8896 st.openrouter_metadata
8897 .get(&model)
8898 .and_then(|m| m.default_reasoning_effort.as_deref()),
8899 &st.sampling_defaults(&model),
8900 ) {
8901 Ok(plan) => plan,
8902 Err(err) => {
8903 return with_request_id(&env.id, bad_request(&err, None));
8904 }
8905 };
8906 plan.request.cache_ns = cache_ns;
8907 plan.request.request_id = env.id.clone();
8908 plan.request.wire_deadline = Some(deadline.at.into_std());
8909 if let Err((message, param)) = apply_model_request_limits(
8910 &mut plan.request,
8911 st.openrouter_metadata.get(&model),
8912 st.caps.get(&model),
8913 ) {
8914 return with_request_id(&env.id, bad_request(&message, Some(param)));
8915 }
8916 if let Err(msg) = nonstream_deadline_gate(
8919 &plan.request,
8920 stream,
8921 deadline,
8922 declared_max_tokens,
8923 st.budget_tokenizers
8924 .as_ref()
8925 .and_then(|t| t.get(&model))
8926 .map(Arc::as_ref),
8927 ) {
8928 return with_request_id(
8929 &env.id,
8930 error_response_coded(
8931 StatusCode::BAD_REQUEST,
8932 &msg,
8933 "invalid_request_error",
8934 Some("max_tokens"),
8935 Some("nonstream_deadline_infeasible"),
8936 ),
8937 );
8938 }
8939 plan.vision_memory = match reserve_vision_memory(&plan) {
8940 Ok(permit) => permit,
8941 Err(err) => {
8942 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
8943 }
8944 };
8945 if draining() {
8948 let receipt = start_request_receipt(
8949 &st,
8950 &env,
8951 &tenant,
8952 &model,
8953 "/v1/chat/completions",
8954 lane,
8955 stream,
8956 effective_max_tokens(&plan.request),
8957 None,
8958 None,
8959 );
8960 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
8961 }
8962 let budget = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
8963 Ok(budget) => budget,
8964 Err(rejection) => {
8965 let (response, error_code) = rejection.into_response();
8966 let receipt = start_request_receipt(
8967 &st,
8968 &env,
8969 &tenant,
8970 &model,
8971 "/v1/chat/completions",
8972 lane,
8973 stream,
8974 effective_max_tokens(&plan.request),
8975 None,
8976 None,
8977 );
8978 return ledger_rejected(receipt, response, error_code, &env.id);
8979 }
8980 };
8981 let receipt = start_request_receipt(
8982 &st,
8983 &env,
8984 &tenant,
8985 &model,
8986 "/v1/chat/completions",
8987 lane,
8988 stream,
8989 effective_max_tokens(&plan.request),
8990 budget.reserved_ctx,
8991 budget.permit,
8992 );
8993 let receipt = if let Some(prompt) = capture_prompt {
8994 arm_capture(receipt, move || prompt)
8995 } else {
8996 receipt
8997 };
8998 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
9002 Ok(slot) => slot,
9003 Err(resp) => {
9004 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
9005 }
9006 };
9007 if let Some(admission) = body_admission.as_mut() {
9008 admission.release();
9009 }
9010 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
9013 Ok(guard) => guard,
9014 Err((resp, outcome)) => {
9015 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
9016 }
9017 };
9018 if let Err(err) = decode_pending_vision(&mut plan) {
9023 return ledger_rejected(
9024 receipt,
9025 rl.attach(bad_request(&err, Some("messages"))),
9026 "invalid_request_error",
9027 &env.id,
9028 );
9029 }
9030 plan.request.vision_memory = plan.vision_memory.take();
9031 drop(vision_preprocess_permit);
9032 let constraint_ready = if plan.request.grammar.is_some() {
9033 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
9034 plan.request.constraint_ready = Some(ready_tx);
9035 Some(ready_rx)
9036 } else {
9037 None
9038 };
9039 meter_admit(&env, &tenant, &model, lane);
9040 let stop_strings = plan.request.stop_strings.clone();
9041 if let Some(trace) = ttft.as_ref() {
9043 trace.mark_submitted();
9044 }
9045 if st
9046 .cmd_tx
9047 .send(Cmd::Generate(Box::new(plan.request)))
9048 .is_err()
9049 {
9050 drop(pending_admit);
9051 return ledger_rejected(
9052 receipt,
9053 rl.attach(worker_unavailable_response()),
9054 "worker_unavailable",
9055 &env.id,
9056 );
9057 }
9058 pending_admit.commit();
9059 if let Some(ready) = constraint_ready {
9065 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
9066 match tokio::time::timeout(bound, ready).await {
9067 Ok(Ok(Ok(()))) => {}
9068 Ok(Ok(Err(err))) => {
9069 return ledger_rejected(
9070 receipt,
9071 rl.attach(engine_error_response(&err)),
9072 engine_error_code(err.class),
9073 &env.id,
9074 );
9075 }
9076 Ok(Err(_)) => {
9077 return ledger_rejected(
9078 receipt,
9079 rl.attach(worker_unavailable_response()),
9080 "worker_unavailable",
9081 &env.id,
9082 );
9083 }
9084 Err(_) if deadline.remaining().is_zero() => {
9085 return ledger_unbilled(
9086 receipt,
9087 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
9088 "deadline_exceeded",
9089 "deadline_exceeded",
9090 &env.id,
9091 );
9092 }
9093 Err(_) => {
9094 return ledger_rejected(
9095 receipt,
9096 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
9097 "constraint_compile_timeout",
9098 &env.id,
9099 );
9100 }
9101 }
9102 }
9103 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
9105 Ok(Ok(rx)) => rx,
9106 Ok(Err((resp, error_code))) => {
9107 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
9108 }
9109 Err(_) => {
9110 return ledger_unbilled(
9111 receipt,
9112 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
9113 "deadline_exceeded",
9114 "deadline_exceeded",
9115 &env.id,
9116 );
9117 }
9118 };
9119 let resp = if stream {
9120 let rx = match peek_first_token(rx, deadline).await {
9122 Ok(rx) => rx,
9123 Err(()) => {
9124 return ledger_unbilled(
9125 receipt,
9126 rl.attach(deadline_exceeded_response(deadline.ms, true)),
9127 "deadline_exceeded",
9128 "deadline_exceeded",
9129 &env.id,
9130 );
9131 }
9132 };
9133 sse_response_with_receipt(
9134 rx,
9135 model,
9136 true,
9137 plan.parser,
9138 env.clone(),
9139 stop_strings,
9140 Some(guard),
9141 receipt,
9142 )
9143 .into_response()
9144 } else {
9145 let mut receipt = receipt;
9148 let resp = blocking_response_with_receipt(
9149 rx,
9150 model,
9151 true,
9152 stop_strings,
9153 plan.parser,
9154 env.clone(),
9155 &mut receipt,
9156 Some(deadline),
9157 )
9158 .await;
9159 drop(guard); resp.into_response()
9161 };
9162 rl.attach(with_request_id(&env.id, resp))
9163}
9164
9165#[cfg(test)]
9174fn sse_response(
9175 rx: worker::EventReceiver,
9176 model: String,
9177 chat: bool,
9178 parser: Option<ToolStreamParser>,
9179 env: Envelope,
9180 stop_strings: Vec<String>,
9181 guard: Option<InflightGuard>,
9182) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9183 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
9184}
9185
9186#[allow(clippy::too_many_arguments)] fn sse_response_with_receipt(
9188 mut rx: worker::EventReceiver,
9189 model: String,
9190 chat: bool,
9191 mut parser: Option<ToolStreamParser>,
9192 env: Envelope,
9193 stop_strings: Vec<String>,
9194 guard: Option<InflightGuard>,
9195 mut receipt: Option<Box<dyn metering::Receipt>>,
9196) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9197 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
9201 .then(|| StopScrubber::new(stop_strings));
9202 let stream = async_stream::stream! {
9203 let _guard = guard;
9206 let mut call_index: usize = 0;
9207 let mut role_sent = false;
9210 macro_rules! chat_chunk {
9211 ($delta:expr, $finish:expr) => {{
9212 let mut delta = $delta;
9213 if chat && !role_sent {
9214 role_sent = true;
9215 delta["role"] = json!("assistant");
9216 }
9217 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
9218 "choices": [{ "index": 0, "delta": delta,
9219 "finish_reason": $finish }] }))
9220 .to_string()
9221 }};
9222 }
9223 macro_rules! piece_chunks {
9225 ($piece:expr) => {{
9226 let mut payloads: Vec<String> = Vec::new();
9227 match $piece {
9228 Piece::Content(text) => {
9229 let text = match scrub.as_mut() {
9230 Some(sc) => sc.push(&text),
9231 None => text,
9232 };
9233 if !text.is_empty() {
9234 payloads.push(chat_chunk!(json!({ "content": text }),
9235 serde_json::Value::Null));
9236 }
9237 }
9238 Piece::Reasoning(text) => payloads.push(
9242 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
9243 Piece::Call(call) => {
9244 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9245 "index": call_index, "id": call.id, "type": "function",
9246 "function": { "name": call.name, "arguments": "" } }] }),
9247 serde_json::Value::Null));
9248 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9249 "index": call_index,
9250 "function": { "arguments": call.arguments } }] }),
9251 serde_json::Value::Null));
9252 call_index += 1;
9253 }
9254 }
9255 payloads
9256 }};
9257 }
9258 let mut terminal = false;
9262 while let Some(ev) = rx.recv().await {
9263 match ev {
9264 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9266 if let Some(receipt) = receipt.as_mut()
9267 && let Err(err) = receipt.record_prompt_usage(
9268 n_prompt as u64,
9269 n_cached as u64,
9270 )
9271 {
9272 eprintln!(
9273 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9274 env.id
9275 );
9276 let _ = receipt.reject(500, "request_ledger_unavailable");
9279 let payload = request_ledger_error_body().to_string();
9280 if chat || openai_compat() {
9281 yield Ok(SseEvent::default().data(payload));
9282 yield Ok(SseEvent::default().data("[DONE]"));
9283 } else {
9284 yield Ok(SseEvent::default().event("error").data(payload));
9285 }
9286 terminal = true;
9287 break;
9288 }
9289 }
9290 Event::Token { id, text } => {
9291 if let Some(receipt) = receipt.as_mut()
9292 && let Err(err) = receipt.record_completion_token()
9293 {
9294 eprintln!(
9295 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9296 env.id
9297 );
9298 let _ = receipt.reject(500, "request_ledger_unavailable");
9299 let payload = request_ledger_error_body().to_string();
9300 if chat || openai_compat() {
9301 yield Ok(SseEvent::default().data(payload));
9302 yield Ok(SseEvent::default().data("[DONE]"));
9303 } else {
9304 yield Ok(SseEvent::default().event("error").data(payload));
9305 }
9306 terminal = true;
9307 break;
9308 }
9309 if let Some(receipt) = receipt.as_mut() {
9312 receipt.capture_completion_delta(&text);
9313 }
9314 if let Some(p) = parser.as_mut() {
9315 for piece in p.push(&text) {
9316 for payload in piece_chunks!(piece) {
9317 yield Ok(SseEvent::default().data(payload));
9318 }
9319 }
9320 continue;
9321 }
9322 let text = match scrub.as_mut() {
9323 Some(sc) => sc.push(&text),
9324 None => text,
9325 };
9326 if text.is_empty() && scrub.is_some() {
9327 continue; }
9329 let payload = if chat {
9330 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
9331 } else if openai_compat() {
9332 env.stamp(json!({ "object": "text_completion", "model": model,
9333 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
9334 .to_string()
9335 } else {
9336 json!({ "model": model, "id": id, "text": text }).to_string()
9337 };
9338 yield Ok(SseEvent::default().data(payload));
9339 }
9340 Event::TokenSnapshot(_) => {}
9344 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
9345 let mut finish = stop_reason_to_finish(&stop_reason);
9346 if let Some(p) = parser.as_mut() {
9347 for piece in p.finish() {
9348 for payload in piece_chunks!(piece) {
9349 yield Ok(SseEvent::default().data(payload));
9350 }
9351 }
9352 if p.n_calls() > 0 { finish = "tool_calls"; }
9353 }
9354 if let Some(sc) = scrub.as_mut() {
9356 let tail = sc.finish();
9357 if !tail.is_empty() {
9358 let payload = if chat {
9359 chat_chunk!(json!({ "content": tail }),
9360 serde_json::Value::Null)
9361 } else {
9362 env.stamp(json!({ "object": "text_completion",
9363 "model": model,
9364 "choices": [{ "index": 0, "text": tail,
9365 "finish_reason": null }] })).to_string()
9366 };
9367 yield Ok(SseEvent::default().data(payload));
9368 }
9369 }
9370 if let Some(receipt) = receipt.as_mut()
9371 && let Err(err) = receipt.complete(
9372 metering::UsageCounts {
9373 prompt_tokens: n_prompt as u64,
9374 cached_prompt_tokens: n_cached as u64,
9375 completion_tokens: n_tokens as u64,
9376 },
9377 elapsed_s,
9378 )
9379 {
9380 eprintln!(
9381 "[ledger] ERROR: request {} completion receipt failed: {err}",
9382 env.id
9383 );
9384 let _ = receipt.reject(500, "request_ledger_unavailable");
9388 let payload = request_ledger_error_body().to_string();
9389 if chat || openai_compat() {
9390 yield Ok(SseEvent::default().data(payload));
9391 yield Ok(SseEvent::default().data("[DONE]"));
9392 } else {
9393 yield Ok(SseEvent::default().event("error").data(payload));
9394 }
9395 terminal = true;
9396 break;
9397 }
9398 if chat || openai_compat() {
9399 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
9400 let fin = if chat {
9401 let mut v = env.stamp(json!({
9402 "object": "chat.completion.chunk", "model": model,
9403 "choices": [{ "index": 0, "delta": {},
9404 "finish_reason": finish }],
9405 "usage": usage }));
9406 if !role_sent {
9408 v["choices"][0]["delta"]["role"] = json!("assistant");
9409 }
9410 v
9411 } else {
9412 env.stamp(json!({ "object": "text_completion", "model": model,
9413 "choices": [{ "index": 0, "text": "",
9414 "finish_reason": finish }],
9415 "usage": usage }))
9416 }.to_string();
9417 yield Ok(SseEvent::default().data(fin));
9418 yield Ok(SseEvent::default().data("[DONE]"));
9419 } else {
9420 let payload = json!({
9421 "stop_reason": stop_reason, "n_tokens": n_tokens,
9422 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
9423 "elapsed_s": elapsed_s
9424 }).to_string();
9425 yield Ok(SseEvent::default().event("done").data(payload));
9426 }
9427 terminal = true;
9428 break;
9429 }
9430 Event::Error(err) => {
9431 let ledger_error = if let Some(receipt) = receipt.as_mut() {
9441 receipt
9442 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
9443 .err()
9444 } else {
9445 None
9446 };
9447 if let Some(ref ledger_error) = ledger_error {
9448 eprintln!(
9449 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
9450 env.id
9451 );
9452 }
9453 let payload = if ledger_error.is_some() {
9454 request_ledger_error_body().to_string()
9455 } else {
9456 engine_error_body(&err).to_string()
9457 };
9458 if chat || openai_compat() {
9459 yield Ok(SseEvent::default().data(payload));
9462 yield Ok(SseEvent::default().data("[DONE]"));
9463 } else {
9464 yield Ok(SseEvent::default().event("error").data(payload));
9467 }
9468 terminal = true;
9469 break;
9470 }
9471 }
9472 }
9473 if !terminal {
9474 let e = worker::EngineError::overloaded(
9480 "worker closed the stream without completing (worker restart in progress)",
9481 );
9482 if let Some(receipt) = receipt.as_mut()
9483 && let Err(ledger_err) = receipt.reject(
9484 class_http(e.class).0.as_u16(),
9485 engine_error_code(e.class),
9486 )
9487 {
9488 eprintln!(
9489 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9490 env.id
9491 );
9492 }
9493 let payload = engine_error_body(&e).to_string();
9494 if chat || openai_compat() {
9495 yield Ok(SseEvent::default().data(payload));
9496 yield Ok(SseEvent::default().data("[DONE]"));
9497 } else {
9498 yield Ok(SseEvent::default().event("error").data(payload));
9499 }
9500 }
9501 };
9502 Sse::new(stream).keep_alive(
9503 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
9506 )
9507}
9508
9509fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
9511 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
9512 text.truncate(offset);
9513 }
9514}
9515
9516fn partial_stop_suffix(s: &str, tag: &str) -> usize {
9519 let mut best = 0;
9520 for (k, _) in tag.char_indices().skip(1) {
9521 if k <= s.len() && s.ends_with(&tag[..k]) {
9522 best = k;
9523 }
9524 }
9525 best
9526}
9527
9528struct StopScrubber {
9534 stops: Vec<String>,
9535 buf: String,
9536 done: bool,
9537}
9538
9539impl StopScrubber {
9540 fn new(stops: Vec<String>) -> Self {
9541 Self {
9542 stops,
9543 buf: String::new(),
9544 done: false,
9545 }
9546 }
9547
9548 fn push(&mut self, text: &str) -> String {
9550 if self.done {
9551 return String::new();
9552 }
9553 self.buf.push_str(text);
9554 if let Some(i) = self
9555 .stops
9556 .iter()
9557 .filter_map(|s| self.buf.find(s.as_str()))
9558 .min()
9559 {
9560 self.done = true;
9561 let out = self.buf[..i].to_string();
9562 self.buf.clear();
9563 return out;
9564 }
9565 let keep = self
9566 .stops
9567 .iter()
9568 .map(|s| partial_stop_suffix(&self.buf, s))
9569 .max()
9570 .unwrap_or(0);
9571 let emit_to = self.buf.len() - keep;
9572 let out = self.buf[..emit_to].to_string();
9573 self.buf.drain(..emit_to);
9574 out
9575 }
9576
9577 fn finish(&mut self) -> String {
9579 if self.done {
9580 self.buf.clear();
9581 return String::new();
9582 }
9583 std::mem::take(&mut self.buf)
9584 }
9585}
9586
9587#[cfg(test)]
9588async fn blocking_response(
9589 rx: worker::EventReceiver,
9590 model: String,
9591 chat: bool,
9592 stop_strings: Vec<String>,
9593 parser: Option<ToolStreamParser>,
9594 env: Envelope,
9595) -> Response {
9596 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
9597 .await
9598}
9599
9600struct BlockingPayload<'a> {
9604 env: &'a Envelope,
9605 model: String,
9606 chat: bool,
9607 finish: &'static str,
9608 text: String,
9609 reasoning: String,
9610 calls: Vec<ParsedToolCall>,
9611 tokens: Vec<u32>,
9612 stop_reason: String,
9613 n_prompt: usize,
9614 n_tokens: usize,
9615 n_cached: usize,
9616 elapsed_s: f64,
9617 spec: Option<worker::SpecUsage>,
9618 deadline_error: Option<serde_json::Value>,
9624}
9625
9626fn blocking_payload(p: BlockingPayload<'_>) -> Response {
9627 let BlockingPayload {
9628 env,
9629 model,
9630 chat,
9631 finish,
9632 text,
9633 reasoning,
9634 calls,
9635 tokens,
9636 stop_reason,
9637 n_prompt,
9638 n_tokens,
9639 n_cached,
9640 elapsed_s,
9641 spec,
9642 deadline_error,
9643 } = p;
9644 if chat {
9645 let content = if !calls.is_empty() && text.is_empty() {
9647 serde_json::Value::Null
9648 } else {
9649 serde_json::Value::String(text)
9650 };
9651 let mut message = json!({ "role": "assistant", "content": content });
9652 if !reasoning.is_empty() {
9655 message["reasoning"] = json!(reasoning);
9656 message["reasoning_details"] = json!([{
9657 "type": "reasoning.text", "text": reasoning }]);
9658 }
9659 if !calls.is_empty() {
9660 message["tool_calls"] =
9661 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
9662 }
9663 let mut body = json!({
9664 "object": "chat.completion", "model": model,
9665 "choices": [{ "index": 0,
9666 "message": message,
9667 "finish_reason": finish }],
9668 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
9669 });
9670 if let Some(err) = deadline_error {
9671 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
9672 body["error"] = err;
9673 }
9674 return Json(env.stamp(body)).into_response();
9675 }
9676 if openai_compat() {
9677 let mut body = json!({
9678 "object": "text_completion", "model": model,
9679 "choices": [{ "index": 0, "text": text,
9680 "finish_reason": finish }],
9681 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
9682 });
9683 if let Some(err) = deadline_error {
9684 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
9685 body["error"] = err;
9686 }
9687 return Json(env.stamp(body)).into_response();
9688 }
9689 Json(CompletionResp {
9690 model,
9691 text,
9692 tokens,
9693 stop_reason,
9694 error: deadline_error,
9695 n_tokens,
9696 prompt_tokens: n_prompt,
9697 cached_tokens: n_cached,
9698 elapsed_s,
9699 })
9700 .into_response()
9701}
9702
9703#[allow(clippy::too_many_arguments)] async fn blocking_response_with_receipt(
9729 mut rx: worker::EventReceiver,
9730 model: String,
9731 chat: bool,
9732 stop_strings: Vec<String>,
9733 mut parser: Option<ToolStreamParser>,
9734 env: Envelope,
9735 receipt: &mut Option<Box<dyn metering::Receipt>>,
9736 deadline: Option<RequestDeadline>,
9737) -> Response {
9738 let mut text = String::new();
9739 let mut reasoning = String::new();
9740 let mut tokens: Vec<u32> = Vec::new();
9741 let mut calls: Vec<ParsedToolCall> = Vec::new();
9742 let consume = |pieces: Vec<Piece>,
9743 text: &mut String,
9744 reasoning: &mut String,
9745 calls: &mut Vec<ParsedToolCall>| {
9746 for piece in pieces {
9747 match piece {
9748 Piece::Content(t) => text.push_str(&t),
9749 Piece::Reasoning(t) => reasoning.push_str(&t),
9750 Piece::Call(c) => calls.push(c),
9751 }
9752 }
9753 };
9754 let started = std::time::Instant::now();
9756 let mut seen_prompt: usize = 0;
9757 let mut seen_cached: usize = 0;
9758 let mut seen_tokens: usize = 0;
9759 loop {
9760 let ev = match deadline {
9761 Some(d) => tokio::select! {
9762 biased;
9763 ev = rx.recv() => ev,
9764 () = tokio::time::sleep_until(d.at) => {
9765 drop(rx);
9768 if seen_tokens == 0 {
9769 if let Some(receipt) = receipt.as_mut()
9773 && let Err(err) = receipt.settle_unbilled(
9774 "deadline_exceeded",
9775 StatusCode::REQUEST_TIMEOUT.as_u16(),
9776 "deadline_exceeded",
9777 )
9778 {
9779 eprintln!(
9780 "[ledger] ERROR: request {} deadline receipt failed: {err}",
9781 env.id
9782 );
9783 return request_ledger_error_response();
9784 }
9785 return deadline_exceeded_response(d.ms, false);
9786 }
9787 if let Some(p) = parser.as_mut() {
9788 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
9789 }
9790 truncate_at_stop(&mut text, &stop_strings);
9791 let elapsed_s = started.elapsed().as_secs_f64();
9792 if let Some(receipt) = receipt.as_mut()
9795 && let Err(err) = receipt.complete_deadline_partial(
9796 metering::UsageCounts {
9797 prompt_tokens: seen_prompt as u64,
9798 cached_prompt_tokens: seen_cached as u64,
9799 completion_tokens: seen_tokens as u64,
9800 },
9801 elapsed_s,
9802 )
9803 {
9804 eprintln!(
9805 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
9806 env.id
9807 );
9808 let _ = receipt.reject(500, "request_ledger_unavailable");
9809 return request_ledger_error_response();
9810 }
9811 eprintln!(
9812 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
9813 {} ms deadline (prompt {}); non-streaming caller advised to stream",
9814 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
9815 );
9816 let err_obj = json!({
9817 "message": format!(
9818 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
9819 the {} tokens produced before the cut are delivered above and are \
9820 billed. Set \"stream\": true for work this long — a stream's \
9821 deadline bounds only the time to first token — or lower max_tokens.",
9822 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
9823 ),
9824 "code": "deadline_exceeded",
9825 "metadata": { "error_type": "timeout", "provider_name": "memra" }
9826 });
9827 return blocking_payload(BlockingPayload {
9828 env: &env,
9829 model,
9830 chat,
9831 finish: "error",
9832 text,
9833 reasoning,
9834 calls,
9835 tokens,
9836 stop_reason: "Deadline".to_string(),
9837 n_prompt: seen_prompt,
9838 n_tokens: seen_tokens,
9839 n_cached: seen_cached,
9840 elapsed_s,
9841 spec: None,
9842 deadline_error: Some(err_obj),
9843 });
9844 }
9845 },
9846 None => rx.recv().await,
9847 };
9848 let Some(ev) = ev else { break };
9849 match ev {
9850 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9852 if let Some(receipt) = receipt.as_mut()
9853 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
9854 {
9855 eprintln!(
9856 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9857 env.id
9858 );
9859 let _ = receipt.reject(500, "request_ledger_unavailable");
9862 return request_ledger_error_response();
9863 }
9864 seen_prompt = n_prompt;
9865 seen_cached = n_cached;
9866 }
9867 Event::Token { id, text: delta } => {
9868 if let Some(receipt) = receipt.as_mut()
9869 && let Err(err) = receipt.record_completion_token()
9870 {
9871 eprintln!(
9872 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9873 env.id
9874 );
9875 let _ = receipt.reject(500, "request_ledger_unavailable");
9876 return request_ledger_error_response();
9877 }
9878 if let Some(receipt) = receipt.as_mut() {
9880 receipt.capture_completion_delta(&delta);
9881 }
9882 tokens.push(id);
9883 seen_tokens += 1;
9884 match parser.as_mut() {
9885 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
9886 None => text.push_str(&delta),
9887 }
9888 }
9889 Event::TokenSnapshot(ids) => tokens = ids,
9890 Event::Done {
9891 stop_reason,
9892 n_tokens,
9893 n_prompt,
9894 n_cached,
9895 elapsed_s,
9896 spec,
9897 } => {
9898 if let Some(p) = parser.as_mut() {
9899 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
9900 }
9901 truncate_at_stop(&mut text, &stop_strings);
9902 let finish = if calls.is_empty() {
9903 stop_reason_to_finish(&stop_reason)
9904 } else {
9905 "tool_calls"
9906 };
9907 if let Some(receipt) = receipt.as_mut()
9908 && let Err(err) = receipt.complete(
9909 metering::UsageCounts {
9910 prompt_tokens: n_prompt as u64,
9911 cached_prompt_tokens: n_cached as u64,
9912 completion_tokens: n_tokens as u64,
9913 },
9914 elapsed_s,
9915 )
9916 {
9917 eprintln!(
9918 "[ledger] ERROR: request {} completion receipt failed: {err}",
9919 env.id
9920 );
9921 let _ = receipt.reject(500, "request_ledger_unavailable");
9924 return request_ledger_error_response();
9925 }
9926 return blocking_payload(BlockingPayload {
9927 env: &env,
9928 model,
9929 chat,
9930 finish,
9931 text,
9932 reasoning,
9933 calls,
9934 tokens,
9935 stop_reason,
9936 n_prompt,
9937 n_tokens,
9938 n_cached,
9939 elapsed_s,
9940 spec,
9941 deadline_error: None,
9942 });
9943 }
9944 Event::Error(err) => {
9945 if let Some(receipt) = receipt.as_mut()
9949 && let Err(ledger_err) = receipt.reject(
9950 class_http(err.class).0.as_u16(),
9951 engine_error_code(err.class),
9952 )
9953 {
9954 eprintln!(
9955 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
9956 env.id
9957 );
9958 return request_ledger_error_response();
9959 }
9960 return engine_error_response(&err);
9961 }
9962 }
9963 }
9964 let e = worker::EngineError::overloaded(
9969 "worker closed the stream without completing (worker restart in progress)",
9970 );
9971 if let Some(receipt) = receipt.as_mut()
9972 && let Err(ledger_err) =
9973 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
9974 {
9975 eprintln!(
9976 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9977 env.id
9978 );
9979 return request_ledger_error_response();
9980 }
9981 engine_error_response(&e)
9982}
9983
9984#[cfg(test)]
9985mod tests {
9986 use super::*;
9987
9988 #[test]
9994 fn capture_children_are_distinct_ledger_identities_under_the_parent_id() {
9995 let parent = Envelope::new(false);
9996 assert!(parent.id.starts_with("cmpl-"));
9997 let a = parent.capture_child(0);
9998 let b = parent.capture_child(1);
9999 let c = parent.capture_child(2);
10000 assert_eq!(a.id, format!("{}.0", parent.id));
10001 assert_eq!(b.id, format!("{}.1", parent.id));
10002 assert_eq!(c.id, format!("{}.2", parent.id));
10003 assert_ne!(a.id, b.id);
10004 assert_ne!(b.id, c.id);
10005 for child in [&a, &b, &c] {
10006 assert!(
10007 child.id.starts_with(&parent.id),
10008 "child nests under the parent by prefix"
10009 );
10010 assert_ne!(
10011 child.id, parent.id,
10012 "a child never reuses the parent's ledger id"
10013 );
10014 assert_eq!(child.created, parent.created);
10015 }
10016 assert_eq!(parent.capture_child(1).id, b.id);
10019 }
10020
10021 #[derive(Debug, Clone, PartialEq)]
10029 enum MeterEvent {
10030 Reserve {
10031 tenant: String,
10032 principal: Option<String>,
10033 model: String,
10034 },
10035 Open {
10036 request_id: String,
10037 tenant: String,
10038 model: String,
10039 route: &'static str,
10040 stream: bool,
10041 with_permit: bool,
10042 },
10043 PromptUsage {
10044 prompt: u64,
10045 cached: u64,
10046 },
10047 Token,
10048 CapturePrompt(serde_json::Value),
10049 CaptureDelta(String),
10050 Complete {
10051 prompt: u64,
10052 cached: u64,
10053 completion: u64,
10054 },
10055 DeadlinePartial {
10056 prompt: u64,
10057 cached: u64,
10058 completion: u64,
10059 },
10060 Reject {
10061 status: u16,
10062 code: String,
10063 },
10064 Unbilled {
10065 outcome: &'static str,
10066 status: u16,
10067 code: String,
10068 },
10069 Dropped {
10072 prompt: u64,
10073 cached: u64,
10074 completion: u64,
10075 },
10076 }
10077
10078 enum ReserveScript {
10081 Admit { with_permit: bool },
10082 Insufficient,
10083 Blocked,
10084 PrincipalCapped,
10085 }
10086
10087 struct MockMetering {
10088 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10089 limits: bool,
10090 limited: bool,
10091 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
10092 captures: bool,
10093 }
10094
10095 impl MockMetering {
10096 fn admit_all() -> Arc<Self> {
10097 Arc::new(MockMetering {
10098 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10099 limits: false,
10100 limited: true,
10101 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10102 captures: false,
10103 })
10104 }
10105
10106 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
10107 Arc::new(MockMetering {
10108 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10109 limits: true,
10110 limited: true,
10111 reserve_script: std::sync::Mutex::new(script.into()),
10112 captures: false,
10113 })
10114 }
10115
10116 fn capturing() -> Arc<Self> {
10117 Arc::new(MockMetering {
10118 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10119 limits: false,
10120 limited: true,
10121 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10122 captures: true,
10123 })
10124 }
10125
10126 fn events(&self) -> Vec<MeterEvent> {
10127 self.events.lock().unwrap().clone()
10128 }
10129 }
10130
10131 impl metering::Metering for MockMetering {
10132 fn enforces_limits(&self) -> bool {
10133 self.limits
10134 }
10135
10136 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
10137 Ok(self.limited)
10138 }
10139
10140 fn reserve(
10141 &self,
10142 tenant: &str,
10143 principal: Option<&str>,
10144 model: &str,
10145 _prompt_tokens: u64,
10146 _completion_bound: u64,
10147 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
10148 self.events.lock().unwrap().push(MeterEvent::Reserve {
10149 tenant: tenant.into(),
10150 principal: principal.map(str::to_owned),
10151 model: model.into(),
10152 });
10153 match self.reserve_script.lock().unwrap().pop_front() {
10154 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
10155 Some(ReserveScript::Admit { with_permit: true }) => {
10156 Ok(Some(Box::new(()) as metering::Permit))
10157 }
10158 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
10159 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
10160 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
10161 }
10162 }
10163
10164 fn open(
10165 &self,
10166 meta: &metering::RequestMeta<'_>,
10167 permit: Option<metering::Permit>,
10168 ) -> Box<dyn metering::Receipt> {
10169 self.events.lock().unwrap().push(MeterEvent::Open {
10170 request_id: meta.request_id.into(),
10171 tenant: meta.tenant.into(),
10172 model: meta.model.into(),
10173 route: meta.route,
10174 stream: meta.stream,
10175 with_permit: permit.is_some(),
10176 });
10177 Box::new(MockReceipt {
10178 events: self.events.clone(),
10179 wants_capture: self.captures,
10180 prompt: 0,
10181 cached: 0,
10182 completion: 0,
10183 finalized: false,
10184 })
10185 }
10186
10187 fn captures(&self, _tenant: &str) -> bool {
10188 self.captures
10189 }
10190
10191 fn limits_health(&self) -> Option<metering::LimitsHealth> {
10192 self.limits.then_some(metering::LimitsHealth {
10193 source_reload_failed: 0,
10194 source_reload_consecutive: 0,
10195 source_available: true,
10196 })
10197 }
10198 }
10199
10200 struct MockReceipt {
10201 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10202 wants_capture: bool,
10203 prompt: u64,
10204 cached: u64,
10205 completion: u64,
10206 finalized: bool,
10207 }
10208
10209 impl metering::Receipt for MockReceipt {
10210 fn wants_capture(&self) -> bool {
10211 self.wants_capture
10212 }
10213
10214 fn arm_capture(&mut self, prompt: serde_json::Value) {
10215 self.events
10216 .lock()
10217 .unwrap()
10218 .push(MeterEvent::CapturePrompt(prompt));
10219 }
10220
10221 fn capture_completion_delta(&mut self, text: &str) {
10222 if self.wants_capture {
10223 self.events
10224 .lock()
10225 .unwrap()
10226 .push(MeterEvent::CaptureDelta(text.into()));
10227 }
10228 }
10229
10230 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
10231 self.prompt = prompt;
10232 self.cached = cached;
10233 self.events
10234 .lock()
10235 .unwrap()
10236 .push(MeterEvent::PromptUsage { prompt, cached });
10237 Ok(())
10238 }
10239
10240 fn record_completion_token(&mut self) -> Result<(), String> {
10241 self.completion += 1;
10242 self.events.lock().unwrap().push(MeterEvent::Token);
10243 Ok(())
10244 }
10245
10246 fn complete(
10247 &mut self,
10248 usage: metering::UsageCounts,
10249 _worker_elapsed_s: f64,
10250 ) -> Result<(), String> {
10251 self.finalized = true;
10252 self.events.lock().unwrap().push(MeterEvent::Complete {
10253 prompt: usage.prompt_tokens,
10254 cached: usage.cached_prompt_tokens,
10255 completion: usage.completion_tokens,
10256 });
10257 Ok(())
10258 }
10259
10260 fn complete_deadline_partial(
10261 &mut self,
10262 usage: metering::UsageCounts,
10263 _worker_elapsed_s: f64,
10264 ) -> Result<(), String> {
10265 self.finalized = true;
10266 self.events
10267 .lock()
10268 .unwrap()
10269 .push(MeterEvent::DeadlinePartial {
10270 prompt: usage.prompt_tokens,
10271 cached: usage.cached_prompt_tokens,
10272 completion: usage.completion_tokens,
10273 });
10274 Ok(())
10275 }
10276
10277 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
10278 self.finalized = true;
10279 self.events.lock().unwrap().push(MeterEvent::Reject {
10280 status,
10281 code: error_code.into(),
10282 });
10283 Ok(())
10284 }
10285
10286 fn settle_unbilled(
10287 &mut self,
10288 outcome: &'static str,
10289 status: u16,
10290 error_code: &str,
10291 ) -> Result<(), String> {
10292 self.finalized = true;
10293 self.events.lock().unwrap().push(MeterEvent::Unbilled {
10294 outcome,
10295 status,
10296 code: error_code.into(),
10297 });
10298 Ok(())
10299 }
10300 }
10301
10302 impl Drop for MockReceipt {
10303 fn drop(&mut self) {
10304 if !self.finalized {
10305 self.events.lock().unwrap().push(MeterEvent::Dropped {
10306 prompt: self.prompt,
10307 cached: self.cached,
10308 completion: self.completion,
10309 });
10310 }
10311 }
10312 }
10313
10314 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
10319
10320 fn gate_env_lock() -> std::sync::MutexGuard<'static, ()> {
10327 let guard = GATE_ENV_LOCK.lock().unwrap_or_else(|poisoned| {
10328 GATE_ENV_LOCK.clear_poison();
10331 poisoned.into_inner()
10332 });
10333 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10334 guard
10335 }
10336
10337 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
10340 let req: CompletionReq = serde_json::from_value(json!({
10341 "model": "qwen/qwen3.8-27b",
10342 "prompt_ids": vec![7u32; prompt_ids],
10343 }))
10344 .unwrap();
10345 let (tx, _rx) = worker::event_channel();
10346 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
10347 request.params.max_new = max_new;
10348 request
10349 }
10350
10351 #[test]
10358 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
10359 let prompt = 30_278u64;
10360 let deadline_ms = TIMEOUT_MS_DEFAULT;
10361 let margin = |max_new: u64| {
10362 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
10363 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
10364 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
10365 };
10366 for allowed in [64u64, 2048, 4096, 5120, 6144] {
10367 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
10368 }
10369 for refused in [8192u64, 16384, 262_144] {
10370 assert!(
10371 !margin(refused),
10372 "{refused} measured as a 408 and must be refused"
10373 );
10374 }
10375 }
10376
10377 #[test]
10378 fn the_gate_names_a_max_tokens_that_actually_fits() {
10379 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
10382 assert!(
10383 fits > 0 && fits < 7_800,
10384 "advice {fits} must fit the measured ceiling"
10385 );
10386 assert_eq!(
10388 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
10389 None
10390 );
10391 }
10392
10393 #[test]
10394 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
10395 let req = gate_request(262_144, 30_000);
10396 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
10397 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
10399 assert!(
10400 err.contains("stream"),
10401 "message must name the streaming alternative: {err}"
10402 );
10403 assert!(
10404 err.contains("max_tokens"),
10405 "message must name the knob: {err}"
10406 );
10407 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
10409 let _l = gate_env_lock(); for off in ["0", "off", "false"] {
10416 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
10417 assert!(
10418 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
10419 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
10420 );
10421 }
10422 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
10423 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
10424 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10425 assert!(
10426 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
10427 "unset means ON (the documented default)"
10428 );
10429 }
10430
10431 #[test]
10438 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
10439 let strip = |src: &str| -> String {
10446 src.lines()
10447 .map(|line| match line.find("//") {
10448 Some(i) => line[..i].to_string(),
10449 None => line.to_string(),
10450 })
10451 .collect::<Vec<_>>()
10452 .join("\n")
10453 };
10454 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
10456 let start = src
10457 .find(signature)
10458 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
10459 let rest = &src[start + signature.len()..];
10460 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
10461 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
10462 &rest[..end]
10463 }
10464 let main_src = strip(include_str!("lib.rs"));
10465 let surfaces_src = strip(include_str!("surfaces.rs"));
10466 for (surface, src, signature) in [
10467 (
10468 "/v1/completions",
10469 &main_src,
10470 "async fn completions_with_admission(",
10471 ),
10472 (
10473 "/v1/chat/completions",
10474 &main_src,
10475 "async fn chat_completions_with_admission(",
10476 ),
10477 (
10478 "/v1/messages + /v1/responses (shared admission)",
10479 &surfaces_src,
10480 "pub(crate) async fn admit_translated(",
10481 ),
10482 ] {
10483 let handler = body(src, signature);
10484 assert!(
10485 handler.contains("nonstream_deadline_gate("),
10486 "{surface} must CALL the feasibility gate inside {signature}"
10487 );
10488 let limits = handler
10491 .find("apply_model_request_limits(")
10492 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
10493 let gate = handler.find("nonstream_deadline_gate(").unwrap();
10494 assert!(
10495 limits < gate,
10496 "{surface}: the gate must run after apply_model_request_limits"
10497 );
10498 }
10499 }
10500
10501 #[test]
10505 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
10506 let err = json!({"code": "deadline_exceeded",
10507 "metadata": {"error_type": "timeout"}});
10508 let cut = CompletionResp {
10509 model: "m".into(),
10510 text: "partial".into(),
10511 tokens: vec![1, 2],
10512 stop_reason: "Deadline".into(),
10513 error: Some(err.clone()),
10514 n_tokens: 2,
10515 prompt_tokens: 9,
10516 cached_tokens: 0,
10517 elapsed_s: 1.0,
10518 };
10519 let v = serde_json::to_value(&cut).unwrap();
10520 assert_eq!(v["stop_reason"], "Deadline");
10521 assert_eq!(v["error"]["code"], "deadline_exceeded");
10522 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
10523 let whole = CompletionResp {
10525 error: None,
10526 stop_reason: "Eos".into(),
10527 ..cut
10528 };
10529 let v = serde_json::to_value(&whole).unwrap();
10530 assert!(
10531 v.get("error").is_none(),
10532 "a complete response must not grow an error key: {v}"
10533 );
10534 }
10535
10536 #[test]
10537 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
10538 let _l = gate_env_lock();
10539 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
10543 assert!(
10544 nonstream_deadline_gate(
10545 &req,
10546 false,
10547 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10548 false,
10549 None,
10550 )
10551 .is_ok(),
10552 "an omitted max_tokens is never gated — context is its only limit"
10553 );
10554 let resolved = gate_request(32_768, 30_000);
10559 assert!(
10560 nonstream_deadline_gate(
10561 &resolved,
10562 false,
10563 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10564 false,
10565 None,
10566 )
10567 .is_ok(),
10568 "a resolved-but-undeclared cap is not the caller's number to be refused over"
10569 );
10570 assert!(
10572 nonstream_deadline_gate(
10573 &resolved,
10574 false,
10575 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10576 true,
10577 None,
10578 )
10579 .is_err()
10580 );
10581 }
10582
10583 #[test]
10584 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
10585 let req = gate_request(64, 1234);
10586 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
10587 let mut text = gate_request(64, 0);
10588 text.prompt_ids.clear();
10589 text.prompt_text = "x".repeat(6_000);
10590 assert_eq!(
10591 prompt_tokens_estimate(&text, None),
10592 1_000,
10593 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
10594 that would have succeeded"
10595 );
10596 }
10597
10598 #[test]
10599 fn vision_memory_reservation_is_bounded_and_released() {
10600 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
10601 let Err(capacity) = try_reserve_vision_memory(1) else {
10602 panic!("a full process vision budget admitted another request");
10603 };
10604 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
10605 let response = vision_memory_error_response(capacity, Some("messages"));
10606 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
10607 assert_eq!(response.headers()["retry-after"], "5");
10608 assert_eq!(response.headers()["retry-after-ms"], "5000");
10609 drop(permit);
10610 assert!(try_reserve_vision_memory(1).is_ok());
10611 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
10612 panic!("an over-limit vision request was admitted");
10613 };
10614 assert!(matches!(request, VisionMemoryError::Request(_)));
10615 let response = vision_memory_error_response(request, Some("messages"));
10616 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
10617 assert_eq!(response.headers()["x-should-retry"], "false");
10618 let _ = try_reserve_vision_memory(1);
10619 }
10620
10621 #[test]
10622 fn header_auth_gate_covers_only_inference_dialects() {
10623 for path in [
10624 "/v1/auth/check",
10625 "/v1/completions",
10626 "/v1/chat/completions",
10627 "/v1/messages",
10628 "/v1/responses",
10629 "/v1/embeddings",
10630 "/v1/rerank",
10631 ] {
10632 assert!(protected_inference_path(path), "{path}");
10633 }
10634 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
10635 assert!(!protected_inference_path(path), "{path}");
10636 }
10637 }
10638 #[tokio::test]
10645 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
10646 use crate::metering::Metering as _;
10647 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
10648
10649 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
10650 let (tx, rx) = worker::event_channel();
10651 tx.send(Event::PromptUsage {
10652 n_prompt: 7,
10653 n_cached: 0,
10654 })
10655 .unwrap();
10656 tx.send(Event::Token {
10657 id: 1,
10658 text: "Hel".into(),
10659 })
10660 .unwrap();
10661 tx.send(Event::Token {
10662 id: 2,
10663 text: "lo".into(),
10664 })
10665 .unwrap();
10666 tx.send(Event::Done {
10667 stop_reason: "eos".into(),
10668 n_tokens: 2,
10669 n_prompt: 7,
10670 n_cached: 0,
10671 elapsed_s: 0.05,
10672 spec: None,
10673 })
10674 .unwrap();
10675 drop(tx);
10676 let mut receipt = receipt;
10677 blocking_response_with_receipt(
10678 rx,
10679 "m".into(),
10680 true,
10681 Vec::new(),
10682 None,
10683 Envelope::new(true),
10684 &mut receipt,
10685 None,
10686 )
10687 .await
10688 };
10689
10690 let plain = MockMetering::admit_all();
10692 let receipt = plain.open(
10693 &metering::RequestMeta {
10694 request_id: "cap-unmarked",
10695 tenant: "unmarked",
10696 principal: None,
10697 model: "m",
10698 route: "/v1/chat/completions",
10699 lane: "interactive",
10700 stream: false,
10701 max_tokens: None,
10702 reserved_ctx: None,
10703 },
10704 None,
10705 );
10706 let response = drive(Some(receipt)).await;
10707 assert_eq!(response.status(), StatusCode::OK);
10708 assert!(
10709 !plain.events().iter().any(|e| matches!(
10710 e,
10711 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
10712 )),
10713 "an unarmed receipt must see no capture traffic: {:?}",
10714 plain.events()
10715 );
10716
10717 let capturing = MockMetering::capturing();
10720 let mut receipt = capturing.open(
10721 &metering::RequestMeta {
10722 request_id: "cap-marked",
10723 tenant: "marked",
10724 principal: None,
10725 model: "m",
10726 route: "/v1/chat/completions",
10727 lane: "interactive",
10728 stream: false,
10729 max_tokens: None,
10730 reserved_ctx: None,
10731 },
10732 None,
10733 );
10734 assert!(receipt.wants_capture());
10735 receipt.arm_capture(prompt.clone());
10736 let response = drive(Some(receipt)).await;
10737 assert_eq!(response.status(), StatusCode::OK);
10738 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
10739 .await
10740 .unwrap();
10741 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
10742 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
10743
10744 let events = capturing.events();
10745 assert!(
10746 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
10747 "prompt must arm byte-exact: {events:?}"
10748 );
10749 let completion: String = events
10750 .iter()
10751 .filter_map(|e| match e {
10752 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
10753 _ => None,
10754 })
10755 .collect();
10756 assert_eq!(
10757 completion, "Hello",
10758 "the deltas must reassemble the served completion byte-exact: {events:?}"
10759 );
10760 assert!(
10761 events.contains(&MeterEvent::Complete {
10762 prompt: 7,
10763 cached: 0,
10764 completion: 2,
10765 }),
10766 "worker-truth usage settles alongside the capture: {events:?}"
10767 );
10768 }
10769
10770 fn tool_caps() -> ModelCaps {
10771 ModelCaps {
10772 tools_branch: true,
10773 qwen_think: true,
10774 think_switch: true,
10775 chat_ok: true,
10776 ..Default::default()
10777 }
10778 }
10779
10780 fn ladder_caps() -> ModelCaps {
10785 ModelCaps {
10786 qwen_effort: true,
10787 ..tool_caps()
10788 }
10789 }
10790
10791 fn gemma_tool_caps() -> ModelCaps {
10792 ModelCaps {
10793 tools_branch: true,
10794 gemma_think: true,
10795 chat_ok: true,
10796 instruct_type: Some("gemma".into()),
10797 ..Default::default()
10798 }
10799 }
10800
10801 fn hy3_tool_caps() -> ModelCaps {
10802 ModelCaps {
10803 tools_branch: true,
10804 hy3: true,
10805 chat_ok: true,
10806 effort_levels: true,
10807 instruct_type: Some("hy3".into()),
10808 ..Default::default()
10809 }
10810 }
10811
10812 fn gemma_template(kind: &str) -> String {
10813 let file = match kind {
10814 "qat" => "qat-trunk-template.jinja",
10815 _ => "official-tooluse-template.jinja",
10816 };
10817 let path = format!(
10818 "{}/../../research/gemma4-tools-20260817/{file}",
10819 env!("CARGO_MANIFEST_DIR")
10820 );
10821 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
10822 }
10823
10824 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
10829 let tools_arr = request
10830 .get("tools")
10831 .and_then(|t| t.as_array())
10832 .cloned()
10833 .unwrap_or_default();
10834 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
10835 (Vec::new(), Vec::new(), HashMap::new())
10836 } else {
10837 prepare_tools(&tools_arr).unwrap()
10838 };
10839 let effort = request
10840 .get("reasoning_effort")
10841 .and_then(|v| v.as_str())
10842 .map(String::from);
10843 let (think, _lvl, _explicit) =
10844 parse_think(&effort, &None, None, None, None, false).unwrap();
10845
10846 let mut turns: Vec<TmplTurn> = Vec::new();
10847 for msg in request["messages"].as_array().unwrap() {
10848 let role = msg["role"].as_str().unwrap();
10849 let role = if role == "developer" { "system" } else { role };
10850 let content =
10851 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
10852 let tool_calls = msg
10853 .get("tool_calls")
10854 .and_then(|a| a.as_array())
10855 .map(|a| {
10856 a.iter()
10857 .map(|tc| {
10858 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
10859 render_req_tool_call(&rtc).unwrap()
10860 })
10861 .collect()
10862 })
10863 .unwrap_or_default();
10864 let tool_responses = msg
10865 .get("tool_responses")
10866 .and_then(|a| a.as_array())
10867 .map(|a| {
10868 a.iter()
10869 .map(|tr| {
10870 (
10871 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
10872 json_to_val(&tr["response"]),
10873 )
10874 })
10875 .collect()
10876 })
10877 .unwrap_or_default();
10878 turns.push(TmplTurn {
10879 role: role.to_string(),
10880 content,
10881 tool_calls,
10882 reasoning: msg
10883 .get("reasoning")
10884 .and_then(|r| r.as_str())
10885 .map(String::from)
10886 .filter(|s| !s.is_empty()),
10887 tool_call_id: msg
10888 .get("tool_call_id")
10889 .and_then(|s| s.as_str())
10890 .map(String::from),
10891 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
10892 tool_responses,
10893 task: None,
10894 tools: Vec::new(),
10895 });
10896 }
10897 chat::apply_chat_template_tools_ex(
10898 Some(template),
10899 &turns,
10900 true,
10901 &tools_json,
10902 &tools_struct,
10903 think,
10904 None,
10905 None,
10906 )
10907 .unwrap()
10908 }
10909
10910 #[test]
10914 fn gemma4_tools_fixtures_match_the_official_jinja() {
10915 let dir = format!(
10916 "{}/../../research/gemma4-tools-20260817/fixtures",
10917 env!("CARGO_MANIFEST_DIR")
10918 );
10919 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10920 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10921 .map(|e| e.unwrap().path())
10922 .filter(|p| p.is_dir())
10923 .collect();
10924 entries.sort();
10925 assert!(
10926 entries.len() >= 14,
10927 "expected >=14 fixtures, found {}",
10928 entries.len()
10929 );
10930 let (mut official, mut qat) = (0u32, 0u32);
10931 for d in entries {
10932 let input: serde_json::Value =
10933 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10934 .unwrap();
10935 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10936 let kind = input
10937 .get("template")
10938 .and_then(|t| t.as_str())
10939 .unwrap_or("official");
10940 match kind {
10941 "qat" => qat += 1,
10942 _ => official += 1,
10943 }
10944 let tmpl = gemma_template(kind);
10945 let got = render_fixture(&input["request"], &tmpl);
10946 assert_eq!(
10947 got, expected,
10948 "fixture {:?} diverged from the jinja oracle",
10949 d
10950 );
10951 }
10952 assert!(
10953 official >= 12 && qat >= 2,
10954 "coverage: {official} official, {qat} qat"
10955 );
10956 }
10957
10958 #[test]
10964 fn gemma4_tools_flow_through_build_chat_request() {
10965 let tmpl = gemma_template("official");
10966 for name in [
10967 "01-system-tools-basic",
10968 "04-single-call-cycle",
10969 "07-multi-cycle-agentic",
10970 ] {
10971 let path = format!(
10972 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
10973 env!("CARGO_MANIFEST_DIR")
10974 );
10975 let input: serde_json::Value =
10976 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
10977 let expected_path = format!(
10978 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
10979 env!("CARGO_MANIFEST_DIR")
10980 );
10981 let expected = std::fs::read_to_string(&expected_path).unwrap();
10982 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
10983 let (tx, _rx) = worker::event_channel();
10984 let plan = build_chat_request(
10985 req,
10986 Some(&gemma_tool_caps()),
10987 tx,
10988 lanes::Lane::Interactive,
10989 None,
10990 )
10991 .unwrap();
10992 let got = chat::apply_chat_template_tools_ex(
10993 Some(&tmpl),
10994 &plan.request.chat_turns,
10995 true,
10996 &plan.request.tools_json,
10997 &plan.request.tools_struct,
10998 plan.request.think,
10999 plan.request.reasoning_effort.as_deref(),
11000 None,
11001 )
11002 .unwrap();
11003 assert_eq!(got, expected, "pipeline render diverged for {name}");
11004 }
11005 }
11006
11007 fn glm5_template() -> String {
11016 let path = format!(
11017 "{}/../../research/glm53-flash-bringup-20260827/chat_template.jinja",
11018 env!("CARGO_MANIFEST_DIR")
11019 );
11020 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
11021 }
11022
11023 fn glm5_caps() -> ModelCaps {
11028 ModelCaps {
11029 tools_branch: true,
11030 qwen_think: true,
11031 think_switch: false,
11032 chat_ok: true,
11033 context_length: 1_048_576,
11034 tokenizer: "glm4".into(),
11035 instruct_type: Some("glm".into()),
11036 effort_levels: true,
11037 glm5: true,
11038 ..Default::default()
11039 }
11040 }
11041
11042 fn glm5_render(body: serde_json::Value) -> Result<String, String> {
11044 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11045 let (tx, _rx) = worker::event_channel();
11046 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)?;
11047 chat::apply_chat_template_tools_ex(
11048 Some(&glm5_template()),
11049 &plan.request.chat_turns,
11050 true,
11051 &plan.request.tools_json,
11052 &plan.request.tools_struct,
11053 plan.request.think,
11054 plan.request.reasoning_effort.as_deref(),
11055 None,
11056 )
11057 }
11058
11059 #[test]
11064 fn glm5_fixtures_match_the_vendor_jinja() {
11065 let dir = format!(
11066 "{}/../../research/glm53-flash-bringup-20260827/surface-fixtures",
11067 env!("CARGO_MANIFEST_DIR")
11068 );
11069 let mut entries: Vec<_> = std::fs::read_dir(&dir)
11070 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
11071 .map(|e| e.unwrap().path())
11072 .filter(|p| p.is_dir())
11073 .collect();
11074 entries.sort();
11075 assert!(
11076 entries.len() >= 22,
11077 "expected >=22 fixtures, found {}",
11078 entries.len()
11079 );
11080 for d in entries {
11081 let input: serde_json::Value =
11082 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11083 .unwrap();
11084 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11085 let got = glm5_render(input["request"].clone())
11086 .unwrap_or_else(|e| panic!("fixture {d:?} refused: {e}"));
11087 assert_eq!(
11088 got, expected,
11089 "fixture {d:?} diverged from the jinja oracle"
11090 );
11091 }
11092 }
11093
11094 #[test]
11100 fn glm5_never_renders_chatml() {
11101 let tmpl = glm5_template();
11102 assert!(tmpl.contains("<think>") && tmpl.contains("add_generation_prompt"));
11104 assert!(tmpl.contains("<tools>"));
11105 assert!(chat::template_is_glm5(&tmpl));
11106 for body in [
11107 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11108 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11109 "tools": [{"type": "function", "function": {"name": "f",
11110 "parameters": {"type": "object", "properties": {}}}}]}),
11111 ] {
11112 let got = glm5_render(body).unwrap();
11113 assert!(
11114 !got.contains("<|im_start|>") && !got.contains("<|im_end|>"),
11115 "glm5 rendered ChatML frames: {got:?}"
11116 );
11117 assert!(
11118 got.starts_with("[gMASK]<sop><|system|>Reasoning Effort: "),
11119 "{got:?}"
11120 );
11121 assert!(got.ends_with("<|assistant|><think>"), "{got:?}");
11122 }
11123 }
11124
11125 #[test]
11129 fn glm5_reasoning_effort_renders_and_keeps_its_max_tier() {
11130 for (sent, line) in [
11131 (None, "Max"),
11132 (Some("low"), "Low"),
11133 (Some("medium"), "High"),
11138 (Some("high"), "High"),
11139 (Some("xhigh"), "Max"),
11140 (Some("max"), "Max"),
11141 (Some("ultra"), "Max"),
11142 ] {
11143 let mut body = json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]});
11144 if let Some(v) = sent {
11145 body["reasoning_effort"] = json!(v);
11146 }
11147 let got = glm5_render(body).unwrap();
11148 assert!(
11149 got.starts_with(&format!("[gMASK]<sop><|system|>Reasoning Effort: {line}<|")),
11150 "reasoning_effort {sent:?} should render {line:?}: {got:?}"
11151 );
11152 }
11153 let sampler_of = |v: &str| {
11156 let req: ChatCompletionReq = serde_json::from_value(
11157 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11160 "reasoning_effort": v, "seed": 7}),
11161 )
11162 .unwrap();
11163 let (tx, _rx) = worker::event_channel();
11164 let plan =
11165 build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11166 .unwrap();
11167 format!("{:?}", plan.request.sampler_cfg)
11168 };
11169 assert_eq!(sampler_of("low"), sampler_of("max"));
11170 assert_eq!(canonical_effort_for("max", true), Some("max"));
11172 assert_eq!(canonical_effort_for("xhigh", true), Some("max"));
11173 assert_eq!(canonical_effort_for("max", false), Some("high"));
11174 }
11175
11176 #[test]
11180 fn glm5_refuses_what_its_template_cannot_honour() {
11181 for (value, needle) in [
11182 ("none", "cannot disable reasoning"),
11183 ("minimal", "cannot disable reasoning"),
11184 ("bogus", "bad reasoning_effort"),
11185 ] {
11186 let err = glm5_render(json!({"model": "m",
11187 "messages": [{"role": "user", "content": "hi"}],
11188 "reasoning_effort": value}))
11189 .err()
11190 .unwrap_or_else(|| panic!("reasoning_effort {value:?} must be refused"));
11191 assert!(err.contains(needle), "{value}: {err}");
11192 }
11193 }
11194
11195 #[test]
11199 fn one_glm5_request_renders_identical_bytes_on_all_three_surfaces() {
11200 let chat = json!({
11207 "model": "m",
11208 "reasoning_effort": "high",
11209 "messages": [
11210 {"role": "user", "content": "Weather in Paris and Rome?"},
11211 {"role": "assistant", "content": null,
11212 "tool_calls": [
11213 {"id": "c1", "type": "function",
11214 "function": {"name": "get_weather",
11215 "arguments": "{\"city\": \"Paris\"}"}},
11216 {"id": "c2", "type": "function",
11217 "function": {"name": "get_weather",
11218 "arguments": "{\"city\": \"Rome\"}"}}]},
11219 {"role": "tool", "tool_call_id": "c2", "content": "rome:27"},
11220 {"role": "tool", "tool_call_id": "c1", "content": "paris:21"}
11221 ],
11222 "tools": [{"type": "function", "function": {
11223 "name": "get_weather", "description": "Get the current weather for a city",
11224 "parameters": {"type": "object",
11225 "properties": {"city": {"type": "string"}},
11226 "required": ["city"]}}}]
11227 });
11228 let responses = responses_api::translate(&json!({
11229 "model": "m",
11230 "reasoning": {"effort": "high"},
11231 "input": [
11232 {"type": "message", "role": "user",
11233 "content": [{"type": "input_text", "text": "Weather in Paris and Rome?"}]},
11234 {"type": "function_call", "call_id": "c1", "name": "get_weather",
11235 "arguments": "{\"city\": \"Paris\"}"},
11236 {"type": "function_call", "call_id": "c2", "name": "get_weather",
11237 "arguments": "{\"city\": \"Rome\"}"},
11238 {"type": "function_call_output", "call_id": "c2", "output": "rome:27"},
11239 {"type": "function_call_output", "call_id": "c1", "output": "paris:21"}
11240 ],
11241 "tools": [{"type": "function", "name": "get_weather",
11242 "description": "Get the current weather for a city",
11243 "parameters": {"type": "object",
11244 "properties": {"city": {"type": "string"}},
11245 "required": ["city"]}}]
11246 }))
11247 .expect("/v1/responses translate");
11248 let messages = anthropic::translate(&json!({
11249 "model": "m",
11250 "max_tokens": 256,
11251 "output_config": {"effort": "high"},
11252 "messages": [
11253 {"role": "user", "content": "Weather in Paris and Rome?"},
11254 {"role": "assistant", "content": [
11255 {"type": "tool_use", "id": "c1", "name": "get_weather",
11256 "input": {"city": "Paris"}},
11257 {"type": "tool_use", "id": "c2", "name": "get_weather",
11258 "input": {"city": "Rome"}}]},
11259 {"role": "user", "content": [
11260 {"type": "tool_result", "tool_use_id": "c2", "content": "rome:27"},
11261 {"type": "tool_result", "tool_use_id": "c1", "content": "paris:21"}]}
11262 ],
11263 "tools": [{"name": "get_weather",
11264 "description": "Get the current weather for a city",
11265 "input_schema": {"type": "object",
11266 "properties": {"city": {"type": "string"}},
11267 "required": ["city"]}}]
11268 }))
11269 .expect("/v1/messages translate");
11270 let want = glm5_render(chat).expect("chat");
11271 assert!(
11273 want.contains(
11274 "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value>\
11275 </tool_call><tool_call>get_weather<arg_key>city</arg_key>\
11276 <arg_value>Rome</arg_value></tool_call>"
11277 ),
11278 "{want:?}"
11279 );
11280 assert!(
11284 want.contains(
11285 "<|observation|><tool_response>paris:21</tool_response>\
11286 <tool_response>rome:27</tool_response>"
11287 ),
11288 "{want:?}"
11289 );
11290 assert!(
11291 want.contains("<|system|>Reasoning Effort: High"),
11292 "{want:?}"
11293 );
11294 for (surface, body) in [
11295 ("/v1/responses", responses),
11296 ("/v1/messages", messages.clone()),
11297 ] {
11298 let got = glm5_render(body).unwrap_or_else(|e| panic!("{surface}: {e}"));
11299 assert_eq!(
11300 got, want,
11301 "{surface} rendered DIFFERENT glm5 prompt bytes than /v1/chat/completions"
11302 );
11303 }
11304 let mut idless = messages;
11310 for m in idless["messages"].as_array_mut().unwrap() {
11311 if m["role"] == "tool" {
11312 m.as_object_mut().unwrap().remove("tool_call_id");
11313 }
11314 }
11315 let got = glm5_render(idless).expect("id-less render");
11316 assert_ne!(
11317 got, want,
11318 "dropping tool_call_id must change the rendered order — this test cannot detect \
11319 a surface that loses ids otherwise"
11320 );
11321 assert!(
11322 got.contains(
11323 "<|observation|><tool_response>rome:27</tool_response>\
11324 <tool_response>paris:21</tool_response>"
11325 ),
11326 "{got:?}"
11327 );
11328 }
11329
11330 #[test]
11333 fn glm5_chat_arms_the_native_tool_parser() {
11334 let req: ChatCompletionReq = serde_json::from_value(json!({
11335 "model": "m", "messages": [{"role": "user", "content": "weather?"}],
11336 "tools": [{"type": "function", "function": {"name": "get_weather",
11337 "parameters": {"type": "object",
11338 "properties": {"city": {"type": "string"}}}}}]}))
11339 .unwrap();
11340 let (tx, _rx) = worker::event_channel();
11341 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11342 .unwrap();
11343 let mut parser = plan.parser.expect("glm5 tools request must carry a parser");
11344 let pieces = parser.push(
11345 "reasoning here</think><tool_call>get_weather<arg_key>city</arg_key>\
11346 <arg_value>Paris</arg_value></tool_call>",
11347 );
11348 let calls: Vec<_> = pieces
11349 .iter()
11350 .filter_map(|p| match p {
11351 toolcall::Piece::Call(c) => Some((c.name.as_str(), c.arguments.as_str())),
11352 _ => None,
11353 })
11354 .collect();
11355 assert_eq!(
11356 calls,
11357 vec![("get_weather", r#"{"city":"Paris"}"#)],
11358 "{pieces:?}"
11359 );
11360 assert!(
11361 pieces
11362 .iter()
11363 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "reasoning here")),
11364 "{pieces:?}"
11365 );
11366 assert!(
11368 !pieces
11369 .iter()
11370 .any(|p| matches!(p, toolcall::Piece::Content(_))),
11371 "{pieces:?}"
11372 );
11373 let req: ChatCompletionReq = serde_json::from_value(
11377 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11378 )
11379 .unwrap();
11380 let (tx, _rx) = worker::event_channel();
11381 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11382 .unwrap();
11383 let mut parser = plan
11384 .parser
11385 .expect("glm5 non-tools request must still split reasoning");
11386 let pieces = parser.push("weighing it</think>The answer.");
11387 assert!(
11388 pieces
11389 .iter()
11390 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "weighing it")),
11391 "{pieces:?}"
11392 );
11393 assert!(
11394 pieces
11395 .iter()
11396 .any(|p| matches!(p, toolcall::Piece::Content(c) if c == "The answer.")),
11397 "{pieces:?}"
11398 );
11399 }
11400
11401 #[test]
11408 fn glm5_plain_fast_path_never_drops_replayed_reasoning() {
11409 let with_reasoning = vec![
11410 chat::Turn {
11411 role: "user".into(),
11412 content: "a".into(),
11413 ..Default::default()
11414 },
11415 chat::Turn {
11416 role: "assistant".into(),
11417 content: "A".into(),
11418 reasoning: Some("I considered a.".into()),
11419 ..Default::default()
11420 },
11421 chat::Turn {
11422 role: "user".into(),
11423 content: "b".into(),
11424 ..Default::default()
11425 },
11426 ];
11427 assert!(!worker::plain_chat_render_path(
11429 &[],
11430 &chat::ThinkMode::Default,
11431 None,
11432 &with_reasoning,
11433 false,
11434 ));
11435 let plain_turns: Vec<chat::Turn> = with_reasoning
11438 .iter()
11439 .cloned()
11440 .map(|mut t| {
11441 t.reasoning = None;
11442 t
11443 })
11444 .collect();
11445 assert!(worker::plain_chat_render_path(
11446 &[],
11447 &chat::ThinkMode::Default,
11448 None,
11449 &plain_turns,
11450 false,
11451 ));
11452 let tmpl = glm5_template();
11455 let via_tools = chat::apply_chat_template_tools_ex(
11456 Some(&tmpl),
11457 &with_reasoning,
11458 true,
11459 &[],
11460 &[],
11461 chat::ThinkMode::Default,
11462 None,
11463 None,
11464 )
11465 .unwrap();
11466 let msgs: Vec<(&str, &str)> = with_reasoning
11467 .iter()
11468 .map(|t| (t.role.as_str(), t.content.as_str()))
11469 .collect();
11470 let via_plain = chat::apply_chat_template_str(Some(&tmpl), &msgs, true);
11471 assert!(
11472 via_tools.contains("<think>I considered a.</think>"),
11473 "{via_tools:?}"
11474 );
11475 assert_ne!(via_tools, via_plain);
11476 let plain_msgs: Vec<(&str, &str)> = plain_turns
11479 .iter()
11480 .map(|t| (t.role.as_str(), t.content.as_str()))
11481 .collect();
11482 assert_eq!(
11483 chat::apply_chat_template_tools_ex(
11484 Some(&tmpl),
11485 &plain_turns,
11486 true,
11487 &[],
11488 &[],
11489 chat::ThinkMode::Default,
11490 None,
11491 None,
11492 )
11493 .unwrap(),
11494 chat::apply_chat_template_str(Some(&tmpl), &plain_msgs, true)
11495 );
11496 }
11497
11498 #[test]
11502 fn glm5_model_row_does_not_claim_structured_output() {
11503 let caps = glm5_caps();
11504 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), None);
11505 assert_eq!(row["capabilities"]["structured_output"], json!(false));
11506 assert_eq!(row["capabilities"]["tools"], json!(true));
11507 assert_eq!(row["capabilities"]["reasoning"], json!(true));
11508 let err = glm5_render(json!({"model": "m",
11510 "messages": [{"role": "user", "content": "hi"}],
11511 "response_format": {"type": "json_object"}}))
11512 .expect_err("response_format must be refused on a switchless think template");
11513 assert!(
11517 err.contains("neither an enable_thinking switch nor a recognizable"),
11518 "{err}"
11519 );
11520 let switchable = model_entry_v1("q", Some(&tool_caps()), None);
11522 assert_eq!(switchable["capabilities"]["structured_output"], json!(true));
11523 let glm_params = openrouter_supported_parameters(Some(&caps), None, true);
11526 assert!(
11527 glm_params.get("structured_outputs").is_none(),
11528 "{glm_params}"
11529 );
11530 let step_like = ModelCaps {
11536 chat_ok: true,
11537 qwen_think: true,
11538 think_switch: false,
11539 think_close: vec![128799],
11540 ..caps.clone()
11541 };
11542 let step_row = model_entry_v1("stepfun/step-3.7-flash", Some(&step_like), None);
11543 assert_eq!(step_row["capabilities"]["structured_output"], json!(true));
11544 let step_params = openrouter_supported_parameters(Some(&step_like), None, true);
11545 assert!(
11546 step_params.get("structured_outputs").is_some(),
11547 "{step_params}"
11548 );
11549 assert!(glm_params.get("json_mode").is_none(), "{glm_params}");
11550 assert!(glm_params.get("tools").is_some(), "{glm_params}");
11551 assert_eq!(
11554 glm_params.get("reasoning_effort"),
11555 Some(&json!({ "type": "enum", "values": ["low", "high", "max"] })),
11556 "{glm_params}"
11557 );
11558 let qwen_params = openrouter_supported_parameters(Some(&tool_caps()), None, true);
11559 assert!(
11560 qwen_params.get("structured_outputs").is_some(),
11561 "{qwen_params}"
11562 );
11563 assert!(qwen_params.get("json_mode").is_some(), "{qwen_params}");
11564 assert!(
11565 qwen_params.get("reasoning_effort").is_none(),
11566 "{qwen_params}"
11567 );
11568 }
11569
11570 #[test]
11578 fn catalog_context_claim_is_capped_by_the_deployment_envelope() {
11579 let caps = glm5_caps();
11580 assert_eq!(caps.context_length, 1_048_576);
11581 let metadata = OpenRouterModelMetadata {
11582 max_prompt_length: Some(126_976),
11583 max_output_length: Some(4_096),
11584 ..Default::default()
11585 };
11586 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
11588 assert_eq!(row["context_length"], json!(131_072));
11589 let or_row = model_entry_openrouter("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
11590 assert_eq!(
11591 or_row["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
11592 json!(131_072)
11593 );
11594 assert_eq!(
11595 published_context_length(Some(&caps), Some(&metadata)),
11596 Some(131_072)
11597 );
11598 assert_eq!(published_context_length(Some(&caps), None), Some(1_048_576));
11600 let half = OpenRouterModelMetadata {
11601 max_output_length: Some(4_096),
11602 ..Default::default()
11603 };
11604 assert_eq!(
11605 published_context_length(Some(&caps), Some(&half)),
11606 Some(1_048_576)
11607 );
11608 let wide = OpenRouterModelMetadata {
11610 max_prompt_length: Some(2_000_000),
11611 max_output_length: Some(2_000_000),
11612 ..Default::default()
11613 };
11614 assert_eq!(
11615 published_context_length(Some(&caps), Some(&wide)),
11616 Some(1_048_576)
11617 );
11618 }
11619
11620 fn dsv4_sentinel() -> String {
11628 let path = format!(
11629 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
11630 env!("CARGO_MANIFEST_DIR")
11631 );
11632 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
11633 }
11634
11635 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
11640 let role = msg["role"].as_str().unwrap().to_string();
11641 let content =
11642 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
11643 let reasoning = msg
11644 .get("reasoning")
11645 .or_else(|| msg.get("reasoning_content"))
11646 .and_then(|r| r.as_str())
11647 .map(String::from)
11648 .filter(|s| !s.is_empty());
11649 let tool_calls = msg
11650 .get("tool_calls")
11651 .and_then(|a| a.as_array())
11652 .map(|a| {
11653 a.iter()
11654 .map(|tc| {
11655 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
11656 render_req_tool_call(&rtc).unwrap()
11657 })
11658 .collect()
11659 })
11660 .unwrap_or_default();
11661 let tools = msg
11662 .get("tools")
11663 .and_then(|a| a.as_array())
11664 .map(|a| {
11665 a.iter()
11666 .filter_map(|t| t.get("function").map(json_to_val))
11667 .collect()
11668 })
11669 .unwrap_or_default();
11670 TmplTurn {
11671 role,
11672 content,
11673 tool_calls,
11674 reasoning,
11675 tool_call_id: msg
11676 .get("tool_call_id")
11677 .and_then(|s| s.as_str())
11678 .map(String::from),
11679 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
11680 tool_responses: Vec::new(),
11681 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
11682 tools,
11683 }
11684 }
11685
11686 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
11687 v.and_then(|t| t.as_array())
11688 .map(|a| {
11689 a.iter()
11690 .filter_map(|t| t.get("function").map(json_to_val))
11691 .collect()
11692 })
11693 .unwrap_or_default()
11694 }
11695
11696 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
11700 let dir = format!(
11701 "{}/../../research/dsv4-template-20260818/{subdir}",
11702 env!("CARGO_MANIFEST_DIR")
11703 );
11704 let tmpl = dsv4_sentinel();
11705 let mut entries: Vec<_> = std::fs::read_dir(&dir)
11706 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
11707 .map(|e| e.unwrap().path())
11708 .filter(|p| p.is_dir())
11709 .collect();
11710 entries.sort();
11711 assert!(
11712 entries.len() >= min_fixtures,
11713 "expected >={min_fixtures} fixtures, found {}",
11714 entries.len()
11715 );
11716 for d in &entries {
11717 let input: serde_json::Value =
11718 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11719 .unwrap();
11720 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11721 let turns: Vec<TmplTurn> = input["turns"]
11722 .as_array()
11723 .unwrap()
11724 .iter()
11725 .map(dsv4_turn)
11726 .collect();
11727 let think = match input["think"].as_str().unwrap() {
11728 "chat" => ThinkMode::NoThink,
11729 _ => ThinkMode::Think,
11730 };
11731 let effort = input
11732 .get("reasoning_effort")
11733 .and_then(|v| v.as_str())
11734 .map(String::from);
11735 let req_tools = dsv4_req_tools(input.get("req_tools"));
11736 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
11737 let got = chat::apply_chat_template_tools_ex(
11738 Some(&tmpl),
11739 &turns,
11740 agp,
11741 &[],
11742 &req_tools,
11743 think,
11744 effort.as_deref(),
11745 Some(encoding),
11746 )
11747 .unwrap();
11748 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
11749 }
11750 }
11751
11752 #[test]
11753 fn dsv4_template_fixtures_match_the_oracle() {
11754 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
11755 }
11756
11757 #[test]
11763 fn dsv4_0731_fixtures_match_the_oracle() {
11764 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
11765 }
11766
11767 #[test]
11768 fn dsv4_artifact_fixtures_are_byte_identical() {
11769 let base = format!(
11773 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
11774 env!("CARGO_MANIFEST_DIR")
11775 );
11776 let tmpl = dsv4_sentinel();
11777 for (n, think) in [
11778 (1u32, ThinkMode::Think),
11779 (2, ThinkMode::Think),
11780 (3, ThinkMode::Think),
11781 (4, ThinkMode::NoThink),
11782 ] {
11783 let td: serde_json::Value = serde_json::from_str(
11784 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
11785 )
11786 .unwrap();
11787 let (messages, tools) = if td.is_object() {
11788 (td["messages"].clone(), td.get("tools").cloned())
11789 } else {
11790 (td.clone(), None)
11791 };
11792 let mut turns: Vec<TmplTurn> = Vec::new();
11793 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
11794 let mut t = dsv4_turn(msg);
11795 if i == 0
11796 && let Some(tl) = &tools
11797 {
11798 t.tools = tl
11799 .as_array()
11800 .unwrap()
11801 .iter()
11802 .filter_map(|x| x.get("function").map(json_to_val))
11803 .collect();
11804 }
11805 turns.push(t);
11806 }
11807 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
11808 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
11812 let got = chat::apply_chat_template_tools_ex(
11813 Some(&tmpl),
11814 &turns,
11815 true,
11816 &[],
11817 &[],
11818 think,
11819 None,
11820 Some(encoding),
11821 )
11822 .unwrap();
11823 assert_eq!(
11824 got, expected,
11825 "artifact fixture {n} diverged from the oracle under {encoding:?}"
11826 );
11827 }
11828 }
11829 }
11830
11831 #[test]
11832 fn dsv4_default_thinkmode_renders_thinking() {
11833 let tmpl = dsv4_sentinel();
11836 let turns = vec![TmplTurn {
11837 role: "user".into(),
11838 content: "Hi".into(),
11839 ..Default::default()
11840 }];
11841 let dflt = chat::apply_chat_template_tools_ex(
11842 Some(&tmpl),
11843 &turns,
11844 true,
11845 &[],
11846 &[],
11847 ThinkMode::Default,
11848 None,
11849 None,
11850 )
11851 .unwrap();
11852 let think = chat::apply_chat_template_tools_ex(
11853 Some(&tmpl),
11854 &turns,
11855 true,
11856 &[],
11857 &[],
11858 ThinkMode::Think,
11859 None,
11860 None,
11861 )
11862 .unwrap();
11863 assert_eq!(dflt, think);
11864 assert!(
11865 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
11866 "{dflt:?}"
11867 );
11868 let chat_mode = chat::apply_chat_template_tools_ex(
11869 Some(&tmpl),
11870 &turns,
11871 true,
11872 &[],
11873 &[],
11874 ThinkMode::NoThink,
11875 None,
11876 None,
11877 )
11878 .unwrap();
11879 assert!(
11880 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
11881 "{chat_mode:?}"
11882 );
11883 }
11884
11885 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
11890 let base = format!(
11891 "{}/../../research/dsv4-template-20260818",
11892 env!("CARGO_MANIFEST_DIR")
11893 );
11894 let refdir = std::path::Path::new(&base).join("ref");
11895 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
11896 .expect("load dsv4 tokenizer from ref dir");
11897 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
11898 let banked: serde_json::Value = serde_json::from_str(
11899 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
11900 .unwrap(),
11901 )
11902 .unwrap();
11903 let obj = banked.as_object().unwrap();
11904 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
11905 for (name, ids_v) in obj {
11906 let rendered =
11907 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
11908 let want: Vec<u32> = ids_v
11909 .as_array()
11910 .unwrap()
11911 .iter()
11912 .map(|v| v.as_u64().unwrap() as u32)
11913 .collect();
11914 let got = tok.encode(&rendered, true);
11915 assert_eq!(got, want, "tokenization diverged for {name}");
11916 }
11917 }
11918
11919 #[test]
11920 fn dsv4_tokenization_crosscheck_matches_official_ids() {
11921 dsv4_run_tokenization_crosscheck("fixtures");
11922 }
11923
11924 #[test]
11928 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
11929 dsv4_run_tokenization_crosscheck("fixtures-0731");
11930 }
11931
11932 #[test]
11933 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
11934 let base = format!(
11944 "{}/../../research/dsv4-template-20260818",
11945 env!("CARGO_MANIFEST_DIR")
11946 );
11947 let refdir = std::path::Path::new(&base).join("ref");
11948 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
11949 .expect("load dsv4 tokenizer from ref dir");
11950 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
11951 let tmpl = dsv4_sentinel();
11952 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
11953 {"type": "function", "function": {
11954 "name": "get_data",
11955 "description": "Fetch a blob",
11956 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
11957 "required": ["key"]}
11958 }}
11959 ])));
11960
11961 let cases: Vec<(&str, String)> = vec![
11962 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
11964 ("space-131k", " ".repeat(131_072)),
11965 ("digit-131k", "7".repeat(131_072)),
11966 (
11967 "mixed-runs",
11968 format!(
11969 "{}{}{}{}",
11970 "Z".repeat(65_536),
11971 " ".repeat(65_536),
11972 "7".repeat(65_536),
11973 "\n".repeat(65_536)
11974 ),
11975 ),
11976 ("cjk-64k", "中".repeat(65_536)),
11977 ("accented-letter-64k", "é".repeat(65_536)),
11978 ];
11979 for (name, blob) in &cases {
11980 let msgs = serde_json::json!([
11981 {"role": "system", "content": "You are a tool-using assistant."},
11982 {"role": "user", "content": "Fetch the blob."},
11983 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
11984 "tool_calls": [{"id": "call_001", "type": "function",
11985 "function": {"name": "get_data",
11986 "arguments": "{\"key\": \"blob\"}"}}]},
11987 {"role": "tool", "tool_call_id": "call_001", "content": blob}
11988 ]);
11989 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
11990 let rendered = chat::apply_chat_template_tools_ex(
11991 Some(&tmpl),
11992 &turns,
11993 true,
11994 &[],
11995 &req_tools,
11996 ThinkMode::Think,
11997 None,
11998 None,
11999 )
12000 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
12001 assert!(
12002 rendered.contains(blob.as_str()),
12003 "{name}: tool result missing from render"
12004 );
12005 let t0 = std::time::Instant::now();
12006 let ids = tok.encode(&rendered, true);
12007 let encode_dt = t0.elapsed();
12008 assert!(!ids.is_empty(), "{name}: empty encode");
12009 let back = tok.decode(&ids);
12010 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
12011 assert!(
12015 encode_dt < std::time::Duration::from_secs(60),
12016 "{name}: encode took {encode_dt:?}"
12017 );
12018 if *name == "ascii-letter-131k"
12022 && let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR")
12023 {
12024 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
12025 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
12026 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
12027 }
12028 }
12029 }
12030
12031 #[test]
12032 fn models_v1_entry_advertises_thinking_support() {
12033 let step_caps = ModelCaps {
12036 effort_levels: true,
12037 ..tool_caps()
12038 };
12039 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
12040 assert_eq!(entry["capabilities"]["reasoning"], true);
12041 assert_eq!(entry["capabilities"]["tools"], true);
12042
12043 let plain = ModelCaps {
12045 chat_ok: true,
12046 ..Default::default()
12047 };
12048 let entry = model_entry_v1("plain", Some(&plain), None);
12049 assert_eq!(entry["capabilities"]["reasoning"], false);
12050 assert_eq!(entry["capabilities"]["tools"], false);
12051 let entry = model_entry_v1("unknown", None, None);
12053 assert_eq!(entry["capabilities"]["reasoning"], false);
12054 assert_eq!(entry["capabilities"]["streaming"], true);
12055 }
12056
12057 #[test]
12058 fn chat_request_preserves_turns_and_openai_stop_forms() {
12059 let payload = serde_json::json!({
12060 "model": "plain_quant",
12061 "messages": [
12062 {"role": "system", "content": "rules"},
12063 {"role": "developer", "content": "dev rules"},
12064 {"role": "user", "content": "task"},
12065 {"role": "assistant", "content": "work"}
12066 ],
12067 "max_tokens": 64,
12068 "temperature": 0.0,
12069 "stop": "<stop>"
12070 });
12071 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12072 let (tx, _rx) = worker::event_channel();
12073 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
12074 let request = plan.request;
12075 assert!(
12076 plan.parser.is_none(),
12077 "no tools -> no parser (isolation contract)"
12078 );
12079 assert!(request.tools_json.is_empty());
12080 assert_eq!(request.think, ThinkMode::Default);
12081 assert_eq!(request.model, "plain_quant");
12082 assert_eq!(request.params.max_new, 64);
12083 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12085 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
12086 }))
12087 .unwrap();
12088 let (tx, _rx) = worker::event_channel();
12089 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
12090 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
12091 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12093 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12094 "max_completion_tokens": 7
12095 }))
12096 .unwrap();
12097 let (tx, _rx) = worker::event_channel();
12098 assert_eq!(
12099 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12100 .unwrap()
12101 .request
12102 .params
12103 .max_new,
12104 7
12105 );
12106 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12108 "model": "plain_quant", "prompt": "task"
12109 }))
12110 .unwrap();
12111 let (tx, _rx) = worker::event_channel();
12112 assert_eq!(
12113 build_request(&req, tx, lanes::Lane::Interactive, None)
12114 .params
12115 .max_new,
12116 worker::MAX_NEW_CTX_BOUNDED
12117 );
12118 let turns: Vec<(String, String)> = request
12119 .chat_turns
12120 .iter()
12121 .map(|t| (t.role.clone(), t.content.clone()))
12122 .collect();
12123 assert_eq!(
12124 turns,
12125 vec![
12126 ("system".into(), "rules".into()),
12127 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
12129 ("assistant".into(), "work".into()),
12130 ]
12131 );
12132 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
12133 assert_eq!(request.stop_strings, vec!["<stop>"]);
12134
12135 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12136 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12137 "stop": ["a", "b"]
12138 }))
12139 .unwrap();
12140 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
12141
12142 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12146 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12147 "stop": ["", "real", ""]
12148 }))
12149 .unwrap();
12150 assert_eq!(req.stop.into_vec(), vec!["real"]);
12151 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12152 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12153 "stop": ""
12154 }))
12155 .unwrap();
12156 assert!(req.stop.into_vec().is_empty());
12157
12158 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12159 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12160 "stop": null
12161 }))
12162 .unwrap();
12163 assert!(req.stop.into_vec().is_empty());
12164 }
12165
12166 #[test]
12167 fn stop_sequence_limits_bound_count_individual_and_aggregate_work() {
12168 let at_limit = StopSequences::Many(vec!["x".repeat(256); MAX_STOP_SEQUENCES]);
12169 assert!(at_limit.validate().is_ok());
12170 assert!(
12171 StopSequences::Many(vec![String::new(); MAX_STOP_SEQUENCES + 1])
12172 .validate()
12173 .unwrap_err()
12174 .contains("at most")
12175 );
12176 assert!(
12177 StopSequences::One("x".repeat(MAX_STOP_SEQUENCE_BYTES + 1))
12178 .validate()
12179 .unwrap_err()
12180 .contains("each stop")
12181 );
12182 assert!(
12183 StopSequences::Many(vec!["x".repeat(300); MAX_STOP_SEQUENCES])
12184 .validate()
12185 .unwrap_err()
12186 .contains("total at most")
12187 );
12188 }
12189
12190 #[tokio::test]
12191 async fn chat_response_has_openai_message_shape() {
12192 let (tx, rx) = worker::event_channel();
12193 tx.send(Event::Token {
12194 id: 1,
12195 text: "hello".into(),
12196 })
12197 .unwrap();
12198 tx.send(Event::Done {
12199 stop_reason: "Eos".into(),
12200 n_tokens: 1,
12201 n_prompt: 42,
12202 n_cached: 30,
12203 elapsed_s: 0.5,
12204 spec: None,
12205 })
12206 .unwrap();
12207 drop(tx);
12208 let response = blocking_response(
12209 rx,
12210 "plain_quant".into(),
12211 true,
12212 Vec::new(),
12213 None,
12214 Envelope::new(true),
12215 )
12216 .await;
12217 assert_eq!(response.status(), StatusCode::OK);
12218 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12219 .await
12220 .unwrap();
12221 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12222 assert_eq!(payload["object"], "chat.completion");
12223 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
12225 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
12226 let fingerprint = payload["system_fingerprint"].as_str().unwrap();
12230 assert!(
12231 build_id::fingerprint_is_well_formed(fingerprint),
12232 "system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
12233 );
12234 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
12235 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
12236 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
12237 assert_eq!(payload["usage"]["prompt_tokens"], 42);
12239 assert_eq!(payload["usage"]["completion_tokens"], 1);
12240 assert_eq!(payload["usage"]["total_tokens"], 43);
12241 assert_eq!(
12242 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12243 30
12244 );
12245 assert!(payload["usage"].get("spec").is_none());
12248 }
12249
12250 #[tokio::test]
12251 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
12252 let (tx, rx) = worker::event_channel();
12253 tx.send(Event::Token {
12255 id: 4,
12256 text: "hello".into(),
12257 })
12258 .unwrap();
12259 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
12260 tx.send(Event::Done {
12261 stop_reason: "MaxNew".into(),
12262 n_tokens: 4,
12263 n_prompt: 2,
12264 n_cached: 0,
12265 elapsed_s: 0.5,
12266 spec: None,
12267 })
12268 .unwrap();
12269 drop(tx);
12270
12271 let response = blocking_response(
12272 rx,
12273 "plain_quant".into(),
12274 false,
12275 Vec::new(),
12276 None,
12277 Envelope::new(false),
12278 )
12279 .await;
12280 assert_eq!(response.status(), StatusCode::OK);
12281 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12282 .await
12283 .unwrap();
12284 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12285 assert_eq!(payload["text"], "hello");
12286 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
12287 assert_eq!(payload["n_tokens"], 4);
12288 }
12289
12290 #[tokio::test]
12293 async fn chat_usage_carries_spec_acceptance_summary() {
12294 let (tx, rx) = worker::event_channel();
12295 tx.send(Event::Token {
12296 id: 1,
12297 text: "hello".into(),
12298 })
12299 .unwrap();
12300 tx.send(Event::Done {
12301 stop_reason: "Eos".into(),
12302 n_tokens: 1,
12303 n_prompt: 42,
12304 n_cached: 0,
12305 elapsed_s: 0.5,
12306 spec: Some(worker::SpecUsage {
12307 rounds: 10,
12308 drafted: 30,
12309 accepted: 21,
12310 }),
12311 })
12312 .unwrap();
12313 drop(tx);
12314 let response = blocking_response(
12315 rx,
12316 "plain_quant".into(),
12317 true,
12318 Vec::new(),
12319 None,
12320 Envelope::new(true),
12321 )
12322 .await;
12323 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12324 .await
12325 .unwrap();
12326 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12327 let sp = &payload["usage"]["spec"];
12328 assert_eq!(sp["rounds"], 10);
12329 assert_eq!(sp["drafted"], 30);
12330 assert_eq!(sp["accepted"], 21);
12331 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
12332 assert_eq!(payload["usage"]["total_tokens"], 43);
12334 }
12335
12336 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
12337 let mut payload = serde_json::json!({
12338 "model": "m",
12339 "messages": [{"role": "user", "content": "Weather in Paris?"}],
12340 "tools": [{"type": "function", "function": {
12341 "name": "get_weather",
12342 "description": "Get current weather",
12343 "parameters": {"type": "object",
12344 "properties": {"city": {"type": "string"},
12345 "days": {"type": "integer"}},
12346 "required": ["city"]}}}],
12347 });
12348 if let Some(obj) = extra.as_object() {
12349 for (k, v) in obj {
12350 payload[k] = v.clone();
12351 }
12352 }
12353 serde_json::from_value(payload).unwrap()
12354 }
12355
12356 #[test]
12360 fn glm5_vision_decode_is_deferred_and_grid_pinned() {
12361 let (tx, _rx) = worker::event_channel();
12362 let req: ChatCompletionReq = serde_json::from_value(json!({
12363 "model": "m", "messages": [{"role": "user", "content": "hi"}],
12364 }))
12365 .unwrap();
12366 let mut plan = build_chat_request(
12367 req,
12368 Some(&ModelCaps {
12369 chat_ok: true,
12370 ..Default::default()
12371 }),
12372 tx,
12373 lanes::Lane::Interactive,
12374 None,
12375 )
12376 .unwrap();
12377 let bmp = |w: u32, h: u32| -> Vec<u8> {
12380 let row = (w * 3).div_ceil(4) * 4;
12381 let size = 54 + row * h;
12382 let mut b = vec![0x42u8, 0x4d];
12383 b.extend_from_slice(&size.to_le_bytes());
12384 b.extend_from_slice(&[0; 4]);
12385 b.extend_from_slice(&54u32.to_le_bytes());
12386 b.extend_from_slice(&40u32.to_le_bytes());
12387 b.extend_from_slice(&w.to_le_bytes());
12388 b.extend_from_slice(&h.to_le_bytes());
12389 b.extend_from_slice(&1u16.to_le_bytes());
12390 b.extend_from_slice(&24u16.to_le_bytes());
12391 b.extend_from_slice(&[0u8; 24]);
12392 b.extend(std::iter::repeat_n(0x7fu8, (row * h) as usize));
12393 b
12394 };
12395 let bytes = bmp(112, 112);
12396 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes).unwrap();
12397 assert_eq!((gh, gw), (8, 8), "identity resize grid");
12398 assert_eq!(memra_engine::vision_glm5::n_merged_for_grid(gh, gw), 16);
12399 plan.pending_glm5.push(PendingGlm5Image {
12400 bytes: bytes.clone(),
12401 gh,
12402 gw,
12403 });
12404 decode_pending_vision(&mut plan).unwrap();
12405 assert_eq!(plan.request.glm5_images.len(), 1);
12406 let unit = &plan.request.glm5_images[0];
12407 assert_eq!((unit.gh, unit.gw), (gh, gw));
12408 assert_eq!(
12409 unit.patches.len(),
12410 gh * gw * memra_engine::vision_glm5::G5V_PATCH_IN
12411 );
12412 plan.request.glm5_images.clear();
12414 plan.pending_glm5.push(PendingGlm5Image {
12415 bytes,
12416 gh: gh + 2,
12417 gw,
12418 });
12419 let err = decode_pending_vision(&mut plan).unwrap_err();
12420 assert!(err.contains("header-planned"), "got: {err}");
12421 }
12422
12423 #[test]
12424 fn vision_decode_is_deferred_and_grid_pinned() {
12425 let (tx, _rx) = worker::event_channel();
12430 let req: ChatCompletionReq = serde_json::from_value(json!({
12431 "model": "m", "messages": [{"role": "user", "content": "hi"}],
12432 }))
12433 .unwrap();
12434 let mut plan = build_chat_request(
12435 req,
12436 Some(&ModelCaps {
12437 chat_ok: true,
12438 ..Default::default()
12439 }),
12440 tx,
12441 lanes::Lane::Interactive,
12442 None,
12443 )
12444 .unwrap();
12445 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
12449 let mut b = Vec::new();
12450 b.extend_from_slice(b"BM");
12451 b.extend_from_slice(&54u32.to_le_bytes());
12452 b.extend_from_slice(&0u32.to_le_bytes());
12453 b.extend_from_slice(&54u32.to_le_bytes());
12454 b.extend_from_slice(&40u32.to_le_bytes());
12455 b.extend_from_slice(&w.to_le_bytes());
12456 b.extend_from_slice(&h.to_le_bytes());
12457 b.extend_from_slice(&1u16.to_le_bytes());
12458 b.extend_from_slice(&24u16.to_le_bytes());
12459 b.extend_from_slice(&[0u8; 24]);
12460 if with_pixels {
12461 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
12462 }
12463 b
12464 };
12465 let bytes = bmp(64, 64, true);
12466 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
12467 plan.pending_images.push(PendingVisionUnit::Still {
12468 bytes: bytes.clone(),
12469 gh,
12470 gw,
12471 });
12472 decode_pending_vision(&mut plan).unwrap();
12473 assert_eq!(plan.request.images.len(), 1);
12474 assert_eq!(
12475 (
12476 plan.request.images[0].prep.gh,
12477 plan.request.images[0].prep.gw
12478 ),
12479 (gh, gw),
12480 "decoded grid must equal the header-planned grid the pad run was rendered from"
12481 );
12482 plan.request.images.clear();
12484 plan.pending_images.push(PendingVisionUnit::Still {
12485 bytes,
12486 gh: gh + 2,
12487 gw,
12488 });
12489 let err = decode_pending_vision(&mut plan).unwrap_err();
12490 assert!(err.contains("header-planned"), "got: {err}");
12491 let bomb = bmp(16_000, 16_000, false);
12494 plan.pending_images.clear();
12495 plan.pending_images.push(PendingVisionUnit::Still {
12496 bytes: bomb,
12497 gh: 2,
12498 gw: 2,
12499 });
12500 let err = decode_pending_vision(&mut plan).unwrap_err();
12501 assert!(err.contains("exceeds the decode budget"), "got: {err}");
12502 }
12503
12504 #[test]
12505 fn tools_request_renders_client_key_order_and_arms_parser() {
12506 let (tx, _rx) = worker::event_channel();
12507 let plan = build_chat_request(
12508 weather_request(json!({})),
12509 Some(&tool_caps()),
12510 tx,
12511 lanes::Lane::Interactive,
12512 None,
12513 )
12514 .unwrap();
12515 assert!(plan.parser.is_some());
12516 assert_eq!(plan.request.tools_json.len(), 1);
12517 assert_eq!(
12519 plan.request.tools_json[0],
12520 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
12521 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
12522 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
12523 \"integer\"}}, \"required\": [\"city\"]}}}"
12524 );
12525 }
12526
12527 #[test]
12528 fn hy3_tools_and_reasoning_flow_through_the_real_chat_plan() {
12529 let (tx, _rx) = worker::event_channel();
12530 let plan = build_chat_request(
12531 weather_request(json!({"reasoning_effort": "high"})),
12532 Some(&hy3_tool_caps()),
12533 tx,
12534 lanes::Lane::Interactive,
12535 None,
12536 )
12537 .unwrap();
12538 assert_eq!(plan.request.think, ThinkMode::Think);
12539 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12540 assert!(
12541 plan.request
12542 .stop_strings
12543 .iter()
12544 .any(|stop| stop == "</tool_calls:opensource>")
12545 );
12546 let rendered = chat::apply_chat_template_tools_ex(
12547 Some("... hy_User ... <tools> ..."),
12548 &plan.request.chat_turns,
12549 true,
12550 &plan.request.tools_json,
12551 &plan.request.tools_struct,
12552 plan.request.think,
12553 plan.request.reasoning_effort.as_deref(),
12554 None,
12555 )
12556 .unwrap();
12557 assert!(rendered.contains("<tool_calls:opensource>"));
12558 assert!(rendered.ends_with("<think:opensource>"));
12559
12560 let mut parser = plan.parser.expect("HY3 tools arm its native parser");
12561 let pieces = parser.push(concat!(
12562 "Need weather.</think:opensource>",
12563 "<tool_calls:opensource><tool_call:opensource>get_weather",
12564 "<tool_sep:opensource>\n<arg_key:opensource>city</arg_key:opensource>\n",
12565 "<arg_value:opensource>Paris</arg_value:opensource>\n",
12566 "</tool_call:opensource></tool_calls:opensource>",
12567 ));
12568 assert!(pieces.contains(&Piece::Reasoning("Need weather.".into())));
12569 assert!(pieces.iter().any(|piece| matches!(piece, Piece::Call(call)
12570 if call.name == "get_weather" && call.arguments == r#"{"city":"Paris"}"#)));
12571 }
12572
12573 #[test]
12574 fn tool_choice_none_strips_tools_and_parser() {
12575 let (tx, _rx) = worker::event_channel();
12576 let plan = build_chat_request(
12577 weather_request(json!({"tool_choice": "none"})),
12578 Some(&tool_caps()),
12579 tx,
12580 lanes::Lane::Interactive,
12581 None,
12582 )
12583 .unwrap();
12584 let mut p = plan
12587 .parser
12588 .expect("think-open chat arms the reasoning splitter");
12589 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
12590 assert_eq!(
12591 pieces,
12592 vec![
12593 Piece::Reasoning("x".into()),
12594 Piece::Content("<tool_call> stays prose".into()),
12595 ]
12596 );
12597 assert!(plan.request.tools_json.is_empty());
12598 let (tx, _rx) = worker::event_channel();
12600 assert!(
12601 build_chat_request(
12602 weather_request(json!({"tool_choice": "required"})),
12603 Some(&tool_caps()),
12604 tx,
12605 lanes::Lane::Interactive,
12606 None
12607 )
12608 .is_err()
12609 );
12610 let (tx, _rx) = worker::event_channel();
12611 assert!(
12612 build_chat_request(
12613 weather_request(json!({"tool_choice":
12614 {"type": "function", "function": {"name": "get_weather"}}})),
12615 Some(&tool_caps()),
12616 tx,
12617 lanes::Lane::Interactive,
12618 None
12619 )
12620 .is_err()
12621 );
12622 }
12623
12624 #[test]
12625 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
12626 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
12627 let _ = std::fs::remove_dir_all(&root);
12628
12629 let st = root.join("st_single");
12631 std::fs::create_dir_all(&st).unwrap();
12632 std::fs::write(st.join("config.json"), "{}").unwrap();
12633 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
12634 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
12635
12636 let sh = root.join("st_sharded");
12638 std::fs::create_dir_all(&sh).unwrap();
12639 std::fs::write(sh.join("config.json"), "{}").unwrap();
12640 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
12641 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
12642
12643 let rp = root.join("repack");
12645 std::fs::create_dir_all(&rp).unwrap();
12646 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
12647 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
12648
12649 let bogus = root.join("bogus");
12651 std::fs::create_dir_all(&bogus).unwrap();
12652 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
12653 assert!(
12654 err.contains("model.safetensors"),
12655 "error should say what is missing: {err}"
12656 );
12657 assert!(
12658 err.contains("manifest.json"),
12659 "error should mention the repack form: {err}"
12660 );
12661
12662 let nc = root.join("no_config");
12664 std::fs::create_dir_all(&nc).unwrap();
12665 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
12666 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
12667 assert!(
12668 err.contains("config.json"),
12669 "error should name config.json: {err}"
12670 );
12671
12672 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
12674 assert!(err.contains("does not exist"), "{err}");
12675
12676 let f = root.join("model.gguf");
12678 std::fs::write(&f, b"g").unwrap();
12679 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
12680
12681 let _ = std::fs::remove_dir_all(&root);
12682 }
12683
12684 #[test]
12685 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
12686 let caps = ModelCaps {
12689 tools_branch: false,
12690 qwen_think: false,
12691 think_switch: false,
12692 chat_ok: false,
12693 ..Default::default()
12694 };
12695 let payload = serde_json::json!({
12696 "model": "st_model",
12697 "messages": [{"role": "user", "content": "hello"}],
12698 });
12699 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12700 let (tx, _rx) = worker::event_channel();
12701 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
12702 Err(e) => e,
12703 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
12704 };
12705 assert!(
12706 err.contains("no chat template"),
12707 "message should name the cause: {err}"
12708 );
12709 assert!(
12710 err.contains("/v1/completions"),
12711 "message should point at the raw-prompt escape hatch: {err}"
12712 );
12713 }
12714
12715 #[test]
12716 fn tools_on_model_without_tools_branch_is_rejected() {
12717 let (tx, _rx) = worker::event_channel();
12718 let caps = ModelCaps {
12719 chat_ok: true,
12720 ..Default::default()
12721 };
12722 assert!(
12723 build_chat_request(
12724 weather_request(json!({})),
12725 Some(&caps),
12726 tx,
12727 lanes::Lane::Interactive,
12728 None
12729 )
12730 .is_err()
12731 );
12732 let (tx, _rx) = worker::event_channel();
12733 assert!(
12734 build_chat_request(
12735 weather_request(json!({})),
12736 None,
12737 tx,
12738 lanes::Lane::Interactive,
12739 None
12740 )
12741 .is_err()
12742 );
12743 }
12744
12745 #[test]
12746 fn reasoning_effort_maps_to_think_switch() {
12747 for (extra, want) in [
12753 (json!({}), ThinkMode::Default),
12754 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
12755 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
12756 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
12757 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
12758 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
12759 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
12760 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
12761 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
12762 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
12766 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
12767 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
12768 (
12772 json!({"reasoning": {"enabled": true, "effort": "none"}}),
12773 ThinkMode::Think,
12774 ),
12775 (
12776 json!({"reasoning": {"enabled": false, "effort": "high"}}),
12777 ThinkMode::NoThink,
12778 ),
12779 ] {
12780 let (tx, _rx) = worker::event_channel();
12781 let plan = build_chat_request(
12782 weather_request(extra.clone()),
12783 Some(&ladder_caps()),
12788 tx,
12789 lanes::Lane::Interactive,
12790 None,
12791 )
12792 .unwrap();
12793 assert_eq!(plan.request.think, want, "extra={extra}");
12794 }
12795 for extra in [
12799 json!({"reasoning_effort": "extreme"}),
12800 json!({"reasoning": {"effort": "banana"}}),
12801 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
12802 json!({"reasoning": {"enabled": true, "effort": ""}}),
12803 ] {
12804 let (tx, _rx) = worker::event_channel();
12805 assert!(
12806 build_chat_request(
12807 weather_request(extra.clone()),
12808 Some(&tool_caps()),
12809 tx,
12810 lanes::Lane::Interactive,
12811 None
12812 )
12813 .is_err(),
12814 "extra={extra} must be rejected by the one allowlist"
12815 );
12816 }
12817 for (raw, want) in [
12820 ("none", Some("none")),
12821 ("minimal", Some("minimal")),
12822 ("low", Some("low")),
12823 ("medium", Some("medium")),
12824 ("high", Some("high")),
12825 ("xhigh", Some("high")),
12826 ("max", Some("high")),
12827 ("ultra", Some("high")),
12828 ("banana", None),
12829 ("", None),
12830 ("HIGH", None),
12831 ] {
12832 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
12833 }
12834 for (raw, want) in [
12837 ("none", Some("none")),
12838 ("minimal", Some("minimal")),
12839 ("low", Some("low")),
12840 ("medium", Some("medium")),
12841 ("high", Some("high")),
12842 ("xhigh", Some("max")),
12843 ("max", Some("max")),
12844 ("ultra", Some("max")),
12845 ("banana", None),
12846 ("", None),
12847 ("MAX", None),
12848 ] {
12849 assert_eq!(
12850 canonical_effort_for(raw, true),
12851 want,
12852 "canonical_effort_for({raw:?}, dsv4)"
12853 );
12854 }
12855 }
12856
12857 #[test]
12858 fn dsv4_reasoning_effort_max_survives_canonicalization() {
12859 let dsv4_caps = ModelCaps {
12865 chat_ok: true,
12866 dsv4: true,
12867 ..Default::default()
12868 };
12869 let build = |caps: &ModelCaps, effort: &str| {
12870 let (tx, _rx) = worker::event_channel();
12871 let req: ChatCompletionReq = serde_json::from_value(json!({
12872 "model": "m",
12873 "messages": [{"role": "user", "content": "hi"}],
12874 "reasoning_effort": effort,
12875 }))
12876 .unwrap();
12877 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
12878 };
12879 for raw in ["max", "xhigh", "ultra"] {
12880 let plan = build(&dsv4_caps, raw).unwrap();
12881 assert_eq!(
12882 plan.request.reasoning_effort.as_deref(),
12883 Some("max"),
12884 "dsv4 {raw:?} must reach the renderer as the max rung"
12885 );
12886 assert_eq!(plan.request.think, chat::ThinkMode::Think);
12887 }
12888 let plan = build(&dsv4_caps, "high").unwrap();
12890 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12891 let step_caps = ModelCaps {
12893 chat_ok: true,
12894 effort_levels: true,
12895 ..Default::default()
12896 };
12897 let plan = build(&step_caps, "max").unwrap();
12898 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12899 }
12900
12901 #[test]
12902 fn default_reasoning_effort_flips_only_the_unset_request() {
12903 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
12910 let (tx, _rx) = worker::event_channel();
12911 build_chat_request_with_trace(
12912 weather_request(extra),
12913 Some(&ladder_caps()),
12914 tx,
12915 lanes::Lane::Interactive,
12916 None,
12917 None,
12918 default_effort,
12919 &ModelSamplingDefaults::default(),
12920 )
12921 .unwrap()
12922 };
12923 for (extra, want) in [
12924 (json!({}), ThinkMode::Think),
12926 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
12930 (json!({"include_reasoning": false}), ThinkMode::NoThink),
12931 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
12933 (json!({"include_reasoning": true}), ThinkMode::Think),
12934 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
12936 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
12937 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
12938 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
12940 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
12941 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
12942 ] {
12943 let plan = build(extra.clone(), Some("high"));
12944 assert_eq!(plan.request.think, want, "extra={extra}");
12945 }
12946 assert_eq!(
12948 build(json!({}), Some("none")).request.think,
12949 ThinkMode::NoThink
12950 );
12951 assert_eq!(
12952 build(json!({"reasoning_effort": "high"}), Some("none"))
12953 .request
12954 .think,
12955 ThinkMode::Think
12956 );
12957 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
12961 }
12962
12963 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
12969 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
12970 {%- else %}'<think>\\n'{%- endif %}";
12971
12972 #[test]
12973 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
12974 let build = |extra: serde_json::Value| {
12980 let (tx, _rx) = worker::event_channel();
12981 build_chat_request(
12982 weather_request(extra),
12983 Some(&tool_caps()),
12984 tx,
12985 lanes::Lane::Interactive,
12986 None,
12987 )
12988 };
12989 for (extra, want) in [
12990 (json!({"enable_thinking": false}), ThinkMode::NoThink),
12991 (json!({"enable_thinking": true}), ThinkMode::Think),
12992 (
12993 json!({"chat_template_kwargs": {"enable_thinking": false}}),
12994 ThinkMode::NoThink,
12995 ),
12996 (
12997 json!({"chat_template_kwargs": {"enable_thinking": true}}),
12998 ThinkMode::Think,
12999 ),
13000 (
13003 json!({"enable_thinking": false, "reasoning_effort": "high"}),
13004 ThinkMode::NoThink,
13005 ),
13006 (
13008 json!({"enable_thinking": false,
13009 "chat_template_kwargs": {"enable_thinking": false}}),
13010 ThinkMode::NoThink,
13011 ),
13012 ] {
13013 let plan = build(extra.clone()).unwrap_or_else(|e| {
13014 panic!("{extra} must be accepted and honored, got 400: {e}");
13015 });
13016 assert_eq!(
13017 plan.request.think, want,
13018 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
13019 );
13020 }
13021 let render = |extra: serde_json::Value| -> String {
13024 let plan = build(extra).unwrap();
13025 chat::apply_chat_template_tools_ex(
13026 Some(SWITCHED_QWEN_TMPL),
13027 &plan.request.chat_turns,
13028 true,
13029 &plan.request.tools_json,
13030 &plan.request.tools_struct,
13031 plan.request.think,
13032 plan.request.reasoning_effort.as_deref(),
13033 None,
13034 )
13035 .unwrap()
13036 };
13037 let off = render(json!({"enable_thinking": false}));
13038 assert!(
13039 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
13040 "enable_thinking:false must render the CLOSED think pair: {off:?}"
13041 );
13042 let on = render(json!({}));
13043 assert!(
13044 on.ends_with("<|im_start|>assistant\n<think>\n"),
13045 "an unset request must still render the template's OPEN think tail: {on:?}"
13046 );
13047 assert_eq!(
13048 off,
13049 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
13050 "both vLLM spellings must render byte-identically"
13051 );
13052 assert_eq!(
13053 off,
13054 render(json!({"reasoning_effort": "none"})),
13055 "the vLLM spelling must render byte-identically to the OpenAI spelling"
13056 );
13057 }
13058
13059 #[test]
13060 fn unknown_chat_template_kwarg_refuses_by_name() {
13061 let build = |extra: serde_json::Value| {
13064 let (tx, _rx) = worker::event_channel();
13065 build_chat_request(
13066 weather_request(extra),
13067 Some(&tool_caps()),
13068 tx,
13069 lanes::Lane::Interactive,
13070 None,
13071 )
13072 };
13073 let refusal = |extra: serde_json::Value, why: &str| -> String {
13074 build(extra).err().unwrap_or_else(|| panic!("{why}"))
13075 };
13076 let err = refusal(
13077 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
13078 "an unimplementable template kwarg must not be accepted",
13079 );
13080 assert!(
13081 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
13082 "the refusal must name the offending key AND the supported one: {err}"
13083 );
13084 let err = refusal(
13085 json!({"chat_template_kwargs": "enable_thinking=false"}),
13086 "a non-object chat_template_kwargs must not be accepted",
13087 );
13088 assert!(
13089 err.contains("must be an object"),
13090 "refusal must say what shape is expected: {err}"
13091 );
13092 let err = refusal(
13093 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
13094 "a stringly-typed switch must not be accepted",
13095 );
13096 assert!(
13097 err.contains("true or false"),
13098 "refusal must name the expected type: {err}"
13099 );
13100 let plan = build(json!({"chat_template_kwargs": null}))
13102 .expect("null chat_template_kwargs is the unset case");
13103 assert_eq!(plan.request.think, ThinkMode::Default);
13104 }
13105
13106 const Q38_TMPL: &str =
13121 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
13122
13123 fn render_with(
13126 tmpl: &str,
13127 caps: &ModelCaps,
13128 extra: serde_json::Value,
13129 default_effort: Option<&str>,
13130 ) -> Result<String, String> {
13131 let mut payload = serde_json::json!({
13132 "model": "m",
13133 "messages": [{"role": "user", "content": "hi"}],
13134 });
13135 if let Some(obj) = extra.as_object() {
13136 for (k, v) in obj {
13137 payload[k] = v.clone();
13138 }
13139 }
13140 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13141 let (tx, _rx) = worker::event_channel();
13142 let plan = build_chat_request_with_trace(
13143 req,
13144 Some(caps),
13145 tx,
13146 lanes::Lane::Interactive,
13147 None,
13148 None,
13149 default_effort,
13150 &ModelSamplingDefaults::default(),
13151 )?;
13152 Ok(chat::apply_chat_template_tools_ex(
13153 Some(tmpl),
13154 &plan.request.chat_turns,
13155 true,
13156 &plan.request.tools_json,
13157 &plan.request.tools_struct,
13158 plan.request.think,
13159 plan.request.reasoning_effort.as_deref(),
13160 None,
13161 )
13162 .unwrap())
13163 }
13164
13165 #[test]
13166 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
13167 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
13174 let xhigh = "Reasoning effort is set to xhigh.";
13175 let low = "Reasoning effort is set to low.";
13176 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
13178 assert!(
13179 r(json!({"reasoning_effort": "high"}))
13180 .unwrap()
13181 .contains(xhigh)
13182 );
13183 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
13186 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
13187 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
13189 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
13190 assert_ne!(low_p, high_p);
13191 assert_ne!(low_p, medium);
13192 assert_ne!(high_p, medium);
13193 for alias in ["xhigh", "max", "ultra"] {
13196 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
13197 }
13198 assert_eq!(r(json!({})).unwrap(), high_p);
13201 assert_eq!(
13204 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
13205 medium
13206 );
13207 let off = r(json!({"reasoning_effort": "none"})).unwrap();
13210 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
13211 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
13212 }
13213
13214 #[test]
13215 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
13216 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
13229focused, moving directly to the conclusion without unnecessary elaboration.";
13230 let expected = format!(
13231 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
13232 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
13233 );
13234 let after_fix = render_with(
13236 Q38_TMPL,
13237 &ladder_caps(),
13238 json!({"reasoning_effort": "low"}),
13239 None,
13240 )
13241 .unwrap();
13242 assert_eq!(
13243 after_fix, expected,
13244 "the shipped prompt for reasoning_effort:\"low\""
13245 );
13246 const ORNITH_TMPL: &str = include_str!(
13249 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
13250 );
13251 let on_deployed_binary = render_with(
13252 ORNITH_TMPL,
13253 &tool_caps(),
13254 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
13255 {"role": "user", "content": "hi"}]}),
13256 None,
13257 )
13258 .unwrap();
13259 assert_eq!(
13260 on_deployed_binary, expected,
13261 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
13262 level, or its reasoning-volume numbers do not describe the shipped prompt"
13263 );
13264 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
13267 assert!(
13268 !ladderless_unset.contains("Reasoning effort is set to"),
13269 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
13270 );
13271 assert_eq!(
13272 ladderless_unset,
13273 render_with(
13274 Q38_TMPL,
13275 &ladder_caps(),
13276 json!({"reasoning_effort": "medium"}),
13277 None
13278 )
13279 .unwrap(),
13280 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
13281 );
13282 }
13283
13284 #[test]
13285 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
13286 let off = render_with(
13293 Q38_TMPL,
13294 &ladder_caps(),
13295 json!({"reasoning_effort": "none"}),
13296 None,
13297 )
13298 .unwrap();
13299 for extra in [
13300 json!({"include_reasoning": false}),
13301 json!({"reasoning": {"exclude": true}}),
13302 ] {
13303 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
13304 assert!(
13305 got.ends_with("<think>\n\n</think>\n\n"),
13306 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
13307 );
13308 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
13309 }
13310 for extra in [
13315 json!({"enable_thinking": true, "include_reasoning": false}),
13316 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
13317 json!({"reasoning": {"enabled": true, "exclude": true}}),
13318 ] {
13319 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
13320 .err()
13321 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
13322 assert!(e.contains("contradictory"), "{extra}: {e}");
13323 assert!(
13324 e.contains("include_reasoning") || e.contains("exclude"),
13325 "{extra}: the refusal must name the suppression field the caller sent: {e}"
13326 );
13327 }
13328 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
13331 for extra in [
13332 json!({"include_reasoning": true}),
13333 json!({"reasoning": {"exclude": false}}),
13334 ] {
13335 assert_eq!(
13336 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
13337 dflt,
13338 "{extra} must not perturb the model's default"
13339 );
13340 }
13341 let switchless = ModelCaps {
13345 think_switch: false,
13346 ..tool_caps()
13347 };
13348 let err = render_with(
13349 Q38_TMPL,
13350 &switchless,
13351 json!({"include_reasoning": false}),
13352 None,
13353 )
13354 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
13355 assert!(err.contains("cannot disable reasoning"), "{err}");
13356 }
13357
13358 #[test]
13359 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
13360 let build = |extra: serde_json::Value| {
13361 let (tx, _rx) = worker::event_channel();
13362 build_chat_request(
13363 weather_request(extra),
13364 Some(&ladder_caps()),
13365 tx,
13366 lanes::Lane::Interactive,
13367 None,
13368 )
13369 };
13370 let err = |extra: serde_json::Value, why: &str| -> String {
13371 build(extra).err().unwrap_or_else(|| panic!("{why}"))
13372 };
13373 let e = err(
13377 json!({"reasoning": {"max_tokens": 1024}}),
13378 "reasoning.max_tokens must not be accepted-and-ignored",
13379 );
13380 assert!(e.contains("reasoning.max_tokens"), "{e}");
13381 assert!(e.contains("ONE output budget"), "{e}");
13382 for extra in [
13387 json!({"reasoning": {"max_tokens": null}}),
13388 json!({"reasoning": {"banana": null}}),
13389 ] {
13390 let e = err(
13391 extra.clone(),
13392 "a null-valued unhonourable key must still refuse",
13393 );
13394 assert!(
13395 e.contains("max_tokens") || e.contains("banana"),
13396 "{extra}: {e}"
13397 );
13398 }
13399 let e = err(
13401 json!({"reasoning": {"budget": 5}}),
13402 "an unknown reasoning key must not be accepted",
13403 );
13404 assert!(
13405 e.contains("reasoning.budget") && e.contains("enabled"),
13406 "{e}"
13407 );
13408 for (extra, want) in [
13412 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
13413 (json!({"reasoning": {"exclude": 1}}), "true or false"),
13414 (json!({"reasoning": {"effort": 3}}), "must be a string"),
13415 ] {
13416 let e = err(
13417 extra.clone(),
13418 "a wrong-typed reasoning key must not be ignored",
13419 );
13420 assert!(e.contains(want), "{extra}: {e}");
13421 }
13422 for extra in [
13427 json!({"reasoning": {"enabled": true}}),
13428 json!({"reasoning": {"effort": "low"}}),
13429 json!({"reasoning": {"exclude": false}}),
13430 json!({"reasoning": null}),
13431 json!({"reasoning": {"effort": null}}),
13432 json!({"reasoning": {"enabled": null, "exclude": null}}),
13433 ] {
13434 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
13435 }
13436 }
13437
13438 #[test]
13439 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
13440 const ORNITH_TMPL: &str = include_str!(
13448 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
13449 );
13450 let explicit_on = render_with(
13454 ORNITH_TMPL,
13455 &tool_caps(),
13456 json!({"reasoning": {"enabled": true}}),
13457 None,
13458 )
13459 .unwrap();
13460 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
13461 for extra in [
13462 json!({"reasoning_effort": "low"}),
13463 json!({"reasoning_effort": "medium"}),
13464 json!({"reasoning_effort": "high"}),
13465 json!({"reasoning_effort": "xhigh"}),
13467 json!({"reasoning": {"effort": "xhigh"}}),
13468 ] {
13469 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
13470 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
13471 assert_eq!(
13472 got, explicit_on,
13473 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
13474 documented translation, not a decorative accept"
13475 );
13476 }
13477 for extra in [
13479 json!({}),
13480 json!({"reasoning_effort": "none"}),
13481 json!({"reasoning_effort": "minimal"}),
13482 json!({"enable_thinking": false}),
13483 ] {
13484 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
13485 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
13486 }
13487 let minimal = render_with(
13490 ORNITH_TMPL,
13491 &tool_caps(),
13492 json!({"reasoning_effort": "minimal"}),
13493 None,
13494 )
13495 .unwrap();
13496 assert!(
13497 minimal.ends_with("<think>\n\n</think>\n\n"),
13498 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
13499 );
13500 let ladder_low = render_with(
13503 Q38_TMPL,
13504 &ladder_caps(),
13505 json!({"reasoning_effort": "low"}),
13506 None,
13507 )
13508 .unwrap();
13509 assert!(
13510 ladder_low.contains("Reasoning effort is set to low."),
13511 "{ladder_low:?}"
13512 );
13513 assert_ne!(
13514 ladder_low,
13515 render_with(
13516 Q38_TMPL,
13517 &ladder_caps(),
13518 json!({"reasoning_effort": "high"}),
13519 None
13520 )
13521 .unwrap(),
13522 "the ladder model's rungs stay distinct prompts"
13523 );
13524 }
13525
13526 #[test]
13527 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
13528 let render_chat = |body: serde_json::Value| -> Result<String, String> {
13539 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13540 let (tx, _rx) = worker::event_channel();
13541 let plan = build_chat_request(
13542 req,
13543 Some(&ladder_caps()),
13544 tx,
13545 lanes::Lane::Interactive,
13546 None,
13547 )?;
13548 Ok(chat::apply_chat_template_tools_ex(
13549 Some(Q38_TMPL),
13550 &plan.request.chat_turns,
13551 true,
13552 &plan.request.tools_json,
13553 &plan.request.tools_struct,
13554 plan.request.think,
13555 plan.request.reasoning_effort.as_deref(),
13556 None,
13557 )
13558 .unwrap())
13559 };
13560 for (intent, chat_body, responses_body, messages_body) in [
13565 (
13566 "reasoning OFF",
13567 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13568 "reasoning_effort": "none"}),
13569 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
13570 json!({"model": "m", "max_tokens": 16,
13571 "messages": [{"role": "user", "content": "hi"}],
13572 "thinking": {"type": "disabled"}}),
13573 ),
13574 (
13575 "reasoning ON at the top rung",
13576 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13577 "reasoning_effort": "xhigh"}),
13578 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
13579 json!({"model": "m", "max_tokens": 16,
13580 "messages": [{"role": "user", "content": "hi"}],
13581 "output_config": {"effort": "xhigh"}}),
13582 ),
13583 (
13584 "reasoning ON at the bottom rung",
13585 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13586 "reasoning_effort": "low"}),
13587 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
13588 json!({"model": "m", "max_tokens": 16,
13589 "messages": [{"role": "user", "content": "hi"}],
13590 "output_config": {"effort": "low"}}),
13591 ),
13592 (
13593 "the model's own default",
13594 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
13595 json!({"model": "m", "input": "hi"}),
13596 json!({"model": "m", "max_tokens": 16,
13597 "messages": [{"role": "user", "content": "hi"}]}),
13598 ),
13599 ] {
13600 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
13601 let via_responses = responses_api::translate(&responses_body)
13602 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
13603 let via_messages = anthropic::translate(&messages_body)
13604 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
13605 for (surface, translated) in [
13606 ("/v1/responses", via_responses),
13607 ("/v1/messages", via_messages),
13608 ] {
13609 let got = render_chat(translated)
13610 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
13611 assert_eq!(
13612 got, chat,
13613 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
13614 /v1/chat/completions — the parameter is honoured on one format and not \
13615 the other"
13616 );
13617 }
13618 }
13619 let switchless = ModelCaps {
13622 think_switch: false,
13623 ..ladder_caps()
13624 };
13625 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
13626 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13627 let (tx, _rx) = worker::event_channel();
13628 let plan =
13629 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
13630 Ok(format!("{:?}", plan.request.think))
13631 };
13632 for (surface, body) in [
13633 (
13634 "/v1/responses",
13635 responses_api::translate(&json!({
13636 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
13637 .unwrap(),
13638 ),
13639 (
13640 "/v1/messages",
13641 anthropic::translate(&json!({
13642 "model": "m", "max_tokens": 16,
13643 "messages": [{"role": "user", "content": "hi"}],
13644 "thinking": {"type": "disabled"}}))
13645 .unwrap(),
13646 ),
13647 ] {
13648 let err = render_switchless(body)
13649 .err()
13650 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
13651 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
13652 }
13653 }
13654
13655 #[test]
13656 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
13657 let build = |extra: serde_json::Value| {
13665 let (tx, _rx) = worker::event_channel();
13666 build_chat_request(
13667 weather_request(extra),
13668 Some(&ladder_caps()),
13669 tx,
13670 lanes::Lane::Interactive,
13671 None,
13672 )
13673 };
13674 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
13675 .expect("preserve_thinking:true is the vendor default the renderer implements");
13676 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
13677 .err()
13678 .expect("preserve_thinking:false (the strip arm) must refuse");
13679 assert!(e.contains("preserve_thinking"), "{e}");
13680 assert!(e.contains("strip"), "{e}");
13681 assert_eq!(
13684 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
13685 .unwrap()
13686 .request
13687 .think,
13688 ThinkMode::NoThink
13689 );
13690 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
13692 .err()
13693 .expect("a stringly-typed preserve_thinking must not be accepted");
13694 assert!(e.contains("true or false"), "{e}");
13695 }
13696
13697 #[test]
13698 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
13699 let dsv4_caps = ModelCaps {
13705 qwen_think: true,
13706 think_switch: false,
13707 dsv4: true,
13708 ..tool_caps()
13709 };
13710 for extra in [
13711 json!({"reasoning_effort": "none"}),
13712 json!({"reasoning": {"enabled": false}}),
13713 json!({"enable_thinking": false}),
13714 json!({"include_reasoning": false}),
13715 ] {
13716 let (tx, _rx) = worker::event_channel();
13717 let plan = build_chat_request(
13718 weather_request(extra.clone()),
13719 Some(&dsv4_caps),
13720 tx,
13721 lanes::Lane::Interactive,
13722 None,
13723 )
13724 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
13725 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
13726 }
13727 }
13728
13729 #[test]
13730 fn contradictory_think_switches_refuse_instead_of_picking_one() {
13731 let build = |extra: serde_json::Value| {
13734 let (tx, _rx) = worker::event_channel();
13735 build_chat_request(
13736 weather_request(extra),
13737 Some(&tool_caps()),
13738 tx,
13739 lanes::Lane::Interactive,
13740 None,
13741 )
13742 };
13743 for extra in [
13744 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
13745 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
13746 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
13747 ] {
13748 match build(extra.clone()) {
13749 Err(err) => assert!(
13750 err.contains("contradictory"),
13751 "the refusal must say the switches contradict: {err}"
13752 ),
13753 Ok(plan) => panic!(
13754 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
13755 plan.request.think
13756 ),
13757 }
13758 }
13759 for extra in [
13761 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
13762 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
13763 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
13764 ] {
13765 build(extra.clone())
13766 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
13767 }
13768 }
13769
13770 #[test]
13771 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
13772 let switchless = ModelCaps {
13777 tools_branch: true,
13778 qwen_think: true,
13779 think_switch: false,
13780 chat_ok: true,
13781 ..Default::default()
13782 };
13783 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
13784 let (tx, _rx) = worker::event_channel();
13785 build_chat_request_with_trace(
13786 weather_request(extra),
13787 Some(caps),
13788 tx,
13789 lanes::Lane::Interactive,
13790 None,
13791 None,
13792 default_effort,
13793 &ModelSamplingDefaults::default(),
13794 )
13795 };
13796 for extra in [
13797 json!({"reasoning_effort": "none"}),
13798 json!({"reasoning_effort": "minimal"}),
13799 json!({"reasoning": {"enabled": false}}),
13800 json!({"enable_thinking": false}),
13801 json!({"chat_template_kwargs": {"enable_thinking": false}}),
13802 ] {
13803 let err = build(extra.clone(), &switchless, None)
13804 .err()
13805 .unwrap_or_else(|| {
13806 panic!(
13807 "{extra} on a switchless think template must not be accepted-and-ignored"
13808 )
13809 });
13810 assert!(
13811 err.contains("cannot disable reasoning"),
13812 "the refusal must say the model cannot disable reasoning: {err}"
13813 );
13814 }
13815 for (extra, default_effort) in [
13819 (json!({}), None),
13820 (json!({"reasoning_effort": "high"}), None),
13823 (json!({"reasoning": {"enabled": true}}), None),
13824 (json!({"enable_thinking": true}), None),
13825 (json!({}), Some("none")),
13826 (json!({}), Some("minimal")),
13827 (json!({}), Some("high")),
13828 ] {
13829 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
13830 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
13831 });
13832 }
13833 assert_eq!(
13836 build(json!({"enable_thinking": false}), &tool_caps(), None)
13837 .unwrap()
13838 .request
13839 .think,
13840 ThinkMode::NoThink
13841 );
13842 }
13843
13844 #[test]
13845 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
13846 let gemma_caps = ModelCaps {
13853 tools_branch: true,
13854 chat_ok: true,
13855 gemma_think: true,
13856 instruct_type: Some("gemma".into()),
13857 ..Default::default()
13858 };
13859 let render =
13860 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
13861 let mut payload = serde_json::json!({
13862 "model": "google/gemma-4-31b-it",
13863 "messages": [{"role": "user", "content": "Weather in Paris?"}],
13864 });
13865 if let Some(obj) = extra.as_object() {
13866 for (k, v) in obj {
13867 payload[k] = v.clone();
13868 }
13869 }
13870 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13871 let (tx, _rx) = worker::event_channel();
13872 let plan = build_chat_request_with_trace(
13873 req,
13874 Some(&gemma_caps),
13875 tx,
13876 lanes::Lane::Interactive,
13877 None,
13878 None,
13879 default_effort,
13880 &ModelSamplingDefaults::default(),
13881 )
13882 .unwrap();
13883 chat::apply_chat_template_tools_ex(
13884 Some(tmpl),
13885 &plan.request.chat_turns,
13886 true,
13887 &plan.request.tools_json,
13888 &plan.request.tools_struct,
13889 plan.request.think,
13890 plan.request.reasoning_effort.as_deref(),
13891 None, )
13893 .unwrap()
13894 };
13895 let official = gemma_template("official");
13896 let unset_with_knob = render(&official, json!({}), Some("high"));
13897 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
13898 assert_eq!(
13899 unset_with_knob, explicit_on,
13900 "knob render must be byte-identical to the explicit think-on render"
13901 );
13902 assert!(
13903 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
13904 "think-on injects the <|think|> system token: {unset_with_knob:?}"
13905 );
13906 assert!(
13907 unset_with_knob.ends_with("<|turn>model\n"),
13908 "think-on generation turn is OPEN: {unset_with_knob:?}"
13909 );
13910 let explicit_off_with_knob =
13914 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
13915 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
13916 assert_eq!(explicit_off_with_knob, explicit_off);
13917 assert!(
13918 !explicit_off_with_knob.contains("<|think|>")
13919 && explicit_off_with_knob.ends_with("<|turn>model\n"),
13920 "explicit off keeps the official template's thinking-off bytes: \
13921 {explicit_off_with_knob:?}"
13922 );
13923 let unset_no_knob = render(&official, json!({}), None);
13925 assert_eq!(
13926 unset_no_knob, explicit_off,
13927 "knobless unset stays the template's own thinking-off default"
13928 );
13929 assert_ne!(unset_no_knob, unset_with_knob);
13930 let qat = gemma_template("qat");
13933 assert!(
13934 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
13935 "QAT knobless unset keeps the closed-channel default"
13936 );
13937 assert_eq!(
13938 render(&qat, json!({}), Some("high")),
13939 render(&qat, json!({"reasoning_effort": "high"}), None),
13940 "QAT knob render must equal the explicit think-on render"
13941 );
13942 }
13943
13944 #[test]
13945 fn default_reasoning_effort_is_validated_at_metadata_load() {
13946 let parsed = OpenRouterMetadataFile::from_toml(
13948 r#"
13949[models.g]
13950default_reasoning_effort = "high"
13951"#,
13952 )
13953 .unwrap();
13954 assert_eq!(
13955 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
13956 Some("high")
13957 );
13958 let err = OpenRouterMetadataFile::from_toml(
13959 r#"
13960[models.g]
13961default_reasoning_effort = "always"
13962"#,
13963 )
13964 .unwrap_err();
13965 assert!(err.contains("default_reasoning_effort"), "{err}");
13966 }
13967
13968 #[test]
13969 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
13970 let effort_caps = ModelCaps {
13982 effort_levels: true,
13983 think_switch: false,
13984 ..tool_caps()
13985 };
13986 for (extra, want) in [
13987 (json!({}), None),
13988 (json!({"reasoning_effort": "low"}), Some("low")),
13989 (json!({"reasoning_effort": "medium"}), Some("medium")),
13990 (json!({"reasoning_effort": "high"}), Some("high")),
13991 (json!({"reasoning": {"effort": "high"}}), Some("high")),
13992 (json!({"reasoning_effort": "xhigh"}), Some("high")),
13994 (json!({"reasoning": {"effort": "max"}}), Some("high")),
13995 ] {
13996 let (tx, _rx) = worker::event_channel();
13997 let plan = build_chat_request(
13998 weather_request(extra.clone()),
13999 Some(&effort_caps),
14000 tx,
14001 lanes::Lane::Interactive,
14002 None,
14003 )
14004 .unwrap();
14005 assert_eq!(
14006 plan.request.reasoning_effort.as_deref(),
14007 want,
14008 "extra={extra}"
14009 );
14010 }
14011 for extra in [
14018 json!({"reasoning_effort": "none"}),
14019 json!({"reasoning_effort": "minimal"}),
14020 json!({"reasoning": {"enabled": false}}),
14021 json!({"enable_thinking": false}),
14022 json!({"include_reasoning": false}),
14023 ] {
14024 let (tx, _rx) = worker::event_channel();
14025 let err = build_chat_request(
14026 weather_request(extra.clone()),
14027 Some(&effort_caps),
14028 tx,
14029 lanes::Lane::Interactive,
14030 None,
14031 )
14032 .err()
14033 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
14034 assert!(
14035 err.contains("cannot disable reasoning"),
14036 "extra={extra}: {err}"
14037 );
14038 }
14039 for extra in [
14046 json!({"reasoning_effort": "high"}),
14047 json!({"reasoning": {"effort": "low"}}),
14048 ] {
14049 let (tx, _rx) = worker::event_channel();
14050 let plan = build_chat_request(
14051 weather_request(extra.clone()),
14052 Some(&tool_caps()),
14053 tx,
14054 lanes::Lane::Interactive,
14055 None,
14056 )
14057 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
14058 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
14059 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
14060 }
14061 let (tx, _rx) = worker::event_channel();
14063 let plan = build_chat_request(
14064 weather_request(json!({})),
14065 Some(&tool_caps()),
14066 tx,
14067 lanes::Lane::Interactive,
14068 None,
14069 )
14070 .unwrap();
14071 assert_eq!(plan.request.reasoning_effort, None);
14072 }
14073
14074 #[test]
14075 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
14076 let payload = serde_json::json!({
14077 "model": "m",
14078 "messages": [
14079 {"role": "user", "content": "Weather in Paris?"},
14080 {"role": "assistant", "content": null, "tool_calls": [
14081 {"id": "call_x", "type": "function", "function": {
14082 "name": "get_weather",
14083 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
14084 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
14085 ],
14086 });
14087 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
14088 let (tx, _rx) = worker::event_channel();
14089 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
14090 .unwrap();
14091 let turns = &plan.request.chat_turns;
14092 assert_eq!(turns[1].tool_calls.len(), 1);
14093 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
14094 assert_eq!(
14095 turns[1].tool_calls[0].params,
14096 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
14097 );
14098 assert_eq!(turns[2].role, "tool");
14099 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
14100 let mut p = plan
14103 .parser
14104 .expect("think-open chat arms the reasoning splitter");
14105 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
14106 assert_eq!(
14107 pieces,
14108 vec![
14109 Piece::Reasoning("thought".into()),
14110 Piece::Content("answer <tool_call> is prose here".into()),
14111 ]
14112 );
14113 }
14114
14115 #[tokio::test]
14116 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
14117 let (tx, rx) = worker::event_channel();
14118 tx.send(Event::Token {
14119 id: 1,
14120 text: "plan</think>\n\n".into(),
14121 })
14122 .unwrap();
14123 tx.send(Event::Token {
14124 id: 2,
14125 text: "<tool_call>\n<function=get_weather>\n\
14126<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
14127 .into(),
14128 })
14129 .unwrap();
14130 tx.send(Event::Done {
14131 stop_reason: "Eos".into(),
14132 n_tokens: 2,
14133 n_prompt: 40,
14134 n_cached: 0,
14135 elapsed_s: 0.5,
14136 spec: None,
14137 })
14138 .unwrap();
14139 drop(tx);
14140 let parser = ToolStreamParser::new(HashMap::new(), true);
14141 let response = blocking_response(
14142 rx,
14143 "m".into(),
14144 true,
14145 Vec::new(),
14146 Some(parser),
14147 Envelope::new(true),
14148 )
14149 .await;
14150 assert_eq!(response.status(), StatusCode::OK);
14151 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
14152 .await
14153 .unwrap();
14154 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14155 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
14156 assert_eq!(
14159 payload["choices"][0]["message"]["content"],
14160 serde_json::Value::Null
14161 );
14162 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
14163 assert_eq!(
14164 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
14165 "plan"
14166 );
14167 let call = &payload["choices"][0]["message"]["tool_calls"][0];
14168 assert_eq!(call["type"], "function");
14169 assert_eq!(call["function"]["name"], "get_weather");
14170 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
14171 assert_eq!(payload["usage"]["prompt_tokens"], 40);
14174 assert_eq!(payload["usage"]["completion_tokens"], 2);
14175 assert_eq!(payload["usage"]["total_tokens"], 42);
14176 assert_eq!(
14177 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
14178 0
14179 );
14180 }
14181
14182 #[test]
14183 fn cache_salt_plumbs_to_the_worker_namespace() {
14184 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14186 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
14187 }))
14188 .unwrap();
14189 let (tx, _rx) = worker::event_channel();
14190 assert_eq!(
14191 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14192 "tenant-a"
14193 );
14194
14195 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14196 "model": "m", "messages": [{"role": "user", "content": "task"}],
14197 "cache_salt": "tenant-b"
14198 }))
14199 .unwrap();
14200 let (tx, _rx) = worker::event_channel();
14201 assert_eq!(
14202 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14203 .unwrap()
14204 .request
14205 .cache_ns,
14206 "tenant-b"
14207 );
14208
14209 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14211 "model": "m", "prompt": "task"
14212 }))
14213 .unwrap();
14214 let (tx, _rx) = worker::event_channel();
14215 assert_eq!(
14216 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14217 ""
14218 );
14219 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14220 "model": "m", "messages": [{"role": "user", "content": "task"}]
14221 }))
14222 .unwrap();
14223 let (tx, _rx) = worker::event_channel();
14224 assert_eq!(
14225 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14226 .unwrap()
14227 .request
14228 .cache_ns,
14229 ""
14230 );
14231 }
14232
14233 #[test]
14234 fn cache_salt_validation_rejects_oversized_value() {
14235 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
14236 assert_eq!(
14237 validate_cache_namespace(&salt, false),
14238 Err("cache_salt must be at most 64 bytes")
14239 );
14240 }
14241
14242 #[test]
14243 fn cache_salt_validation_rejects_reserved_open_namespace() {
14244 let salt = Some("t:acme\u{1f}private".to_string());
14245 assert_eq!(
14246 validate_cache_namespace(&salt, false),
14247 Err("cache_salt must not use the reserved t: prefix without a keyring")
14248 );
14249 }
14250
14251 #[test]
14252 fn cache_salt_validation_accepts_normal_value() {
14253 let raw = "tenant-A_7.c2VjcmV0LXNjb3Bl+/=";
14254 let salt = Some(raw.to_string());
14255 assert_eq!(validate_cache_namespace(&salt, false).unwrap(), raw);
14256 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
14257 let max_raw = "a".repeat(CACHE_SALT_MAX_BYTES);
14258 let max = Some(max_raw.clone());
14259 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max_raw);
14260 }
14261
14262 #[test]
14263 fn cache_salt_validation_rejects_unsupported_characters() {
14264 let salt = Some("tenant salt".to_string());
14265 assert_eq!(
14266 validate_cache_namespace(&salt, false),
14267 Err("cache_salt contains unsupported characters")
14268 );
14269 }
14270
14271 #[test]
14272 fn affinity_key_honors_both_client_conventions_in_priority_order() {
14273 use axum::http::HeaderMap;
14274 let hdr = |v: &str| {
14275 let mut h = HeaderMap::new();
14276 h.insert("x-session-id", v.parse().unwrap());
14277 h
14278 };
14279 let empty = HeaderMap::new();
14280 let s = |v: &str| Some(v.to_string());
14281 assert_eq!(
14283 affinity_key(&s("explicit"), &None, &empty).unwrap(),
14284 s("explicit")
14285 );
14286 assert_eq!(
14287 affinity_key(&None, &s("openai-user"), &empty).unwrap(),
14288 s("openai-user")
14289 );
14290 assert_eq!(
14291 affinity_key(&None, &None, &hdr("hdr-id")).unwrap(),
14292 s("hdr-id")
14293 );
14294 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")).unwrap(), s("a"));
14297 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")).unwrap(), s("b"));
14298 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")).unwrap(), None);
14301 assert_eq!(affinity_key(&s(""), &s("real"), &empty).unwrap(), s("real"));
14302 assert_eq!(
14304 affinity_key(&s(" padded "), &None, &empty).unwrap(),
14305 s("padded")
14306 );
14307 assert_eq!(affinity_key(&None, &None, &empty).unwrap(), None);
14309 assert!(
14310 affinity_key(
14311 &s(&"x".repeat(MAX_CLIENT_IDENTIFIER_BYTES + 1)),
14312 &None,
14313 &empty,
14314 )
14315 .unwrap_err()
14316 .contains("at most")
14317 );
14318 assert!(
14319 affinity_key(&s("forged\nlog"), &None, &empty)
14320 .unwrap_err()
14321 .contains("control")
14322 );
14323 }
14324
14325 #[test]
14326 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
14327 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14328 "model": "m", "prompt": "task", "session_id": "conv-1"
14329 }))
14330 .unwrap();
14331 let (tx, _rx) = worker::event_channel();
14332 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14333 assert_eq!(
14334 build_request(&req, tx, lanes::Lane::Interactive, key)
14335 .affinity
14336 .as_deref(),
14337 Some("conv-1")
14338 );
14339 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14341 "model": "m", "messages": [{"role": "user", "content": "task"}],
14342 "user": "conv-2"
14343 }))
14344 .unwrap();
14345 let (tx, _rx) = worker::event_channel();
14346 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14347 assert_eq!(
14348 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
14349 .unwrap()
14350 .request
14351 .affinity
14352 .as_deref(),
14353 Some("conv-2")
14354 );
14355 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14357 "model": "m", "prompt": "task"
14358 }))
14359 .unwrap();
14360 let (tx, _rx) = worker::event_channel();
14361 assert!(
14362 build_request(&req, tx, lanes::Lane::Interactive, None)
14363 .affinity
14364 .is_none()
14365 );
14366 }
14367
14368 async fn sse_data_lines(resp: Response) -> Vec<String> {
14370 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14371 .await
14372 .unwrap();
14373 String::from_utf8(bytes.to_vec())
14374 .unwrap()
14375 .lines()
14376 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
14377 .collect()
14378 }
14379
14380 #[tokio::test]
14381 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
14382 let feed = |think: bool| {
14388 let (tx, rx) = worker::event_channel();
14389 let body = if think {
14390 "a plan</think>\n\nanswer"
14391 } else {
14392 "answer"
14393 };
14394 tx.send(Event::Token {
14395 id: 1,
14396 text: body.into(),
14397 })
14398 .unwrap();
14399 tx.send(Event::Done {
14400 stop_reason: "Eos".into(),
14401 n_tokens: 3,
14402 n_prompt: 10,
14403 n_cached: 0,
14404 elapsed_s: 0.1,
14405 spec: None,
14406 })
14407 .unwrap();
14408 drop(tx);
14409 rx
14410 };
14411 let resp = blocking_response(
14413 feed(true),
14414 "m".into(),
14415 true,
14416 Vec::new(),
14417 Some(ToolStreamParser::reasoning_only()),
14418 Envelope::new(true),
14419 )
14420 .await;
14421 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14422 .await
14423 .unwrap();
14424 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14425 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
14426 assert_eq!(
14427 v["choices"][0]["message"]["reasoning_details"][0]["text"],
14428 "a plan"
14429 );
14430 assert_eq!(v["choices"][0]["message"]["content"], "answer");
14431 let resp = blocking_response(
14434 feed(false),
14435 "m".into(),
14436 true,
14437 Vec::new(),
14438 None,
14439 Envelope::new(true),
14440 )
14441 .await;
14442 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14443 .await
14444 .unwrap();
14445 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14446 assert!(
14447 v["choices"][0]["message"].get("reasoning").is_none(),
14448 "a reasoning-off response must carry no reasoning field: {v}"
14449 );
14450 assert_eq!(v["choices"][0]["message"]["content"], "answer");
14451 let resp = sse_response(
14453 feed(true),
14454 "m".into(),
14455 true,
14456 Some(ToolStreamParser::reasoning_only()),
14457 Envelope::new(true),
14458 Vec::new(),
14459 None,
14460 )
14461 .into_response();
14462 let lines = sse_data_lines(resp).await;
14463 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14464 .iter()
14465 .map(|l| serde_json::from_str(l).unwrap())
14466 .collect();
14467 let reasoning: String = chunks
14468 .iter()
14469 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
14470 .collect();
14471 assert_eq!(
14472 reasoning, "a plan",
14473 "think text must stream as delta.reasoning"
14474 );
14475 let content: String = chunks
14476 .iter()
14477 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
14478 .collect();
14479 assert_eq!(content, "answer", "content must exclude the think segment");
14480 let resp = sse_response(
14482 feed(false),
14483 "m".into(),
14484 true,
14485 None,
14486 Envelope::new(true),
14487 Vec::new(),
14488 None,
14489 )
14490 .into_response();
14491 let lines = sse_data_lines(resp).await;
14492 for l in &lines[..lines.len() - 1] {
14493 let c: serde_json::Value = serde_json::from_str(l).unwrap();
14494 assert!(
14495 c["choices"][0]["delta"].get("reasoning").is_none(),
14496 "a reasoning-off stream must carry no reasoning deltas: {c}"
14497 );
14498 }
14499 }
14500
14501 #[tokio::test]
14502 async fn stream_chunks_carry_envelope_and_first_delta_role() {
14503 let (tx, rx) = worker::event_channel();
14504 tx.send(Event::Token {
14505 id: 1,
14506 text: "he".into(),
14507 })
14508 .unwrap();
14509 tx.send(Event::Token {
14510 id: 2,
14511 text: "llo".into(),
14512 })
14513 .unwrap();
14514 tx.send(Event::Done {
14515 stop_reason: "Eos".into(),
14516 n_tokens: 2,
14517 n_prompt: 10,
14518 n_cached: 0,
14519 elapsed_s: 0.1,
14520 spec: None,
14521 })
14522 .unwrap();
14523 drop(tx);
14524 let resp = sse_response(
14525 rx,
14526 "m".into(),
14527 true,
14528 None,
14529 Envelope::new(true),
14530 Vec::new(),
14531 None,
14532 )
14533 .into_response();
14534 let lines = sse_data_lines(resp).await;
14535 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
14536 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14537 .iter()
14538 .map(|l| serde_json::from_str(l).unwrap())
14539 .collect();
14540 let id = chunks[0]["id"].as_str().unwrap().to_string();
14542 assert!(id.starts_with("chatcmpl-"));
14543 for c in &chunks {
14544 assert_eq!(c["id"], id.as_str());
14545 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
14546 let fingerprint = c["system_fingerprint"].as_str().unwrap();
14547 assert!(
14548 build_id::fingerprint_is_well_formed(fingerprint),
14549 "chunk system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
14550 );
14551 assert_eq!(c["object"], "chat.completion.chunk");
14552 }
14553 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
14555 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
14556 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
14557 let fin = chunks.last().unwrap();
14559 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
14560 assert_eq!(fin["usage"]["prompt_tokens"], 10);
14561 }
14562
14563 #[tokio::test]
14564 async fn stream_token_events_equal_usage_on_every_finish_path() {
14565 for (stop_reason, expected_finish) in [
14566 ("Eos", "stop"),
14567 ("Callback", "stop"),
14568 ("MaxNew", "length"),
14569 ("ContextFull", "length"),
14570 ] {
14571 let (tx, rx) = worker::event_channel();
14572 tx.send(Event::Token {
14575 id: 248_046,
14576 text: String::new(),
14577 })
14578 .unwrap();
14579 tx.send(Event::Done {
14580 stop_reason: stop_reason.into(),
14581 n_tokens: 1,
14582 n_prompt: 8,
14583 n_cached: 8,
14584 elapsed_s: 0.1,
14585 spec: None,
14586 })
14587 .unwrap();
14588 drop(tx);
14589
14590 let resp = sse_response(
14591 rx,
14592 "m".into(),
14593 true,
14594 None,
14595 Envelope::new(true),
14596 Vec::new(),
14597 None,
14598 )
14599 .into_response();
14600 let lines = sse_data_lines(resp).await;
14601 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
14602 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14603 .iter()
14604 .map(|line| serde_json::from_str(line).unwrap())
14605 .collect();
14606 let token_events = chunks
14607 .iter()
14608 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
14609 .count();
14610 let terminal = chunks.last().unwrap();
14611 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
14612 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
14613 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
14614 }
14615 }
14616
14617 #[tokio::test]
14618 async fn stream_excludes_stop_text_like_non_stream_does() {
14619 let (tx, rx) = worker::event_channel();
14623 tx.send(Event::Token {
14624 id: 1,
14625 text: "answer\nPro".into(),
14626 })
14627 .unwrap();
14628 tx.send(Event::Token {
14629 id: 2,
14630 text: "blem: leaked prompt".into(),
14631 })
14632 .unwrap();
14633 tx.send(Event::Done {
14634 stop_reason: "Callback".into(),
14635 n_tokens: 2,
14636 n_prompt: 8,
14637 n_cached: 0,
14638 elapsed_s: 0.1,
14639 spec: None,
14640 })
14641 .unwrap();
14642 drop(tx);
14643 let resp = sse_response(
14644 rx,
14645 "m".into(),
14646 true,
14647 None,
14648 Envelope::new(true),
14649 vec!["Problem:".into()],
14650 None,
14651 )
14652 .into_response();
14653 let lines = sse_data_lines(resp).await;
14654 let content: String = lines
14655 .iter()
14656 .filter(|l| *l != "[DONE]")
14657 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
14658 .filter_map(|c| {
14659 c["choices"][0]["delta"]["content"]
14660 .as_str()
14661 .map(str::to_string)
14662 })
14663 .collect();
14664 assert_eq!(content, "answer\n");
14665
14666 let (tx, rx) = worker::event_channel();
14668 tx.send(Event::Token {
14669 id: 1,
14670 text: "ends in Pro".into(),
14671 })
14672 .unwrap();
14673 tx.send(Event::Done {
14674 stop_reason: "Eos".into(),
14675 n_tokens: 1,
14676 n_prompt: 8,
14677 n_cached: 0,
14678 elapsed_s: 0.1,
14679 spec: None,
14680 })
14681 .unwrap();
14682 drop(tx);
14683 let resp = sse_response(
14684 rx,
14685 "m".into(),
14686 true,
14687 None,
14688 Envelope::new(true),
14689 vec!["Problem:".into()],
14690 None,
14691 )
14692 .into_response();
14693 let lines = sse_data_lines(resp).await;
14694 let content: String = lines
14695 .iter()
14696 .filter(|l| *l != "[DONE]")
14697 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
14698 .filter_map(|c| {
14699 c["choices"][0]["delta"]["content"]
14700 .as_str()
14701 .map(str::to_string)
14702 })
14703 .collect();
14704 assert_eq!(content, "ends in Pro");
14705 }
14706
14707 #[tokio::test]
14708 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
14709 let (tx, rx) = worker::event_channel();
14710 tx.send(Event::Error(worker::EngineError::engine("boom")))
14711 .unwrap();
14712 drop(tx);
14713 let resp = sse_response(
14714 rx,
14715 "m".into(),
14716 true,
14717 None,
14718 Envelope::new(true),
14719 Vec::new(),
14720 None,
14721 )
14722 .into_response();
14723 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14724 .await
14725 .unwrap();
14726 let body = String::from_utf8(bytes.to_vec()).unwrap();
14727 assert!(
14729 !body.contains("event: error"),
14730 "named SSE event leaked: {body}"
14731 );
14732 let lines: Vec<&str> = body
14733 .lines()
14734 .filter_map(|l| l.strip_prefix("data: "))
14735 .collect();
14736 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
14737 assert_eq!(err["error"]["message"], "boom");
14738 assert_eq!(err["error"]["type"], "server_error");
14739 assert_eq!(err["error"]["code"], "engine_error");
14740 assert_eq!(lines.last(), Some(&"[DONE]"));
14741 }
14742
14743 #[test]
14744 fn ttft_sse_marker_ignores_keepalive_comments() {
14745 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
14746 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
14747 assert!(is_sse_data_frame(
14748 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
14749 ));
14750 }
14751
14752 #[tokio::test]
14753 async fn error_bodies_use_the_openai_object_shape() {
14754 let (tx, rx) = worker::event_channel();
14755 tx.send(Event::Error(worker::EngineError::model_not_found(
14756 "unknown model \"x\"",
14757 )))
14758 .unwrap();
14759 drop(tx);
14760 let response =
14761 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
14762 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
14763 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
14764 .await
14765 .unwrap();
14766 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14767 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
14769 assert_eq!(payload["error"]["type"], "invalid_request_error");
14770 assert_eq!(payload["error"]["param"], "model");
14771 assert_eq!(payload["error"]["code"], "model_not_found");
14772 }
14773
14774 fn retry_after(resp: &Response) -> Option<String> {
14781 resp.headers()
14782 .get(axum::http::header::RETRY_AFTER)
14783 .and_then(|v| v.to_str().ok())
14784 .map(str::to_string)
14785 }
14786
14787 async fn body_value(resp: Response) -> serde_json::Value {
14790 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14791 .await
14792 .expect("body");
14793 serde_json::from_slice(&bytes).expect("json body")
14794 }
14795
14796 async fn chat_completion_admitted(st: &AppState, req: serde_json::Value) -> Response {
14812 let mut last_shed = serde_json::Value::Null;
14813 for _ in 0..50 {
14814 let resp = chat_completions(
14815 State(st.clone()),
14816 HeaderMap::new(),
14817 None,
14818 Json(serde_json::from_value(req.clone()).unwrap()),
14819 )
14820 .await;
14821 if resp.status() != StatusCode::TOO_MANY_REQUESTS {
14822 return resp;
14823 }
14824 let body = body_value(resp).await;
14825 let code = body["error"]["code"].as_str().unwrap_or_default();
14826 assert!(
14827 code.starts_with("shed_"),
14828 "only a contention shed may be retried; any other 429 is a finding: {body}"
14829 );
14830 last_shed = body;
14831 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
14832 }
14833 panic!(
14836 "still shed after 50 attempts — either load the retry budget cannot absorb \
14837 or a shed that no longer clears; last refusal: {last_shed}"
14838 );
14839 }
14840
14841 #[test]
14842 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
14843 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
14845 assert_eq!(
14846 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
14847 TIMEOUT_MS_DEFAULT
14848 );
14849 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
14853 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
14854 }
14855 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
14857 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
14858 assert!(err.contains("timeout_ms"), "{err}");
14859 assert!(
14860 err.contains(&TIMEOUT_MS_MIN.to_string())
14861 && err.contains(&TIMEOUT_MS_MAX.to_string()),
14862 "the message must state the range: {err}"
14863 );
14864 assert!(
14865 err.contains("stream"),
14866 "the message must point at streaming for longer work: {err}"
14867 );
14868 }
14869 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
14871 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
14872 assert!(
14873 err.contains("timeout_ms") && err.contains("stream"),
14874 "{err}"
14875 );
14876 }
14877 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
14879 }
14880
14881 #[tokio::test]
14884 #[allow(clippy::await_holding_lock)] async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
14886 let _l = drain_lock();
14887 let st = fake_worker_state();
14888
14889 let comp = completions(
14890 State(st.clone()),
14891 HeaderMap::new(),
14892 None,
14893 Json(
14894 serde_json::from_value(json!({
14895 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
14896 .unwrap(),
14897 ),
14898 )
14899 .await;
14900 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
14901 let chat = chat_completions(
14902 State(st.clone()),
14903 HeaderMap::new(),
14904 None,
14905 Json(
14906 serde_json::from_value(json!({
14907 "model": "m", "messages": [{"role": "user", "content": "t"}],
14908 "timeout_ms": 90_001}))
14909 .unwrap(),
14910 ),
14911 )
14912 .await;
14913 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
14914 let resp_api = responses_api::responses(
14915 State(st.clone()),
14916 HeaderMap::new(),
14917 None,
14918 axum::body::Bytes::from(
14919 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
14920 ),
14921 )
14922 .await;
14923 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
14924 let msgs = anthropic::messages(
14925 State(st.clone()),
14926 HeaderMap::new(),
14927 None,
14928 axum::body::Bytes::from(
14929 json!({"model": "m", "max_tokens": 16,
14930 "messages": [{"role": "user", "content": "t"}],
14931 "timeout_ms": 90_001})
14932 .to_string(),
14933 ),
14934 )
14935 .await;
14936 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
14937
14938 for (surface, resp) in [
14940 ("/v1/completions", comp),
14941 ("/v1/chat/completions", chat),
14942 ("/v1/responses", resp_api),
14943 ] {
14944 let body = body_value(resp).await;
14945 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
14946 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
14947 let m = body["error"]["message"].as_str().unwrap();
14948 assert!(
14949 m.contains("90000") && m.contains("stream"),
14950 "{surface}: {m}"
14951 );
14952 }
14953 let body = body_value(msgs).await;
14955 assert_eq!(body["error"]["type"], "invalid_request_error");
14956 let m = body["error"]["message"].as_str().unwrap();
14957 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
14958 }
14959
14960 #[tokio::test]
14963 #[allow(clippy::await_holding_lock)] async fn a_non_integer_timeout_ms_is_a_named_400() {
14965 let _l = drain_lock();
14966 let st = fake_worker_state();
14967 let resp = chat_completions(
14968 State(st),
14969 HeaderMap::new(),
14970 None,
14971 Json(
14972 serde_json::from_value(json!({
14973 "model": "m", "messages": [{"role": "user", "content": "t"}],
14974 "timeout_ms": "30s"}))
14975 .unwrap(),
14976 ),
14977 )
14978 .await;
14979 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
14980 let body = body_value(resp).await;
14981 assert_eq!(body["error"]["param"], "timeout_ms");
14982 }
14983
14984 #[tokio::test]
14990 #[allow(clippy::await_holding_lock)] async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
14992 let _l = drain_lock();
14993 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
14998 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
14999 let worker_cancel = cancel_seen.clone();
15000 let health = health::WorkerHealth::new();
15001 let h = health.clone();
15002 std::thread::spawn(move || {
15003 h.mark_ready();
15004 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15005 worker::release_pending_admit();
15006 worker::release_admission_reservation(req.lane);
15007 let _ = req.tx.send(Event::PromptUsage {
15008 n_prompt: 1,
15009 n_cached: 0,
15010 });
15011 let _ = req.tx.send(Event::Token {
15012 id: 1,
15013 text: "partial".into(),
15014 });
15015 for _ in 0..5_000 {
15018 if req.tx.is_closed() {
15019 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
15020 break;
15021 }
15022 std::thread::sleep(std::time::Duration::from_millis(1));
15023 }
15024 }
15025 });
15026 for _ in 0..2_000 {
15027 if health.live().is_ok() {
15028 break;
15029 }
15030 std::thread::sleep(std::time::Duration::from_millis(1));
15031 }
15032 let mut st = fake_worker_state();
15033 st.cmd_tx = cmd_tx;
15034 st.health = health;
15035 let mock = MockMetering::admit_all();
15036 st.metering = Some(mock.clone());
15037
15038 let resp = chat_completion_admitted(
15039 &st,
15040 json!({
15041 "model": "m", "messages": [{"role": "user", "content": "t"}],
15042 "timeout_ms": 1_000}),
15043 )
15044 .await;
15045
15046 assert_eq!(resp.status(), StatusCode::OK);
15051 let body = body_value(resp).await;
15052 assert!(
15053 body["choices"][0]["message"]["content"]
15054 .as_str()
15055 .unwrap()
15056 .contains("partial"),
15057 "the tokens generated before the cut must be delivered: {body}"
15058 );
15059 assert_eq!(body["choices"][0]["finish_reason"], "error");
15063 assert_eq!(
15064 body["choices"][0]["native_finish_reason"],
15065 "deadline_exceeded"
15066 );
15067 assert_eq!(body["error"]["code"], "deadline_exceeded");
15068 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
15069 let message = body["error"]["message"].as_str().unwrap();
15070 assert!(
15071 message.contains("1000") && message.contains("stream"),
15072 "the partial must name the deadline and the streaming alternative: {message}"
15073 );
15074 assert_eq!(body["usage"]["completion_tokens"], 1);
15075
15076 let mut cancelled = false;
15081 for _ in 0..500 {
15082 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
15083 cancelled = true;
15084 break;
15085 }
15086 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
15087 }
15088 assert!(
15089 cancelled,
15090 "the deadline must CANCEL generation (worker's event channel closed)"
15091 );
15092
15093 let events = mock.events();
15098 assert!(
15099 events.contains(&MeterEvent::DeadlinePartial {
15100 prompt: 1,
15101 cached: 0,
15102 completion: 1,
15103 }),
15104 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
15105 );
15106 assert!(
15107 !events
15108 .iter()
15109 .any(|e| matches!(e, MeterEvent::Complete { .. })),
15110 "a deadline cut must stay distinguishable from a full answer: {events:?}"
15111 );
15112 }
15113
15114 #[tokio::test]
15118 #[allow(clippy::await_holding_lock)] async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
15120 let _l = drain_lock();
15121 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15122 let health = health::WorkerHealth::new();
15123 let h = health.clone();
15124 std::thread::spawn(move || {
15125 h.mark_ready();
15126 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15129 worker::release_pending_admit();
15130 worker::release_admission_reservation(req.lane);
15131 let _ = req.tx.send(Event::PromptUsage {
15132 n_prompt: 1,
15133 n_cached: 0,
15134 });
15135 for _ in 0..5_000 {
15136 if req.tx.is_closed() {
15137 break;
15138 }
15139 std::thread::sleep(std::time::Duration::from_millis(1));
15140 }
15141 }
15142 });
15143 for _ in 0..2_000 {
15144 if health.live().is_ok() {
15145 break;
15146 }
15147 std::thread::sleep(std::time::Duration::from_millis(1));
15148 }
15149 let mut st = fake_worker_state();
15150 st.cmd_tx = cmd_tx;
15151 st.health = health;
15152 let mock = MockMetering::admit_all();
15153 st.metering = Some(mock.clone());
15154 let resp = chat_completion_admitted(
15155 &st,
15156 json!({
15157 "model": "m", "messages": [{"role": "user", "content": "t"}],
15158 "timeout_ms": 1_000}),
15159 )
15160 .await;
15161 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15162 assert!(resp.headers().get("x-should-retry").is_none());
15164 assert_eq!(retry_after(&resp), None);
15165 let body = body_value(resp).await;
15166 assert_eq!(body["error"]["code"], "deadline_exceeded");
15167 assert!(
15168 body["error"]["message"]
15169 .as_str()
15170 .unwrap()
15171 .contains("not billed"),
15172 "the zero-token 408 keeps the billing promise: {body}"
15173 );
15174 let events = mock.events();
15175 assert!(
15176 events.contains(&MeterEvent::Unbilled {
15177 outcome: "deadline_exceeded",
15178 status: 408,
15179 code: "deadline_exceeded".into(),
15180 }),
15181 "the named zero-debit census outcome, not the generic reject — every sibling \
15182 deadline path settles this one: {events:?}"
15183 );
15184 }
15185
15186 #[tokio::test]
15189 #[allow(clippy::await_holding_lock)] async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
15191 let _l = drain_lock();
15192 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15194 let health = health::WorkerHealth::new();
15195 let h = health.clone();
15196 std::thread::spawn(move || {
15197 h.mark_ready();
15198 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15199 worker::release_pending_admit();
15200 worker::release_admission_reservation(req.lane);
15201 let _ = req.tx.send(Event::PromptUsage {
15202 n_prompt: 1,
15203 n_cached: 0,
15204 });
15205 while !req.tx.is_closed() {
15206 std::thread::sleep(std::time::Duration::from_millis(1));
15207 }
15208 }
15209 });
15210 for _ in 0..2_000 {
15211 if health.live().is_ok() {
15212 break;
15213 }
15214 std::thread::sleep(std::time::Duration::from_millis(1));
15215 }
15216 let mut st = fake_worker_state();
15217 st.cmd_tx = cmd_tx;
15218 st.health = health;
15219 let mock = MockMetering::admit_all();
15220 st.metering = Some(mock.clone());
15221
15222 let resp = chat_completion_admitted(
15223 &st,
15224 json!({
15225 "model": "m", "messages": [{"role": "user", "content": "t"}],
15226 "stream": true, "timeout_ms": 1_000}),
15227 )
15228 .await;
15229 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15232 let body = body_value(resp).await;
15233 assert_eq!(body["error"]["code"], "deadline_exceeded");
15234 assert!(
15235 body["error"]["message"]
15236 .as_str()
15237 .unwrap()
15238 .contains("first token"),
15239 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
15240 );
15241 let events = mock.events();
15242 assert!(
15243 events.contains(&MeterEvent::Unbilled {
15244 outcome: "deadline_exceeded",
15245 status: 408,
15246 code: "deadline_exceeded".into(),
15247 }),
15248 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
15249 );
15250 }
15251
15252 #[tokio::test]
15256 #[allow(clippy::await_holding_lock)] async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
15258 let _l = drain_lock();
15259 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
15262 let mock = MockMetering::admit_all();
15263 st.metering = Some(mock.clone());
15264 let resp = chat_completion_admitted(
15265 &st,
15266 json!({
15267 "model": "m", "messages": [{"role": "user", "content": "t"}],
15268 "stream": true, "timeout_ms": 1_000}),
15269 )
15270 .await;
15271 assert_eq!(
15272 resp.status(),
15273 StatusCode::OK,
15274 "TTFT was met — 200 is correct"
15275 );
15276 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15277 .await
15278 .expect("the stream must run to completion past the deadline");
15279 let text = String::from_utf8(bytes.to_vec()).unwrap();
15280 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
15281 let events = mock.events();
15282 assert!(
15283 events
15284 .iter()
15285 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
15286 "a stream past its deadline after first token still settles as COMPLETE with \
15287 all four tokens: {events:?}"
15288 );
15289 }
15290
15291 fn admission_counters_guard() -> std::sync::MutexGuard<'static, ()> {
15298 static COUNTERS: std::sync::Mutex<()> = std::sync::Mutex::new(());
15299 COUNTERS
15300 .lock()
15301 .unwrap_or_else(|poisoned| poisoned.into_inner())
15302 }
15303
15304 struct CounterRestore<'a>(&'a std::sync::atomic::AtomicUsize, usize);
15310 impl Drop for CounterRestore<'_> {
15311 fn drop(&mut self) {
15312 self.0.store(self.1, std::sync::atomic::Ordering::Release);
15313 }
15314 }
15315
15316 #[allow(clippy::result_large_err)] fn reserve_interactive_through_contention(
15327 st: &AppState,
15328 rl: &RateLimit,
15329 deadline_ms: u64,
15330 ) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
15331 let reserve = || {
15332 reserve_pending_admit(
15333 st,
15334 lanes::Lane::Interactive,
15335 rl,
15336 RequestDeadline::starting_now(deadline_ms),
15337 )
15338 };
15339 let mut g = reserve();
15340 for _ in 0..50 {
15341 match &g {
15342 Ok(_) => break,
15343 Err((_, "shed_deadline")) => {
15344 std::thread::sleep(std::time::Duration::from_millis(10));
15345 g = reserve();
15346 }
15347 Err((_, outcome)) => panic!("unexpected refusal: {outcome}"),
15348 }
15349 }
15350 g
15351 }
15352
15353 #[test]
15356 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
15357 let _counters = admission_counters_guard();
15358 let st = fake_worker_state();
15359 let lane = lanes::Lane::Interactive;
15360 let cap = lane_cap(lane);
15361 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15362 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
15363 let _restore = CounterRestore(counter, prev);
15364 let rl = RateLimit {
15365 limit: cap,
15366 remaining: 0,
15367 reset_s: 1,
15368 };
15369 let (resp, outcome) = reserve_pending_admit(
15370 &st,
15371 lane,
15372 &rl,
15373 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15374 )
15375 .map(|_| ())
15376 .expect_err("a backlog at the bound must shed");
15377 assert_eq!(outcome, "shed_queue");
15378 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15379 assert!(
15380 retry_after(&resp).is_some(),
15381 "a shed must carry Retry-After so the router's spill can act on it"
15382 );
15383 let stamped = rl.attach(resp);
15385 for h in [
15386 "x-ratelimit-limit",
15387 "x-ratelimit-remaining",
15388 "x-ratelimit-reset",
15389 ] {
15390 assert!(stamped.headers().get(h).is_some(), "missing {h}");
15391 }
15392 }
15393
15394 #[test]
15398 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
15399 let _counters = admission_counters_guard();
15400 let st = fake_worker_state();
15401 let lane = lanes::Lane::Interactive;
15402 let cap = lane_cap(lane);
15403 {
15404 let mut m = st.metrics.lock().unwrap();
15405 m.completed = 10;
15406 m.tokens_out = 1_000;
15407 m.step_p50_ms = 10.0; }
15409 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15410 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15412 let rl = RateLimit {
15413 limit: cap,
15414 remaining: 0,
15415 reset_s: 1,
15416 };
15417 let admitted = reserve_pending_admit(
15419 &st,
15420 lane,
15421 &rl,
15422 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15423 );
15424 assert!(
15425 admitted.is_ok(),
15426 "a request whose deadline covers the estimate must be admitted"
15427 );
15428 drop(admitted); let (resp, outcome) = reserve_pending_admit(
15431 &st,
15432 lane,
15433 &rl,
15434 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15435 )
15436 .map(|_| ())
15437 .expect_err("a deadline shorter than the estimated wait must shed");
15438 assert_eq!(outcome, "shed_deadline");
15439 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15440 assert!(retry_after(&resp).is_some());
15441 }
15442
15443 #[test]
15446 fn deadline_shed_is_interactive_only_and_silent_with_free_slots() {
15447 let _counters = admission_counters_guard();
15448 let st = fake_worker_state();
15449 let cap = lane_cap(lanes::Lane::Interactive);
15450 {
15451 let mut m = st.metrics.lock().unwrap();
15452 m.completed = 10;
15453 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
15455 }
15456 let free = RateLimit {
15458 limit: cap,
15459 remaining: 1,
15460 reset_s: 0,
15461 };
15462 let g = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
15467 assert!(
15468 g.is_ok(),
15469 "free capacity must admit regardless of the estimate"
15470 );
15471 drop(g);
15472 let full = RateLimit {
15475 limit: cap,
15476 remaining: 0,
15477 reset_s: 5,
15478 };
15479 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
15480 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15481 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15483 let g = reserve_pending_admit(
15484 &st,
15485 lane,
15486 &full,
15487 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15488 );
15489 assert!(
15490 g.is_ok(),
15491 "{lane:?} must not be deadline-shed by the interactive gate"
15492 );
15493 drop(g);
15494 }
15495 }
15496
15497 #[test]
15504 fn a_saturated_queue_with_free_http_slots_queues_silently_without_a_ceiling() {
15505 let _counters = admission_counters_guard();
15506 let st = fake_worker_state();
15507 let lane = lanes::Lane::Interactive;
15508 let cap = lane_cap(lane);
15509 {
15510 let mut m = st.metrics.lock().unwrap();
15511 m.completed = 10;
15512 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
15515 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15516 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15518 let free = RateLimit {
15521 limit: cap,
15522 remaining: 1,
15523 reset_s: 0,
15524 };
15525 let g = reserve_pending_admit(
15526 &st,
15527 lane,
15528 &free,
15529 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15530 );
15531 assert!(
15532 g.is_ok(),
15533 "flag off: a ~20 s projected wait whose deadline can absorb it queues \
15534 silently (no 429) - the darklanes#5 defect shape, preserved by default"
15535 );
15536 drop(g);
15537 }
15538
15539 #[test]
15544 fn the_queue_wait_ceiling_sheds_with_429_retry_after_and_the_ratelimit_trio() {
15545 let _counters = admission_counters_guard();
15546 let st = fake_worker_state();
15547 let lane = lanes::Lane::Interactive;
15548 let cap = lane_cap(lane);
15549 {
15550 let mut m = st.metrics.lock().unwrap();
15551 m.completed = 10;
15552 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
15555 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15556 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15558 let free = RateLimit {
15559 limit: cap,
15560 remaining: 1,
15561 reset_s: 0,
15562 };
15563 let (resp, outcome) = reserve_pending_admit_with_ceiling(
15564 &st,
15565 lane,
15566 &free,
15567 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15568 5, )
15570 .map(|_| ())
15571 .expect_err("a projected wait past the ceiling must shed");
15572 assert_eq!(outcome, "shed_queue_wait");
15573 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15574 assert_eq!(
15575 retry_after(&resp).as_deref(),
15576 Some("20"),
15577 "Retry-After must carry the estimate (~10 s/wave x 2 waves)"
15578 );
15579 assert_eq!(
15580 resp.headers()
15581 .get("retry-after-ms")
15582 .and_then(|v| v.to_str().ok()),
15583 Some("20000"),
15584 "the ms twin must match"
15585 );
15586 let stamped = free.attach(resp);
15587 for h in [
15588 "x-ratelimit-limit",
15589 "x-ratelimit-remaining",
15590 "x-ratelimit-reset",
15591 ] {
15592 assert!(stamped.headers().get(h).is_some(), "missing {h}");
15593 }
15594 }
15595
15596 #[test]
15600 fn the_queue_wait_ceiling_admits_under_it_and_never_touches_dark_lanes() {
15601 let _counters = admission_counters_guard();
15602 let st = fake_worker_state();
15603 let lane = lanes::Lane::Interactive;
15604 let cap = lane_cap(lane);
15605 {
15606 let mut m = st.metrics.lock().unwrap();
15607 m.completed = 10;
15608 m.tokens_out = 1_000;
15609 m.step_p50_ms = 100.0; }
15611 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15612 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel);
15613 let _restore = CounterRestore(counter, prev);
15614 let free = RateLimit {
15615 limit: cap,
15616 remaining: 1,
15617 reset_s: 0,
15618 };
15619 let g = reserve_pending_admit_with_ceiling(
15620 &st,
15621 lane,
15622 &free,
15623 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15624 60, );
15626 assert!(
15627 g.is_ok(),
15628 "an estimate under the ceiling must admit and queue as before"
15629 );
15630 drop(g);
15631 let full = RateLimit {
15633 limit: cap,
15634 remaining: 0,
15635 reset_s: 5,
15636 };
15637 for dark in [lanes::Lane::Judge, lanes::Lane::Harvest] {
15638 let counter = &worker::ADMISSION_RESERVATIONS[dark.idx()];
15639 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15641 let g = reserve_pending_admit_with_ceiling(
15642 &st,
15643 dark,
15644 &full,
15645 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15646 1,
15647 );
15648 assert!(
15649 g.is_ok(),
15650 "{dark:?} must not be shed by the interactive queue-wait ceiling"
15651 );
15652 drop(g);
15653 }
15654 }
15655
15656 #[test]
15660 fn the_queue_wait_ceiling_leaves_the_existing_shed_arms_first_and_unchanged() {
15661 let _counters = admission_counters_guard();
15662 let st = fake_worker_state();
15663 let lane = lanes::Lane::Interactive;
15664 let cap = lane_cap(lane);
15665 {
15666 let mut m = st.metrics.lock().unwrap();
15667 m.completed = 10;
15668 m.tokens_out = 1_000;
15669 m.step_p50_ms = 100.0;
15670 }
15671 let rl = RateLimit {
15672 limit: cap,
15673 remaining: 0,
15674 reset_s: 1,
15675 };
15676 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15677 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
15679 let _restore = CounterRestore(counter, prev);
15680 assert!(matches!(
15681 reserve_pending_admit_with_ceiling(
15682 &st,
15683 lane,
15684 &rl,
15685 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15686 1,
15687 ),
15688 Err((_, "shed_queue"))
15689 ));
15690 counter.store(cap, std::sync::atomic::Ordering::Release);
15692 assert!(matches!(
15693 reserve_pending_admit_with_ceiling(
15694 &st,
15695 lane,
15696 &rl,
15697 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15698 1,
15699 ),
15700 Err((_, "shed_deadline"))
15701 ));
15702 }
15703
15704 #[test]
15709 fn the_queue_wait_ceiling_is_wired_through_the_production_wrapper() {
15710 let src = include_str!("lib.rs");
15711 let code: String = src
15712 .lines()
15713 .map(|l| l.split("//").next().unwrap_or(""))
15714 .collect::<Vec<_>>()
15715 .join("\n");
15716 let start = code
15717 .find("pub(crate) fn reserve_pending_admit(")
15718 .expect("the production wrapper exists");
15719 let rest = &code[start..];
15720 let end = rest.find("\nfn ").unwrap_or(rest.len());
15721 let wrapper = &rest[..end];
15722 assert!(
15723 wrapper.contains(
15724 "reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())"
15725 ),
15726 "every production ingress must judge the ceiling the env read armed"
15727 );
15728 }
15729
15730 #[test]
15731 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
15732 let _counters = admission_counters_guard();
15733 let st = fake_worker_state();
15734 let cap = lane_cap(lanes::Lane::Interactive);
15735 let bound = max_queue_depth(cap);
15736 assert!(bound > 0, "the queue bound must admit at least one request");
15737 let rl = RateLimit {
15738 limit: cap,
15739 remaining: 0,
15740 reset_s: 1,
15741 };
15742 let _ = worker::PENDING_ADMITS.fetch_update(
15743 std::sync::atomic::Ordering::AcqRel,
15744 std::sync::atomic::Ordering::Acquire,
15745 |_| Some(0),
15746 );
15747 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
15748 let _restore = CounterRestore(counter, 0);
15749 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
15750 let guard = reserve_pending_admit(
15751 &st,
15752 lanes::Lane::Interactive,
15753 &rl,
15754 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15755 )
15756 .expect("the final queue slot should be reservable");
15757 assert_eq!(
15758 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
15759 1
15760 );
15761 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
15762 drop(guard);
15763 assert_eq!(
15764 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
15765 0
15766 );
15767 assert_eq!(
15768 counter.load(std::sync::atomic::Ordering::Acquire),
15769 bound - 1
15770 );
15771
15772 counter.store(bound, std::sync::atomic::Ordering::Release);
15773 let rejected = reserve_pending_admit(
15774 &st,
15775 lanes::Lane::Interactive,
15776 &rl,
15777 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15778 );
15779 assert!(matches!(rejected, Err((_, "shed_queue"))));
15780 }
15781
15782 #[test]
15783 fn admission_reservations_are_lane_scoped() {
15784 let _counters = admission_counters_guard();
15785 let st = fake_worker_state();
15786 let harvest = lanes::Lane::Harvest;
15787 let interactive = lanes::Lane::Interactive;
15788 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
15789 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
15790 let _restore = CounterRestore(harvest_counter, 0);
15791 harvest_counter.store(
15792 max_queue_depth(lane_cap(harvest)),
15793 std::sync::atomic::Ordering::Release,
15794 );
15795 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
15796 let free = RateLimit {
15797 limit: lane_cap(interactive),
15798 remaining: 1,
15799 reset_s: 0,
15800 };
15801 let guard = reserve_pending_admit(
15812 &st,
15813 interactive,
15814 &free,
15815 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15816 )
15817 .expect("a full harvest queue must not consume interactive capacity");
15818 drop(guard);
15819 let tight = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
15820 assert!(
15821 tight.is_ok(),
15822 "a full harvest queue must not deadline-shed a tight interactive request \
15823 (a backlog that outlasts the retry budget here is a cross-lane leak, not \
15824 contention)"
15825 );
15826 drop(tight);
15827 let harvest_rl = RateLimit {
15828 limit: lane_cap(harvest),
15829 remaining: 0,
15830 reset_s: 1,
15831 };
15832 assert!(matches!(
15833 reserve_pending_admit(
15834 &st,
15835 harvest,
15836 &harvest_rl,
15837 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
15838 ),
15839 Err((_, "shed_queue"))
15840 ));
15841 }
15842
15843 #[test]
15844 fn taxonomy_maps_every_class_to_its_status_and_code() {
15845 use worker::{EngineError as E, ErrClass as C};
15846 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
15847 (
15848 E::invalid_param("bad json", "response_format"),
15849 StatusCode::BAD_REQUEST,
15850 "invalid_request_error",
15851 "",
15852 ),
15853 (
15854 E::context_length("prompt (9000 tok) >= context cap (8192)"),
15855 StatusCode::BAD_REQUEST,
15856 "invalid_request_error",
15857 "context_length_exceeded",
15858 ),
15859 (
15860 E::model_not_found("unknown model \"nope\""),
15861 StatusCode::BAD_REQUEST,
15862 "invalid_request_error",
15863 "model_not_found",
15864 ),
15865 (
15866 E::rate_limit("lane judge is at capacity, retry"),
15867 StatusCode::TOO_MANY_REQUESTS,
15868 "rate_limit_error",
15869 "rate_limit_exceeded",
15870 ),
15871 (
15872 E::overloaded("no VRAM for a new session"),
15873 StatusCode::SERVICE_UNAVAILABLE,
15874 "server_error",
15875 "overloaded",
15876 ),
15877 (
15878 E::engine("graph step failed: launch error"),
15879 StatusCode::INTERNAL_SERVER_ERROR,
15880 "server_error",
15881 "engine_error",
15882 ),
15883 ];
15884 for (err, want_status, want_type, want_code) in cases {
15885 let (status, etype, code) = class_http(err.class);
15886 assert_eq!(status, want_status, "{:?}", err);
15887 assert_eq!(etype, want_type, "{:?}", err);
15888 if !want_code.is_empty() {
15889 assert_eq!(code, Some(want_code), "{:?}", err);
15890 }
15891 let body = engine_error_body(&err);
15893 assert_eq!(body["error"]["message"], err.message);
15894 assert_eq!(body["error"]["type"], want_type);
15895 }
15896 for c in [
15898 C::InvalidRequest,
15899 C::ContextLength,
15900 C::ModelNotFound,
15901 C::RateLimit,
15902 C::Overloaded,
15903 C::Engine,
15904 ] {
15905 let (s, t, _) = class_http(c);
15906 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
15907 assert!(!t.is_empty());
15908 }
15909 }
15910
15911 #[test]
15912 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
15913 let e = worker::EngineError::engine(
15918 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
15919 );
15920 let resp = engine_error_response(&e);
15921 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
15922 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
15923 }
15924
15925 #[test]
15926 fn retry_headers_follow_the_sdk_contract() {
15927 for e in [
15931 worker::EngineError::rate_limit("shed"),
15932 worker::EngineError::overloaded("no VRAM"),
15933 ] {
15934 let resp = engine_error_response(&e);
15935 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
15936 let secs: u64 = ra
15937 .parse()
15938 .expect("Retry-After must be integer delay-seconds");
15939 assert!(
15940 secs > 0 && secs <= 60,
15941 "Retry-After {secs}s outside the honored window"
15942 );
15943 let ms = resp
15944 .headers()
15945 .get("retry-after-ms")
15946 .unwrap()
15947 .to_str()
15948 .unwrap();
15949 assert_eq!(
15950 ms.parse::<u64>().unwrap(),
15951 secs * 1000,
15952 "the two headers disagree"
15953 );
15954 assert!(
15955 resp.headers().get("x-should-retry").is_none(),
15956 "a retryable class must not say x-should-retry: false"
15957 );
15958 }
15959 }
15960
15961 #[tokio::test]
15968 async fn admit_predict_reject_matches_shed_contract() {
15969 let shed = retry_contract_response(
15971 (
15972 StatusCode::TOO_MANY_REQUESTS,
15973 Json(error_body(
15974 "interactive queue is at its bound",
15975 "rate_limit_error",
15976 None,
15977 Some("shed_queue"),
15978 )),
15979 )
15980 .into_response(),
15981 Some(7),
15982 );
15983 let predict = engine_error_response(&worker::EngineError::rate_limit_after(
15986 "predicted KV-to-completion exceeds the box budget; retry",
15987 7,
15988 ));
15989 assert_eq!(shed.status(), predict.status());
15990 for header in ["retry-after", "retry-after-ms"] {
15991 assert_eq!(
15992 shed.headers().get(header),
15993 predict.headers().get(header),
15994 "header {header} must be byte-identical to the shed contract"
15995 );
15996 }
15997 let shed_body: serde_json::Value = serde_json::from_slice(
15998 &axum::body::to_bytes(shed.into_body(), usize::MAX)
15999 .await
16000 .unwrap(),
16001 )
16002 .unwrap();
16003 let predict_body: serde_json::Value = serde_json::from_slice(
16004 &axum::body::to_bytes(predict.into_body(), usize::MAX)
16005 .await
16006 .unwrap(),
16007 )
16008 .unwrap();
16009 assert_eq!(shed_body["error"]["type"], predict_body["error"]["type"]);
16010 assert_eq!(predict_body["error"]["type"], "rate_limit_error");
16011 let shed_keys: Vec<&String> = shed_body["error"].as_object().unwrap().keys().collect();
16012 let predict_keys: Vec<&String> =
16013 predict_body["error"].as_object().unwrap().keys().collect();
16014 assert_eq!(shed_keys, predict_keys, "same body schema, key for key");
16015 assert_eq!(predict_body["error"]["code"], "rate_limit_exceeded");
16016
16017 let clamped = engine_error_response(&worker::EngineError::rate_limit_after("m", 400));
16019 assert_eq!(retry_after(&clamped).as_deref(), Some("60"));
16020 let plain = engine_error_response(&worker::EngineError::rate_limit("m"));
16022 assert_eq!(retry_after(&plain).as_deref(), Some("2"));
16023 }
16024
16025 #[tokio::test]
16026 #[allow(clippy::await_holding_lock)] async fn command_send_failure_obeys_the_retry_contract() {
16028 let _l = drain_lock();
16029 let mut st = fake_worker_state();
16030 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
16031 drop(cmd_rx);
16032 st.cmd_tx = cmd_tx;
16033
16034 let completion = completions(
16035 State(st.clone()),
16036 axum::http::HeaderMap::new(),
16037 None,
16038 Json(
16039 serde_json::from_value(serde_json::json!({
16040 "model": "m", "prompt": "test"
16041 }))
16042 .unwrap(),
16043 ),
16044 )
16045 .await;
16046 let chat = chat_completions(
16047 State(st),
16048 axum::http::HeaderMap::new(),
16049 None,
16050 Json(
16051 serde_json::from_value(serde_json::json!({
16052 "model": "m", "messages": [{"role": "user", "content": "test"}]
16053 }))
16054 .unwrap(),
16055 ),
16056 )
16057 .await;
16058
16059 for resp in [completion, chat] {
16060 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16061 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16062 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
16063 assert_ne!(
16064 resp.headers()
16065 .get("x-should-retry")
16066 .and_then(|v| v.to_str().ok()),
16067 Some("false")
16068 );
16069 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16070 .await
16071 .unwrap();
16072 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16073 assert_eq!(payload["error"]["type"], "server_error");
16074 assert_eq!(payload["error"]["code"], "overloaded");
16075 }
16076 }
16077
16078 #[test]
16079 fn unfixable_client_errors_say_x_should_retry_false() {
16080 for e in [
16083 worker::EngineError::model_not_found("unknown model \"x\""),
16084 worker::EngineError::context_length("prompt too long"),
16085 worker::EngineError::invalid_param("bad", "messages"),
16086 ] {
16087 let resp = engine_error_response(&e);
16088 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16089 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16090 assert!(
16091 retry_after(&resp).is_none(),
16092 "a 400 must not promise a retry window"
16093 );
16094 }
16095 }
16096
16097 #[tokio::test]
16098 async fn a_closed_worker_channel_is_503_not_500() {
16099 let (tx, rx) = worker::event_channel();
16103 drop(tx);
16104 let resp =
16105 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
16106 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16107 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
16108 }
16109
16110 #[tokio::test]
16111 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
16112 let (tx, rx) = worker::event_channel();
16115 tx.send(Event::Error(worker::EngineError::rate_limit(
16116 "lane judge shed: interactive p99 over budget, retry",
16117 )))
16118 .unwrap();
16119 let (resp, error_code) = peek_admission(rx)
16120 .await
16121 .expect_err("a shed must not be forwarded into the stream");
16122 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16123 assert_eq!(error_code, "rate_limit_exceeded");
16124 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16125 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16126 .await
16127 .unwrap();
16128 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16129 assert!(
16130 payload["error"].is_object(),
16131 "bare-string error body: {payload}"
16132 );
16133 assert_eq!(payload["error"]["type"], "rate_limit_error");
16134 assert!(
16135 payload["error"]["message"]
16136 .as_str()
16137 .unwrap()
16138 .contains("shed")
16139 );
16140 }
16141
16142 #[tokio::test]
16143 async fn interactive_admission_error_is_a_preheader_429() {
16144 let (tx, rx) = worker::event_channel();
16147 tx.send(Event::Error(worker::EngineError::rate_limit(
16148 "KV capacity unavailable",
16149 )))
16150 .unwrap();
16151 let (resp, error_code) = peek_admission(rx)
16152 .await
16153 .expect_err("admission error must stay pre-header");
16154 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16155 assert_eq!(error_code, "rate_limit_exceeded");
16156 }
16157
16158 #[tokio::test]
16159 async fn admission_peek_preserves_context_error_for_the_ledger() {
16160 let (tx, rx) = worker::event_channel();
16161 tx.send(Event::Error(worker::EngineError::context_length(
16162 "prompt exceeds configured model maximum",
16163 )))
16164 .unwrap();
16165 let (resp, error_code) = peek_admission(rx)
16166 .await
16167 .expect_err("context rejection must stay pre-header");
16168 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16169 assert_eq!(error_code, "context_length_exceeded");
16170 }
16171
16172 #[tokio::test]
16173 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
16174 let (tx, rx) = worker::event_channel();
16175 tx.send(Event::PromptUsage {
16176 n_prompt: 262_143,
16177 n_cached: 0,
16178 })
16179 .unwrap();
16180 let mut replay = peek_admission(rx).await.expect("successful admission");
16181 assert!(matches!(
16182 replay.recv().await,
16183 Some(Event::PromptUsage {
16184 n_prompt: 262_143,
16185 n_cached: 0
16186 }),
16187 ));
16188 }
16189
16190 #[test]
16191 fn penalties_plumb_from_http_to_sampler_config() {
16192 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
16195 "model": "m", "messages": [{"role": "user", "content": "task"}],
16196 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
16197 }))
16198 .unwrap();
16199 let (tx, _rx) = worker::event_channel();
16200 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16201 .unwrap()
16202 .request
16203 .sampler_cfg;
16204 assert_eq!(cfg.penalty_freq, 0.5);
16205 assert_eq!(cfg.penalty_present, 0.25);
16206 assert_eq!(cfg.penalty_repeat, 1.1);
16207 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16208
16209 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16210 "model": "m", "prompt": "task", "frequency_penalty": 1.5
16211 }))
16212 .unwrap();
16213 let (tx, _rx) = worker::event_channel();
16214 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16215 assert_eq!(cfg.penalty_freq, 1.5);
16216 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16217
16218 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16220 "model": "m", "prompt": "task"
16221 }))
16222 .unwrap();
16223 let (tx, _rx) = worker::event_channel();
16224 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16225 assert_eq!(cfg.penalty_last_n, 0);
16226 assert_eq!(cfg.penalty_repeat, 1.0);
16227 }
16228
16229 #[test]
16230 fn omitted_temperature_is_openai_default_not_greedy() {
16231 let chat_temp = |body: serde_json::Value| {
16247 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16248 let (tx, _rx) = worker::event_channel();
16249 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16250 .unwrap()
16251 .request
16252 .sampler_cfg
16253 .temperature
16254 };
16255 let comp_temp = |body: serde_json::Value| {
16256 let req: CompletionReq = serde_json::from_value(body).unwrap();
16257 let (tx, _rx) = worker::event_channel();
16258 build_request(&req, tx, lanes::Lane::Interactive, None)
16259 .sampler_cfg
16260 .temperature
16261 };
16262
16263 assert_eq!(
16265 chat_temp(serde_json::json!({
16266 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
16267 1.0,
16268 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
16269 );
16270 assert_eq!(
16271 comp_temp(serde_json::json!({
16272 "model": "m", "prompt": "t"})),
16273 1.0,
16274 "omitted completions temperature must be the OpenAI 1.0 default"
16275 );
16276
16277 assert_eq!(
16279 chat_temp(serde_json::json!({
16280 "model": "m", "messages": [{"role": "user", "content": "t"}],
16281 "temperature": 0.0})),
16282 0.0,
16283 "explicit temperature 0 must stay greedy"
16284 );
16285 assert_eq!(
16286 comp_temp(serde_json::json!({
16287 "model": "m", "prompt": "t", "temperature": 0})),
16288 0.0,
16289 "explicit temperature 0 must stay greedy"
16290 );
16291 assert!(
16293 memra_engine::sampler::Sampler::new(sampler_config(
16294 0.0,
16295 0,
16296 1.0,
16297 0.0,
16298 0.0,
16299 0.0,
16300 1.0,
16301 Some(0)
16302 ))
16303 .is_greedy()
16304 );
16305 assert!(
16306 !memra_engine::sampler::Sampler::new(sampler_config(
16307 1.0,
16308 0,
16309 1.0,
16310 0.0,
16311 0.0,
16312 0.0,
16313 1.0,
16314 Some(0)
16315 ))
16316 .is_greedy()
16317 );
16318
16319 assert_eq!(
16321 chat_temp(serde_json::json!({
16322 "model": "m", "messages": [{"role": "user", "content": "t"}],
16323 "temperature": 0.7})),
16324 0.7
16325 );
16326
16327 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16331 "model": "m", "prompt": "t"}))
16332 .unwrap();
16333 let (tx, _rx) = worker::event_channel();
16334 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16335 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
16336 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
16337 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
16338 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
16339 assert!(
16344 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
16345 "the omitted-temperature default must ride sampled spec's pure-temp regime"
16346 );
16347 }
16348
16349 #[test]
16350 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
16351 let caps = ModelCaps {
16352 chat_temperature_default: Some(0.5),
16353 chat_top_p_default: Some(0.9),
16354 chat_ok: true,
16355 ..Default::default()
16356 };
16357 let cfg = |extra: serde_json::Value| {
16358 let mut body = serde_json::json!({
16359 "model": "step35",
16360 "messages": [{"role": "user", "content": "task"}]
16361 });
16362 body.as_object_mut()
16363 .unwrap()
16364 .extend(extra.as_object().unwrap().clone());
16365 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16366 let (tx, _rx) = worker::event_channel();
16367 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
16368 .unwrap()
16369 .request
16370 .sampler_cfg
16371 };
16372
16373 let omitted = cfg(serde_json::json!({}));
16374 assert_eq!(omitted.temperature, 0.5);
16375 assert_eq!(omitted.top_p, 0.9);
16376
16377 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
16378 assert_eq!(explicit_temp.temperature, 0.7);
16379 assert_eq!(
16380 explicit_temp.top_p, 0.9,
16381 "omitting top_p must retain StepFun's nucleus default"
16382 );
16383
16384 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
16385 assert_eq!(
16386 explicit.temperature, 0.0,
16387 "explicit greedy must remain authoritative"
16388 );
16389 assert_eq!(
16390 explicit.top_p, 1.0,
16391 "explicit untruncated sampling must remain authoritative"
16392 );
16393 }
16394
16395 fn qwen38_vendor_defaults() -> SamplingDefaults {
16399 SamplingDefaults {
16400 temperature: Some(1.0),
16401 top_p: Some(0.95),
16402 top_k: Some(20),
16403 min_p: Some(0.0),
16404 presence_penalty: Some(0.0),
16405 repetition_penalty: Some(1.0),
16406 frequency_penalty: None,
16407 }
16408 }
16409
16410 fn gemma4_vendor_defaults() -> SamplingDefaults {
16415 SamplingDefaults {
16416 temperature: Some(1.0),
16417 top_p: Some(0.95),
16418 top_k: Some(64),
16419 ..Default::default()
16420 }
16421 }
16422
16423 #[test]
16424 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
16425 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
16430 let chat = |extra: serde_json::Value| {
16431 let mut body = serde_json::json!({
16432 "model": "google/gemma-4-31b-it",
16433 "messages": [{"role": "user", "content": "task"}],
16434 "seed": 7
16436 });
16437 body.as_object_mut()
16438 .unwrap()
16439 .extend(extra.as_object().unwrap().clone());
16440 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16441 let (tx, _rx) = worker::event_channel();
16442 build_chat_request_with_trace(
16443 req,
16444 Some(&ModelCaps {
16445 chat_ok: true,
16446 ..Default::default()
16447 }),
16448 tx,
16449 lanes::Lane::Interactive,
16450 None,
16451 None,
16452 None,
16453 &d,
16454 )
16455 .unwrap()
16456 .request
16457 .sampler_cfg
16458 };
16459
16460 let omitted = chat(serde_json::json!({}));
16462 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
16463 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
16464 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
16465 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
16467 assert_eq!(omitted.penalty_repeat, 1.0);
16468 assert_eq!(omitted.penalty_freq, 0.0);
16469 assert_eq!(omitted.penalty_present, 0.0);
16470 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
16471 assert!(
16472 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
16473 "the vendor default must NOT be greedy — that is the whole point of the lane"
16474 );
16475
16476 let greedy = chat(serde_json::json!({"temperature": 0}));
16479 assert_eq!(
16480 greedy.temperature, 0.0,
16481 "explicit temperature 0 stays greedy"
16482 );
16483 assert!(
16484 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
16485 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
16486 spec/graph exactness arms"
16487 );
16488
16489 let one_field = chat(serde_json::json!({"top_k": 3}));
16491 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
16492 assert_eq!(
16493 one_field.temperature, 1.0,
16494 "omitting temperature still takes the vendor value"
16495 );
16496 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
16497
16498 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
16501 assert_eq!(
16502 disabled.top_k, 0,
16503 "an explicit top_k 0 means KEEP ALL, not 'unset'"
16504 );
16505 assert_eq!(
16506 disabled.top_p, 1.0,
16507 "an explicit top_p 1.0 means untruncated"
16508 );
16509
16510 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
16512 assert_eq!(penal.penalty_present, 1.5);
16513 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16514 }
16515
16516 #[test]
16517 fn vendor_sampling_defaults_are_identical_on_every_surface() {
16518 let d = qwen38_vendor_defaults();
16530 let md = ModelSamplingDefaults::single(d);
16531 let comp = |extra: serde_json::Value| {
16532 let mut body = serde_json::json!({
16533 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
16534 body.as_object_mut()
16535 .unwrap()
16536 .extend(extra.as_object().unwrap().clone());
16537 let req: CompletionReq = serde_json::from_value(body).unwrap();
16538 let (tx, _rx) = worker::event_channel();
16539 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
16540 };
16541 let chat = |extra: serde_json::Value| {
16542 let mut body = serde_json::json!({
16543 "model": "qwen/qwen3.8-27b",
16544 "messages": [{"role": "user", "content": "task"}],
16545 "seed": 11 });
16546 body.as_object_mut()
16547 .unwrap()
16548 .extend(extra.as_object().unwrap().clone());
16549 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16550 let (tx, _rx) = worker::event_channel();
16551 build_chat_request_with_trace(
16552 req,
16553 Some(&ModelCaps {
16554 chat_ok: true,
16555 ..Default::default()
16556 }),
16557 tx,
16558 lanes::Lane::Interactive,
16559 None,
16560 None,
16561 None,
16562 &md,
16563 )
16564 .unwrap()
16565 .request
16566 .sampler_cfg
16567 };
16568
16569 for extra in [
16570 serde_json::json!({}),
16571 serde_json::json!({"temperature": 0}),
16572 serde_json::json!({"temperature": 0.0}),
16573 serde_json::json!({"temperature": 0.7}),
16574 serde_json::json!({"top_p": 1.0}),
16575 serde_json::json!({"top_k": 0}),
16576 serde_json::json!({"min_p": 0.05}),
16577 serde_json::json!({"repetition_penalty": 1.1}),
16578 serde_json::json!({"frequency_penalty": 0.5}),
16579 serde_json::json!({"presence_penalty": 1.5}),
16580 serde_json::json!({
16581 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
16582 "frequency_penalty": 0.1, "presence_penalty": 0.2,
16583 "repetition_penalty": 1.05 }),
16584 ] {
16585 let c = comp(extra.clone());
16586 let h = chat(extra.clone());
16587 assert_eq!(
16588 (
16589 c.temperature,
16590 c.top_p,
16591 c.top_k,
16592 c.min_p,
16593 c.penalty_repeat,
16594 c.penalty_freq,
16595 c.penalty_present,
16596 c.penalty_last_n,
16597 c.seed
16598 ),
16599 (
16600 h.temperature,
16601 h.top_p,
16602 h.top_k,
16603 h.min_p,
16604 h.penalty_repeat,
16605 h.penalty_freq,
16606 h.penalty_present,
16607 h.penalty_last_n,
16608 h.seed
16609 ),
16610 "/v1/completions and /v1/chat/completions disagree on {extra} — \
16611 standard-surface-law violation"
16612 );
16613 }
16614
16615 let omitted = comp(serde_json::json!({}));
16617 assert_eq!(
16618 omitted.temperature, 1.0,
16619 "qwen3.8 card thinking temperature"
16620 );
16621 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
16622 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
16623 assert!(
16625 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
16626 .is_greedy()
16627 );
16628 }
16629
16630 #[tokio::test]
16643 #[allow(clippy::await_holding_lock)] async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
16645 let _l = drain_lock();
16646 let step_caps = ModelCaps {
16647 chat_ok: true,
16648 chat_temperature_default: Some(0.5),
16649 chat_top_p_default: Some(0.9),
16650 ..Default::default()
16651 };
16652 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
16653 let st = fake_worker_state_full(
16654 1,
16655 std::time::Duration::ZERO,
16656 HashMap::from([("m".to_string(), step_caps)]),
16657 Some(cfg_tx),
16658 );
16659 let fields = |saw: &WorkerSaw| {
16664 let c = &saw.sampler_cfg;
16665 (
16666 c.temperature,
16667 c.top_p,
16668 c.top_k,
16669 c.min_p,
16670 c.penalty_repeat,
16671 c.penalty_freq,
16672 c.penalty_present,
16673 c.penalty_last_n,
16674 )
16675 };
16676 let worker_saw = |surface: &str| {
16677 cfg_rx
16678 .recv_timeout(std::time::Duration::from_secs(10))
16679 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
16680 };
16681
16682 let resp = completions(
16683 State(st.clone()),
16684 axum::http::HeaderMap::new(),
16685 None,
16686 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
16687 )
16688 .await;
16689 assert_eq!(
16690 resp.status(),
16691 StatusCode::OK,
16692 "/v1/completions rejected the omitted-sampling request"
16693 );
16694 let comp = worker_saw("/v1/completions");
16695
16696 let resp = chat_completions(
16697 State(st.clone()),
16698 axum::http::HeaderMap::new(),
16699 None,
16700 Json(
16701 serde_json::from_value(serde_json::json!({
16702 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
16703 .unwrap(),
16704 ),
16705 )
16706 .await;
16707 assert_eq!(
16708 resp.status(),
16709 StatusCode::OK,
16710 "/v1/chat/completions rejected the omitted-sampling request"
16711 );
16712 let chat = worker_saw("/v1/chat/completions");
16713
16714 let resp = anthropic::messages(
16715 State(st.clone()),
16716 axum::http::HeaderMap::new(),
16717 None,
16718 axum::body::Bytes::from(
16719 serde_json::json!({
16720 "model": "m", "max_tokens": 16,
16721 "messages": [{"role": "user", "content": "t"}]})
16722 .to_string(),
16723 ),
16724 )
16725 .await;
16726 assert_eq!(
16727 resp.status(),
16728 StatusCode::OK,
16729 "/v1/messages rejected the omitted-sampling request"
16730 );
16731 let msg = worker_saw("/v1/messages");
16732
16733 let resp = responses_api::responses(
16734 State(st.clone()),
16735 axum::http::HeaderMap::new(),
16736 None,
16737 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
16738 )
16739 .await;
16740 assert_eq!(
16741 resp.status(),
16742 StatusCode::OK,
16743 "/v1/responses rejected the omitted-sampling request"
16744 );
16745 let rsp = worker_saw("/v1/responses");
16746
16747 for (surface, cfg) in [
16748 ("/v1/completions", &comp),
16749 ("/v1/messages", &msg),
16750 ("/v1/responses", &rsp),
16751 ] {
16752 assert_eq!(
16753 fields(cfg),
16754 fields(&chat),
16755 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
16756 for the same omitted-sampling request — standard-surface-law violation \
16757 (hermes d991b51699218285)"
16758 );
16759 }
16760 assert_eq!(
16763 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
16764 (0.5, 0.9),
16765 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
16766 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
16767 );
16768 }
16769
16770 #[tokio::test]
16780 #[allow(clippy::await_holding_lock)] async fn same_effort_value_resolves_identically_on_every_surface() {
16782 let _l = drain_lock();
16783 let caps = ModelCaps {
16786 chat_ok: true,
16787 effort_levels: true,
16788 ..Default::default()
16789 };
16790 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
16791 let st = fake_worker_state_full(
16792 1,
16793 std::time::Duration::ZERO,
16794 HashMap::from([("m".to_string(), caps)]),
16795 Some(saw_tx),
16796 );
16797 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
16798 match surface {
16799 "/v1/chat/completions" => {
16800 chat_completions(
16801 State(st),
16802 axum::http::HeaderMap::new(),
16803 None,
16804 Json(
16805 serde_json::from_value(serde_json::json!({
16806 "model": "m", "max_tokens": 8,
16807 "reasoning_effort": effort,
16808 "messages": [{"role": "user", "content": "t"}]}))
16809 .unwrap(),
16810 ),
16811 )
16812 .await
16813 }
16814 "/v1/responses" => {
16815 responses_api::responses(
16816 State(st),
16817 axum::http::HeaderMap::new(),
16818 None,
16819 axum::body::Bytes::from(
16820 serde_json::json!({
16821 "model": "m", "max_output_tokens": 8, "input": "t",
16822 "reasoning": {"effort": effort}})
16823 .to_string(),
16824 ),
16825 )
16826 .await
16827 }
16828 "/v1/messages" => {
16829 anthropic::messages(
16830 State(st),
16831 axum::http::HeaderMap::new(),
16832 None,
16833 axum::body::Bytes::from(
16834 serde_json::json!({
16835 "model": "m", "max_tokens": 8,
16836 "messages": [{"role": "user", "content": "t"}],
16837 "output_config": {"effort": effort}})
16838 .to_string(),
16839 ),
16840 )
16841 .await
16842 }
16843 other => panic!("unknown surface {other}"),
16844 }
16845 };
16846 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
16847
16848 for (effort, want_think, want_level) in [
16851 ("none", ThinkMode::NoThink, Some("low")),
16852 ("minimal", ThinkMode::NoThink, Some("low")),
16853 ("low", ThinkMode::Think, Some("low")),
16854 ("medium", ThinkMode::Think, Some("medium")),
16855 ("high", ThinkMode::Think, Some("high")),
16856 ("xhigh", ThinkMode::Think, Some("high")),
16858 ] {
16859 for surface in SURFACES {
16860 let resp = send(st.clone(), surface, effort).await;
16861 assert_eq!(
16862 resp.status(),
16863 StatusCode::OK,
16864 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
16865 diverged again (issue #31)"
16866 );
16867 let saw = saw_rx
16868 .recv_timeout(std::time::Duration::from_secs(10))
16869 .unwrap_or_else(|_| {
16870 panic!("{surface}: effort {effort:?} request never reached the worker")
16871 });
16872 assert_eq!(
16873 (saw.think, saw.reasoning_effort.as_deref()),
16874 (want_think, want_level),
16875 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
16876 reasoning surface — the parameter was dropped or remapped before \
16877 parse_think (issue #31 regression)"
16878 );
16879 }
16880 }
16881
16882 for effort in ["bogus", "banana", ""] {
16885 for surface in SURFACES {
16886 let resp = send(st.clone(), surface, effort).await;
16887 assert_eq!(
16888 resp.status(),
16889 StatusCode::BAD_REQUEST,
16890 "{surface} accepted effort {effort:?} — silent-accept regression \
16891 (issue #31: the value never reached parse_think's allowlist)"
16892 );
16893 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
16895 .await
16896 .unwrap();
16897 let v: serde_json::Value = serde_json::from_slice(&body)
16898 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
16899 match surface {
16900 "/v1/messages" => {
16901 assert_eq!(v["type"], "error", "{surface} error envelope");
16902 assert_eq!(
16903 v["error"]["type"], "invalid_request_error",
16904 "{surface} error type"
16905 );
16906 }
16907 _ => {
16908 assert!(
16909 v["error"]["message"].is_string(),
16910 "{surface} OpenAI-shaped error body: {v}"
16911 );
16912 }
16913 }
16914 }
16915 }
16916
16917 let resp = anthropic::messages(
16921 State(st.clone()),
16922 axum::http::HeaderMap::new(),
16923 None,
16924 axum::body::Bytes::from(
16925 serde_json::json!({
16926 "model": "m", "max_tokens": 8,
16927 "messages": [{"role": "user", "content": "t"}],
16928 "thinking": {"type": "enabled"},
16929 "output_config": {"effort": "none"}})
16930 .to_string(),
16931 ),
16932 )
16933 .await;
16934 assert_eq!(resp.status(), StatusCode::OK);
16935 let saw = saw_rx
16936 .recv_timeout(std::time::Duration::from_secs(10))
16937 .expect("thinking+effort request never reached the worker");
16938 assert_eq!(
16939 saw.think,
16940 ThinkMode::Think,
16941 "thinking.type (the documented Anthropic lever) must win the switch over \
16942 output_config.effort"
16943 );
16944 let resp = anthropic::messages(
16945 State(st.clone()),
16946 axum::http::HeaderMap::new(),
16947 None,
16948 axum::body::Bytes::from(
16949 serde_json::json!({
16950 "model": "m", "max_tokens": 8,
16951 "messages": [{"role": "user", "content": "t"}],
16952 "thinking": {"type": "enabled"},
16953 "output_config": {"effort": "banana"}})
16954 .to_string(),
16955 ),
16956 )
16957 .await;
16958 assert_eq!(
16959 resp.status(),
16960 StatusCode::BAD_REQUEST,
16961 "an invalid effort must 400 even next to an explicit thinking.type — \
16962 precedence must not re-open the silent-accept hole"
16963 );
16964 }
16965
16966 #[test]
16967 fn vendor_sampling_defaults_are_boot_validated() {
16968 let parsed = OpenRouterMetadataFile::from_toml(
16971 r#"
16972[models.g]
16973default_temperature = 1.0
16974default_top_p = 0.95
16975default_top_k = 64
16976default_min_p = 0.0
16977default_presence_penalty = 0.0
16978default_frequency_penalty = 0.0
16979default_repetition_penalty = 1.0
16980"#,
16981 )
16982 .unwrap();
16983 let g = parsed.get("g").unwrap();
16984 assert_eq!(g.default_temperature, Some(1.0));
16985 assert_eq!(g.default_top_p, Some(0.95));
16986 assert_eq!(g.default_top_k, Some(64));
16987
16988 let err = OpenRouterMetadataFile::from_toml(
16992 r#"
16993[models.g]
16994default_temperature = 0.0
16995"#,
16996 )
16997 .unwrap_err();
16998 assert!(err.contains("default_temperature"), "{err}");
16999 assert!(
17000 err.contains("greedy"),
17001 "the refusal must say WHY a zero default is refused: {err}"
17002 );
17003
17004 for bad in [
17005 "default_temperature = 2.5",
17006 "default_temperature = -1.0",
17007 "default_top_p = 0.0",
17008 "default_top_p = 1.5",
17009 "default_min_p = 1.0",
17010 "default_min_p = -0.1",
17011 "default_presence_penalty = 3.0",
17012 "default_frequency_penalty = -2.5",
17013 "default_repetition_penalty = 0.0",
17014 ] {
17015 let err =
17016 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
17017 let key = bad.split(' ').next().unwrap();
17018 assert!(err.contains(key), "{bad} must be refused by name: {err}");
17019 }
17020
17021 let err = OpenRouterMetadataFile::from_toml(
17025 r#"
17026[models.g]
17027default_temperture = 1.0
17028"#,
17029 )
17030 .unwrap_err();
17031 assert!(
17032 err.contains("unknown field"),
17033 "an unknown key must be fatal, which is what makes binary-first ordering \
17034 mandatory: {err}"
17035 );
17036 }
17037
17038 #[test]
17039 fn non_thinking_sampling_arm_is_boot_validated() {
17040 let parsed = OpenRouterMetadataFile::from_toml(
17044 r#"
17045[models.q]
17046default_temperature = 1.0
17047default_top_p = 0.95
17048default_top_k = 20
17049
17050[models.q.non_thinking_sampling]
17051temperature = 0.7
17052top_p = 0.8
17053top_k = 20
17054presence_penalty = 1.5
17055"#,
17056 )
17057 .unwrap();
17058 let arm = parsed
17059 .get("q")
17060 .unwrap()
17061 .non_thinking_sampling
17062 .as_ref()
17063 .unwrap();
17064 assert_eq!(arm.temperature, Some(0.7));
17065 assert_eq!(arm.top_p, Some(0.8));
17066 assert_eq!(arm.top_k, Some(20));
17067 assert_eq!(arm.presence_penalty, Some(1.5));
17068 assert_eq!(
17069 arm.min_p, None,
17070 "undeclared arm fields stay undeclared, never invented"
17071 );
17072
17073 let err = OpenRouterMetadataFile::from_toml(
17077 r#"
17078[models.q]
17079[models.q.non_thinking_sampling]
17080temperature = 0.0
17081"#,
17082 )
17083 .unwrap_err();
17084 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
17085 assert!(err.contains("greedy"), "{err}");
17086
17087 let err = OpenRouterMetadataFile::from_toml(
17090 r#"
17091[models.q]
17092[models.q.non_thinking_sampling]
17093"#,
17094 )
17095 .unwrap_err();
17096 assert!(err.contains("non_thinking_sampling"), "{err}");
17097 assert!(err.contains("declare"), "{err}");
17098
17099 for bad in [
17101 "temperature = 2.5",
17102 "top_p = 0.0",
17103 "top_p = 1.5",
17104 "min_p = 1.0",
17105 "presence_penalty = 3.0",
17106 "frequency_penalty = -2.5",
17107 "repetition_penalty = 0.0",
17108 ] {
17109 let err = OpenRouterMetadataFile::from_toml(&format!(
17110 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
17111 ))
17112 .unwrap_err();
17113 let key = bad.split(' ').next().unwrap();
17114 assert!(
17115 err.contains(&format!("non_thinking_sampling.{key}")),
17116 "the refusal for {bad:?} must name the nested key: {err}"
17117 );
17118 }
17119
17120 let err = OpenRouterMetadataFile::from_toml(
17124 r#"
17125[models.q]
17126[models.q.non_thinking_sampling]
17127temperture = 0.7
17128"#,
17129 )
17130 .unwrap_err();
17131 assert!(err.contains("unknown field"), "{err}");
17132 }
17133
17134 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
17139 SamplingDefaults {
17140 temperature: Some(0.7),
17141 top_p: Some(0.8),
17142 top_k: Some(20),
17143 presence_penalty: Some(1.5),
17144 ..Default::default()
17145 }
17146 }
17147
17148 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
17149 ModelSamplingDefaults {
17150 thinking: qwen38_vendor_defaults(),
17151 non_thinking: Some(qwen38_non_thinking_defaults()),
17152 }
17153 }
17154
17155 fn qwen38_caps() -> ModelCaps {
17159 ModelCaps {
17160 chat_ok: true,
17161 qwen_think: true,
17162 think_switch: true,
17163 ..Default::default()
17164 }
17165 }
17166
17167 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
17170 (
17171 c.temperature,
17172 c.top_p,
17173 c.top_k,
17174 c.min_p,
17175 c.penalty_present,
17176 c.penalty_freq,
17177 c.penalty_repeat,
17178 c.penalty_last_n,
17179 c.seed,
17180 )
17181 }
17182
17183 fn build_with_arms(
17184 defaults: &ModelSamplingDefaults,
17185 caps: &ModelCaps,
17186 default_effort: Option<&str>,
17187 extra: serde_json::Value,
17188 ) -> Request {
17189 let mut body = serde_json::json!({
17190 "model": "m",
17191 "messages": [{"role": "user", "content": "task"}],
17192 "seed": 3
17194 });
17195 body.as_object_mut()
17196 .unwrap()
17197 .extend(extra.as_object().unwrap().clone());
17198 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17199 let (tx, _rx) = worker::event_channel();
17200 build_chat_request_with_trace(
17201 req,
17202 Some(caps),
17203 tx,
17204 lanes::Lane::Interactive,
17205 None,
17206 None,
17207 default_effort,
17208 defaults,
17209 )
17210 .unwrap()
17211 .request
17212 }
17213
17214 #[test]
17215 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
17216 let two_arm = qwen38_two_arm_defaults();
17221 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
17222 let caps = qwen38_caps();
17223
17224 let off_spellings = [
17226 serde_json::json!({"reasoning_effort": "none"}),
17227 serde_json::json!({"enable_thinking": false}),
17228 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
17229 serde_json::json!({"reasoning": {"enabled": false}}),
17230 ];
17231 for extra in &off_spellings {
17232 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
17233 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
17234 let c = &r.sampler_cfg;
17235 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
17236 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
17237 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
17238 assert_eq!(
17239 c.penalty_present, 1.5,
17240 "{extra}: non-thinking presence_penalty"
17241 );
17242 assert_eq!(
17243 c.penalty_last_n,
17244 memra_engine::spec::PEN_WINDOW_MAX,
17245 "{extra}: the arm's presence penalty uses the cross-path history window"
17246 );
17247 assert_eq!(
17248 c.min_p, 0.0,
17249 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
17250 );
17251
17252 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
17255 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
17256 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
17257 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
17258 assert_eq!(
17259 s.sampler_cfg.penalty_present, 0.0,
17260 "{extra}: single-arm model"
17261 );
17262 }
17263
17264 for extra in [
17267 serde_json::json!({}),
17268 serde_json::json!({"enable_thinking": true}),
17269 serde_json::json!({"reasoning_effort": "high"}),
17270 serde_json::json!({"reasoning": {"enabled": true}}),
17271 ] {
17272 for defaults in [&two_arm, &single_arm] {
17273 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
17274 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
17275 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
17276 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
17277 assert_eq!(
17278 c.penalty_present, 0.0,
17279 "{extra}: thinking arm has no presence"
17280 );
17281 }
17282 }
17283
17284 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
17287 assert_eq!(
17288 c.temperature, 0.7,
17289 "deployment-default off = non-thinking arm"
17290 );
17291 let c = build_with_arms(
17293 &two_arm,
17294 &caps,
17295 Some("none"),
17296 serde_json::json!({"enable_thinking": true}),
17297 )
17298 .sampler_cfg;
17299 assert_eq!(
17300 c.temperature, 1.0,
17301 "explicit ON beats the deployment default"
17302 );
17303
17304 let c = build_with_arms(
17306 &two_arm,
17307 &caps,
17308 None,
17309 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
17310 )
17311 .sampler_cfg;
17312 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
17313 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
17314 let c = build_with_arms(
17315 &two_arm,
17316 &caps,
17317 None,
17318 serde_json::json!({
17319 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
17320 )
17321 .sampler_cfg;
17322 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
17323 assert_eq!(
17324 c.penalty_present, 0.0,
17325 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
17326 is a value, not an absence"
17327 );
17328 assert_eq!(
17329 c.penalty_last_n, 0,
17330 "all penalties off => no history window"
17331 );
17332 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
17333
17334 let c = build_with_arms(
17337 &two_arm,
17338 &caps,
17339 None,
17340 serde_json::json!({"enable_thinking": false, "temperature": 0}),
17341 )
17342 .sampler_cfg;
17343 assert!(
17344 memra_engine::sampler::Sampler::new(c).is_greedy(),
17345 "explicit temperature 0 must stay greedy on the non-thinking arm"
17346 );
17347
17348 let c = build_with_arms(
17351 &single_arm,
17352 &caps,
17353 None,
17354 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
17355 )
17356 .sampler_cfg;
17357 assert_eq!(c.temperature, 0.55);
17358 assert_eq!(
17359 c.top_p, 0.95,
17360 "single-arm model: unset top_p takes its one arm"
17361 );
17362 }
17363
17364 #[test]
17365 fn sampling_arms_never_blend_field_by_field() {
17366 let parsed = OpenRouterMetadataFile::from_toml(
17371 r#"
17372[models.m]
17373default_temperature = 1.0
17374default_min_p = 0.05
17375
17376[models.m.non_thinking_sampling]
17377temperature = 0.6
17378"#,
17379 )
17380 .unwrap();
17381 let caps = ModelCaps {
17382 chat_temperature_default: Some(0.5),
17383 chat_top_p_default: Some(0.9),
17384 ..Default::default()
17385 };
17386 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
17387 let client = ClientSampling {
17388 seed: Some(1),
17389 ..Default::default()
17390 };
17391
17392 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
17393 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
17394 assert_eq!(
17395 off.min_p, 0.0,
17396 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
17397 );
17398 assert_eq!(
17399 off.top_p, 1.0,
17400 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
17401 );
17402
17403 for mode in [ThinkMode::Default, ThinkMode::Think] {
17405 let on = resolve_sampler_config(client, d.for_mode(mode));
17406 assert_eq!(on.temperature, 1.0);
17407 assert_eq!(on.min_p, 0.05);
17408 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
17409 }
17410 }
17411
17412 #[test]
17413 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
17414 let caps = qwen38_caps();
17423 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
17424 let two_arm = qwen38_two_arm_defaults();
17425
17426 let bodies = [
17427 serde_json::json!({}),
17428 serde_json::json!({"enable_thinking": true}),
17429 serde_json::json!({"reasoning_effort": "high"}),
17430 serde_json::json!({"reasoning_effort": "none"}),
17431 serde_json::json!({"enable_thinking": false}),
17432 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
17433 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
17434 serde_json::json!({"enable_thinking": false, "temperature": 0}),
17435 ];
17436 for extra in &bodies {
17437 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
17439 let mut client = ClientSampling {
17440 seed: Some(3),
17441 ..Default::default()
17442 };
17443 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
17444 client.temperature = Some(t as f32);
17445 }
17446 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
17447 client.top_p = Some(p as f32);
17448 }
17449 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
17450 assert_eq!(
17451 sampler_key(&r.sampler_cfg),
17452 sampler_key(&pre_arm),
17453 "{extra}: single-arm model diverged from the pre-arm resolution law"
17454 );
17455
17456 if r.think != ThinkMode::NoThink {
17459 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
17460 assert_eq!(t.think, r.think, "{extra}");
17461 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
17462 assert_eq!(
17463 sampler_key(&t.sampler_cfg),
17464 sampler_key(&r.sampler_cfg),
17465 "{extra}: a thinking-on request must not feel the non-thinking arm"
17466 );
17467 }
17468 }
17469 }
17470
17471 #[test]
17472 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
17473 let r = build_with_arms(
17479 &qwen38_two_arm_defaults(),
17480 &qwen38_caps(),
17481 None,
17482 serde_json::json!({"response_format": {"type": "json_object"}}),
17483 );
17484 assert_eq!(
17485 r.think,
17486 ThinkMode::NoThink,
17487 "constraint forces the switch off"
17488 );
17489 assert_eq!(
17490 r.sampler_cfg.temperature, 0.7,
17491 "and the arm follows the real mode"
17492 );
17493 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
17494 }
17495
17496 #[test]
17497 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
17498 let caps = ModelCaps {
17503 chat_temperature_default: Some(0.5),
17504 chat_top_p_default: Some(0.9),
17505 chat_ok: true,
17506 ..Default::default()
17507 };
17508 let metadata = OpenRouterModelMetadata {
17509 default_temperature: Some(1.0),
17510 default_top_p: Some(0.95),
17511 default_top_k: Some(64),
17512 ..Default::default()
17513 };
17514
17515 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
17516 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
17517 assert_eq!(caps_only.top_p, Some(0.9));
17518 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
17519
17520 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
17521 assert_eq!(
17522 both.temperature,
17523 Some(1.0),
17524 "metadata outranks the arch cap"
17525 );
17526 assert_eq!(both.top_p, Some(0.95));
17527 assert_eq!(both.top_k, Some(64));
17528
17529 let partial = SamplingDefaults::resolve(
17531 Some(&OpenRouterModelMetadata {
17532 default_temperature: Some(0.7),
17533 ..Default::default()
17534 }),
17535 Some(&caps),
17536 );
17537 assert_eq!(partial.temperature, Some(0.7));
17538 assert_eq!(
17539 partial.top_p,
17540 Some(0.9),
17541 "an undeclared metadata field must fall through to the cap, not to 1.0"
17542 );
17543
17544 assert_eq!(
17546 SamplingDefaults::resolve(None, None),
17547 SamplingDefaults::default()
17548 );
17549 }
17550
17551 #[test]
17552 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
17553 let resolved = |d: &SamplingDefaults| {
17566 resolve_sampler_config(
17567 ClientSampling {
17568 seed: Some(1),
17569 ..Default::default()
17570 },
17571 d,
17572 )
17573 };
17574
17575 assert!(
17577 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
17578 .is_spec_sampling(),
17579 "the API-standard default must stay in the fast pure-temp regime"
17580 );
17581
17582 for (name, d) in [
17583 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
17584 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
17585 ] {
17586 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
17587 assert!(
17588 !sampler.is_greedy(),
17589 "{name}: vendor default must not be greedy"
17590 );
17591 assert!(
17592 !sampler.is_spec_sampling(),
17593 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
17594 starts passing, either the vendor numbers changed or the in-graph draft \
17595 learned filters, and the perf note in docs/SERVING.md needs revisiting"
17596 );
17597 }
17598
17599 let opted_out = resolve_sampler_config(
17601 ClientSampling {
17602 top_p: Some(1.0),
17603 top_k: Some(0),
17604 seed: Some(1),
17605 ..Default::default()
17606 },
17607 &qwen38_vendor_defaults(),
17608 );
17609 assert!(
17610 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
17611 "explicitly disabling the filters must restore the pure-temp regime"
17612 );
17613 }
17614
17615 #[test]
17616 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
17617 let comp_seed = |body: serde_json::Value| {
17624 let req: CompletionReq = serde_json::from_value(body).unwrap();
17625 let (tx, _rx) = worker::event_channel();
17626 build_request(&req, tx, lanes::Lane::Interactive, None)
17627 .sampler_cfg
17628 .seed
17629 };
17630 let chat_seed = |body: serde_json::Value| {
17631 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17632 let (tx, _rx) = worker::event_channel();
17633 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17634 .unwrap()
17635 .request
17636 .sampler_cfg
17637 .seed
17638 };
17639
17640 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
17643 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
17644 let c = chat_seed(serde_json::json!({
17645 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
17646 assert_ne!(
17647 a, 0,
17648 "omitted seed must not be the pinned 0 that caused the loop"
17649 );
17650 assert_ne!(b, 0);
17651 assert_ne!(c, 0);
17652 assert_ne!(
17653 a, b,
17654 "two seed-omitting requests must get DIFFERENT streams"
17655 );
17656 assert_ne!(a, c);
17657
17658 assert_eq!(
17661 comp_seed(serde_json::json!({
17662 "model": "m", "prompt": "t", "seed": 0})),
17663 0,
17664 "explicit seed 0 must stay 0 — the determinism gates depend on it"
17665 );
17666 assert_eq!(
17667 comp_seed(serde_json::json!({
17668 "model": "m", "prompt": "t", "seed": 12345})),
17669 12345
17670 );
17671 assert_eq!(
17672 chat_seed(serde_json::json!({
17673 "model": "m", "messages": [{"role": "user", "content": "t"}],
17674 "seed": 777})),
17675 777
17676 );
17677 assert_eq!(
17679 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
17680 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
17681 );
17682
17683 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
17686 assert_eq!(
17687 seeds.len(),
17688 256,
17689 "fresh_seed must not collide across rapid calls"
17690 );
17691 assert!(!seeds.contains(&0));
17692 }
17693
17694 #[test]
17695 fn response_format_builds_grammar_only_when_present() {
17696 let mk = |rf: Option<serde_json::Value>| {
17700 let mut body = serde_json::json!({
17701 "model": "m", "messages": [{"role": "user", "content": "t"}]});
17702 if let Some(rf) = rf {
17703 body["response_format"] = rf;
17704 }
17705 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17706 let (tx, _rx) = worker::event_channel();
17707 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17708 };
17709 assert!(mk(None).unwrap().request.grammar.is_none());
17710 assert!(
17711 mk(Some(serde_json::json!({"type": "text"})))
17712 .unwrap()
17713 .request
17714 .grammar
17715 .is_none()
17716 );
17717 assert!(matches!(
17718 mk(Some(serde_json::json!({"type": "json_object"})))
17719 .unwrap()
17720 .request
17721 .grammar,
17722 Some(constrained::GrammarSpec::JsonObject)
17723 ));
17724 assert!(matches!(
17725 mk(Some(serde_json::json!({"type": "json_schema",
17726 "json_schema": {"schema": {"type": "object"}}})))
17727 .unwrap()
17728 .request
17729 .grammar,
17730 Some(constrained::GrammarSpec::JsonSchema(_))
17731 ));
17732 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
17734 }
17735
17736 #[test]
17746 fn response_format_think_table_switch_postthink_refusal() {
17747 let mk = |caps: &ModelCaps| {
17748 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17749 "model": "m", "messages": [{"role": "user", "content": "t"}],
17750 "response_format": {"type": "json_object"}}))
17751 .unwrap();
17752 let (tx, _rx) = worker::event_channel();
17753 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
17754 };
17755 let switch = ModelCaps {
17757 chat_ok: true,
17758 qwen_think: true,
17759 think_switch: true,
17760 ..Default::default()
17761 };
17762 let plan = mk(&switch).unwrap();
17763 assert_eq!(
17764 plan.request.think,
17765 memra_tokenizer::chat::ThinkMode::NoThink,
17766 "switch-carrying template must keep the grammar-from-token-1 path"
17767 );
17768 assert!(plan.request.grammar.is_some());
17769
17770 let postthink = ModelCaps {
17772 chat_ok: true,
17773 qwen_think: true,
17774 think_switch: false,
17775 think_close: vec![128799],
17776 ..Default::default()
17777 };
17778 let plan = mk(&postthink).unwrap();
17779 assert_ne!(
17780 plan.request.think,
17781 memra_tokenizer::chat::ThinkMode::NoThink,
17782 "post-think constrained request must keep the think channel ON"
17783 );
17784 assert!(plan.request.grammar.is_some());
17785
17786 let no_contract = ModelCaps {
17788 chat_ok: true,
17789 qwen_think: true,
17790 think_switch: false,
17791 think_close: Vec::new(),
17792 ..Default::default()
17793 };
17794 let err = match mk(&no_contract) {
17795 Err(err) => err,
17796 Ok(_) => panic!("think-forced template with no close contract must refuse"),
17797 };
17798 assert!(
17799 err.contains("think-close"),
17800 "refusal must name the missing close contract: {err}"
17801 );
17802 }
17803
17804 #[test]
17805 fn unsupported_semantic_params_are_named_rejections() {
17806 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17808 "model": "m", "messages": [{"role": "user", "content": "t"}],
17809 "response_format": {"type": "json_object"}
17810 }))
17811 .unwrap();
17812 assert!(req.response_format.is_some());
17813 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17814 "model": "m", "messages": [{"role": "user", "content": "t"}],
17815 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
17816 "user": "u-1", "stream_options": {"include_usage": true}
17817 }))
17818 .unwrap();
17819 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
17821 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
17822 assert_eq!(req.n, Some(1));
17823 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
17825 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
17826 assert_eq!(param, "logit_bias");
17827 assert_eq!(msg, "logit_bias is not supported (why)");
17828 }
17829
17830 #[test]
17831 fn completions_accept_openai_stop_forms() {
17832 for (value, expected) in [
17833 (serde_json::json!("Problem:"), vec!["Problem:"]),
17834 (
17835 serde_json::json!(["Question:", "Problem:"]),
17836 vec!["Question:", "Problem:"],
17837 ),
17838 (serde_json::Value::Null, Vec::<&str>::new()),
17839 ] {
17840 let req: CompletionReq = serde_json::from_value(serde_json::json!({
17841 "model": "plain_quant", "prompt": "task", "stop": value
17842 }))
17843 .unwrap();
17844 assert_eq!(req.stop.into_vec(), expected);
17845 }
17846 }
17847
17848 fn fake_worker_state() -> AppState {
17855 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
17856 }
17857
17858 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
17859 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
17860 }
17861
17862 struct WorkerSaw {
17867 sampler_cfg: SamplerConfig,
17868 think: ThinkMode,
17869 reasoning_effort: Option<String>,
17870 }
17871
17872 fn fake_worker_state_full(
17879 steps: usize,
17880 step_delay: std::time::Duration,
17881 caps: HashMap<String, ModelCaps>,
17882 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
17883 ) -> AppState {
17884 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
17885 let health = health::WorkerHealth::new();
17886 let h = health.clone();
17887 std::thread::spawn(move || {
17888 h.mark_ready();
17889 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
17890 if let Some(tx) = &saw_tx {
17891 let _ = tx.send(WorkerSaw {
17892 sampler_cfg: req.sampler_cfg.clone(),
17893 think: req.think,
17894 reasoning_effort: req.reasoning_effort.clone(),
17895 });
17896 }
17897 worker::release_pending_admit();
17901 worker::release_admission_reservation(req.lane);
17902 h.beat_busy();
17903 if let Some(ready) = req.constraint_ready.take() {
17904 let _ = ready.send(Ok(()));
17905 }
17906 let _ = req.tx.send(Event::PromptUsage {
17907 n_prompt: 1,
17908 n_cached: 0,
17909 });
17910 if let Some(spec) = req.capture.as_ref() {
17915 let _ = req.tx.send(Event::PromptCapture {
17916 hidden: spec.hidden.then(|| vec![1.0, 0.0]),
17917 logits: if spec.logit_pieces.is_empty() {
17918 Vec::new()
17919 } else {
17920 vec![2.0, 0.0]
17921 },
17922 });
17923 }
17924 for step in 0..steps {
17925 h.beat_busy();
17926 let text = if steps == 1 { "ok" } else { "x" };
17927 let _ = req.tx.send(Event::Token {
17928 id: step as u32 + 1,
17929 text: text.into(),
17930 });
17931 if !step_delay.is_zero() {
17932 std::thread::sleep(step_delay);
17933 }
17934 }
17935 let _ = req.tx.send(Event::Done {
17936 stop_reason: "Eos".into(),
17937 n_tokens: steps,
17938 n_prompt: 1,
17939 n_cached: 0,
17940 elapsed_s: 0.01,
17941 spec: None,
17942 });
17943 h.set_phase(health::PHASE_IDLE);
17944 }
17945 });
17946 for _ in 0..2000 {
17949 if health.live().is_ok() {
17950 break;
17951 }
17952 std::thread::sleep(std::time::Duration::from_millis(1));
17953 }
17954 AppState {
17955 cmd_tx,
17956 models: Arc::new(vec!["m".into()]),
17957 caps: Arc::new(caps),
17958 openrouter_metadata: Arc::new(HashMap::new()),
17959 provider_metadata: Arc::new(None),
17960 metering: None,
17961
17962 budget_tokenizers: None,
17963 api_auth: ApiAuth::default(),
17964 metrics_auth: MetricsAuth::default(),
17965 metrics: SharedMetrics::default(),
17966 inflight: Arc::new(Default::default()),
17967 tenant_inflight: Arc::new(Default::default()),
17968 health,
17969 bg: None,
17970 }
17971 }
17972
17973 #[tokio::test]
17974 #[allow(clippy::await_holding_lock)] async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
17976 let _l = drain_lock();
17977 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
17978 let normal_state = st.clone();
17979 let normal = tokio::spawn(async move {
17980 chat_completions(
17981 State(normal_state),
17982 axum::http::HeaderMap::new(),
17983 None,
17984 Json(
17985 serde_json::from_value(serde_json::json!({
17986 "model": "m",
17987 "messages": [{"role": "user", "content": "keep decoding"}],
17988 }))
17989 .unwrap(),
17990 ),
17991 )
17992 .await
17993 });
17994 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
17995
17996 let mut deep = serde_json::json!({"type": "string"});
17997 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
17998 deep = serde_json::json!({"allOf": [deep]});
17999 }
18000 let bad = chat_completions(
18001 State(st.clone()),
18002 axum::http::HeaderMap::new(),
18003 None,
18004 Json(
18005 serde_json::from_value(serde_json::json!({
18006 "model": "m",
18007 "messages": [{"role": "user", "content": "bad schema"}],
18008 "response_format": {
18009 "type": "json_schema",
18010 "json_schema": {"schema": deep},
18011 },
18012 }))
18013 .unwrap(),
18014 ),
18015 )
18016 .await;
18017 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
18018 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
18019 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
18020 .await
18021 .unwrap();
18022 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18023 assert!(
18024 payload["error"]["message"]
18025 .as_str()
18026 .unwrap()
18027 .contains("maximum nesting depth")
18028 );
18029 assert!(
18030 !normal.is_finished(),
18031 "bad schema stalled or replaced the normal decode"
18032 );
18033
18034 let normal_response = normal.await.unwrap();
18035 assert_eq!(normal_response.status(), StatusCode::OK);
18036 let snapshot = st.health.snapshot();
18037 assert!(
18038 st.health.live().is_ok(),
18039 "normal decode left health stalled"
18040 );
18041 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
18042 }
18043
18044 #[tokio::test]
18045 #[allow(clippy::await_holding_lock)] async fn valid_response_format_preflight_preserves_generation() {
18047 let _l = drain_lock();
18048 let response = chat_completions(
18049 State(fake_worker_state()),
18050 axum::http::HeaderMap::new(),
18051 None,
18052 Json(
18053 serde_json::from_value(serde_json::json!({
18054 "model": "m",
18055 "messages": [{"role": "user", "content": "valid schema"}],
18056 "response_format": {"type": "json_object"},
18057 }))
18058 .unwrap(),
18059 ),
18060 )
18061 .await;
18062 assert_eq!(response.status(), StatusCode::OK);
18063 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18064 .await
18065 .unwrap();
18066 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18067 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
18068 }
18069
18070 #[tokio::test]
18071 #[allow(clippy::await_holding_lock)] async fn unknown_model_refuses_model_not_found_before_admission() {
18073 let _l = drain_lock();
18074 let response = chat_completions(
18079 State(fake_worker_state()),
18080 axum::http::HeaderMap::new(),
18081 None,
18082 Json(
18083 serde_json::from_value(serde_json::json!({
18084 "model": "qwen/qwen3.8-27b-typo",
18085 "messages": [{"role": "user", "content": "hi"}],
18086 }))
18087 .unwrap(),
18088 ),
18089 )
18090 .await;
18091 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
18092 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18093 .await
18094 .unwrap();
18095 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18096 assert_eq!(payload["error"]["code"], "model_not_found");
18097 assert_eq!(payload["error"]["type"], "invalid_request_error");
18098
18099 let response = completions(
18101 State(fake_worker_state()),
18102 axum::http::HeaderMap::new(),
18103 None,
18104 Json(
18105 serde_json::from_value(serde_json::json!({
18106 "model": "nope",
18107 "prompt": "hi",
18108 }))
18109 .unwrap(),
18110 ),
18111 )
18112 .await;
18113 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
18114 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18115 .await
18116 .unwrap();
18117 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18118 assert_eq!(payload["error"]["code"], "model_not_found");
18119 }
18120
18121 const METRICS_KEY_ACME: &str = "completion-acme-secret";
18122 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
18123
18124 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
18125 let spec = format!(
18126 "acme:{},blue:{}",
18127 auth::sha256_hex(METRICS_KEY_ACME),
18128 auth::sha256_hex(METRICS_KEY_BLUE),
18129 );
18130 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
18131 let mut st = fake_worker_state();
18132 st.api_auth.keyring = Some(keyring);
18133 st.metrics_auth = MetricsAuth::new(
18134 true,
18135 st.api_auth.configured(),
18136 metrics_token.map(str::to_string),
18137 );
18138 {
18139 let mut metrics = st.metrics.lock().unwrap();
18140 metrics.admitted = 17;
18141 metrics.prompt_tokens_in = 400;
18142 metrics.cached_tokens_in = 60;
18143 metrics.prefix_hits = 2;
18144 metrics.prefix_misses = 3;
18145 metrics.prefix_inserts = 5;
18146 metrics.prefix_evictions = 7;
18147 metrics.prefix_skips_budget = 9;
18148 metrics.prefix_skips_pinned = 10;
18149 metrics.prefix_hit_tokens = 11;
18150 metrics.lcp_hist[4] = 13;
18151 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
18152 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
18153 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
18154 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
18155 metrics.prefix_entries = 29;
18156 metrics.prefix_bytes = 31;
18157 metrics.active_sessions = 3;
18158 metrics.queued_requests = 5;
18159 metrics.admission_inflight.insert("m".into(), 4);
18160 metrics
18161 .admission_booked_bytes
18162 .insert("m".into(), 41_000_000);
18163 metrics.continuation_pool_entries = 7;
18164 metrics.spec_pool_entries = 11;
18165 metrics.cuda_driver_free_bytes = 13;
18166 metrics.cuda_pool_reserved_bytes = 17;
18167 metrics.cuda_pool_used_bytes = 19;
18168 metrics.cuda_pool_cached_bytes = 23;
18169 metrics.batch_size_last = 37;
18170 metrics.spec.insert(
18171 "m".into(),
18172 memra_engine::spec::SpecTelemetry {
18173 rounds: 2,
18174 drafted: 6,
18175 accepted: 4,
18176 ..Default::default()
18177 },
18178 );
18179 let mut spec_window = memra_engine::spec::SpecTelemetry {
18180 rounds: 4,
18181 drafted: 12,
18182 accepted: 6,
18183 ..Default::default()
18184 };
18185 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
18186 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
18187 metrics.spec_window.insert("m".into(), spec_window);
18188 metrics.constraint_compiler_fail_closed.insert(
18189 "m".into(),
18190 Arc::new(std::sync::atomic::AtomicBool::new(true)),
18191 );
18192 }
18193 st
18194 }
18195
18196 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
18197 let mut headers = HeaderMap::new();
18198 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
18199 let response = get_metrics(State(st), headers).await;
18200 assert_eq!(response.status(), StatusCode::OK);
18201 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18202 .await
18203 .unwrap();
18204 serde_json::from_slice(&bytes).unwrap()
18205 }
18206
18207 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
18208 let mut headers = HeaderMap::new();
18209 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
18210 let response = yield_metrics(State(st), headers).await;
18211 assert_eq!(response.status(), StatusCode::OK);
18212 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18213 .await
18214 .unwrap();
18215 serde_json::from_slice(&bytes).unwrap()
18216 }
18217
18218 #[test]
18219 fn exposed_open_bind_is_refused_before_server_start() {
18220 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
18221 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
18222
18223 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
18224 assert!(err.contains("refusing unauthenticated non-loopback bind"));
18225 assert!(err.contains("MEMRA_API_KEY"));
18226 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
18227 assert!(validate_bind_security("[::]:8000", false, false).is_err());
18228
18229 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
18230 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
18231 }
18232
18233 #[tokio::test]
18234 async fn keyed_metrics_require_and_accept_api_bearer() {
18235 let mut st = fake_worker_state();
18236 st.api_auth.single_key = Some(Arc::from("completion-secret"));
18237 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
18238
18239 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
18240 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
18241 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
18242 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
18243
18244 let mut headers = HeaderMap::new();
18245 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
18246 assert_eq!(
18247 get_metrics(State(st.clone()), headers.clone())
18248 .await
18249 .status(),
18250 StatusCode::OK,
18251 );
18252 let body = metrics_json(st.clone(), "completion-secret").await;
18253 assert!(
18254 body.get("admitted").is_some(),
18255 "the legacy single-key domain keeps cumulative counters",
18256 );
18257 assert!(
18258 body.get("active_sessions").is_none(),
18259 "a static completion key is not an operator metrics principal",
18260 );
18261 assert_eq!(
18262 yield_metrics(State(st), headers).await.status(),
18263 StatusCode::OK
18264 );
18265 }
18266
18267 #[tokio::test]
18268 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
18269 let st = multi_key_metrics_state(None);
18270 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
18271 assert_eq!(
18272 body.as_object().unwrap().len(),
18273 2,
18274 "completion metrics must contain only tenant-scoped rows",
18275 );
18276 let tenants = body["tenants"].as_object().unwrap();
18277 assert_eq!(tenants.len(), 1);
18278 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
18279 assert!(!tenants.contains_key("t:blue"));
18280 let adsd = body["adsd_suspect_total"].as_object().unwrap();
18281 assert_eq!(adsd.len(), 1);
18282 assert_eq!(adsd["t:acme"], 1);
18283 assert!(!adsd.contains_key("t:blue"));
18284
18285 let mut headers = HeaderMap::new();
18286 headers.insert(
18287 "authorization",
18288 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
18289 );
18290 assert_eq!(
18291 yield_metrics(State(st), headers).await.status(),
18292 StatusCode::FORBIDDEN,
18293 "the process-wide yield view requires an operator metrics token",
18294 );
18295 }
18296
18297 #[tokio::test]
18298 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
18299 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
18300 for operator_only in [
18301 "prefix_cache_entries",
18302 "prefix_cache_bytes",
18303 "prefix_cache_skips_budget",
18304 "prefix_cache_skips_pinned",
18305 "active_sessions",
18306 "queued_requests",
18307 "admission_inflight",
18308 "admission_booked_bytes",
18309 "continuation_pool_entries",
18310 "spec_pool_entries",
18311 "cuda_driver_free_bytes",
18312 "cuda_pool_reserved_bytes",
18313 "cuda_pool_used_bytes",
18314 "cuda_pool_cached_bytes",
18315 "constraint_compiler_fail_closed",
18316 "serve_idle_seconds",
18317 "spec",
18318 "spec_tau",
18319 "spec_accept_by_position",
18320 "dual_pp",
18321 "pp_wave",
18322 "peer_probe_bypassed",
18323 "peer_probe_boundary_copies",
18324 "peer_probe_runtime_reprobes",
18325 "peer_probe_runtime_failures",
18326 "peer_probe_deferred_total",
18327 "peer_probe_integrity_degraded",
18328 "peer_probe_degraded_to_host_bounce",
18329 ] {
18330 assert!(
18331 body.get(operator_only).is_none(),
18332 "tenant metrics must not expose operator field {operator_only}",
18333 );
18334 }
18335 }
18336
18337 #[test]
18338 fn populated_spec_acceptance_metrics_are_operator_only() {
18339 for scope in [
18340 MetricsScope::CompletionDomain,
18341 MetricsScope::Tenant("t:acme".into()),
18342 ] {
18343 let mut body = json!({});
18344 insert_spec_acceptance_metrics(&mut body, &scope, || {
18345 panic!("tenant scope evaluated the process-wide spec snapshot")
18346 });
18347 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
18348 assert!(
18349 body.get("spec_accept_by_position").is_none(),
18350 "{scope:?} leaked the accept histogram"
18351 );
18352 }
18353
18354 let mut telemetry = memra_engine::spec::SpecTelemetry {
18355 rounds: 4,
18356 drafted: 12,
18357 accepted: 6,
18358 ..Default::default()
18359 };
18360 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
18361 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
18362 let mut body = json!({});
18363 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
18364 HashMap::from([("model-a".to_string(), telemetry)])
18365 });
18366 assert_eq!(body["spec_tau"]["model-a"], 1.5);
18367 let histogram = &body["spec_accept_by_position"]["model-a"];
18368 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
18369 assert_eq!(histogram["rounds"], 4);
18370 assert_eq!(histogram["offered"], json!([4, 4, 4]));
18371 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
18372 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
18373 }
18374
18375 #[test]
18376 fn populated_dual_pp_metrics_are_operator_only() {
18377 let populated = DualPpMetricsSnapshot {
18378 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
18379 stage_samples: [1, 1, 1, 1],
18380 dropped_timing_samples: 0,
18381 overlaps: 17,
18382 slot_pairs: 19,
18383 slot_uses: [19, 19],
18384 slot_collisions: 0,
18385 };
18386 for scope in [
18387 MetricsScope::CompletionDomain,
18388 MetricsScope::Tenant("t:acme".into()),
18389 ] {
18390 let mut body = json!({});
18391 insert_dual_pp_metrics(&mut body, &scope, || populated);
18392 assert!(
18393 body.get("dual_pp").is_none(),
18394 "{scope:?} leaked dual PP topology"
18395 );
18396 }
18397
18398 let mut body = json!({});
18399 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
18400 assert_eq!(body["dual_pp"]["overlaps"], 17);
18401 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
18402 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
18403 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
18404 assert_eq!(
18405 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
18406 1.0
18407 );
18408 }
18409
18410 #[test]
18411 fn populated_pp_wave_metrics_are_operator_only() {
18412 let populated = PpWaveMetricsSnapshot {
18413 ticks: 11,
18414 cells: 96,
18415 overlaps: 37,
18416 };
18417 for scope in [
18418 MetricsScope::CompletionDomain,
18419 MetricsScope::Tenant("t:acme".into()),
18420 ] {
18421 let mut body = json!({});
18422 insert_pp_wave_metrics(&mut body, &scope, || populated);
18423 assert!(
18424 body.get("pp_wave").is_none(),
18425 "{scope:?} leaked PP wave topology"
18426 );
18427 }
18428
18429 let mut body = json!({});
18430 insert_pp_wave_metrics(&mut body, &MetricsScope::All, || populated);
18431 assert_eq!(body["pp_wave"]["ticks"], 11);
18432 assert_eq!(body["pp_wave"]["cells"], 96);
18433 assert_eq!(body["pp_wave"]["overlaps"], 37);
18434 }
18435
18436 #[test]
18437 fn peer_probe_metrics_are_operator_only() {
18438 let populated = memra_engine::pp::PeerProbeMetrics {
18439 bypassed: 1,
18440 boundary_copies: 8_192,
18441 runtime_probes: 1,
18442 runtime_failures: 0,
18443 deferred_total: 4,
18444 integrity_degraded: true,
18445 degraded_to_host_bounce: true,
18446 };
18447 for scope in [
18448 MetricsScope::CompletionDomain,
18449 MetricsScope::Tenant("t:acme".into()),
18450 ] {
18451 let mut body = json!({});
18452 insert_peer_probe_metrics(&mut body, &scope, || populated);
18453 assert!(body.get("peer_probe_bypassed").is_none());
18454 }
18455
18456 let mut body = json!({});
18457 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
18458 assert_eq!(body["peer_probe_bypassed"], 1);
18459 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
18460 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
18461 assert_eq!(body["peer_probe_runtime_failures"], 0);
18462 assert_eq!(body["peer_probe_deferred_total"], 4);
18463 assert_eq!(body["peer_probe_integrity_degraded"], true);
18464 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
18465 }
18466
18467 #[tokio::test]
18468 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
18469 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
18470 for operator_only in [
18471 "lcp_histogram",
18472 "cache_hit_token_ratio",
18473 "prefix_cache_hits",
18474 "prefix_cache_misses",
18475 "prefix_cache_inserts",
18476 "prefix_cache_evictions",
18477 "prefix_cache_skips_budget",
18478 "prefix_cache_skips_pinned",
18479 "prefix_cache_hit_tokens",
18480 ] {
18481 assert!(
18482 tenant_body.get(operator_only).is_none(),
18483 "tenant metrics must not expose global prefix field {operator_only}",
18484 );
18485 }
18486 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
18487 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
18488 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
18489 assert_eq!(
18490 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
18491 0.4
18492 );
18493
18494 let operator_body = metrics_json(
18495 multi_key_metrics_state(Some("scrape-secret")),
18496 "scrape-secret",
18497 )
18498 .await;
18499 assert_eq!(operator_body["prefix_cache_hits"], 2);
18500 assert_eq!(operator_body["prefix_cache_misses"], 3);
18501 assert_eq!(operator_body["prefix_cache_inserts"], 5);
18502 assert_eq!(operator_body["prefix_cache_evictions"], 7);
18503 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
18504 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
18505 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
18506 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
18507 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
18508 }
18509
18510 #[tokio::test]
18511 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
18512 let st = multi_key_metrics_state(Some("scrape-secret"));
18513 let mut completion_headers = HeaderMap::new();
18514 completion_headers.insert(
18515 "authorization",
18516 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
18517 );
18518 assert_eq!(
18519 get_metrics(State(st.clone()), completion_headers.clone())
18520 .await
18521 .status(),
18522 StatusCode::FORBIDDEN,
18523 );
18524 assert_eq!(
18525 yield_metrics(State(st.clone()), completion_headers)
18526 .await
18527 .status(),
18528 StatusCode::FORBIDDEN,
18529 );
18530
18531 let body = metrics_json(st.clone(), "scrape-secret").await;
18532 let tenants = body["tenants"].as_object().unwrap();
18533 assert_eq!(tenants.len(), 2);
18534 assert!(tenants.contains_key("t:acme"));
18535 assert!(tenants.contains_key("t:blue"));
18536 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
18537 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
18538 assert_eq!(body["active_sessions"], 3);
18539 assert_eq!(body["queued_requests"], 5);
18540 assert_eq!(body["admission_inflight"]["m"], 4);
18542 assert_eq!(body["admission_booked_bytes"]["m"], 41_000_000);
18543 assert_eq!(body["prefix_cache_bytes"], 31);
18544 assert_eq!(body["cuda_driver_free_bytes"], 13);
18545 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
18546 assert_eq!(body["spec"]["m"]["drafted"], 6);
18547 assert_eq!(body["spec_tau"]["m"], 1.5);
18548 assert_eq!(
18549 body["spec_accept_by_position"]["m"]["accepted"],
18550 json!([3, 2, 1])
18551 );
18552 let yield_body = yield_metrics_json(st, "scrape-secret").await;
18553 assert_eq!(yield_body["batch_size_last"], 37);
18554 }
18555
18556 #[tokio::test]
18557 async fn metrics_token_protects_public_override_without_api_keys() {
18558 let mut st = fake_worker_state();
18559 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
18560
18561 assert_eq!(
18562 get_metrics(State(st.clone()), HeaderMap::new())
18563 .await
18564 .status(),
18565 StatusCode::UNAUTHORIZED,
18566 );
18567 let mut headers = HeaderMap::new();
18568 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
18569 assert_eq!(
18570 get_metrics(State(st.clone()), headers.clone())
18571 .await
18572 .status(),
18573 StatusCode::OK,
18574 );
18575 assert_eq!(
18576 yield_metrics(State(st), headers).await.status(),
18577 StatusCode::OK
18578 );
18579 }
18580
18581 #[tokio::test]
18582 async fn no_key_loopback_metrics_remain_open_for_development() {
18583 let mut st = fake_worker_state();
18584 st.metrics_auth = MetricsAuth::new(true, false, None);
18585 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
18586 assert_eq!(response.status(), StatusCode::OK);
18587 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18588 .await
18589 .unwrap();
18590 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18591 assert!(
18592 body.get("active_sessions").is_some(),
18593 "no-key loopback development keeps full operator visibility",
18594 );
18595 assert_eq!(
18596 yield_metrics(State(st), HeaderMap::new()).await.status(),
18597 StatusCode::OK,
18598 );
18599 }
18600
18601 #[test]
18602 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
18603 let metrics = SharedMetrics::default();
18604 let rl = RateLimit::compute(4, 1, &metrics);
18606 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
18607 let rl = RateLimit::compute(4, 3, &metrics);
18608 assert_eq!(rl.remaining, 1);
18609 let rl = RateLimit::compute(4, 4, &metrics);
18611 assert_eq!(rl.remaining, 0);
18612 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
18613 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
18615 let m = worker::Metrics {
18617 completed: 2,
18618 tokens_out: 200,
18619 step_p50_ms: 20.0,
18620 ..Default::default()
18621 };
18622 assert_eq!(reset_estimate_s(&m), 2); }
18624
18625 #[test]
18626 fn inflight_guard_counts_up_and_frees_on_drop() {
18627 let counts: InflightCounts = Arc::new(Default::default());
18628 let tenants: TenantGauge = Arc::new(Default::default());
18629 let (g1, n1, t1) = InflightGuard::try_acquire(
18630 counts.clone(),
18631 lanes::Lane::Interactive,
18632 tenants.clone(),
18633 "acme",
18634 None,
18635 )
18636 .unwrap();
18637 let (g2, n2, t2) = InflightGuard::try_acquire(
18638 counts.clone(),
18639 lanes::Lane::Interactive,
18640 tenants.clone(),
18641 "acme",
18642 None,
18643 )
18644 .unwrap();
18645 assert_eq!((n1, n2), (1, 2));
18646 assert_eq!((t1, t2), (1, 2));
18648 let (gj, nj, tj) = InflightGuard::try_acquire(
18650 counts.clone(),
18651 lanes::Lane::Judge,
18652 tenants.clone(),
18653 "blue",
18654 None,
18655 )
18656 .unwrap();
18657 assert_eq!((nj, tj), (1, 1));
18658 drop(g1);
18659 drop(gj);
18660 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
18661 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
18662 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
18663 assert!(tenants.lock().unwrap().get("blue").is_none());
18665 drop(g2);
18666 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
18667 assert!(tenants.lock().unwrap().is_empty());
18668 }
18669
18670 #[test]
18671 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
18672 let counts: InflightCounts = Arc::new(Default::default());
18673 let tenants: TenantGauge = Arc::new(Default::default());
18674 let start = Arc::new(std::sync::Barrier::new(3));
18675 let attempted = Arc::new(std::sync::Barrier::new(3));
18676 let mut joins = Vec::new();
18677 for _ in 0..2 {
18678 let counts = counts.clone();
18679 let tenants = tenants.clone();
18680 let start = start.clone();
18681 let attempted = attempted.clone();
18682 joins.push(std::thread::spawn(move || {
18683 start.wait();
18684 let result = InflightGuard::try_acquire(
18685 counts,
18686 lanes::Lane::Interactive,
18687 tenants,
18688 "preview_001",
18689 Some(1),
18690 );
18691 let won = result.is_ok();
18692 attempted.wait(); drop(result);
18694 won
18695 }));
18696 }
18697 start.wait();
18698 attempted.wait();
18699 let wins = joins
18700 .into_iter()
18701 .map(|join| join.join().unwrap())
18702 .filter(|won| *won)
18703 .count();
18704 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
18705 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
18706 assert!(tenants.lock().unwrap().is_empty());
18707 }
18708
18709 #[tokio::test]
18710 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
18711 let st = fake_worker_state();
18712 let tenant = auth::TenantCtx {
18713 tenant: "preview_001".into(),
18714 lane_class: auth::LaneClass::Interactive,
18715 rate_limit: Some(1),
18716 key_prefix: None,
18717 };
18718 let first_env = Envelope::new(true);
18719 let (guard, first_rl) =
18720 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
18721 Ok(slot) => slot,
18722 Err(_) => panic!("the first request must acquire the tenant slot"),
18723 };
18724 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
18725
18726 let second_env = Envelope::new(true);
18727 let response =
18728 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
18729 Err(response) => response,
18730 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
18731 };
18732 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
18733 assert_eq!(response.headers()["retry-after"], "2");
18734 assert_eq!(response.headers()["retry-after-ms"], "2000");
18735 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
18736 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
18737 assert_eq!(response.headers()["x-request-id"], second_env.id);
18738 assert_eq!(
18739 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18740 1,
18741 "rejected request must not consume a lane slot"
18742 );
18743 assert_eq!(
18744 st.tenant_inflight
18745 .lock()
18746 .unwrap()
18747 .get("preview_001")
18748 .copied(),
18749 Some(1),
18750 "rejected request must not increment the tenant gauge"
18751 );
18752 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18753 .await
18754 .unwrap();
18755 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18756 assert_eq!(payload["error"]["type"], "rate_limit_error");
18757 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
18758 assert!(
18759 payload["error"]["message"]
18760 .as_str()
18761 .unwrap()
18762 .contains("concurrent request limit")
18763 );
18764
18765 drop(guard);
18766 let _ = InflightGuard::try_acquire(
18767 st.inflight.clone(),
18768 lanes::Lane::Interactive,
18769 st.tenant_inflight.clone(),
18770 "preview_001",
18771 Some(1),
18772 )
18773 .expect("slot must reopen after the in-flight request completes");
18774 }
18775
18776 #[test]
18777 fn tenant_rate_limit_override_is_min_with_global_cap() {
18778 let metrics = SharedMetrics::default();
18779 let unlimited = auth::TenantCtx::default_tenant();
18780 let capped = auth::TenantCtx {
18781 tenant: "acme".into(),
18782 lane_class: auth::LaneClass::Interactive,
18783 rate_limit: Some(2),
18784 key_prefix: None,
18785 };
18786 let global = lane_cap(lanes::Lane::Interactive);
18787 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
18789 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
18790 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
18792 assert_eq!((rl.limit, rl.remaining), (2, 1));
18793 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
18794 assert_eq!(rl.remaining, 0);
18795 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
18796 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
18800 assert_eq!(rl.remaining, 0);
18801 let wide = auth::TenantCtx {
18802 rate_limit: Some(global + 100),
18803 ..capped.clone()
18804 };
18805 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
18806 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
18807 }
18808
18809 #[test]
18810 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
18811 let batch = auth::TenantCtx {
18812 tenant: "bulk".into(),
18813 lane_class: auth::LaneClass::Batch,
18814 rate_limit: None,
18815 key_prefix: None,
18816 };
18817 let interactive = auth::TenantCtx::default_tenant();
18818 let hdr = |v: Option<&str>| {
18819 let mut h = axum::http::HeaderMap::new();
18820 if let Some(v) = v {
18821 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
18822 }
18823 h
18824 };
18825 assert_eq!(
18827 lane_for_tenant(&hdr(None), &interactive).unwrap(),
18828 lanes::Lane::Interactive
18829 );
18830 assert_eq!(
18831 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
18832 lanes::Lane::Judge
18833 );
18834 assert_eq!(
18836 lane_for_tenant(&hdr(None), &batch).unwrap(),
18837 lanes::Lane::Harvest
18838 );
18839 assert_eq!(
18840 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
18841 lanes::Lane::Judge
18842 );
18843 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
18844 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
18845 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
18847 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
18848 }
18849
18850 #[tokio::test]
18851 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
18852 let hdr = |v: &str| {
18857 let mut h = axum::http::HeaderMap::new();
18858 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
18859 h
18860 };
18861 let body = |resp: Response| async move {
18862 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18863 .await
18864 .unwrap();
18865 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
18866 };
18867
18868 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
18869 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
18870 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
18871 let payload = body(resp).await;
18872 assert!(
18873 payload["error"].is_object(),
18874 "bare-string error body: {payload}"
18875 );
18876 assert_eq!(payload["error"]["type"], "invalid_request_error");
18877 assert_eq!(payload["error"]["param"], "x-lane");
18878 assert_eq!(payload["error"]["code"], "invalid_lane");
18879
18880 let batch = auth::TenantCtx {
18881 tenant: "bulk".into(),
18882 lane_class: auth::LaneClass::Batch,
18883 rate_limit: None,
18884 key_prefix: None,
18885 };
18886 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
18887 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
18888 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
18889 let payload = body(resp).await;
18890 assert_eq!(payload["error"]["type"], "authentication_error");
18891 assert_eq!(payload["error"]["param"], "x-lane");
18892 }
18893
18894 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
18897
18898 fn drain_lock() -> std::sync::MutexGuard<'static, ()> {
18911 let guard = DRAIN_LOCK.lock().unwrap_or_else(|poisoned| {
18912 DRAIN_LOCK.clear_poison();
18915 poisoned.into_inner()
18916 });
18917 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18918 guard
18919 }
18920
18921 struct DrainingRestore;
18927 impl Drop for DrainingRestore {
18928 fn drop(&mut self) {
18929 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18930 }
18931 }
18932
18933 #[tokio::test]
18934 #[allow(clippy::await_holding_lock)] async fn responses_carry_rate_limit_headers_and_slot_frees() {
18936 let _l = drain_lock();
18937 let st = fake_worker_state();
18938 let resp = chat_completions(
18941 State(st.clone()),
18942 axum::http::HeaderMap::new(),
18943 None,
18944 Json(
18945 serde_json::from_value(serde_json::json!({
18946 "model": "m", "messages": [{"role": "user", "content": "t"}]
18947 }))
18948 .unwrap(),
18949 ),
18950 )
18951 .await;
18952 assert_eq!(resp.status(), StatusCode::OK);
18953 let h = resp.headers();
18954 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
18955 let remaining: usize = h["x-ratelimit-remaining"]
18956 .to_str()
18957 .unwrap()
18958 .parse()
18959 .unwrap();
18960 assert_eq!(remaining, limit - 1);
18961 assert_eq!(h["x-ratelimit-reset"], "0");
18962 assert_eq!(
18963 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18964 0,
18965 "slot must free at completion"
18966 );
18967 let resp = completions(
18970 State(st.clone()),
18971 axum::http::HeaderMap::new(),
18972 None,
18973 Json(
18974 serde_json::from_value(serde_json::json!({
18975 "model": "m", "prompt": "t", "stream": true
18976 }))
18977 .unwrap(),
18978 ),
18979 )
18980 .await;
18981 assert_eq!(resp.status(), StatusCode::OK);
18982 assert!(resp.headers().contains_key("x-ratelimit-limit"));
18983 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
18984 assert!(resp.headers().contains_key("x-ratelimit-reset"));
18985 assert_eq!(
18986 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18987 1,
18988 "stream in flight holds the slot"
18989 );
18990 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
18991 .await
18992 .unwrap();
18993 assert_eq!(
18994 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18995 0,
18996 "slot must free when the stream completes"
18997 );
18998 }
18999
19000 #[tokio::test]
19006 #[allow(clippy::await_holding_lock)] async fn multi_item_capture_requests_open_one_receipt_per_item_under_child_ids() {
19008 let _l = drain_lock();
19009 let mut st = fake_worker_state();
19010 let mock = MockMetering::admit_all();
19011 st.metering = Some(mock.clone());
19012
19013 let resp = embed_api::embeddings_admitted(
19014 State(st.clone()),
19015 HeaderMap::new(),
19016 AdmittedJson(
19017 serde_json::from_value(json!({"model": "m", "input": ["a", "bb", "ccc"]})).unwrap(),
19018 BodyAdmissionLease(None),
19019 ),
19020 )
19021 .await;
19022 assert_eq!(resp.status(), StatusCode::OK);
19023 let parent = resp.headers()["x-request-id"].to_str().unwrap().to_string();
19024 assert!(
19025 !parent.contains('.'),
19026 "the caller sees the parent id: {parent}"
19027 );
19028 let body: serde_json::Value = serde_json::from_slice(
19029 &axum::body::to_bytes(resp.into_body(), usize::MAX)
19030 .await
19031 .unwrap(),
19032 )
19033 .unwrap();
19034 assert_eq!(body["data"].as_array().map(Vec::len), Some(3));
19035 let events = mock.events();
19036 let opened: Vec<(String, &'static str)> = events
19037 .iter()
19038 .filter_map(|e| match e {
19039 MeterEvent::Open {
19040 request_id, route, ..
19041 } => Some((request_id.clone(), *route)),
19042 _ => None,
19043 })
19044 .collect();
19045 assert_eq!(
19046 opened,
19047 vec![
19048 (format!("{parent}.0"), "/v1/embeddings"),
19049 (format!("{parent}.1"), "/v1/embeddings"),
19050 (format!("{parent}.2"), "/v1/embeddings"),
19051 ],
19052 "one receipt per input, each under its own child id: {events:?}"
19053 );
19054 assert_eq!(
19055 events
19056 .iter()
19057 .filter(|e| matches!(e, MeterEvent::Complete { .. }))
19058 .count(),
19059 3,
19060 "every input settles its own receipt: {events:?}"
19061 );
19062
19063 let resp = embed_api::rerank_admitted(
19064 State(st),
19065 HeaderMap::new(),
19066 AdmittedJson(
19067 serde_json::from_value(
19068 json!({"model": "m", "query": "q", "documents": ["d0", "d1"]}),
19069 )
19070 .unwrap(),
19071 BodyAdmissionLease(None),
19072 ),
19073 )
19074 .await;
19075 assert_eq!(resp.status(), StatusCode::OK);
19076 let parent = resp.headers()["x-request-id"].to_str().unwrap().to_string();
19077 let opened: Vec<String> = mock
19078 .events()
19079 .into_iter()
19080 .skip(events.len())
19081 .filter_map(|e| match e {
19082 MeterEvent::Open {
19083 request_id,
19084 route: "/v1/rerank",
19085 ..
19086 } => Some(request_id),
19087 _ => None,
19088 })
19089 .collect();
19090 assert_eq!(opened, vec![format!("{parent}.0"), format!("{parent}.1")]);
19091 }
19092
19093 #[tokio::test]
19094 #[allow(clippy::await_holding_lock)] async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
19096 let _l = drain_lock();
19097 let mut st = fake_worker_state();
19098 let mock = MockMetering::admit_all();
19099 st.metering = Some(mock.clone());
19100
19101 let nonstream = chat_completions(
19102 State(st.clone()),
19103 HeaderMap::new(),
19104 None,
19105 Json(
19106 serde_json::from_value(json!({
19107 "model": "m",
19108 "messages": [{"role": "user", "content": "t"}],
19109 }))
19110 .unwrap(),
19111 ),
19112 )
19113 .await;
19114 assert_eq!(nonstream.status(), StatusCode::OK);
19115 let nonstream_id = nonstream.headers()["x-request-id"]
19116 .to_str()
19117 .unwrap()
19118 .to_string();
19119
19120 let stream = completions(
19121 State(st),
19122 HeaderMap::new(),
19123 None,
19124 Json(
19125 serde_json::from_value(json!({
19126 "model": "m",
19127 "prompt": "t",
19128 "stream": true,
19129 }))
19130 .unwrap(),
19131 ),
19132 )
19133 .await;
19134 assert_eq!(stream.status(), StatusCode::OK);
19135 let stream_id = stream.headers()["x-request-id"]
19136 .to_str()
19137 .unwrap()
19138 .to_string();
19139 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
19140 .await
19141 .unwrap();
19142
19143 let events = mock.events();
19147 let opened: Vec<&str> = events
19148 .iter()
19149 .filter_map(|e| match e {
19150 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
19151 _ => None,
19152 })
19153 .collect();
19154 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
19155 let completes = events
19156 .iter()
19157 .filter(|e| {
19158 matches!(
19159 e,
19160 MeterEvent::Complete {
19161 prompt: 1,
19162 cached: 0,
19163 completion: 1,
19164 }
19165 )
19166 })
19167 .count();
19168 assert_eq!(
19169 completes, 2,
19170 "both surfaces settle complete with worker-truth usage: {events:?}"
19171 );
19172 }
19173
19174 #[tokio::test]
19175 #[allow(clippy::await_holding_lock)] async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
19177 let _l = drain_lock();
19178 let mock = MockMetering::with_limits(vec![
19184 ReserveScript::Insufficient,
19185 ReserveScript::Admit { with_permit: false },
19186 ReserveScript::Blocked,
19187 ReserveScript::Admit { with_permit: true },
19188 ]);
19189 let mut st = fake_worker_state();
19190 st.metering = Some(mock.clone());
19191
19192 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
19194 assert_eq!(metrics.status(), StatusCode::OK);
19195 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
19196 .await
19197 .unwrap();
19198 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
19199 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
19200 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
19201 assert_eq!(metrics_body["budget_source_available"], true);
19202
19203 let request = || {
19204 Json(
19205 serde_json::from_value::<CompletionReq>(json!({
19206 "model": "m",
19207 "prompt_ids": [1],
19208 "max_tokens": 1,
19209 }))
19210 .unwrap(),
19211 )
19212 };
19213
19214 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19215 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
19216 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
19217 .await
19218 .unwrap();
19219 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
19220 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
19221 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
19222
19223 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19224 assert_eq!(included.status(), StatusCode::OK);
19225
19226 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19229 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
19230
19231 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19232 assert_eq!(admitted.status(), StatusCode::OK);
19233
19234 let events = mock.events();
19235 let terminal: Vec<&MeterEvent> = events
19236 .iter()
19237 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
19238 .collect();
19239 assert_eq!(
19240 terminal.len(),
19241 4,
19242 "four requests, four terminal settles: {events:?}"
19243 );
19244 assert!(matches!(
19245 terminal[0],
19246 MeterEvent::Reject { status: 402, .. }
19247 ));
19248 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
19249 assert!(matches!(
19250 terminal[2],
19251 MeterEvent::Reject { status: 402, .. }
19252 ));
19253 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
19254 let permits: Vec<bool> = events
19256 .iter()
19257 .filter_map(|e| match e {
19258 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
19259 _ => None,
19260 })
19261 .collect();
19262 assert_eq!(
19263 permits,
19264 vec![false, false, false, true],
19265 "the permit rides the receipt exactly when reserve minted one: {events:?}"
19266 );
19267 }
19268
19269 #[tokio::test]
19273 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
19274 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
19275 let mut st = fake_worker_state();
19276 st.metering = Some(mock.clone());
19277 let tenant = auth::TenantCtx {
19278 tenant: "acme".into(),
19279 lane_class: auth::LaneClass::Interactive,
19280 rate_limit: None,
19281 key_prefix: Some("mk-acme-testprefix00".into()),
19282 };
19283 let mut request = gate_request(1, 1);
19284 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
19285 .expect_err("a capped key must be refused at admission");
19286 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
19287 let (response, outcome) = rejection.into_response();
19288 assert_eq!(outcome, "key_spend_cap_reached");
19289 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
19290 let body = body_value(response).await;
19291 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
19292 assert!(
19293 body["error"]["message"].as_str().unwrap().contains("cap"),
19294 "the 402 must point at the KEY's cap, not tenant credit: {body}"
19295 );
19296 let events = mock.events();
19297 assert!(
19298 events.contains(&MeterEvent::Reserve {
19299 tenant: "acme".into(),
19300 principal: Some("mk-acme-testprefix00".into()),
19301 model: "qwen/qwen3.8-27b".into(),
19302 }),
19303 "the key prefix must reach reserve: {events:?}"
19304 );
19305 }
19306
19307 #[tokio::test]
19308 #[allow(clippy::await_holding_lock)] async fn streaming_client_disconnect_records_partial_usage_and_cost() {
19310 let _l = drain_lock();
19311 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
19312 let mock = MockMetering::admit_all();
19313 st.metering = Some(mock.clone());
19314
19315 let response = completions(
19316 State(st),
19317 HeaderMap::new(),
19318 None,
19319 Json(
19320 serde_json::from_value(json!({
19321 "model": "m",
19322 "prompt": "disconnect after one delta",
19323 "stream": true,
19324 }))
19325 .unwrap(),
19326 ),
19327 )
19328 .await;
19329 assert_eq!(response.status(), StatusCode::OK);
19330 let request_id = response.headers()["x-request-id"]
19331 .to_str()
19332 .unwrap()
19333 .to_string();
19334 let mut body = Box::pin(response.into_body().into_data_stream());
19335 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
19336 .await
19337 .expect("stream ended before first delta")
19338 .expect("stream body failed");
19339 assert!(
19340 is_sse_data_frame(&first),
19341 "first frame was not SSE data: {first:?}"
19342 );
19343 drop(body);
19344
19345 let mut dropped = None;
19348 for _ in 0..500 {
19349 if let Some(event) = mock
19350 .events()
19351 .into_iter()
19352 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
19353 {
19354 dropped = Some(event);
19355 break;
19356 }
19357 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
19358 }
19359 let events = mock.events();
19360 assert!(
19361 events
19362 .iter()
19363 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
19364 "the receipt was opened under the caller-visible request id: {events:?}"
19365 );
19366 assert_eq!(
19367 dropped,
19368 Some(MeterEvent::Dropped {
19369 prompt: 1,
19370 cached: 0,
19371 completion: 1,
19372 }),
19373 "a client disconnect must leave the partial counts on the dropped receipt \
19374 (the implementation prices that drop): {events:?}"
19375 );
19376 }
19377
19378 #[tokio::test]
19379 #[allow(clippy::await_holding_lock)] async fn draining_rejects_new_requests_with_503_and_retry_after() {
19381 let _l = drain_lock();
19382 let st = fake_worker_state();
19383 let _down = DrainingRestore;
19386 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
19387 let resp = chat_completions(
19389 State(st.clone()),
19390 axum::http::HeaderMap::new(),
19391 None,
19392 Json(
19393 serde_json::from_value(serde_json::json!({
19394 "model": "m", "messages": [{"role": "user", "content": "t"}]
19395 }))
19396 .unwrap(),
19397 ),
19398 )
19399 .await;
19400 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19401 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
19406 let ra_s: u64 = ra
19407 .parse()
19408 .expect("Retry-After must be integer delay-seconds");
19409 assert!(
19410 ra_s > 0 && ra_s <= 60,
19411 "Retry-After {ra_s}s is outside the honored window"
19412 );
19413 let ra_ms: u64 = resp.headers()["retry-after-ms"]
19414 .to_str()
19415 .unwrap()
19416 .parse()
19417 .unwrap();
19418 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
19419 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19420 .await
19421 .unwrap();
19422 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19423 assert!(
19424 payload["error"]["message"]
19425 .as_str()
19426 .unwrap()
19427 .contains("draining")
19428 );
19429 assert_eq!(payload["error"]["type"], "server_error");
19430 assert_eq!(payload["error"]["code"], "draining");
19431 let resp = completions(
19432 State(st.clone()),
19433 axum::http::HeaderMap::new(),
19434 None,
19435 Json(
19436 serde_json::from_value(serde_json::json!({
19437 "model": "m", "prompt": "t"
19438 }))
19439 .unwrap(),
19440 ),
19441 )
19442 .await;
19443 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19444 assert!(resp.headers().contains_key("retry-after"));
19445 assert_eq!(
19446 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
19447 0,
19448 "rejected requests must not hold slots"
19449 );
19450 let resp = health_live(State(st.clone())).await.into_response();
19453 assert_eq!(
19454 resp.status(),
19455 StatusCode::OK,
19456 "a drain must not look like a liveness fault"
19457 );
19458 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19459 .await
19460 .unwrap();
19461 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19462 assert_eq!(payload["status"], "draining");
19463 let resp = health_ready(State(st.clone())).await.into_response();
19465 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19466 let retry_s = drain_deadline_s().clamp(1, 60);
19467 let retry_s_text = retry_s.to_string();
19468 let retry_ms_text = (retry_s * 1000).to_string();
19469 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
19470 assert_eq!(
19471 resp.headers().get("retry-after-ms").unwrap(),
19472 retry_ms_text.as_str()
19473 );
19474 assert_ne!(
19475 resp.headers()
19476 .get("x-should-retry")
19477 .and_then(|v| v.to_str().ok()),
19478 Some("false")
19479 );
19480 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19481 .await
19482 .unwrap();
19483 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19484 assert_eq!(payload["status"], "not_ready");
19485 assert!(payload["detail"].as_str().unwrap().contains("draining"));
19486 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
19487 let resp = chat_completions(
19489 State(st.clone()),
19490 axum::http::HeaderMap::new(),
19491 None,
19492 Json(
19493 serde_json::from_value(serde_json::json!({
19494 "model": "m", "messages": [{"role": "user", "content": "t"}]
19495 }))
19496 .unwrap(),
19497 ),
19498 )
19499 .await;
19500 assert_eq!(resp.status(), StatusCode::OK);
19501 }
19502
19503 #[tokio::test]
19506 #[allow(clippy::await_holding_lock)] async fn health_is_green_only_while_the_worker_is_alive() {
19508 let _l = drain_lock();
19511 let st = fake_worker_state();
19512 let resp = health_live(State(st.clone())).await.into_response();
19515 assert_eq!(resp.status(), StatusCode::OK);
19516 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19517 .await
19518 .unwrap();
19519 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19520 assert_eq!(payload["status"], "ok");
19521 assert_eq!(payload["worker"]["phase"], "idle");
19522 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
19523 let ready = health_ready(State(st.clone())).await.into_response();
19524 assert_eq!(ready.status(), StatusCode::OK);
19525
19526 st.health.mark_dead("worker thread panicked: test-injected");
19530 let resp = health_live(State(st.clone())).await.into_response();
19531 assert_eq!(
19532 resp.status(),
19533 StatusCode::SERVICE_UNAVAILABLE,
19534 "a dead worker MUST NOT report a healthy liveness"
19535 );
19536 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19537 .await
19538 .unwrap();
19539 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19540 assert_eq!(payload["status"], "unhealthy");
19541 assert!(
19543 payload["detail"]
19544 .as_str()
19545 .unwrap()
19546 .contains("test-injected"),
19547 "cause not surfaced: {payload}"
19548 );
19549 let ready = health_ready(State(st.clone())).await.into_response();
19550 assert_eq!(
19551 ready.status(),
19552 StatusCode::SERVICE_UNAVAILABLE,
19553 "dead is also not ready"
19554 );
19555
19556 st.health.mark_ready();
19559 assert_eq!(
19560 health_live(State(st.clone()))
19561 .await
19562 .into_response()
19563 .status(),
19564 StatusCode::OK,
19565 "mark_ready must clear the latch (a successful respawn)"
19566 );
19567 }
19568
19569 #[tokio::test]
19570 #[allow(clippy::await_holding_lock)] async fn readyz_peer_probe_integrity_is_present_and_advisory() {
19572 let _l = drain_lock();
19573 let st = fake_worker_state();
19574
19575 let ready = health_ready(State(st.clone())).await.into_response();
19576 assert_eq!(ready.status(), StatusCode::OK);
19577 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
19578 .await
19579 .unwrap();
19580 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19581 assert_eq!(payload["peer_probe_integrity"], "ok");
19582
19583 st.health.note_peer_probe_deferral(2, false);
19584 let deferred = health_ready(State(st.clone())).await.into_response();
19585 assert_eq!(deferred.status(), StatusCode::OK);
19586 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
19587 .await
19588 .unwrap();
19589 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19590 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
19591
19592 st.health.note_peer_probe_deferral(4, true);
19593 let degraded = health_ready(State(st.clone())).await.into_response();
19594 assert_eq!(
19595 degraded.status(),
19596 StatusCode::OK,
19597 "peer degradation is advisory while plain serving remains healthy"
19598 );
19599 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
19600 .await
19601 .unwrap();
19602 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19603 assert_eq!(payload["peer_probe_integrity"], "degraded");
19604
19605 st.health.mark_dead("test-injected worker failure");
19606 let unready = health_ready(State(st)).await.into_response();
19607 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
19608 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
19609 .await
19610 .unwrap();
19611 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19612 assert_eq!(
19613 payload["peer_probe_integrity"], "degraded",
19614 "the advisory field must also survive an unrelated readiness failure"
19615 );
19616 }
19617
19618 #[tokio::test]
19619 #[allow(clippy::await_holding_lock)] async fn liveness_failure_obeys_the_retry_contract() {
19621 let _l = drain_lock();
19626 let st = fake_worker_state();
19627 st.health
19628 .mark_dead("worker thread panicked: retry-contract-test");
19629
19630 let resp = health_live(State(st)).await.into_response();
19631 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19632 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
19633 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
19634 assert_ne!(
19635 resp.headers()
19636 .get("x-should-retry")
19637 .and_then(|v| v.to_str().ok()),
19638 Some("false")
19639 );
19640 }
19641
19642 #[tokio::test]
19643 #[allow(clippy::await_holding_lock)] async fn readiness_failure_obeys_the_retry_contract() {
19645 let _l = drain_lock();
19646 let st = fake_worker_state();
19647 st.health
19648 .mark_dead("worker thread panicked: retry-contract-test");
19649
19650 let resp = health_ready(State(st)).await.into_response();
19651 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19652 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
19653 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
19654 assert_ne!(
19655 resp.headers()
19656 .get("x-should-retry")
19657 .and_then(|v| v.to_str().ok()),
19658 Some("false")
19659 );
19660 }
19661
19662 #[tokio::test]
19663 #[allow(clippy::await_holding_lock)] async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
19665 let _l = drain_lock();
19675 let st = fake_worker_state();
19676 assert_eq!(
19677 health_live(State(st.clone()))
19678 .await
19679 .into_response()
19680 .status(),
19681 StatusCode::OK
19682 );
19683 st.health
19684 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
19685 let resp = health_live(State(st.clone())).await.into_response();
19686 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19687 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19688 .await
19689 .unwrap();
19690 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19691 assert!(
19692 payload["detail"]
19693 .as_str()
19694 .unwrap()
19695 .contains("probe exceeded")
19696 );
19697 st.health.mark_ready();
19700 assert_eq!(
19701 health_live(State(st.clone()))
19702 .await
19703 .into_response()
19704 .status(),
19705 StatusCode::SERVICE_UNAVAILABLE,
19706 "a GPU fault must not be cleared by an in-process respawn"
19707 );
19708 }
19709
19710 #[test]
19711 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
19712 let caps = ModelCaps {
19714 tools_branch: true,
19715 hy3: false,
19716 qwen_think: true,
19717 think_switch: true,
19718 chat_ok: true,
19719 context_length: 262144,
19720 tokenizer: "qwen2".into(),
19721 instruct_type: Some("chatml".into()),
19722 effort_levels: false,
19723 qwen_effort: false,
19724 gemma_think: false,
19725 dsv4: false,
19726 glm5: false,
19727 chat_temperature_default: None,
19728 chat_top_p_default: None,
19729 n_vocab: 151_936,
19730 think_close: Vec::new(),
19731 };
19732 let e = model_entry_v1("main", Some(&caps), None);
19733 assert_eq!(e["id"], "main");
19734 assert_eq!(e["name"], "main");
19735 assert_eq!(e["object"], "model");
19736 assert_eq!(e["context_length"], 262144);
19737 assert!(e["pricing"]["input"].is_null());
19739 assert!(e["pricing"]["output"].is_null());
19740
19741 let meta = OpenRouterModelMetadata {
19744 pricing: OpenRouterPricing {
19745 prompt: Some("0.00000038".into()),
19746 cached_prompt: Some("0.0000002".into()),
19747 completion: Some("0.0000026".into()),
19748 ..Default::default()
19749 },
19750 input_modalities: vec!["image".into(), "video".into()],
19751 max_output_length: Some(32768),
19752 ..Default::default()
19753 };
19754 let e = model_entry_v1("main", Some(&caps), Some(&meta));
19755 assert_eq!(e["pricing"]["currency"], "USD");
19758 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
19759 assert_eq!(e["pricing"]["input"], "0.38");
19760 assert_eq!(e["pricing"]["output"], "2.60");
19761 assert_eq!(e["pricing"]["cached_input"], "0.20");
19762 assert!(e["pricing"]["cache_write"].is_null());
19763 assert_eq!(e["pricing"]["minimum_request"], "0");
19764 assert_eq!(e["owned_by"], "main");
19765 assert_eq!(e["type"], "chat");
19766 assert_eq!(e["max_output_tokens"], 32768);
19767 assert_eq!(e["endpoints"], json!(["chat/completions"]));
19768 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
19769 assert_eq!(e["output_modalities"], json!(["text"]));
19770 assert_eq!(e["capabilities"]["streaming"], true);
19771 assert_eq!(e["capabilities"]["tools"], true);
19772 assert_eq!(e["lifecycle"]["status"], "active");
19773 assert!(e["lifecycle"]["deprecation_at"].is_null());
19774 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
19775 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
19776 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
19780 keys.sort_unstable();
19781 assert_eq!(
19782 keys,
19783 [
19784 "capabilities",
19785 "context_length",
19786 "endpoints",
19787 "id",
19788 "input_modalities",
19789 "lifecycle",
19790 "max_output_tokens",
19791 "name",
19792 "object",
19793 "output_modalities",
19794 "owned_by",
19795 "pricing",
19796 "reliability",
19797 "type",
19798 ],
19799 "unexpected /v1/models entry keys"
19800 );
19801 let mut price_keys: Vec<&str> = e["pricing"]
19802 .as_object()
19803 .unwrap()
19804 .keys()
19805 .map(String::as_str)
19806 .collect();
19807 price_keys.sort_unstable();
19808 assert_eq!(
19809 price_keys,
19810 [
19811 "cache_write",
19812 "cached_input",
19813 "currency",
19814 "input",
19815 "minimum_request",
19816 "output",
19817 "unit",
19818 ],
19819 "unexpected /v1/models pricing keys"
19820 );
19821
19822 let e = model_entry_v1("m", None, None);
19824 assert!(e["context_length"].is_null());
19825 assert!(e["max_output_tokens"].is_null());
19826 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
19828 assert!(e["context_length"].is_null());
19829 }
19830
19831 #[test]
19837 fn catalog_row_follows_the_declared_surface() {
19838 let caps = ModelCaps {
19839 tools_branch: true,
19840 ..Default::default()
19841 };
19842
19843 let embed = OpenRouterModelMetadata {
19844 surface: Some("embedding".into()),
19845 max_output_length: Some(1),
19846 ..Default::default()
19847 };
19848 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
19849 assert_eq!(e["type"], "embedding");
19850 assert_eq!(e["endpoints"], json!(["embeddings"]));
19851 assert_eq!(e["output_modalities"], json!(["embeddings"]));
19852 assert_eq!(e["capabilities"]["streaming"], false);
19853 assert_eq!(
19854 e["capabilities"]["tools"], false,
19855 "an embedder has no tools"
19856 );
19857 assert_eq!(e["capabilities"]["reasoning"], false);
19858 assert_eq!(e["capabilities"]["structured_output"], false);
19859 assert_eq!(e["capabilities"]["prompt_caching"], false);
19860 assert!(
19861 e["max_output_tokens"].is_null(),
19862 "a surface that emits no completion tokens must not advertise a ceiling"
19863 );
19864
19865 let rerank = OpenRouterModelMetadata {
19866 surface: Some("rerank".into()),
19867 ..Default::default()
19868 };
19869 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
19870 assert_eq!(r["type"], "rerank");
19871 assert_eq!(r["endpoints"], json!(["rerank"]));
19872 assert_eq!(r["output_modalities"], json!(["rerank"]));
19873 assert_eq!(r["capabilities"]["tools"], false);
19874 assert_eq!(r["capabilities"]["reasoning"], false);
19875
19876 let chat = OpenRouterModelMetadata {
19879 max_output_length: Some(32768),
19880 ..Default::default()
19881 };
19882 let c = model_entry_v1("main", Some(&caps), Some(&chat));
19883 assert_eq!(c["type"], "chat");
19884 assert_eq!(c["endpoints"], json!(["chat/completions"]));
19885 assert_eq!(c["output_modalities"], json!(["text"]));
19886 assert_eq!(c["capabilities"]["tools"], true);
19887 assert_eq!(c["max_output_tokens"], 32768);
19888 }
19889
19890 #[test]
19893 fn unknown_surface_is_rejected_at_config_load() {
19894 let bad = OpenRouterModelMetadata {
19895 surface: Some("embeddings".into()), ..Default::default()
19897 };
19898 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
19899 .expect_err("an unknown surface must not load");
19900 assert!(err.contains("surface"), "{err}");
19901
19902 for good in ["chat", "embedding", "rerank"] {
19903 let ok = OpenRouterModelMetadata {
19904 surface: Some(good.into()),
19905 ..Default::default()
19906 };
19907 assert!(
19908 validate_openrouter_metadata("m", &ok).is_ok(),
19909 "{good} must load"
19910 );
19911 }
19912 }
19913
19914 #[test]
19915 fn per_million_price_is_exact_decimal_shift() {
19916 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
19918 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
19919 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
19920 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
19921 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
19922 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
19923 assert_eq!(per_million_price("not-a-price"), None);
19924 assert_eq!(per_million_price(""), None);
19925 }
19926
19927 #[test]
19928 fn metadata_provider_block_parses_and_validates() {
19929 let (_, provider) = OpenRouterMetadataFile::parse(
19930 r#"
19931 [provider]
19932 id = "tiyuvta"
19933 status_url = "https://status.tiyuvta.ai"
19934 support_contact = "mailto:support@tiyuvta.ai"
19935 incident_contact = "mailto:incidents@tiyuvta.ai"
19936 regions = ["eu-central"]
19937 "#,
19938 )
19939 .unwrap();
19940 let provider = provider.unwrap();
19941 assert_eq!(provider.id, "tiyuvta");
19942 assert_eq!(provider.regions, vec!["eu-central"]);
19943 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
19945 assert!(err.contains("provider.id"), "{err}");
19946 let err = OpenRouterMetadataFile::parse(
19948 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
19949 )
19950 .unwrap_err();
19951 assert!(err.contains("must be a URI"), "{err}");
19952 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
19954 assert!(provider.is_none());
19955 }
19956
19957 #[test]
19958 fn models_openai_default_body_stays_byte_identical() {
19959 let body = models_openai_body(&["main".into(), "judge".into()]);
19960 let bytes = serde_json::to_vec(&body).unwrap();
19961 assert_eq!(
19962 bytes,
19963 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
19964 );
19965 }
19966
19967 #[test]
19968 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
19969 let loaded = vec![
19971 "qwen/qwen3.6-27b".to_string(),
19972 "qwen/qwen3.6-35b-a3b".to_string(),
19973 ];
19974 assert_eq!(
19975 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
19976 Some("qwen/qwen3.6-35b-a3b"),
19977 );
19978 assert_eq!(
19979 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
19980 Some("qwen/qwen3.6-27b"),
19981 );
19982 assert_eq!(
19984 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
19985 Some("qwen/qwen3.6-35b-a3b"),
19986 );
19987 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
19989 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
19990 assert_eq!(canonical_model_id(&loaded, ""), None);
19991 }
19992
19993 #[test]
19994 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
19995 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
19998 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
19999 assert_eq!(
20001 canonical_model_id(&loaded, "a/shared-name").as_deref(),
20002 Some("a/shared-name")
20003 );
20004 assert_eq!(
20005 canonical_model_id(&loaded, "b/shared-name").as_deref(),
20006 Some("b/shared-name")
20007 );
20008 let bare = vec!["solo".to_string()];
20010 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
20011 }
20012
20013 #[test]
20014 fn openrouter_models_entry_serializes_complete_metadata() {
20015 let metadata = OpenRouterMetadataFile::from_toml(
20016 r#"
20017[models.main]
20018hugging_face_id = "Qwen/Qwen3.6-27B"
20019created = 1786032000
20020quantization = "nvfp4"
20021description = "Qwen3.6 27B served by memra."
20022max_prompt_length = 245760
20023max_output_length = 16384
20024default_output_length = 4096
20025is_ready = true
20026is_free = false
20027discount_to_user = 0.1
20028openrouter_slug = "qwen/qwen3.6-27b"
20029datacenters = [{ country_code = "US", region = "us-east" }]
20030zdr = true
20031hipaa = false
20032
20033[models.main.pricing]
20034prompt = "0.000000234"
20035cached_prompt = "0.0000000585"
20036cache_write = "0.000000234"
20037completion = "0.000001872"
20038internal_reasoning = "0.000001872"
20039request = "0.01"
20040
20041[models.main.capacity]
20042prompt_tpm = 1000000
20043cached_prompt_tpm = 2000000
20044completion_tpm = 500000
20045request_rpm = 1000
20046concurrency = 64
20047"#,
20048 )
20049 .unwrap();
20050 let caps = ModelCaps {
20051 tools_branch: true,
20052 qwen_think: true,
20053 think_switch: true,
20054 chat_ok: true,
20055 context_length: 262144,
20056 tokenizer: "qwen2".into(),
20057 instruct_type: Some("chatml".into()),
20058 ..Default::default()
20059 };
20060 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
20061
20062 assert_eq!(entry["schema_version"], "2.4");
20063 assert_eq!(entry["id"], "main");
20064 assert_eq!(entry["name"], "main");
20065 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
20066 assert_eq!(entry["created"], 1786032000u64);
20067 assert_eq!(entry["quantization"], "nvfp4");
20068 assert_eq!(entry["tokenizer"], "qwen2");
20069 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
20070 assert!(
20071 entry.get("object").is_none(),
20072 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
20073 );
20074
20075 let input = &entry["input_modalities"][0];
20076 assert_eq!(input["type"], "text");
20077 assert_eq!(
20078 input["supported_inputs"]["max_context_length"]["value"],
20079 262144
20080 );
20081 assert_eq!(
20082 input["supported_inputs"]["max_prompt_length"]["value"],
20083 245760
20084 );
20085 let input_prices = input["pricing"].as_array().unwrap();
20086 let input_price = |kind: &str| {
20087 input_prices
20088 .iter()
20089 .find(|price| price["type"] == kind)
20090 .unwrap()
20091 };
20092 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
20093 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
20094 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
20095 assert_eq!(input["capacity"][0]["value"], 1000000);
20096 assert_eq!(input["capacity"][1]["value"], 2000000);
20097
20098 let output = &entry["output_modalities"][0];
20099 assert_eq!(output["type"], "text");
20100 assert_eq!(output["max_length"]["value"], 16384);
20101 assert_eq!(output["streaming"], true);
20102 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
20103 assert_eq!(
20104 output["supported_parameters"]["structured_outputs"]["type"],
20105 "boolean"
20106 );
20107 assert_eq!(
20108 output["supported_parameters"]["reasoning"]["type"],
20109 "boolean"
20110 );
20111 assert_eq!(output["pricing"][0]["type"], "completion");
20112 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
20113 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
20114 assert_eq!(output["capacity"][0]["value"], 500000);
20115 assert_eq!(output["capacity"][1]["type"], "concurrency");
20116 assert_eq!(output["capacity"][1]["value"], 64);
20117
20118 assert_eq!(entry["pricing"][0]["type"], "request");
20119 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
20120 assert_eq!(entry["capacity"][0]["value"], 1000);
20121 assert_eq!(entry["is_ready"], true);
20122 assert_eq!(entry["is_free"], false);
20123 assert_eq!(entry["discount_to_user"], 0.1);
20124 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
20125 assert_eq!(entry["datacenters"][0]["country_code"], "US");
20126 assert_eq!(entry["compliance"]["zdr"], true);
20127 assert_eq!(entry["compliance"]["hipaa"], false);
20128 }
20129
20130 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
20135[models."qwen/qwen3.6-35b-a3b"]
20136hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
20137created = 1777260255
20138quantization = "int4"
20139description = "Qwen3.6 35B-A3B fixture entry."
20140max_prompt_length = 262144
20141max_output_length = 262144
20142default_output_length = 8192
20143is_ready = true
20144is_free = false
20145discount_to_user = 0.0
20146openrouter_slug = "qwen/qwen3.6-35b-a3b"
20147zdr = false
20148hipaa = false
20149
20150[[models."qwen/qwen3.6-35b-a3b".datacenters]]
20151country_code = "CA"
20152region = "Ontario"
20153
20154[models."qwen/qwen3.6-35b-a3b".pricing]
20155prompt = "0.0000000931"
20156cached_prompt = "0.0000000652"
20157completion = "0.0000009025"
20158
20159[models."qwen/qwen3.6-35b-a3b".capacity]
20160prompt_tpm = 780000
20161cached_prompt_tpm = 310000
20162completion_tpm = 9600
20163request_rpm = 160
20164concurrency = 16
20165
20166[planned_models."qwen/qwen3.8-27b"]
20167description = "Planned fixture entry; must never be emitted."
20168max_prompt_length = 262144
20169max_output_length = 262144
20170default_output_length = 8192
20171is_ready = false
20172is_free = false
20173discount_to_user = 0.0
20174openrouter_slug = "qwen/qwen3.8-27b"
20175zdr = false
20176hipaa = false
20177
20178[planned_models."qwen/qwen3.8-27b".pricing]
20179prompt = "0.0000002745"
20180cached_prompt = "0.0000001922"
20181completion = "0.0000022800"
20182
20183[planned_models."google/gemma-4-26b-a4b-it"]
20184hugging_face_id = "google/gemma-4-26B-A4B-it"
20185created = 1775227989
20186quantization = "int4"
20187description = "Planned fixture entry; must never be emitted."
20188max_prompt_length = 262144
20189max_output_length = 262144
20190default_output_length = 8192
20191is_ready = false
20192is_free = false
20193discount_to_user = 0.0
20194openrouter_slug = "google/gemma-4-26b-a4b-it"
20195zdr = false
20196hipaa = false
20197
20198[planned_models."google/gemma-4-26b-a4b-it".pricing]
20199prompt = "0.0000000665"
20200cached_prompt = "0.0000000466"
20201completion = "0.0000003230"
20202"#;
20203
20204 #[test]
20205 fn gateway_registry_generates_the_staged_active_shape() {
20206 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
20207 let caps = ModelCaps {
20208 tools_branch: true,
20209 qwen_think: true,
20210 think_switch: true,
20211 chat_ok: true,
20212 context_length: 262144,
20213 tokenizer: "qwen2".into(),
20214 instruct_type: Some("chatml".into()),
20215 ..Default::default()
20216 };
20217 let q35_entry = model_entry_openrouter(
20218 "qwen/qwen3.6-35b-a3b",
20219 Some(&caps),
20220 metadata.get("qwen/qwen3.6-35b-a3b"),
20221 );
20222 assert_eq!(q35_entry["created"], 1777260255u64);
20223 assert_eq!(q35_entry["quantization"], "int4");
20224 assert_eq!(q35_entry["is_ready"], true);
20225 assert_eq!(
20226 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
20227 262144
20228 );
20229 assert_eq!(
20230 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
20231 262144
20232 );
20233 assert_eq!(
20234 q35_entry["output_modalities"][0]["max_length"]["value"],
20235 262144
20236 );
20237 let prices = q35_entry["input_modalities"][0]["pricing"]
20238 .as_array()
20239 .unwrap();
20240 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
20241 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
20242 assert_eq!(
20246 q35_entry["input_modalities"][0]["capacity"][0]["value"],
20247 780000
20248 );
20249 assert_eq!(
20250 q35_entry["input_modalities"][0]["capacity"][1]["value"],
20251 310000
20252 );
20253 assert_eq!(
20254 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
20255 262144
20256 );
20257 assert_eq!(
20258 q35_entry["output_modalities"][0]["capacity"][0]["value"],
20259 9600
20260 );
20261 assert_eq!(
20262 q35_entry["output_modalities"][0]["capacity"][1]["value"],
20263 16
20264 );
20265 assert_eq!(
20266 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
20267 "0.0000009025"
20268 );
20269 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
20271
20272 assert_eq!(
20273 metadata.len(),
20274 1,
20275 "planned models must never enter the active map"
20276 );
20277 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
20278 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
20279 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
20280
20281 let openmodels = model_entry_openmodels(
20282 "qwen/qwen3.6-35b-a3b",
20283 Some(&caps),
20284 metadata.get("qwen/qwen3.6-35b-a3b"),
20285 )
20286 .unwrap();
20287 assert_eq!(openmodels["currency"], "USD");
20288 assert_eq!(openmodels["max_output_length"], 262144);
20289 assert_eq!(openmodels["is_ready"], true);
20290 assert_eq!(openmodels["is_free"], false);
20291 assert_eq!(openmodels["discount_to_user"], 0.0);
20292 }
20293
20294 #[test]
20295 fn gateway_registry_limits_are_live_request_limits() {
20296 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
20297 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
20298 let caps = ModelCaps {
20299 context_length: 262_144,
20300 ..Default::default()
20301 };
20302 let build = |value: serde_json::Value| {
20303 let req: CompletionReq = serde_json::from_value(value).unwrap();
20304 let (tx, _rx) = worker::event_channel();
20305 build_request(&req, tx, lanes::Lane::Interactive, None)
20306 };
20307
20308 let mut omitted = build(json!({
20309 "model": "qwen/qwen3.6-35b-a3b",
20310 "prompt_ids": [1, 2, 3]
20311 }));
20312 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
20313 assert_eq!(omitted.params.max_new, 8_192);
20314 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
20315
20316 let mut field_top = build(json!({
20317 "model": "qwen/qwen3.6-35b-a3b",
20318 "prompt_ids": [1],
20319 "max_tokens": 262144
20320 }));
20321 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
20322 assert_eq!(field_top.params.max_new, 262_144);
20323 assert_eq!(
20324 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
20325 262_044,
20326 "the field-top output request is accepted but bounded by remaining trained context",
20327 );
20328
20329 let mut too_much_output = build(json!({
20330 "model": "qwen/qwen3.6-35b-a3b",
20331 "prompt_ids": [1],
20332 "max_tokens": 262145
20333 }));
20334 let (message, param) =
20335 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
20336 .unwrap_err();
20337 assert_eq!(param, "max_tokens");
20338 assert!(message.contains("262145"));
20339
20340 let mut oversized_allocation = build(json!({
20341 "model": "qwen/qwen3.6-35b-a3b",
20342 "prompt_ids": [1],
20343 "max_tokens": 1,
20344 "max_ctx": 262145
20345 }));
20346 let (_, param) =
20347 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
20348 .unwrap_err();
20349 assert_eq!(param, "max_ctx");
20350 }
20351
20352 #[test]
20353 fn planned_registry_entries_are_validated_but_never_activated() {
20354 let parsed = OpenRouterMetadataFile::from_toml(
20355 r#"
20356[planned_models.future]
20357max_output_length = 262144
20358default_output_length = 8192
20359
20360[planned_models.future.pricing]
20361prompt = "0.0000001"
20362"#,
20363 )
20364 .unwrap();
20365 assert!(parsed.is_empty());
20366
20367 let error = OpenRouterMetadataFile::from_toml(
20368 r#"
20369[planned_models.future]
20370default_output_length = 8192
20371"#,
20372 )
20373 .unwrap_err();
20374 assert!(error.contains("requires max_output_length"));
20375 }
20376
20377 #[test]
20382 fn every_catalog_feed_honours_the_declared_surface() {
20383 let metadata = OpenRouterMetadataFile::from_toml(
20384 r#"
20385[models."qwen/qwen3-embedding-8b"]
20386surface = "embedding"
20387created = 1787961600
20388max_output_length = 1
20389is_ready = true
20390is_free = false
20391discount_to_user = 0.0
20392
20393[models."qwen/qwen3-embedding-8b".pricing]
20394prompt = "0.00000001"
20395cached_prompt = "0.0"
20396completion = "0.0"
20397
20398[models."main"]
20399created = 1787443200
20400max_output_length = 32768
20401is_ready = true
20402is_free = false
20403discount_to_user = 0.0
20404
20405[models."main".pricing]
20406prompt = "0.00000025"
20407cached_prompt = "0.00000009"
20408completion = "0.0000012"
20409"#,
20410 )
20411 .unwrap();
20412 let caps = ModelCaps {
20413 tools_branch: true,
20414 qwen_think: true,
20415 think_switch: true,
20420 chat_ok: true,
20421 context_length: 32768,
20422 ..Default::default()
20423 };
20424 let embed = metadata.get("qwen/qwen3-embedding-8b");
20425 let chat = metadata.get("main");
20426
20427 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
20429 let out = &or["output_modalities"][0];
20430 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
20431 assert!(
20432 out.get("streaming").is_none(),
20433 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
20434 );
20435 let params = &out["supported_parameters"];
20440 assert_eq!(
20441 params.as_object().map(|o| o.len()),
20442 Some(0),
20443 "no completion parameter belongs on an embedder row: {params}"
20444 );
20445 for field in [
20446 "tools",
20447 "tool_choice",
20448 "reasoning",
20449 "max_tokens",
20450 "json_mode",
20451 "structured_outputs",
20452 "stop",
20453 "temperature",
20454 "seed",
20455 ] {
20456 assert!(params[field].is_null(), "{field} leaked onto an embedder");
20457 }
20458 assert!(
20459 out["max_length"].is_null(),
20460 "a surface emitting no completion tokens advertises no ceiling: {out}"
20461 );
20462
20463 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
20465 .expect("openmodels entry builds");
20466 assert_eq!(om["output_modalities"], json!(["embeddings"]));
20467 let features = om["supported_features"].as_array().unwrap();
20468 assert!(
20469 !features
20470 .iter()
20471 .any(|f| f == "tool_calling" || f == "reasoning"),
20472 "chat-only features leaked onto an embedder: {features:?}"
20473 );
20474
20475 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
20477 assert_eq!(v1["type"], "embedding");
20478 assert_eq!(v1["capabilities"]["tools"], false);
20479
20480 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
20482 let out_chat = &or_chat["output_modalities"][0];
20483 assert_eq!(out_chat["type"], "text");
20484 assert_eq!(out_chat["streaming"], true);
20485 assert!(!out_chat["supported_parameters"]["tools"].is_null());
20486 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
20487 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
20488 assert_eq!(out_chat["max_length"]["value"], 32768u64);
20489 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
20490 assert_eq!(om_chat["output_modalities"], json!(["text"]));
20491 assert!(
20492 om_chat["supported_features"]
20493 .as_array()
20494 .unwrap()
20495 .iter()
20496 .any(|f| f == "tool_calling")
20497 );
20498 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
20499 }
20500
20501 #[test]
20508 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
20509 let raw = std::fs::read_to_string(concat!(
20510 env!("CARGO_MANIFEST_DIR"),
20511 "/../../research/gateway-20260812/raw/sources/",
20512 "openrouter-provider-schema-v2.4-20260812.json"
20513 ))
20514 .expect("vendored Provider Monitor 2.4 schema is in-tree");
20515 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
20516 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
20517 .as_array()
20518 .expect("OutputModality is a oneOf");
20519
20520 let metadata = OpenRouterMetadataFile::from_toml(
20521 r#"
20522[models."embed"]
20523surface = "embedding"
20524created = 1787961600
20525max_output_length = 1
20526is_ready = true
20527is_free = false
20528discount_to_user = 0.0
20529
20530[models."embed".pricing]
20531prompt = "0.00000001"
20532cached_prompt = "0.0"
20533completion = "0.0"
20534
20535[models."rr"]
20536surface = "rerank"
20537created = 1787961600
20538max_output_length = 1
20539is_ready = true
20540is_free = false
20541discount_to_user = 0.0
20542
20543[models."rr".pricing]
20544prompt = "0.00000003"
20545cached_prompt = "0.0"
20546completion = "0.0"
20547
20548[models."chatty"]
20549created = 1787443200
20550max_output_length = 32768
20551is_ready = true
20552is_free = false
20553discount_to_user = 0.0
20554
20555[models."chatty".pricing]
20556prompt = "0.00000025"
20557cached_prompt = "0.00000009"
20558completion = "0.0000012"
20559"#,
20560 )
20561 .unwrap();
20562 let caps = ModelCaps {
20563 tools_branch: true,
20564 qwen_think: true,
20565 chat_ok: true,
20566 context_length: 32768,
20567 ..Default::default()
20568 };
20569
20570 for (alias, want_type) in [
20571 ("embed", "embeddings"),
20572 ("rr", "rerank"),
20573 ("chatty", "text"),
20574 ] {
20575 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
20576 let modality = &row["output_modalities"][0];
20577 assert_eq!(modality["type"], want_type, "{alias}: {row}");
20578
20579 let branch = branches
20581 .iter()
20582 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
20583 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
20584 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
20585 .as_object()
20586 .expect("branch properties")
20587 .keys()
20588 .map(String::as_str)
20589 .collect();
20590 for key in modality.as_object().expect("modality object").keys() {
20591 assert!(
20592 allowed.contains(key.as_str()),
20593 "{alias}: {key:?} is not a property of the {want_type:?} branch \
20594 (additionalProperties:false); allowed = {allowed:?}"
20595 );
20596 }
20597 for req in branch["required"].as_array().into_iter().flatten() {
20598 let req = req.as_str().expect("required entry is a string");
20599 assert!(
20600 modality.get(req).is_some(),
20601 "{alias}: required property {req:?} missing from the {want_type:?} branch"
20602 );
20603 }
20604 }
20605 }
20606
20607 #[test]
20608 fn openrouter_models_entry_omits_undeclared_optional_fields() {
20609 let entry = model_entry_openrouter("minimal", None, None);
20610 let object = entry.as_object().unwrap();
20611 for field in [
20612 "hugging_face_id",
20613 "created",
20614 "quantization",
20615 "tokenizer",
20616 "description",
20617 "pricing",
20618 "capacity",
20619 "is_ready",
20620 "is_free",
20621 "discount_to_user",
20622 "openrouter",
20623 "datacenters",
20624 "compliance",
20625 ] {
20626 assert!(
20627 !object.contains_key(field),
20628 "optional field {field} must be absent, not null"
20629 );
20630 }
20631 assert_eq!(entry["schema_version"], "2.4");
20632 assert_eq!(entry["input_modalities"][0]["type"], "text");
20633 assert!(
20634 entry["input_modalities"][0]
20635 .get("supported_inputs")
20636 .is_none()
20637 );
20638 assert!(entry["input_modalities"][0].get("pricing").is_none());
20639 assert!(entry["input_modalities"][0].get("capacity").is_none());
20640 assert_eq!(entry["output_modalities"][0]["type"], "text");
20641 assert_eq!(entry["output_modalities"][0]["streaming"], true);
20642 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
20643 assert!(entry["output_modalities"][0].get("max_length").is_none());
20644 assert!(entry["output_modalities"][0].get("pricing").is_none());
20645 assert!(entry["output_modalities"][0].get("capacity").is_none());
20646 }
20647
20648 #[test]
20649 fn openmodels_entry_serializes_standard_provider_shape() {
20650 let metadata = OpenRouterMetadataFile::from_toml(
20651 r#"
20652[models."qwen/qwen3.6-27b"]
20653created = 1786032000
20654max_output_length = 16384
20655is_ready = true
20656is_free = false
20657discount_to_user = 0.05
20658
20659[models."qwen/qwen3.6-27b".pricing]
20660prompt = "0.000000291"
20661cached_prompt = "0.000000291"
20662completion = "0.000002763"
20663request = "0"
20664"#,
20665 )
20666 .unwrap();
20667 let caps = ModelCaps {
20668 tools_branch: true,
20669 qwen_think: true,
20670 chat_ok: true,
20671 context_length: 262144,
20672 ..Default::default()
20673 };
20674 let entry = model_entry_openmodels(
20675 "qwen/qwen3.6-27b",
20676 Some(&caps),
20677 metadata.get("qwen/qwen3.6-27b"),
20678 )
20679 .unwrap();
20680
20681 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
20682 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
20683 assert_eq!(entry["created"], 1786032000u64);
20684 assert_eq!(entry["input_modalities"], json!(["text"]));
20685 assert_eq!(entry["output_modalities"], json!(["text"]));
20686 assert_eq!(entry["context_length"], 262144u64);
20687 assert_eq!(entry["max_output_length"], 16384u64);
20688 assert_eq!(entry["currency"], "USD");
20689 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
20690 assert_eq!(entry["pricing"]["completion"], "0.000002763");
20691 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
20692 assert_eq!(entry["pricing"]["request"], "0");
20693 assert_eq!(
20694 entry["supported_features"],
20695 json!(["tool_calling", "reasoning"])
20696 );
20697 assert_eq!(entry["is_ready"], true);
20698 assert_eq!(entry["is_free"], false);
20699 assert_eq!(entry["discount_to_user"], 0.05);
20700 assert!(entry.get("schema_version").is_none());
20701 assert!(entry.get("quantization").is_none());
20702 }
20703
20704 #[test]
20705 fn openmodels_entry_rejects_missing_operator_metadata() {
20706 let caps = ModelCaps {
20707 context_length: 262144,
20708 ..Default::default()
20709 };
20710 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
20711 assert_eq!(
20712 error,
20713 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
20714 );
20715 }
20716
20717 #[tokio::test]
20718 async fn blocking_response_excludes_stop_text_across_token_events() {
20719 let (tx, rx) = worker::event_channel();
20720 tx.send(Event::Token {
20721 id: 1,
20722 text: "answer\nPro".into(),
20723 })
20724 .unwrap();
20725 tx.send(Event::Token {
20726 id: 2,
20727 text: "blem: leaked prompt".into(),
20728 })
20729 .unwrap();
20730 tx.send(Event::Done {
20731 stop_reason: "Callback".into(),
20732 n_tokens: 2,
20733 n_prompt: 8,
20734 n_cached: 0,
20735 elapsed_s: 0.5,
20736 spec: None,
20737 })
20738 .unwrap();
20739 drop(tx);
20740 let response = blocking_response(
20741 rx,
20742 "plain_quant".into(),
20743 false,
20744 vec!["Problem:".into()],
20745 None,
20746 Envelope::new(false),
20747 )
20748 .await;
20749 assert_eq!(response.status(), StatusCode::OK);
20750 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
20751 .await
20752 .unwrap();
20753 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20754 assert_eq!(payload["text"], "answer\n");
20755 assert_eq!(payload["stop_reason"], "Callback");
20756 }
20757
20758 #[test]
20762 fn step_walker_expansion_and_separator_law() {
20763 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
20765 let uri = format!("data:image/png;base64,{PNG64}");
20766 let content = serde_json::json!([
20767 {"type": "text", "text": "look at"},
20768 {"type": "text", "text": "this:"},
20769 {"type": "image_url", "image_url": {"url": uri}},
20770 {"type": "text", "text": "what is it?"},
20771 ]);
20772 let mut pending: Vec<PendingStepImage> = Vec::new();
20773 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
20774 let mut expansion = String::from("<im_start>");
20775 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
20776 expansion.push_str("<im_patch>");
20777 }
20778 expansion.push_str("<im_end>");
20779 assert_eq!(out, format!("look at this:{expansion}what is it?"));
20782 assert_eq!(pending.len(), 1);
20783 assert_eq!(pending[0].plan.n_tiles, 0);
20784 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
20785
20786 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
20788 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
20789 let http = serde_json::json!([
20790 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
20791 ]);
20792 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
20793 }
20794}
20795
20796#[cfg(test)]
20802mod build_identity_tests {
20803 use super::{BUILD_GIT_SHA, BUILD_ID_NOTE, BUILD_ID_SRC, SYSTEM_FINGERPRINT, build_id};
20804
20805 #[test]
20807 fn baked_fingerprint_is_real_and_well_formed() {
20808 assert!(!SYSTEM_FINGERPRINT.is_empty());
20809 assert_ne!(SYSTEM_FINGERPRINT, "memra-unknown");
20810 assert!(
20811 !SYSTEM_FINGERPRINT.contains("unknown"),
20812 "fingerprint {SYSTEM_FINGERPRINT:?} still carries the degraded literal"
20813 );
20814 assert!(
20815 build_id::fingerprint_is_well_formed(SYSTEM_FINGERPRINT),
20816 "fingerprint {SYSTEM_FINGERPRINT:?} is not memra-<version>-<12 hex>"
20817 );
20818 assert!(
20821 SYSTEM_FINGERPRINT.starts_with(concat!("memra-", env!("CARGO_PKG_VERSION"), "-")),
20822 "fingerprint {SYSTEM_FINGERPRINT:?} does not name this crate version"
20823 );
20824 }
20825
20826 #[test]
20829 fn the_shape_check_rejects_what_shipped_to_prod() {
20830 assert!(!build_id::fingerprint_is_well_formed("memra-unknown"));
20831 assert!(!build_id::fingerprint_is_well_formed(
20832 "memra-0.123.0-unknown"
20833 ));
20834 let old_form = format!("memra-{}", "0".repeat(12));
20839 assert!(!build_id::fingerprint_is_well_formed(&old_form));
20840 assert!(!build_id::fingerprint_is_well_formed(""));
20841 assert!(!build_id::fingerprint_is_well_formed("memra-"));
20842 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-"));
20843 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-abc"));
20845 assert!(!build_id::fingerprint_is_well_formed(
20846 "memra-0.123.0-ABCDEF012345"
20847 ));
20848 assert!(!build_id::fingerprint_is_well_formed(
20849 "memra-0.123.0-zzzzzzzzzzzz"
20850 ));
20851 assert!(build_id::fingerprint_is_well_formed(
20853 "memra-0.123.0-4b1f9c02d7a3"
20854 ));
20855 }
20856
20857 #[test]
20865 fn build_id_is_rederivable_from_the_source_tree() {
20866 let root = build_id::workspace_root(env!("CARGO_MANIFEST_DIR"));
20867 let scan = root.as_deref().and_then(build_id::content_id);
20868 match scan {
20869 Some(scan) => {
20870 assert_eq!(
20871 BUILD_ID_SRC,
20872 build_id::BUILD_ID_SRC_TREE,
20873 "the source tree is readable, so the baked id must come from it"
20874 );
20875 assert!(BUILD_ID_NOTE.is_empty(), "note set on a non-degraded build");
20876 let expected =
20877 format!(concat!("memra-", env!("CARGO_PKG_VERSION"), "-{}"), scan.id);
20878 assert_eq!(
20879 SYSTEM_FINGERPRINT,
20880 expected,
20881 "baked fingerprint disagrees with a re-derivation over {} files: the id \
20882 is not a pure function of the source tree, or the build script did not \
20883 re-run after an edit",
20884 scan.files.len()
20885 );
20886 assert!(scan.files.len() > 100, "suspiciously small hashed file set");
20887 }
20888 None => {
20889 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
20893 assert!(
20894 !BUILD_ID_NOTE.is_empty(),
20895 "a degraded build must state its reason so the boot WARN can print it"
20896 );
20897 }
20898 }
20899 }
20900
20901 #[test]
20904 fn identity_is_independent_of_git_history() {
20905 let id = SYSTEM_FINGERPRINT.rsplit_once('-').unwrap().1;
20906 assert_ne!(
20907 id, BUILD_GIT_SHA,
20908 "the content id equals the git sha; the identity must not be history, it has to \
20909 survive a rewrite that changes every commit"
20910 );
20911 assert!(
20912 !SYSTEM_FINGERPRINT.contains(BUILD_GIT_SHA),
20913 "the git sha leaked into the customer-visible fingerprint {SYSTEM_FINGERPRINT:?}"
20914 );
20915 assert!(!BUILD_GIT_SHA.is_empty());
20918 }
20919
20920 #[test]
20923 fn content_digest_is_deterministic_and_change_sensitive() {
20924 let a = build_id::degraded_build_id("memra-server", "0.123.0");
20925 let b = build_id::degraded_build_id("memra-server", "0.123.0");
20926 assert_eq!(a, b, "the digest is not deterministic");
20927 assert_eq!(a.len(), build_id::BUILD_ID_HEX);
20928 assert!(
20929 a.chars()
20930 .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
20931 );
20932 assert_ne!(a, build_id::degraded_build_id("memra-server", "0.123.1"));
20933 assert_ne!(a, build_id::degraded_build_id("memra-serve", "r0.123.0"));
20934 assert_eq!(build_id::render_build_id(0).len(), build_id::BUILD_ID_HEX);
20936 assert_eq!(
20937 build_id::render_build_id(0),
20938 "0".repeat(build_id::BUILD_ID_HEX)
20939 );
20940 }
20941
20942 #[test]
20945 fn two_scans_of_one_tree_agree() {
20946 let Some(root) = build_id::workspace_root(env!("CARGO_MANIFEST_DIR")) else {
20947 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
20948 return;
20949 };
20950 let first = build_id::content_id(&root).expect("first scan");
20951 let second = build_id::content_id(&root).expect("second scan");
20952 assert_eq!(first.id, second.id);
20953 assert_eq!(first.files.len(), second.files.len());
20954 }
20955}
20956
20957#[cfg(test)]
20963mod vision_placement_gate_tests {
20964 use super::vision_media_admissible;
20965
20966 #[test]
20967 fn a_media_part_is_admitted_only_when_the_placement_admits() {
20968 assert_eq!(vision_media_admissible(true, "image"), Ok(()));
20969 assert_eq!(vision_media_admissible(true, "video"), Ok(()));
20970 let err = vision_media_admissible(false, "image").unwrap_err();
20971 assert!(
20972 err.starts_with("image input is not enabled on this deployment"),
20973 "same named refusal the armed-off path gives, so clients see one contract: {err}"
20974 );
20975 assert!(
20976 err.contains("placement"),
20977 "the refusal names its cause: {err}"
20978 );
20979 let err = vision_media_admissible(false, "video").unwrap_err();
20980 assert!(
20981 err.starts_with("video input is not enabled on this deployment"),
20982 "{err}"
20983 );
20984 }
20985
20986 fn live_src() -> String {
20987 let src: String = include_str!("lib.rs")
20988 .lines()
20989 .map(|l| match l.find("//") {
20990 Some(i) => &l[..i],
20991 None => l,
20992 })
20993 .collect::<Vec<_>>()
20994 .join("\n");
20995 let end = src
20996 .find("\nmod vision_placement_gate_tests")
20997 .expect("this test module exists");
20998 src[..end].to_string()
20999 }
21000
21001 fn item_body<'a>(live: &'a str, head: &str) -> &'a str {
21003 let start = live
21004 .find(head)
21005 .unwrap_or_else(|| panic!("{head} not found — did it get renamed?"));
21006 let body = &live[start..];
21007 let end = body.find("\n}\n").expect("item body closes");
21008 &body[..end]
21009 }
21010
21011 fn head_of(s: &str, n: usize) -> &str {
21013 match s.char_indices().nth(n) {
21014 Some((i, _)) => &s[..i],
21015 None => s,
21016 }
21017 }
21018
21019 #[test]
21024 fn no_family_switch_reads_the_placement_decision() {
21025 let live = live_src();
21026 for switch in [
21027 "fn vision_enabled()",
21028 "fn gemma_vision_enabled()",
21029 "fn step_vision_enabled()",
21030 ] {
21031 let body = item_body(&live, switch);
21032 assert!(
21033 !body.contains("vision_placement_serving")
21034 && !body.contains("vision_placement_admits"),
21035 "{switch} routes text rendering; it must stay keyed on the operator knobs alone"
21036 );
21037 }
21038 let walker = item_body(&live, "fn content_to_text_vision(");
21039 assert!(
21040 walker.contains(
21041 "if step_vision_enabled() {\n return content_to_text_vision_step(v, step_images);"
21042 ),
21043 "the step walker dispatch is keyed on the armed switch alone"
21044 );
21045 }
21046
21047 #[test]
21050 fn every_media_accepting_arm_passes_the_placement_gate() {
21051 let live = live_src();
21052 let step = item_body(&live, "fn content_to_text_vision_step(");
21053 let arm = step
21054 .split("Some(\"image_url\") => {")
21055 .nth(1)
21056 .expect("the step walker has an image arm");
21057 assert!(
21058 head_of(arm, 120).contains("vision_placement_admits(\"image\")?;"),
21059 "the step image arm must pass the placement gate first: {}",
21060 head_of(arm, 120)
21061 );
21062 let walker = item_body(&live, "fn content_to_text_vision(");
21063 for (head, kind) in [
21064 (
21065 "Some(\"image_url\") if gemma_vision_enabled() => {",
21066 "image",
21067 ),
21068 ("Some(\"image_url\") => {", "image"),
21069 ("Some(\"video_url\") => {", "video"),
21070 ] {
21071 let arm = walker
21072 .split(head)
21073 .nth(1)
21074 .unwrap_or_else(|| panic!("{head} is not an arm of the walker"));
21075 let window = head_of(arm, 400);
21076 assert!(
21077 window.contains(&format!("vision_placement_admits(\"{kind}\")?;")),
21078 "{head} must pass the placement gate before planning anything: {window}"
21079 );
21080 }
21081 assert!(live.contains("GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)"));
21085 let gate = item_body(&live, "fn vision_placement_admits(");
21087 assert!(gate.contains("vision_media_admissible(vision_placement_serving(), kind)"));
21088 }
21089}