1pub mod auth;
57pub(crate) mod constrained;
58pub(crate) mod darklane;
62pub(crate) mod health;
66pub(crate) mod lanes {
67 pub use memra_lanes::*;
68}
69mod admit_predict;
74mod affinity;
81mod anthropic;
87#[allow(dead_code)] mod build_id;
92mod dsv4_serve;
93mod embed_api;
94pub mod metering;
101mod responses_api;
102mod surfaces;
103mod toolcall;
104mod ttft;
105mod worker;
106
107use std::collections::HashMap;
108use std::net::{SocketAddr, ToSocketAddrs};
109use std::sync::mpsc::Sender;
110use std::sync::{Arc, Mutex};
111
112use axum::{
113 Extension, Json, Router,
114 body::Body,
115 extract::{DefaultBodyLimit, FromRequest, Query, Request as AxumRequest, State},
116 http::{
117 HeaderMap, StatusCode,
118 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
119 },
120 middleware::{self, Next},
121 response::{
122 IntoResponse, Response,
123 sse::{Event as SseEvent, Sse},
124 },
125 routing::{get, post},
126};
127use futures_core::Stream as _;
128use serde::de::DeserializeOwned;
129use serde::{Deserialize, Serialize};
130use serde_json::json;
131use tower::ServiceExt as _;
132
133use memra_engine::decode::GenParams;
134use memra_engine::sampler::SamplerConfig;
135use memra_tokenizer::{
136 Tokenizer,
137 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
138};
139use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
140use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
141
142const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
166const MAX_BODY_ADMISSIONS: usize = 4;
167const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
168#[allow(clippy::identity_op)] const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
173const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
174const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
175const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
176const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
177const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
178const MAX_STOP_SEQUENCES: usize = 16;
179const MAX_STOP_SEQUENCE_BYTES: usize = 1_024;
180const MAX_STOP_SEQUENCES_BYTES: usize = 4 * 1_024;
181const MAX_CLIENT_IDENTIFIER_BYTES: usize = 256;
182const MAX_HTTP_CONNECTIONS: usize = 1_024;
183const MAX_HTTP2_STREAMS_PER_CONNECTION: u32 = 128;
184const HTTP1_HEADER_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
185const HTTP_CONNECTION_MAX_LIFETIME: std::time::Duration = std::time::Duration::from_secs(300);
186
187fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
188 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
189 SEMAPHORE
190 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
191 .clone()
192}
193
194fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
195 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
196 SEMAPHORE
197 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
198 .clone()
199}
200
201#[derive(Clone)]
202pub(crate) struct BodyAdmissionGuard {
203 permit: Arc<Mutex<Option<tokio::sync::OwnedSemaphorePermit>>>,
204}
205
206impl BodyAdmissionGuard {
207 fn new(permit: tokio::sync::OwnedSemaphorePermit) -> Self {
208 Self {
209 permit: Arc::new(Mutex::new(Some(permit))),
210 }
211 }
212
213 pub(crate) fn release(&self) {
214 if let Ok(mut permit) = self.permit.lock() {
215 permit.take();
216 }
217 }
218}
219
220pub(crate) struct BodyAdmissionLease(Option<BodyAdmissionGuard>);
221
222impl BodyAdmissionLease {
223 fn release(&mut self) {
224 if let Some(admission) = self.0.take() {
225 admission.release();
226 }
227 }
228
229 pub(crate) fn guard(&self) -> Option<&BodyAdmissionGuard> {
230 self.0.as_ref()
231 }
232}
233
234impl Drop for BodyAdmissionLease {
235 fn drop(&mut self) {
236 self.release();
237 }
238}
239
240pub(crate) struct AdmittedJson<T>(pub(crate) T, pub(crate) BodyAdmissionLease);
241
242#[axum::async_trait]
243impl<S, T> FromRequest<S> for AdmittedJson<T>
244where
245 S: Send + Sync,
246 T: DeserializeOwned,
247{
248 type Rejection = axum::extract::rejection::JsonRejection;
249
250 async fn from_request(req: AxumRequest, state: &S) -> Result<Self, Self::Rejection> {
251 let admission = req.extensions().get::<BodyAdmissionGuard>().cloned();
252 let parsed = Json::<T>::from_request(req, state).await;
253 parsed.map(|Json(value)| Self(value, BodyAdmissionLease(admission)))
254 }
255}
256
257fn declared_body_length(req: &AxumRequest) -> Option<usize> {
258 req.headers()
259 .get(CONTENT_LENGTH)
260 .and_then(|value| value.to_str().ok())
261 .and_then(|value| value.parse().ok())
262}
263
264fn body_requires_admission(req: &AxumRequest) -> bool {
265 if req.headers().contains_key(TRANSFER_ENCODING) {
269 return true;
270 }
271 declared_body_length(req).is_none_or(|length| length > BODY_ADMISSION_BYPASS_BYTES)
272}
273
274fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
278 let Some(length) = declared_body_length(req) else {
279 return BODY_READ_TIMEOUT;
280 };
281 let bytes = length as u64;
282 let extra_seconds =
283 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
284 let seconds = BODY_READ_TIMEOUT
285 .as_secs()
286 .saturating_add(extra_seconds)
287 .min(BODY_READ_TIMEOUT_MAX.as_secs());
288 std::time::Duration::from_secs(seconds)
289}
290
291async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
296 let resp = next.run(req).await;
297 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
298 return resp;
299 }
300 error_response_coded(
301 StatusCode::PAYLOAD_TOO_LARGE,
302 &format!(
303 "request body exceeds the {} MiB limit",
304 MAX_BODY_BYTES / (1024 * 1024)
305 ),
306 "invalid_request_error",
307 None,
308 Some("request_too_large"),
309 )
310}
311
312fn apply_body_limit(app: Router) -> Router {
315 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
316 .layer(middleware::from_fn(shape_payload_too_large))
317}
318
319fn protected_inference_path(path: &str) -> bool {
320 matches!(
321 path,
322 "/v1/auth/check"
323 | "/v1/completions"
324 | "/v1/chat/completions"
325 | "/v1/messages"
326 | "/v1/responses"
327 | "/v1/embeddings"
328 | "/v1/rerank"
329 )
330}
331
332async fn shape_inference_early_response(path: &str, response: Response) -> Response {
336 let request_id = Envelope::new(path != "/v1/completions");
337 if path == "/v1/messages" {
338 anthropic::with_anthropic_request_id(
339 &request_id.id,
340 anthropic::reshape_error(response, &request_id.id).await,
341 )
342 } else {
343 with_request_id(&request_id.id, response)
344 }
345}
346
347async fn authenticate_inference_before_body(
352 State(st): State<AppState>,
353 mut req: AxumRequest,
354 next: Next,
355) -> Response {
356 if !protected_inference_path(req.uri().path()) {
357 return next.run(req).await;
358 }
359 let path = req.uri().path().to_string();
360 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
364 return shape_inference_early_response(
365 &path,
366 error_response_coded(
367 StatusCode::PAYLOAD_TOO_LARGE,
368 &format!(
369 "request body exceeds the {} MiB limit",
370 MAX_BODY_BYTES / (1024 * 1024)
371 ),
372 "invalid_request_error",
373 None,
374 Some("request_too_large"),
375 ),
376 )
377 .await;
378 }
379 let headers = req.headers();
380 let bearer = bearer_token(headers);
381 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
382 let api_key = headers
383 .get("x-api-key")
384 .and_then(|value| value.to_str().ok());
385 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
386 } else {
387 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
388 };
389 if let Err(why) = auth {
390 return shape_inference_early_response(&path, authentication_error(why)).await;
391 }
392 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
400 let body_admission = if body_requires_admission(&req) {
401 body_admission_semaphore()
402 } else {
403 small_body_admission_semaphore()
404 };
405 let body_permit = match body_admission.try_acquire_owned() {
406 Ok(permit) => permit,
407 Err(tokio::sync::TryAcquireError::Closed) => {
408 let response = retry_contract_response(
409 error_response_coded(
410 StatusCode::SERVICE_UNAVAILABLE,
411 "request body admission is unavailable",
412 "server_error",
413 None,
414 Some("body_admission_unavailable"),
415 ),
416 Some(BODY_ADMISSION_RETRY_AFTER_S),
417 );
418 return shape_inference_early_response(&path, response).await;
419 }
420 Err(tokio::sync::TryAcquireError::NoPermits) => {
421 let response = retry_contract_response(
422 error_response_coded(
423 StatusCode::TOO_MANY_REQUESTS,
424 "request body admission is busy",
425 "rate_limit_error",
426 None,
427 Some("body_admission_busy"),
428 ),
429 Some(BODY_ADMISSION_RETRY_AFTER_S),
430 );
431 return shape_inference_early_response(&path, response).await;
432 }
433 };
434 let body_admission_guard = BodyAdmissionGuard::new(body_permit);
439 req.extensions_mut().insert(body_admission_guard.clone());
440 let body = std::mem::replace(req.body_mut(), Body::empty());
441 let mut body = Box::pin(body.into_data_stream());
442 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
443 let body_timed_out_flag = body_timed_out.clone();
444 let guarded_body = async_stream::stream! {
445 loop {
446 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
447 if remaining.is_zero() {
448 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
449 yield Err(std::io::Error::new(
450 std::io::ErrorKind::TimedOut,
451 "request body read deadline exceeded",
452 ));
453 break;
454 }
455 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
456 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
457 Ok(frame) => frame,
458 Err(_) => {
459 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
460 yield Err(std::io::Error::new(
461 std::io::ErrorKind::TimedOut,
462 "request body idle timeout exceeded",
463 ));
464 break;
465 }
466 };
467 match frame {
468 Some(Ok(bytes)) => yield Ok(bytes),
469 Some(Err(error)) => {
470 yield Err(std::io::Error::other(error.to_string()));
471 break;
472 }
473 None => break,
474 }
475 }
476 };
477 *req.body_mut() = Body::from_stream(guarded_body);
478 let response = next.run(req).await;
479 body_admission_guard.release();
480 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
481 let request_id = Envelope::new(path != "/v1/completions");
482 let timeout = error_response_coded(
483 StatusCode::REQUEST_TIMEOUT,
484 "request body read timed out",
485 "invalid_request_error",
486 None,
487 Some("request_body_timeout"),
488 );
489 return if path == "/v1/messages" {
490 anthropic::with_anthropic_request_id(
491 &request_id.id,
492 anthropic::reshape_error(timeout, &request_id.id).await,
493 )
494 } else {
495 with_request_id(&request_id.id, timeout)
496 };
497 }
498 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
499 let request_id = Envelope::new(true);
500 return anthropic::with_anthropic_request_id(
501 &request_id.id,
502 anthropic::reshape_error(response, &request_id.id).await,
503 );
504 }
505 response
506}
507
508#[cfg(test)]
509mod body_limit_tests {
510 use super::*;
511
512 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
513
514 fn test_app() -> Router {
518 let app = Router::new()
519 .route(
520 "/bytes",
521 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
522 )
523 .route(
524 "/json",
525 post(
526 |AdmittedJson(v, _admission): AdmittedJson<serde_json::Value>| async move {
527 v["pad"].as_str().unwrap_or("").len().to_string()
528 },
529 ),
530 );
531 apply_body_limit(app)
532 }
533
534 fn streamed_body(chunks: usize) -> Body {
535 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
538 Body::from_stream(async_stream::stream! {
539 for _ in 0..chunks {
540 yield Ok::<_, std::io::Error>(chunk.clone());
541 }
542 })
543 }
544
545 #[tokio::test]
546 async fn bodies_past_the_old_2mib_default_are_accepted() {
547 for (path, body) in [
550 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
551 (
552 "/json",
553 Body::from(
554 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
555 ),
556 ),
557 ] {
558 let resp = test_app()
559 .oneshot(
560 axum::http::Request::post(path)
561 .header(CONTENT_TYPE, "application/json")
562 .body(body)
563 .unwrap(),
564 )
565 .await
566 .unwrap();
567 assert_eq!(resp.status(), StatusCode::OK, "{path}");
568 }
569 }
570
571 #[tokio::test]
572 async fn body_at_exactly_the_limit_is_accepted() {
573 let resp = test_app()
574 .oneshot(
575 axum::http::Request::post("/bytes")
576 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
577 .unwrap(),
578 )
579 .await
580 .unwrap();
581 assert_eq!(resp.status(), StatusCode::OK);
582 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
583 .await
584 .unwrap();
585 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
586 }
587
588 #[tokio::test]
589 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
590 for path in ["/bytes", "/json"] {
594 let resp = test_app()
595 .oneshot(
596 axum::http::Request::post(path)
597 .header(CONTENT_TYPE, "application/json")
598 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
599 .unwrap(),
600 )
601 .await
602 .unwrap();
603 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
604 assert_eq!(
605 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
606 Some(b"false".as_ref()),
607 "{path}: retrying identical bytes cannot fix a 413"
608 );
609 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
610 .await
611 .unwrap();
612 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
613 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
614 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
615 assert!(
616 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
617 "{path}: message names the limit"
618 );
619 }
620 }
621
622 #[tokio::test]
623 async fn authenticated_body_admission_is_finite() {
624 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
625 let semaphore = body_admission_semaphore();
626 let mut permits = Vec::new();
627 for _ in 0..MAX_BODY_ADMISSIONS {
628 permits.push(semaphore.clone().acquire_owned().await.unwrap());
629 }
630 assert!(
631 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
632 .await
633 .is_err(),
634 "body parser admission must not be unbounded"
635 );
636 drop(permits);
637 assert!(semaphore.acquire().await.is_ok());
638 }
639
640 #[tokio::test]
641 async fn small_body_admission_is_finite_and_separate() {
642 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
643 let large = body_admission_semaphore();
644 let small = small_body_admission_semaphore();
645 let mut small_permits = Vec::new();
646 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
647 small_permits.push(small.clone().acquire_owned().await.unwrap());
648 }
649 assert!(
650 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
651 .await
652 .is_err(),
653 "small body parser admission must be bounded"
654 );
655 assert!(
656 large.clone().try_acquire().is_ok(),
657 "small uploads must not consume large-upload permits"
658 );
659 drop(small_permits);
660 assert!(small.acquire().await.is_ok());
661 }
662
663 #[test]
664 fn small_declared_bodies_bypass_large_upload_admission() {
665 let request = axum::http::Request::post("/v1/chat/completions")
666 .header(CONTENT_LENGTH, "2048")
667 .body(Body::empty())
668 .unwrap();
669 assert!(!body_requires_admission(&request));
670
671 let request = axum::http::Request::post("/v1/chat/completions")
672 .header(
673 CONTENT_LENGTH,
674 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
675 )
676 .body(Body::empty())
677 .unwrap();
678 assert!(body_requires_admission(&request));
679
680 let request = axum::http::Request::post("/v1/chat/completions")
681 .header(CONTENT_LENGTH, "2048")
682 .header(TRANSFER_ENCODING, "chunked")
683 .body(Body::empty())
684 .unwrap();
685 assert!(body_requires_admission(&request));
686 }
687
688 #[test]
689 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
690 let unknown = axum::http::Request::post("/v1/chat/completions")
691 .body(Body::empty())
692 .unwrap();
693 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
694
695 let large = axum::http::Request::post("/v1/chat/completions")
696 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
697 .body(Body::empty())
698 .unwrap();
699 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
700 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
701
702 let absurd = axum::http::Request::post("/v1/chat/completions")
703 .header(CONTENT_LENGTH, u64::MAX.to_string())
704 .body(Body::empty())
705 .unwrap();
706 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
707 }
708
709 #[tokio::test]
710 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
711 let too_large = shape_inference_early_response(
712 "/v1/messages",
713 error_response_coded(
714 StatusCode::PAYLOAD_TOO_LARGE,
715 "request body exceeds the 192 MiB limit",
716 "invalid_request_error",
717 None,
718 Some("request_too_large"),
719 ),
720 )
721 .await;
722 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
723 let house_id = too_large.headers()["x-request-id"].clone();
724 assert_eq!(too_large.headers()["request-id"], house_id);
725 assert_eq!(too_large.headers()["x-should-retry"], "false");
726 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
727 .await
728 .unwrap();
729 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
730 assert_eq!(payload["type"], "error");
731 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
732
733 let busy = shape_inference_early_response(
734 "/v1/chat/completions",
735 retry_contract_response(
736 error_response_coded(
737 StatusCode::TOO_MANY_REQUESTS,
738 "request body admission is busy",
739 "rate_limit_error",
740 None,
741 Some("body_admission_busy"),
742 ),
743 Some(BODY_ADMISSION_RETRY_AFTER_S),
744 ),
745 )
746 .await;
747 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
748 assert!(!busy.headers()["x-request-id"].is_empty());
749 assert_eq!(busy.headers()["retry-after"], "1");
750 assert_eq!(busy.headers()["retry-after-ms"], "1000");
751 assert!(busy.headers().get("x-should-retry").is_none());
752 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
753 .await
754 .unwrap();
755 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
756 assert_eq!(payload["error"]["code"], "body_admission_busy");
757 }
758
759 #[tokio::test]
760 async fn vision_preprocess_admission_is_fail_fast_and_retryable() {
761 let semaphore = Box::leak(Box::new(tokio::sync::Semaphore::new(1)));
762 let held = semaphore.try_acquire().unwrap();
763 let busy = try_vision_preprocess_with(true, semaphore).unwrap_err();
764 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
765 assert_eq!(busy.headers()["retry-after"], "1");
766 drop(held);
767 assert!(
768 try_vision_preprocess_with(true, semaphore)
769 .unwrap()
770 .is_some()
771 );
772 assert!(
773 try_vision_preprocess_with(false, semaphore)
774 .unwrap()
775 .is_none()
776 );
777 }
778
779 #[tokio::test]
780 async fn typed_json_retains_body_admission_until_handler_validation_releases_it() {
781 #[derive(Clone)]
782 struct Signals {
783 parsed: Arc<tokio::sync::Notify>,
784 finish: Arc<tokio::sync::Notify>,
785 semaphore: Arc<tokio::sync::Semaphore>,
786 }
787
788 let semaphore = Arc::new(tokio::sync::Semaphore::new(1));
789 let guard = BodyAdmissionGuard::new(semaphore.clone().try_acquire_owned().unwrap());
790 let signals = Signals {
791 parsed: Arc::new(tokio::sync::Notify::new()),
792 finish: Arc::new(tokio::sync::Notify::new()),
793 semaphore: semaphore.clone(),
794 };
795 let app = Router::new()
796 .route(
797 "/",
798 post(
799 |Extension(signals): Extension<Signals>,
800 AdmittedJson(_, mut admission): AdmittedJson<serde_json::Value>| async move {
801 assert_eq!(
802 signals.semaphore.available_permits(),
803 0,
804 "typed deserialization alone must not release post-parse admission"
805 );
806 admission.release();
807 assert_eq!(signals.semaphore.available_permits(), 1);
808 signals.parsed.notify_one();
809 signals.finish.notified().await;
810 "ok"
811 },
812 ),
813 )
814 .layer(Extension(signals.clone()))
815 .layer(Extension(guard));
816 let response = tokio::spawn(
817 app.oneshot(
818 axum::http::Request::post("/")
819 .header(CONTENT_TYPE, "application/json")
820 .body(Body::from(r#"{"value":1}"#))
821 .unwrap(),
822 ),
823 );
824 signals.parsed.notified().await;
825 assert_eq!(
826 semaphore.available_permits(),
827 1,
828 "validated work must release admission before generation waits"
829 );
830 signals.finish.notify_one();
831 assert_eq!(response.await.unwrap().unwrap().status(), StatusCode::OK);
832 }
833
834 #[tokio::test]
835 async fn transport_closes_stalled_headers_and_caps_connections() {
836 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
837
838 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
839 let address = listener.local_addr().unwrap();
840 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
841 let server = tokio::spawn(serve_bounded_http_with_limits(
842 listener,
843 Router::new().route("/", get(|| async { "ok" })).route(
844 "/slow",
845 get(|| async {
846 tokio::time::sleep(std::time::Duration::from_millis(140)).await;
847 "slow-ok"
848 }),
849 ),
850 async move {
851 let _ = shutdown_rx.await;
852 },
853 std::time::Duration::from_millis(30),
854 1,
855 std::time::Duration::from_millis(80),
856 ));
857
858 let mut stalled = tokio::net::TcpStream::connect(address).await.unwrap();
859 stalled.write_all(b"GET / HT").await.unwrap();
860 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
861 let mut excess = tokio::net::TcpStream::connect(address).await.unwrap();
862 let mut bytes = Vec::new();
863 tokio::time::timeout(
864 std::time::Duration::from_millis(250),
865 excess.read_to_end(&mut bytes),
866 )
867 .await
868 .expect("connection beyond the cap must be closed promptly")
869 .unwrap();
870
871 bytes.clear();
872 tokio::time::timeout(
873 std::time::Duration::from_millis(500),
874 stalled.read_to_end(&mut bytes),
875 )
876 .await
877 .expect("stalled request headers must hit the configured deadline")
878 .unwrap();
879
880 let mut idle = tokio::net::TcpStream::connect(address).await.unwrap();
881 idle.write_all(b"GET / HTTP/1.1\r\nHost: local\r\n\r\n")
882 .await
883 .unwrap();
884 bytes.clear();
885 tokio::time::timeout(
886 std::time::Duration::from_millis(500),
887 idle.read_to_end(&mut bytes),
888 )
889 .await
890 .expect("an idle keep-alive connection must hit the maximum lifetime")
891 .unwrap();
892 assert!(String::from_utf8_lossy(&bytes).contains("200 OK"));
893
894 let mut active = tokio::net::TcpStream::connect(address).await.unwrap();
895 active
896 .write_all(b"GET /slow HTTP/1.1\r\nHost: local\r\n\r\n")
897 .await
898 .unwrap();
899 bytes.clear();
900 tokio::time::timeout(
901 std::time::Duration::from_millis(500),
902 active.read_to_end(&mut bytes),
903 )
904 .await
905 .expect("an active response must finish across the connection age boundary")
906 .unwrap();
907 let active_response = String::from_utf8_lossy(&bytes);
908 assert!(active_response.contains("200 OK"), "{active_response}");
909 assert!(active_response.contains("slow-ok"), "{active_response}");
910
911 let h2_stream = tokio::net::TcpStream::connect(address).await.unwrap();
915 let (mut h2_client, h2_connection) = h2::client::handshake(h2_stream).await.unwrap();
916 let h2_driver = tokio::spawn(h2_connection);
917 let request = axum::http::Request::builder()
918 .uri(format!("http://{address}/"))
919 .body(())
920 .unwrap();
921 let (response, _) = h2_client.send_request(request, true).unwrap();
922 let response = tokio::time::timeout(std::time::Duration::from_millis(500), response)
923 .await
924 .expect("HTTP/2 handshake and response must complete")
925 .expect("HTTP/2 connection must stay alive through the response");
926 assert_eq!(response.status(), StatusCode::OK);
927 drop(h2_client);
928 h2_driver.abort();
929 let _ = h2_driver.await;
930
931 let _ = shutdown_tx.send(());
932 server.await.unwrap().unwrap();
933 }
934}
935
936#[derive(Clone, Default)]
937struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
938
939fn is_sse_data_frame(bytes: &[u8]) -> bool {
940 bytes
941 .windows(b"data:".len())
942 .any(|window| window == b"data:")
943}
944
945async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
946 let trace = ttft::start(req.uri().path());
947 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
948 let response = next.run(req).await;
949 let Some(trace) = trace else {
950 return response;
951 };
952 let is_sse = response
953 .headers()
954 .get(CONTENT_TYPE)
955 .and_then(|value| value.to_str().ok())
956 .is_some_and(|value| value.starts_with("text/event-stream"));
957 if !is_sse {
958 return response;
959 }
960
961 let (parts, body) = response.into_parts();
964 let mut body = Box::pin(body.into_data_stream());
965 let stream = async_stream::stream! {
966 while let Some(frame) =
967 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
968 {
969 if frame
970 .as_ref()
971 .is_ok_and(|bytes| is_sse_data_frame(bytes))
972 {
973 trace.mark_first_sse_byte();
974 }
975 yield frame;
976 }
977 };
978 Response::from_parts(parts, Body::from_stream(stream))
979}
980
981const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
982const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
983
984#[derive(Debug, Clone, Default, Deserialize)]
985#[serde(deny_unknown_fields)]
986struct OpenRouterMetadataFile {
987 #[serde(default)]
988 models: HashMap<String, OpenRouterModelMetadata>,
989 #[serde(default)]
992 planned_models: HashMap<String, OpenRouterModelMetadata>,
993 #[serde(default)]
996 provider: Option<ProviderMetadata>,
997}
998
999#[derive(Debug, Clone, Deserialize)]
1003#[serde(deny_unknown_fields)]
1004struct ProviderMetadata {
1005 id: String,
1006 #[serde(default)]
1007 status_url: Option<String>,
1008 #[serde(default)]
1009 support_contact: Option<String>,
1010 #[serde(default)]
1011 incident_contact: Option<String>,
1012 #[serde(default)]
1013 regions: Vec<String>,
1014}
1015
1016#[derive(Debug, Clone, Default, Deserialize)]
1018#[serde(deny_unknown_fields)]
1019struct LifecycleMetadata {
1020 #[serde(default)]
1021 status: Option<String>,
1022 #[serde(default)]
1023 deprecation_at: Option<String>,
1024 #[serde(default)]
1025 retirement_at: Option<String>,
1026 #[serde(default)]
1027 replacement_model_id: Option<String>,
1028}
1029
1030#[derive(Debug, Clone, Default, Deserialize)]
1032#[serde(deny_unknown_fields)]
1033struct ReliabilityMetadata {
1034 #[serde(default)]
1035 first_token_timeout_seconds: Option<u64>,
1036 #[serde(default)]
1037 completion_timeout_seconds: Option<u64>,
1038 #[serde(default)]
1039 stream_idle_timeout_seconds: Option<u64>,
1040 #[serde(default)]
1041 capacity_scope: Option<String>,
1042}
1043
1044#[derive(Debug, Clone, Default, Deserialize)]
1045#[serde(deny_unknown_fields)]
1046struct OpenRouterModelMetadata {
1047 #[serde(default)]
1049 owned_by: Option<String>,
1050 #[serde(default)]
1051 lifecycle: Option<LifecycleMetadata>,
1052 #[serde(default)]
1053 reliability: Option<ReliabilityMetadata>,
1054 #[serde(default)]
1055 hugging_face_id: Option<String>,
1056 #[serde(default)]
1057 created: Option<u64>,
1058 #[serde(default)]
1059 quantization: Option<String>,
1060 #[serde(default)]
1061 description: Option<String>,
1062 #[serde(default)]
1063 max_prompt_length: Option<u64>,
1064 #[serde(default)]
1065 max_output_length: Option<u64>,
1066 #[serde(default)]
1069 default_output_length: Option<u64>,
1070 #[serde(default)]
1071 pricing: OpenRouterPricing,
1072 #[serde(default)]
1073 capacity: OpenRouterCapacity,
1074 #[serde(default)]
1075 is_ready: Option<bool>,
1076 #[serde(default)]
1077 is_free: Option<bool>,
1078 #[serde(default)]
1079 discount_to_user: Option<f64>,
1080 #[serde(default)]
1081 openrouter_slug: Option<String>,
1082 #[serde(default)]
1083 datacenters: Vec<OpenRouterDatacenter>,
1084 #[serde(default)]
1088 input_modalities: Vec<String>,
1089 #[serde(default)]
1104 surface: Option<String>,
1105 #[serde(default)]
1106 zdr: Option<bool>,
1107 #[serde(default)]
1108 hipaa: Option<bool>,
1109 #[serde(default)]
1119 default_reasoning_effort: Option<String>,
1120 #[serde(default)]
1138 default_temperature: Option<f32>,
1139 #[serde(default)]
1140 default_top_p: Option<f32>,
1141 #[serde(default)]
1143 default_top_k: Option<usize>,
1144 #[serde(default)]
1145 default_min_p: Option<f32>,
1146 #[serde(default)]
1147 default_presence_penalty: Option<f32>,
1148 #[serde(default)]
1149 default_frequency_penalty: Option<f32>,
1150 #[serde(default)]
1152 default_repetition_penalty: Option<f32>,
1153 #[serde(default)]
1174 non_thinking_sampling: Option<SamplingArmMetadata>,
1175}
1176
1177#[derive(Debug, Clone, Default, Deserialize)]
1183#[serde(deny_unknown_fields)]
1184struct SamplingArmMetadata {
1185 #[serde(default)]
1186 temperature: Option<f32>,
1187 #[serde(default)]
1188 top_p: Option<f32>,
1189 #[serde(default)]
1190 top_k: Option<usize>,
1191 #[serde(default)]
1192 min_p: Option<f32>,
1193 #[serde(default)]
1194 presence_penalty: Option<f32>,
1195 #[serde(default)]
1196 frequency_penalty: Option<f32>,
1197 #[serde(default)]
1198 repetition_penalty: Option<f32>,
1199}
1200
1201impl SamplingArmMetadata {
1202 fn is_empty(&self) -> bool {
1203 self.temperature.is_none()
1204 && self.top_p.is_none()
1205 && self.top_k.is_none()
1206 && self.min_p.is_none()
1207 && self.presence_penalty.is_none()
1208 && self.frequency_penalty.is_none()
1209 && self.repetition_penalty.is_none()
1210 }
1211}
1212
1213#[derive(Debug, Clone, Default, Deserialize)]
1214#[serde(deny_unknown_fields)]
1215struct OpenRouterPricing {
1216 #[serde(default)]
1217 prompt: Option<String>,
1218 #[serde(default)]
1219 cached_prompt: Option<String>,
1220 #[serde(default)]
1221 cache_write: Option<String>,
1222 #[serde(default)]
1223 completion: Option<String>,
1224 #[serde(default)]
1225 internal_reasoning: Option<String>,
1226 #[serde(default)]
1227 request: Option<String>,
1228}
1229
1230#[derive(Debug, Clone, Default, Deserialize)]
1231#[serde(deny_unknown_fields)]
1232struct OpenRouterCapacity {
1233 #[serde(default)]
1234 prompt_tpm: Option<u64>,
1235 #[serde(default)]
1236 cached_prompt_tpm: Option<u64>,
1237 #[serde(default)]
1238 completion_tpm: Option<u64>,
1239 #[serde(default)]
1240 request_rpm: Option<u64>,
1241 #[serde(default)]
1242 concurrency: Option<u64>,
1243}
1244
1245#[derive(Debug, Clone, Deserialize, Serialize)]
1246#[serde(deny_unknown_fields)]
1247struct OpenRouterDatacenter {
1248 country_code: String,
1249 #[serde(default, skip_serializing_if = "Option::is_none")]
1250 region: Option<String>,
1251}
1252
1253impl OpenRouterMetadataFile {
1254 fn parse(
1255 text: &str,
1256 ) -> Result<
1257 (
1258 HashMap<String, OpenRouterModelMetadata>,
1259 Option<ProviderMetadata>,
1260 ),
1261 String,
1262 > {
1263 let file: Self =
1264 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
1265 for (alias, metadata) in &file.models {
1266 validate_openrouter_metadata(alias, metadata)?;
1267 }
1268 for (alias, metadata) in &file.planned_models {
1269 validate_openrouter_metadata(alias, metadata)?;
1270 if file.models.contains_key(alias) {
1271 return Err(format!(
1272 "model alias {alias:?} appears in both models and planned_models"
1273 ));
1274 }
1275 }
1276 if let Some(provider) = &file.provider {
1277 if provider.id.is_empty() {
1278 return Err("provider.id must be a non-empty slug".into());
1279 }
1280 for (field, value) in [
1282 ("provider.support_contact", &provider.support_contact),
1283 ("provider.incident_contact", &provider.incident_contact),
1284 ] {
1285 if let Some(value) = value
1286 && !value.contains(':')
1287 {
1288 return Err(format!(
1289 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1290 ));
1291 }
1292 }
1293 }
1294 Ok((file.models, file.provider))
1295 }
1296
1297 #[cfg(test)]
1298 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1299 Self::parse(text).map(|(models, _)| models)
1300 }
1301}
1302
1303fn per_million_price(per_token: &str) -> Option<String> {
1307 if !valid_price_string(per_token) {
1308 return None;
1309 }
1310 let (whole, frac) = match per_token.split_once('.') {
1311 Some((whole, frac)) => (whole, frac),
1312 None => (per_token, ""),
1313 };
1314 let mut digits = format!("{whole}{frac}");
1315 let point = whole.len() + 6;
1316 while digits.len() < point {
1317 digits.push('0');
1318 }
1319 let (int_part, frac_part) = digits.split_at(point);
1320 let int_part = int_part.trim_start_matches('0');
1321 let int_part = if int_part.is_empty() { "0" } else { int_part };
1322 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1323 while frac_out.len() < 2 {
1324 frac_out.push('0');
1325 }
1326 Some(format!("{int_part}.{frac_out}"))
1327}
1328
1329fn valid_price_string(value: &str) -> bool {
1330 let mut parts = value.split('.');
1331 let whole = parts.next().unwrap_or_default();
1332 let fraction = parts.next();
1333 !whole.is_empty()
1334 && whole.bytes().all(|b| b.is_ascii_digit())
1335 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1336 && parts.next().is_none()
1337}
1338
1339fn validate_openrouter_metadata(
1340 alias: &str,
1341 metadata: &OpenRouterModelMetadata,
1342) -> Result<(), String> {
1343 if alias.is_empty() {
1344 return Err("models metadata contains an empty model alias".into());
1345 }
1346 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1349 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1350 {
1351 return Err(format!(
1352 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1353 reasoning_effort level (none|minimal|low|medium|high)"
1354 ));
1355 }
1356 validate_sampling_defaults(alias, metadata)?;
1357 for m in &metadata.input_modalities {
1358 if m != "image" && m != "video" {
1359 return Err(format!(
1360 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1361 ));
1362 }
1363 }
1364 if let Some(sfc) = metadata.surface.as_deref()
1365 && !matches!(sfc, "chat" | "embedding" | "rerank")
1366 {
1367 return Err(format!(
1368 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1369 ));
1370 }
1371 if let Some(q) = metadata.quantization.as_deref()
1372 && !matches!(
1373 q,
1374 "int4"
1375 | "int8"
1376 | "fp4"
1377 | "mxfp4"
1378 | "nvfp4"
1379 | "fp6"
1380 | "fp8"
1381 | "mxfp8"
1382 | "fp16"
1383 | "bf16"
1384 | "fp32"
1385 )
1386 {
1387 return Err(format!(
1388 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1389 ));
1390 }
1391 for (field, value) in [
1392 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1393 (
1394 "pricing.cached_prompt",
1395 metadata.pricing.cached_prompt.as_deref(),
1396 ),
1397 (
1398 "pricing.cache_write",
1399 metadata.pricing.cache_write.as_deref(),
1400 ),
1401 ("pricing.completion", metadata.pricing.completion.as_deref()),
1402 (
1403 "pricing.internal_reasoning",
1404 metadata.pricing.internal_reasoning.as_deref(),
1405 ),
1406 ("pricing.request", metadata.pricing.request.as_deref()),
1407 ] {
1408 if let Some(value) = value
1409 && !valid_price_string(value)
1410 {
1411 return Err(format!(
1412 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1413 ));
1414 }
1415 }
1416 for (field, value) in [
1417 ("created", metadata.created),
1418 ("max_prompt_length", metadata.max_prompt_length),
1419 ("max_output_length", metadata.max_output_length),
1420 ("default_output_length", metadata.default_output_length),
1421 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1422 (
1423 "capacity.cached_prompt_tpm",
1424 metadata.capacity.cached_prompt_tpm,
1425 ),
1426 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1427 ("capacity.request_rpm", metadata.capacity.request_rpm),
1428 ("capacity.concurrency", metadata.capacity.concurrency),
1429 ] {
1430 if let Some(value) = value
1431 && value > JSON_SAFE_INTEGER_MAX
1432 {
1433 return Err(format!(
1434 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1435 ));
1436 }
1437 }
1438 for (field, value) in [
1439 ("max_prompt_length", metadata.max_prompt_length),
1440 ("max_output_length", metadata.max_output_length),
1441 ("default_output_length", metadata.default_output_length),
1442 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1443 (
1444 "capacity.cached_prompt_tpm",
1445 metadata.capacity.cached_prompt_tpm,
1446 ),
1447 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1448 ("capacity.request_rpm", metadata.capacity.request_rpm),
1449 ("capacity.concurrency", metadata.capacity.concurrency),
1450 ] {
1451 if value == Some(0) {
1452 return Err(format!(
1453 "model {alias:?}: {field} must be greater than zero when declared"
1454 ));
1455 }
1456 }
1457 if let (Some(default), Some(maximum)) =
1458 (metadata.default_output_length, metadata.max_output_length)
1459 && default > maximum
1460 {
1461 return Err(format!(
1462 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1463 ));
1464 }
1465 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1466 return Err(format!(
1467 "model {alias:?}: default_output_length requires max_output_length"
1468 ));
1469 }
1470 if let Some(discount) = metadata.discount_to_user
1471 && (!discount.is_finite() || discount >= 1.0)
1472 {
1473 return Err(format!(
1474 "model {alias:?}: discount_to_user must be finite and less than 1"
1475 ));
1476 }
1477 if metadata
1478 .openrouter_slug
1479 .as_deref()
1480 .is_some_and(str::is_empty)
1481 {
1482 return Err(format!(
1483 "model {alias:?}: openrouter_slug must not be empty when declared"
1484 ));
1485 }
1486 for dc in &metadata.datacenters {
1487 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1488 return Err(format!(
1489 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1490 dc.country_code
1491 ));
1492 }
1493 }
1494 Ok(())
1495}
1496
1497fn validate_sampling_defaults(
1512 alias: &str,
1513 metadata: &OpenRouterModelMetadata,
1514) -> Result<(), String> {
1515 validate_sampling_arm(
1516 alias,
1517 &[
1518 "default_temperature",
1519 "default_top_p",
1520 "default_min_p",
1521 "default_presence_penalty",
1522 "default_frequency_penalty",
1523 "default_repetition_penalty",
1524 ],
1525 metadata.default_temperature,
1526 metadata.default_top_p,
1527 metadata.default_min_p,
1528 metadata.default_presence_penalty,
1529 metadata.default_frequency_penalty,
1530 metadata.default_repetition_penalty,
1531 )?;
1532 if let Some(arm) = &metadata.non_thinking_sampling {
1533 if arm.is_empty() {
1537 return Err(format!(
1538 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1539 least one vendor recommendation or delete the table"
1540 ));
1541 }
1542 validate_sampling_arm(
1543 alias,
1544 &[
1545 "non_thinking_sampling.temperature",
1546 "non_thinking_sampling.top_p",
1547 "non_thinking_sampling.min_p",
1548 "non_thinking_sampling.presence_penalty",
1549 "non_thinking_sampling.frequency_penalty",
1550 "non_thinking_sampling.repetition_penalty",
1551 ],
1552 arm.temperature,
1553 arm.top_p,
1554 arm.min_p,
1555 arm.presence_penalty,
1556 arm.frequency_penalty,
1557 arm.repetition_penalty,
1558 )?;
1559 }
1560 Ok(())
1561}
1562
1563#[allow(clippy::too_many_arguments)]
1569fn validate_sampling_arm(
1570 alias: &str,
1571 keys: &[&str; 6],
1572 temperature: Option<f32>,
1573 top_p: Option<f32>,
1574 min_p: Option<f32>,
1575 presence_penalty: Option<f32>,
1576 frequency_penalty: Option<f32>,
1577 repetition_penalty: Option<f32>,
1578) -> Result<(), String> {
1579 if let Some(t) = temperature
1580 && (!t.is_finite() || t <= 0.0 || t > 2.0)
1581 {
1582 return Err(format!(
1583 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1584 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1585 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1586 vendor recommendation, not greedy); clients reach greedy by sending an \
1587 explicit temperature 0.",
1588 keys[0]
1589 ));
1590 }
1591 if let Some(p) = top_p
1592 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1593 {
1594 return Err(format!(
1595 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1596 keys[1]
1597 ));
1598 }
1599 if let Some(m) = min_p
1600 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1601 {
1602 return Err(format!(
1603 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1604 keys[2]
1605 ));
1606 }
1607 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1608 if let Some(v) = value
1609 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1610 {
1611 return Err(format!(
1612 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1613 ));
1614 }
1615 }
1616 if let Some(r) = repetition_penalty
1617 && (!r.is_finite() || r <= 0.0)
1618 {
1619 return Err(format!(
1620 "model {alias:?}: {} {r} must be finite and \
1621 greater than zero (1.0 = off)",
1622 keys[5]
1623 ));
1624 }
1625 Ok(())
1626}
1627
1628fn load_openrouter_metadata(
1629 models: &[(String, String, Option<String>)],
1630) -> Result<
1631 (
1632 HashMap<String, OpenRouterModelMetadata>,
1633 Option<ProviderMetadata>,
1634 ),
1635 String,
1636> {
1637 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1638 Ok(path) => path,
1639 Err(_) => return Ok((HashMap::new(), None)),
1640 };
1641 let p = std::path::Path::new(&path);
1642 if !p.is_file() {
1643 return Err(format!(
1644 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1645 ));
1646 }
1647 let text =
1648 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1649 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1650 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1651 for alias in metadata.keys() {
1652 if !models.iter().any(|(name, _, _)| name == alias) {
1653 return Err(format!(
1654 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1655 ));
1656 }
1657 }
1658 eprintln!(
1659 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1660 metadata.len()
1661 );
1662 Ok((metadata, provider))
1663}
1664
1665#[derive(Clone)]
1666struct AppState {
1667 cmd_tx: Sender<Cmd>,
1668 models: Arc<Vec<String>>,
1669 caps: Arc<HashMap<String, ModelCaps>>,
1670 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1671 provider_metadata: Arc<Option<ProviderMetadata>>,
1673 metering: Option<Arc<dyn metering::Metering>>,
1679 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1682 api_auth: ApiAuth,
1685 metrics_auth: MetricsAuth,
1687 metrics: SharedMetrics,
1688 inflight: InflightCounts,
1692 tenant_inflight: TenantGauge,
1695 health: health::SharedHealth,
1699 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1703}
1704
1705impl AppState {
1706 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1720 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1721 }
1722}
1723
1724#[derive(Clone, Default)]
1725struct ApiAuth {
1726 keyring: Option<&'static auth::KeyStore>,
1727 single_key: Option<Arc<str>>,
1728}
1729
1730impl ApiAuth {
1731 fn from_env() -> Result<ApiAuth, String> {
1732 let single_key = match std::env::var("MEMRA_API_KEY") {
1733 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1734 Ok(key) => Some(Arc::from(key)),
1735 Err(std::env::VarError::NotPresent) => None,
1736 Err(std::env::VarError::NotUnicode(_)) => {
1737 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1738 }
1739 };
1740 Ok(ApiAuth {
1741 keyring: auth::global(),
1742 single_key,
1743 })
1744 }
1745
1746 fn configured(&self) -> bool {
1747 self.keyring.is_some() || self.single_key.is_some()
1748 }
1749}
1750
1751#[derive(Clone, Default)]
1752struct MetricsAuth {
1753 required: bool,
1754 token: Option<Arc<str>>,
1755}
1756
1757impl MetricsAuth {
1758 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1759 let token = token.map(Arc::from);
1760 MetricsAuth {
1761 required: !bind_loopback || api_auth_configured || token.is_some(),
1762 token,
1763 }
1764 }
1765}
1766
1767fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1768 let mut resolved = addr
1769 .to_socket_addrs()
1770 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1771 let first = resolved
1772 .next()
1773 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1774 let mut loopback = first.ip().to_canonical().is_loopback();
1775 for socket in resolved {
1776 loopback &= socket.ip().to_canonical().is_loopback();
1777 }
1778 Ok((first, loopback))
1779}
1780
1781fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1782 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1783}
1784
1785fn validate_bind_security(
1786 addr: &str,
1787 api_auth_configured: bool,
1788 allow_open_bind: bool,
1789) -> Result<bool, String> {
1790 let loopback = bind_is_loopback(addr)?;
1791 if !loopback && !api_auth_configured && !allow_open_bind {
1792 return Err(format!(
1793 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1794 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1795 ));
1796 }
1797 Ok(loopback)
1798}
1799
1800type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1818
1819type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1822
1823struct InflightGuard {
1827 counts: InflightCounts,
1828 idx: usize,
1829 tenants: TenantGauge,
1830 tenant: String,
1831}
1832
1833impl InflightGuard {
1834 fn try_acquire(
1838 counts: InflightCounts,
1839 lane: lanes::Lane,
1840 tenants: TenantGauge,
1841 tenant: &str,
1842 tenant_cap: Option<usize>,
1843 ) -> Result<(Self, usize, usize), usize> {
1844 let idx = lane.idx();
1845 let nt = {
1846 let mut m = tenants.lock().unwrap();
1847 let e = m.entry(tenant.to_string()).or_insert(0);
1848 if tenant_cap.is_some_and(|cap| *e >= cap) {
1849 return Err(*e);
1850 }
1851 *e += 1;
1852 *e
1853 };
1854 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1855 Ok((
1856 InflightGuard {
1857 counts,
1858 idx,
1859 tenants,
1860 tenant: tenant.to_string(),
1861 },
1862 n,
1863 nt,
1864 ))
1865 }
1866}
1867
1868impl Drop for InflightGuard {
1869 fn drop(&mut self) {
1870 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1871 let mut m = self.tenants.lock().unwrap();
1872 if let Some(e) = m.get_mut(&self.tenant) {
1873 *e -= 1;
1874 if *e == 0 {
1875 m.remove(&self.tenant);
1876 }
1877 }
1878 }
1879}
1880
1881fn lane_cap(lane: lanes::Lane) -> usize {
1885 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1886 CAPS.get_or_init(|| {
1887 let batching = std::env::var("MEMRA_SERVE_BATCH")
1888 .map(|v| v != "0")
1889 .unwrap_or(true);
1890 let interactive = if batching {
1891 std::env::var("MEMRA_MAX_SESSIONS")
1892 .ok()
1893 .and_then(|v| v.parse().ok())
1894 .unwrap_or(64)
1895 } else {
1896 worker::MAX_ACTIVE
1897 };
1898 let p = lanes::LanePolicy::from_env();
1899 [interactive, p.max_sessions[1], p.max_sessions[2]]
1900 })[lane.idx()]
1901}
1902
1903fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1906 if m.completed > 0 && m.step_p50_ms > 0.0 {
1907 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1908 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1909 }
1910 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1911 *D.get_or_init(|| {
1912 std::env::var("MEMRA_RL_RESET_S")
1913 .ok()
1914 .and_then(|v| v.parse().ok())
1915 .unwrap_or(2)
1916 })
1917}
1918
1919pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1936pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1937pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1938
1939pub(crate) fn timeout_ms_max() -> u64 {
1949 static V: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1950 *V.get_or_init(|| {
1951 std::env::var("MEMRA_TIMEOUT_MS_MAX")
1952 .ok()
1953 .and_then(|s| s.parse::<u64>().ok())
1954 .filter(|&ms| ms >= TIMEOUT_MS_MIN)
1955 .unwrap_or(TIMEOUT_MS_MAX)
1956 })
1957}
1958
1959pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1963 let max = timeout_ms_max();
1964 let Some(v) = v.filter(|v| !v.is_null()) else {
1965 return Ok(max);
1967 };
1968 let Some(ms) = v.as_u64() else {
1969 return Err(format!(
1970 "timeout_ms must be an integer number of milliseconds in \
1971 {TIMEOUT_MS_MIN}..={max}, got {v}; for work longer than \
1972 {max} ms use \"stream\": true — the deadline then bounds only the \
1973 time to first token and the stream may run as long as it needs"
1974 ));
1975 };
1976 if !(TIMEOUT_MS_MIN..=max).contains(&ms) {
1977 return Err(format!(
1978 "timeout_ms {ms} is outside the accepted range \
1979 {TIMEOUT_MS_MIN}..={max} (milliseconds). {max} is a \
1980 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1981 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1982 lie. For work longer than {max} ms use \"stream\": true — the \
1983 deadline then bounds only the time to first token and the stream may run as \
1984 long as it needs"
1985 ));
1986 }
1987 Ok(ms)
1988}
1989
1990#[derive(Clone, Copy)]
1993pub(crate) struct RequestDeadline {
1994 pub(crate) at: tokio::time::Instant,
1995 pub(crate) ms: u64,
1996}
1997
1998impl RequestDeadline {
1999 pub(crate) fn starting_now(ms: u64) -> Self {
2000 Self {
2001 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
2002 ms,
2003 }
2004 }
2005
2006 pub(crate) fn remaining(&self) -> std::time::Duration {
2007 self.at
2008 .saturating_duration_since(tokio::time::Instant::now())
2009 }
2010}
2011
2012pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
2018 let what = if stream {
2019 "the first token was produced"
2020 } else {
2021 "the response completed"
2022 };
2023 let msg = format!(
2024 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
2025 {what}; generation was cancelled and this request is not billed"
2026 );
2027 error_response_coded(
2028 StatusCode::REQUEST_TIMEOUT,
2029 &msg,
2030 "timeout",
2031 Some("timeout_ms"),
2032 Some("deadline_exceeded"),
2033 )
2034}
2035
2036pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
2076
2077pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
2082
2083pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
2087
2088fn env_flag_on(name: &'static str, default_on: bool) -> bool {
2094 match std::env::var(name) {
2095 Ok(v) => !matches!(
2096 v.trim().to_ascii_lowercase().as_str(),
2097 "0" | "off" | "false"
2098 ),
2099 Err(_) => default_on,
2100 }
2101}
2102
2103pub(crate) fn env_u64(name: &'static str, default: u64) -> u64 {
2106 std::env::var(name)
2107 .ok()
2108 .and_then(|v| v.parse::<u64>().ok())
2109 .filter(|v| *v > 0)
2110 .unwrap_or(default)
2111}
2112
2113const CHARS_PER_TOKEN_FLOOR: usize = 6;
2124
2125pub(crate) fn prompt_tokens_estimate(
2126 request: &worker::Request,
2127 tokenizer: Option<&Tokenizer>,
2128) -> u64 {
2129 if !request.prompt_ids.is_empty() {
2130 return request.prompt_ids.len() as u64;
2131 }
2132 let mut text = String::new();
2133 text.push_str(&request.prompt_text);
2134 for turn in &request.chat_turns {
2135 text.push_str(&turn.content);
2136 }
2137 for tool in &request.tools_json {
2138 text.push_str(tool);
2139 }
2140 if let Some(tokenizer) = tokenizer {
2141 return tokenizer.encode(text.as_str(), false).len() as u64;
2142 }
2143 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
2144}
2145
2146pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
2150 let prefill_ms = prompt_tokens
2151 .saturating_mul(1_000)
2152 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
2153 .unwrap_or(u64::MAX);
2154 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
2155 if decode_ms == 0 {
2156 return None;
2157 }
2158 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
2159}
2160
2161pub(crate) fn nonstream_deadline_gate(
2171 request: &worker::Request,
2172 stream: bool,
2173 deadline: RequestDeadline,
2174 caller_declared_max_tokens: bool,
2175 tokenizer: Option<&Tokenizer>,
2176) -> Result<(), String> {
2177 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
2178 return Ok(());
2179 }
2180 let max_new = request.params.max_new as u64;
2181 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
2188 {
2189 return Ok(());
2190 }
2191 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
2192 let remaining_ms = deadline.remaining().as_millis() as u64;
2193 let prefill_ms = prompt_tokens.saturating_mul(1_000)
2194 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
2195 let decode_ms = max_new.saturating_mul(1_000)
2196 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
2197 let est_ms = prefill_ms.saturating_add(decode_ms);
2198 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
2199 if est_ms <= bound_ms {
2200 return Ok(());
2201 }
2202 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
2203 let advice = match fits {
2204 Some(fits) if fits > 0 => format!(
2205 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
2206 stream's deadline bounds only the time to first token, so it may run as long \
2207 as it needs"
2208 ),
2209 _ => format!(
2210 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
2211 token, so no max_tokens fits: set \"stream\": true"
2212 ),
2213 };
2214 Err(format!(
2215 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
2216 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
2217 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
2218 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
2219 rather than after the deadline: {advice}",
2220 est_ms / 1_000,
2221 ))
2222}
2223
2224fn max_queue_depth(cap: usize) -> usize {
2228 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
2229 D.get_or_init(|| {
2230 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
2231 .ok()
2232 .and_then(|v| v.parse().ok())
2233 })
2234 .unwrap_or(cap.saturating_mul(4))
2235}
2236
2237fn queue_wait_ceiling_s() -> u64 {
2244 static S: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2245 *S.get_or_init(|| {
2246 std::env::var("MEMRA_QUEUE_WAIT_CEILING_S")
2247 .ok()
2248 .and_then(|v| v.parse().ok())
2249 .unwrap_or(0)
2250 })
2251}
2252
2253pub(crate) struct PendingAdmissionGuard {
2274 reserved: bool,
2275 lane: lanes::Lane,
2276}
2277
2278impl PendingAdmissionGuard {
2279 pub(crate) fn commit(mut self) {
2283 self.reserved = false;
2284 std::mem::forget(self);
2285 }
2286}
2287
2288impl Drop for PendingAdmissionGuard {
2289 fn drop(&mut self) {
2290 if self.reserved {
2291 worker::release_pending_admit();
2292 worker::release_admission_reservation(self.lane);
2293 }
2294 }
2295}
2296
2297#[allow(clippy::result_large_err)] pub(crate) fn reserve_pending_admit(
2299 st: &AppState,
2300 lane: lanes::Lane,
2301 rl: &RateLimit,
2302 deadline: RequestDeadline,
2303) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2304 reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())
2305}
2306
2307#[allow(clippy::result_large_err)] fn reserve_pending_admit_with_ceiling(
2312 st: &AppState,
2313 lane: lanes::Lane,
2314 rl: &RateLimit,
2315 deadline: RequestDeadline,
2316 ceiling_s: u64,
2317) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2318 let cap = lane_cap(lane).max(1);
2323 let bound = max_queue_depth(cap);
2324 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2325 loop {
2326 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2327 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2328 let backlog = reservations;
2332 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2333 if backlog >= bound {
2334 let msg = format!(
2335 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2336 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2337 coarse estimate, not a promise)",
2338 lane.as_str()
2339 );
2340 let resp = retry_contract_response(
2341 (
2342 StatusCode::TOO_MANY_REQUESTS,
2343 Json(error_body(
2344 &msg,
2345 "rate_limit_error",
2346 None,
2347 Some("shed_queue"),
2348 )),
2349 )
2350 .into_response(),
2351 Some(est_wait_s),
2352 );
2353 return Err((resp, "shed_queue"));
2354 }
2355 let remaining_ms = deadline.remaining().as_millis() as u64;
2356 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2360 if lane == lanes::Lane::Interactive
2361 && waits_for_capacity
2362 && est_wait_s.saturating_mul(1_000) > remaining_ms
2363 {
2364 let msg = format!(
2365 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2366 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2367 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2368 estimate, not a promise)"
2369 );
2370 let resp = retry_contract_response(
2371 (
2372 StatusCode::TOO_MANY_REQUESTS,
2373 Json(error_body(
2374 &msg,
2375 "rate_limit_error",
2376 None,
2377 Some("shed_deadline"),
2378 )),
2379 )
2380 .into_response(),
2381 Some(est_wait_s),
2382 );
2383 return Err((resp, "shed_deadline"));
2384 }
2385 if lane == lanes::Lane::Interactive
2393 && waits_for_capacity
2394 && ceiling_s > 0
2395 && est_wait_s > ceiling_s
2396 {
2397 let msg = format!(
2398 "estimated queue wait ~{est_wait_s}s exceeds this deployment's queue-wait \
2399 ceiling ({ceiling_s}s); this request was not admitted and is not billed; \
2400 retry after ~{est_wait_s}s (a coarse estimate, not a promise)"
2401 );
2402 let resp = retry_contract_response(
2403 (
2404 StatusCode::TOO_MANY_REQUESTS,
2405 Json(error_body(
2406 &msg,
2407 "rate_limit_error",
2408 None,
2409 Some("shed_queue_wait"),
2410 )),
2411 )
2412 .into_response(),
2413 Some(est_wait_s),
2414 );
2415 return Err((resp, "shed_queue_wait"));
2416 }
2417 if reservations_for_lane
2418 .compare_exchange(
2419 reservations,
2420 reservations.saturating_add(1),
2421 std::sync::atomic::Ordering::AcqRel,
2422 std::sync::atomic::Ordering::Acquire,
2423 )
2424 .is_ok()
2425 {
2426 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2430 return Ok(PendingAdmissionGuard {
2431 reserved: true,
2432 lane,
2433 });
2434 }
2435 }
2436}
2437
2438static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2449
2450fn draining() -> bool {
2451 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2452}
2453
2454fn drain_deadline_s() -> u64 {
2456 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2457 *D.get_or_init(|| {
2458 std::env::var("MEMRA_DRAIN_S")
2459 .ok()
2460 .and_then(|v| v.parse().ok())
2461 .unwrap_or(30)
2462 })
2463}
2464
2465fn drain_response() -> Response {
2475 let resp = (
2476 StatusCode::SERVICE_UNAVAILABLE,
2477 Json(error_body(
2478 "server is draining (shutdown in progress); retry",
2479 "server_error",
2480 None,
2481 Some("draining"),
2482 )),
2483 )
2484 .into_response();
2485 retry_contract_response(resp, Some(drain_deadline_s()))
2486}
2487
2488struct RateLimit {
2490 limit: usize,
2491 remaining: usize,
2492 reset_s: u64,
2493}
2494
2495impl RateLimit {
2496 fn at_admit(
2501 lane: lanes::Lane,
2502 n_inflight: usize,
2503 metrics: &SharedMetrics,
2504 tenant: &auth::TenantCtx,
2505 n_tenant: usize,
2506 ) -> Self {
2507 let global = lane_cap(lane);
2508 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2509 return Self::compute(global, n_inflight, metrics);
2510 };
2511 let headroom = t
2512 .saturating_sub(n_tenant)
2513 .min(global.saturating_sub(n_inflight));
2514 Self::compute(t, t - headroom, metrics)
2516 }
2517
2518 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2519 let remaining = limit.saturating_sub(n_inflight);
2520 let reset_s = if remaining > 0 {
2521 0
2522 } else {
2523 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2524 reset_estimate_s(&m)
2525 };
2526 RateLimit {
2527 limit,
2528 remaining,
2529 reset_s,
2530 }
2531 }
2532
2533 fn attach(&self, mut resp: Response) -> Response {
2535 let h = resp.headers_mut();
2536 for (k, v) in [
2537 ("x-ratelimit-limit", self.limit as u64),
2538 ("x-ratelimit-remaining", self.remaining as u64),
2539 ("x-ratelimit-reset", self.reset_s),
2540 ] {
2541 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2542 h.insert(axum::http::HeaderName::from_static(k), v);
2543 }
2544 }
2545 resp
2546 }
2547}
2548
2549#[allow(clippy::result_large_err)] fn acquire_request_slot(
2554 st: &AppState,
2555 lane: lanes::Lane,
2556 tenant: &auth::TenantCtx,
2557 env: &Envelope,
2558) -> Result<(InflightGuard, RateLimit), Response> {
2559 let global = lane_cap(lane);
2560 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2561 match InflightGuard::try_acquire(
2562 st.inflight.clone(),
2563 lane,
2564 st.tenant_inflight.clone(),
2565 &tenant.tenant,
2566 tenant_cap,
2567 ) {
2568 Ok((guard, n_inflight, n_tenant)) => {
2569 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2570 Ok((guard, rl))
2571 }
2572 Err(n_tenant) => {
2573 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2574 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2575 let error =
2576 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2577 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2578 }
2579 }
2580}
2581
2582#[derive(Deserialize)]
2584struct CompletionReq {
2585 model: String,
2586 #[serde(default)]
2587 prompt: String,
2588 #[serde(default)]
2590 prompt_ids: Vec<u32>,
2591 #[serde(default)]
2594 max_tokens: Option<usize>,
2595 #[serde(default)]
2608 temperature: Option<f32>,
2609 #[serde(default)]
2610 top_p: Option<f32>,
2611 #[serde(default)]
2613 top_k: Option<usize>,
2614 #[serde(default)]
2616 min_p: Option<f32>,
2617 #[serde(default)]
2619 frequency_penalty: Option<f32>,
2620 #[serde(default)]
2621 presence_penalty: Option<f32>,
2622 #[serde(default)]
2624 repetition_penalty: Option<f32>,
2625 #[serde(default)]
2631 seed: Option<u64>,
2632 #[serde(default)]
2633 stop: StopSequences,
2634 #[serde(default)]
2637 logit_bias: Option<serde_json::Value>,
2638 #[serde(default)]
2639 logprobs: Option<serde_json::Value>,
2640 #[serde(default)]
2641 n: Option<usize>,
2642 #[serde(default)]
2643 best_of: Option<usize>,
2644 #[serde(default)]
2646 chat: bool,
2647 #[serde(default)]
2649 stream: bool,
2650 #[serde(default)]
2652 max_ctx: Option<usize>,
2653 #[serde(default)]
2655 trace_id: Option<String>,
2656 #[serde(default)]
2660 cache_salt: Option<String>,
2661 #[serde(default)]
2665 session_id: Option<String>,
2666 #[serde(default)]
2667 user: Option<String>,
2668 #[serde(default)]
2672 timeout_ms: Option<serde_json::Value>,
2673}
2674
2675#[derive(Deserialize)]
2676struct ChatMessage {
2677 role: String,
2678 #[serde(default)]
2680 content: serde_json::Value,
2681 #[serde(default)]
2683 tool_calls: Vec<ReqToolCall>,
2684 #[serde(default)]
2687 tool_call_id: Option<String>,
2688 #[serde(default)]
2691 name: Option<String>,
2692 #[serde(default, alias = "reasoning_content")]
2702 reasoning: Option<String>,
2703}
2704
2705#[derive(Deserialize)]
2706struct ReqToolCall {
2707 #[serde(default)]
2708 #[allow(dead_code)]
2709 id: Option<String>,
2710 function: ReqToolFunction,
2711}
2712
2713#[derive(Deserialize)]
2714struct ReqToolFunction {
2715 name: String,
2716 #[serde(default)]
2718 arguments: serde_json::Value,
2719}
2720
2721#[derive(Clone, Default, Deserialize)]
2722#[serde(untagged)]
2723enum StopSequences {
2724 One(String),
2725 Many(Vec<String>),
2726 #[default]
2727 None,
2728}
2729
2730impl StopSequences {
2731 fn into_vec(self) -> Vec<String> {
2737 let stops = match self {
2738 Self::One(stop) => vec![stop],
2739 Self::Many(stops) => stops,
2740 Self::None => Vec::new(),
2741 };
2742 stops.into_iter().filter(|s| !s.is_empty()).collect()
2743 }
2744
2745 fn validate(&self) -> Result<(), String> {
2746 let stops: &[String] = match self {
2747 Self::One(stop) => std::slice::from_ref(stop),
2748 Self::Many(stops) => stops,
2749 Self::None => &[],
2750 };
2751 if stops.len() > MAX_STOP_SEQUENCES {
2752 return Err(format!(
2753 "stop accepts at most {MAX_STOP_SEQUENCES} sequences"
2754 ));
2755 }
2756 let mut total = 0usize;
2757 for stop in stops {
2758 let bytes = stop.len();
2759 if bytes > MAX_STOP_SEQUENCE_BYTES {
2760 return Err(format!(
2761 "each stop sequence must be at most {MAX_STOP_SEQUENCE_BYTES} UTF-8 bytes"
2762 ));
2763 }
2764 total = total
2765 .checked_add(bytes)
2766 .ok_or_else(|| "stop sequence byte count overflowed".to_string())?;
2767 }
2768 if total > MAX_STOP_SEQUENCES_BYTES {
2769 return Err(format!(
2770 "stop sequences must total at most {MAX_STOP_SEQUENCES_BYTES} UTF-8 bytes"
2771 ));
2772 }
2773 Ok(())
2774 }
2775}
2776
2777#[derive(Deserialize)]
2783struct ChatCompletionReq {
2784 model: String,
2785 messages: Vec<ChatMessage>,
2786 #[serde(default, alias = "max_completion_tokens")]
2789 max_tokens: Option<usize>,
2790 #[serde(default)]
2793 temperature: Option<f32>,
2794 #[serde(default)]
2795 top_p: Option<f32>,
2796 #[serde(default)]
2800 top_k: Option<usize>,
2801 #[serde(default)]
2803 min_p: Option<f32>,
2804 #[serde(default)]
2806 frequency_penalty: Option<f32>,
2807 #[serde(default)]
2808 presence_penalty: Option<f32>,
2809 #[serde(default)]
2811 repetition_penalty: Option<f32>,
2812 #[serde(default)]
2814 seed: Option<u64>,
2815 #[serde(default)]
2816 stop: StopSequences,
2817 #[serde(default)]
2818 stream: bool,
2819 #[serde(default)]
2820 max_ctx: Option<usize>,
2821 #[serde(default)]
2826 response_format: Option<serde_json::Value>,
2827 #[serde(default)]
2828 logit_bias: Option<serde_json::Value>,
2829 #[serde(default)]
2830 logprobs: Option<serde_json::Value>,
2831 #[serde(default)]
2832 top_logprobs: Option<usize>,
2833 #[serde(default)]
2834 n: Option<usize>,
2835 #[serde(default)]
2837 tools: Vec<serde_json::Value>,
2838 #[serde(default)]
2840 tool_choice: Option<serde_json::Value>,
2841 #[serde(default)]
2847 reasoning_effort: Option<String>,
2848 #[serde(default)]
2857 reasoning: Option<serde_json::Value>,
2858 #[serde(default)]
2873 include_reasoning: Option<bool>,
2874 #[serde(default)]
2883 enable_thinking: Option<bool>,
2884 #[serde(default)]
2890 chat_template_kwargs: Option<serde_json::Value>,
2891 #[serde(default)]
2895 cache_salt: Option<String>,
2896 #[serde(default)]
2898 session_id: Option<String>,
2899 #[serde(default)]
2900 user: Option<String>,
2901 #[serde(default)]
2906 timeout_ms: Option<serde_json::Value>,
2907}
2908fn one() -> f32 {
2909 1.0
2910}
2911fn default_temperature() -> f32 {
2917 1.0
2918}
2919
2920#[derive(Debug, Clone, Copy, Default, PartialEq)]
2940struct SamplingDefaults {
2941 temperature: Option<f32>,
2942 top_p: Option<f32>,
2943 top_k: Option<usize>,
2944 min_p: Option<f32>,
2945 frequency_penalty: Option<f32>,
2946 presence_penalty: Option<f32>,
2947 repetition_penalty: Option<f32>,
2948}
2949
2950impl SamplingDefaults {
2951 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2954 SamplingDefaults {
2955 temperature: metadata
2956 .and_then(|m| m.default_temperature)
2957 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2958 top_p: metadata
2959 .and_then(|m| m.default_top_p)
2960 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2961 top_k: metadata.and_then(|m| m.default_top_k),
2962 min_p: metadata.and_then(|m| m.default_min_p),
2963 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2964 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2965 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2966 }
2967 }
2968}
2969
2970#[derive(Debug, Clone, Copy, Default, PartialEq)]
2983struct ModelSamplingDefaults {
2984 thinking: SamplingDefaults,
2985 non_thinking: Option<SamplingDefaults>,
2986}
2987
2988impl ModelSamplingDefaults {
2989 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2990 ModelSamplingDefaults {
2991 thinking: SamplingDefaults::resolve(metadata, caps),
2992 non_thinking: metadata
2997 .and_then(|m| m.non_thinking_sampling.as_ref())
2998 .map(|arm| SamplingDefaults {
2999 temperature: arm.temperature,
3000 top_p: arm.top_p,
3001 top_k: arm.top_k,
3002 min_p: arm.min_p,
3003 frequency_penalty: arm.frequency_penalty,
3004 presence_penalty: arm.presence_penalty,
3005 repetition_penalty: arm.repetition_penalty,
3006 }),
3007 }
3008 }
3009
3010 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
3023 match (think, &self.non_thinking) {
3024 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
3025 _ => &self.thinking,
3026 }
3027 }
3028
3029 #[cfg(test)] fn single(thinking: SamplingDefaults) -> Self {
3033 ModelSamplingDefaults {
3034 thinking,
3035 non_thinking: None,
3036 }
3037 }
3038}
3039
3040#[derive(Debug, Clone, Copy, Default)]
3046struct ClientSampling {
3047 temperature: Option<f32>,
3048 top_p: Option<f32>,
3049 top_k: Option<usize>,
3050 min_p: Option<f32>,
3051 frequency_penalty: Option<f32>,
3052 presence_penalty: Option<f32>,
3053 repetition_penalty: Option<f32>,
3054 seed: Option<u64>,
3055}
3056
3057impl From<&CompletionReq> for ClientSampling {
3058 fn from(r: &CompletionReq) -> Self {
3059 ClientSampling {
3060 temperature: r.temperature,
3061 top_p: r.top_p,
3062 top_k: r.top_k,
3063 min_p: r.min_p,
3064 frequency_penalty: r.frequency_penalty,
3065 presence_penalty: r.presence_penalty,
3066 repetition_penalty: r.repetition_penalty,
3067 seed: r.seed,
3068 }
3069 }
3070}
3071
3072impl From<&ChatCompletionReq> for ClientSampling {
3073 fn from(r: &ChatCompletionReq) -> Self {
3074 ClientSampling {
3075 temperature: r.temperature,
3076 top_p: r.top_p,
3077 top_k: r.top_k,
3078 min_p: r.min_p,
3079 frequency_penalty: r.frequency_penalty,
3080 presence_penalty: r.presence_penalty,
3081 repetition_penalty: r.repetition_penalty,
3082 seed: r.seed,
3083 }
3084 }
3085}
3086
3087fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
3093 sampler_config(
3094 client
3095 .temperature
3096 .or(defaults.temperature)
3097 .unwrap_or_else(default_temperature),
3098 client.top_k.or(defaults.top_k).unwrap_or(0),
3099 client.top_p.or(defaults.top_p).unwrap_or_else(one),
3100 client.min_p.or(defaults.min_p).unwrap_or(0.0),
3101 client
3102 .frequency_penalty
3103 .or(defaults.frequency_penalty)
3104 .unwrap_or(0.0),
3105 client
3106 .presence_penalty
3107 .or(defaults.presence_penalty)
3108 .unwrap_or(0.0),
3109 client
3110 .repetition_penalty
3111 .or(defaults.repetition_penalty)
3112 .unwrap_or_else(one),
3113 client.seed,
3114 )
3115}
3116
3117#[derive(Serialize)]
3118struct CompletionResp {
3119 model: String,
3120 text: String,
3121 tokens: Vec<u32>,
3122 stop_reason: String,
3126 #[serde(default, skip_serializing_if = "Option::is_none")]
3131 error: Option<serde_json::Value>,
3132 n_tokens: usize,
3133 prompt_tokens: usize,
3136 cached_tokens: usize,
3137 elapsed_s: f64,
3138}
3139
3140fn usage_json(
3148 n_prompt: usize,
3149 n_tokens: usize,
3150 n_cached: usize,
3151 elapsed_s: f64,
3152 spec: Option<worker::SpecUsage>,
3153) -> serde_json::Value {
3154 let mut u = json!({
3155 "prompt_tokens": n_prompt,
3156 "completion_tokens": n_tokens,
3157 "total_tokens": n_prompt + n_tokens,
3158 "prompt_tokens_details": { "cached_tokens": n_cached },
3159 "elapsed_s": elapsed_s,
3160 });
3161 if let Some(sp) = spec {
3162 u["spec"] = json!({
3163 "rounds": sp.rounds,
3164 "drafted": sp.drafted,
3165 "accepted": sp.accepted,
3166 "acceptance_rate": if sp.drafted > 0 {
3167 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
3168 });
3169 }
3170 u
3171}
3172
3173pub const SYSTEM_FINGERPRINT: &str = concat!(
3205 "memra-",
3206 env!("CARGO_PKG_VERSION"),
3207 "-",
3208 env!("MEMRA_BUILD_ID")
3209);
3210
3211pub const BUILD_ID_SRC: &str = env!("MEMRA_BUILD_ID_SRC");
3213
3214pub const BUILD_ID_NOTE: &str = env!("MEMRA_BUILD_ID_NOTE");
3216
3217pub const BUILD_GIT_SHA: &str = env!("MEMRA_BUILD_SHA");
3221
3222pub fn build_identity_line() -> String {
3225 format!("[server] build: {SYSTEM_FINGERPRINT} (id: {BUILD_ID_SRC}, git: {BUILD_GIT_SHA})")
3226}
3227
3228fn gen_hex128() -> String {
3231 use std::hash::{BuildHasher, Hasher};
3232 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3233 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3234 let t = std::time::SystemTime::now()
3235 .duration_since(std::time::UNIX_EPOCH)
3236 .map(|d| d.as_nanos() as u64)
3237 .unwrap_or(0);
3238 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
3239 h1.write_u64(n);
3240 h1.write_u64(t);
3241 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
3242 h2.write_u64(t.rotate_left(17));
3243 h2.write_u64(n);
3244 format!("{:016x}{:016x}", h1.finish(), h2.finish())
3245}
3246
3247#[derive(Clone)]
3250struct Envelope {
3251 id: String,
3252 created: u64,
3253}
3254
3255impl Envelope {
3256 fn new(chat: bool) -> Self {
3257 Envelope {
3258 id: format!(
3259 "{}-{}",
3260 if chat { "chatcmpl" } else { "cmpl" },
3261 gen_hex128()
3262 ),
3263 created: std::time::SystemTime::now()
3264 .duration_since(std::time::UNIX_EPOCH)
3265 .map(|d| d.as_secs())
3266 .unwrap_or(0),
3267 }
3268 }
3269
3270 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
3272 v["id"] = json!(self.id);
3273 v["created"] = json!(self.created);
3274 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
3275 v
3276 }
3277}
3278
3279fn with_request_id(id: &str, mut resp: Response) -> Response {
3281 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
3282 resp.headers_mut()
3283 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
3284 }
3285 resp
3286}
3287
3288fn openai_compat() -> bool {
3296 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3297 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
3298 Ok("openai") => true,
3299 Ok(_) => false,
3300 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
3301 })
3302}
3303
3304fn cache_namespace(cache_salt: &Option<String>) -> String {
3316 cache_salt.clone().unwrap_or_default()
3317}
3318
3319const CACHE_SALT_MAX_BYTES: usize = 64;
3320
3321fn validate_cache_namespace(
3322 cache_salt: &Option<String>,
3323 keyring_configured: bool,
3324) -> Result<String, &'static str> {
3325 let raw = cache_namespace(cache_salt);
3326 if raw.len() > CACHE_SALT_MAX_BYTES {
3327 return Err("cache_salt must be at most 64 bytes");
3328 }
3329 if !keyring_configured && raw.starts_with("t:") {
3330 return Err("cache_salt must not use the reserved t: prefix without a keyring");
3331 }
3332 if !raw
3333 .bytes()
3334 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
3335 {
3336 return Err("cache_salt contains unsupported characters");
3337 }
3338 Ok(raw)
3339}
3340
3341fn affinity_key(
3357 session_id: &Option<String>,
3358 user: &Option<String>,
3359 headers: &axum::http::HeaderMap,
3360) -> Result<Option<String>, String> {
3361 let clean = |s: &str| -> Result<Option<String>, String> {
3362 let t = s.trim();
3363 if t.is_empty() {
3364 Ok(None)
3365 } else if t.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3366 Err(format!(
3367 "session identity must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3368 ))
3369 } else if t.chars().any(char::is_control) {
3370 Err("session identity must not contain control characters".into())
3371 } else {
3372 Ok(Some(t.to_string()))
3373 }
3374 };
3375 if let Some(value) = session_id.as_deref()
3376 && let Some(value) = clean(value)?
3377 {
3378 return Ok(Some(value));
3379 }
3380 if let Some(value) = user.as_deref()
3381 && let Some(value) = clean(value)?
3382 {
3383 return Ok(Some(value));
3384 }
3385 match headers.get("x-session-id") {
3386 Some(value) => clean(
3387 value
3388 .to_str()
3389 .map_err(|_| "x-session-id must contain visible ASCII or UTF-8 text")?,
3390 ),
3391 None => Ok(None),
3392 }
3393}
3394
3395fn validate_client_identifier(value: Option<&str>, name: &str) -> Result<(), String> {
3396 let Some(value) = value else {
3397 return Ok(());
3398 };
3399 if value.len() > MAX_CLIENT_IDENTIFIER_BYTES {
3400 return Err(format!(
3401 "{name} must be at most {MAX_CLIENT_IDENTIFIER_BYTES} UTF-8 bytes"
3402 ));
3403 }
3404 if value.chars().any(char::is_control) {
3405 return Err(format!("{name} must not contain control characters"));
3406 }
3407 Ok(())
3408}
3409
3410fn error_body(
3415 message: &str,
3416 etype: &str,
3417 param: Option<&str>,
3418 code: Option<&str>,
3419) -> serde_json::Value {
3420 json!({ "error": {
3421 "message": message,
3422 "type": etype,
3423 "param": param,
3424 "code": code,
3425 } })
3426}
3427
3428fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3429 error_response_coded(status, message, etype, param, None)
3430}
3431
3432fn error_response_coded(
3437 status: StatusCode,
3438 message: &str,
3439 etype: &str,
3440 param: Option<&str>,
3441 code: Option<&str>,
3442) -> Response {
3443 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3444 if status.is_client_error()
3445 && status != StatusCode::TOO_MANY_REQUESTS
3446 && status != StatusCode::REQUEST_TIMEOUT
3447 && status != StatusCode::CONFLICT
3448 {
3449 resp.headers_mut().insert(
3450 "x-should-retry",
3451 axum::http::HeaderValue::from_static("false"),
3452 );
3453 }
3454 resp
3455}
3456
3457fn bad_request(message: &str, param: Option<&str>) -> Response {
3458 error_response(
3459 StatusCode::BAD_REQUEST,
3460 message,
3461 "invalid_request_error",
3462 param,
3463 )
3464}
3465
3466const 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>) {
3495 use worker::ErrClass as C;
3496 match class {
3497 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3498 C::ContextLength => (
3499 StatusCode::BAD_REQUEST,
3500 "invalid_request_error",
3501 Some("context_length_exceeded"),
3502 ),
3503 C::ModelNotFound => (
3504 StatusCode::BAD_REQUEST,
3505 "invalid_request_error",
3506 Some("model_not_found"),
3507 ),
3508 C::RateLimit => (
3509 StatusCode::TOO_MANY_REQUESTS,
3510 "rate_limit_error",
3511 Some("rate_limit_exceeded"),
3512 ),
3513 C::Overloaded => (
3514 StatusCode::SERVICE_UNAVAILABLE,
3515 "server_error",
3516 Some("overloaded"),
3517 ),
3518 C::Engine => (
3519 StatusCode::INTERNAL_SERVER_ERROR,
3520 "server_error",
3521 Some("engine_error"),
3522 ),
3523 }
3524}
3525
3526fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3528 use worker::ErrClass as C;
3529 match class {
3530 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3531 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3532 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3536 }
3537}
3538
3539fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3542 let (_, etype, code) = class_http(e.class);
3543 error_body(&e.message, etype, e.param, code)
3544}
3545
3546fn engine_error_response(e: &worker::EngineError) -> Response {
3552 engine_error_response_with_retry_after(
3553 e,
3554 e.retry_after_s.or_else(|| class_retry_after_s(e.class)),
3555 )
3556}
3557
3558fn engine_error_response_with_retry_after(
3559 e: &worker::EngineError,
3560 retry_after_s: Option<u64>,
3561) -> Response {
3562 let (status, _, _) = class_http(e.class);
3563 let resp = (status, Json(engine_error_body(e))).into_response();
3564 retry_contract_response(resp, retry_after_s)
3565}
3566
3567fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3569 let status = resp.status();
3570 let h = resp.headers_mut();
3571 match retry_after_s {
3572 Some(secs) => {
3573 let secs = secs.clamp(1, 60);
3575 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3576 h.insert(axum::http::header::RETRY_AFTER, v);
3577 }
3578 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3579 h.insert("retry-after-ms", v);
3580 }
3581 }
3582 None if status.is_client_error() => {
3583 h.insert(
3586 "x-should-retry",
3587 axum::http::HeaderValue::from_static("false"),
3588 );
3589 }
3590 None => {}
3591 }
3592 resp
3593}
3594
3595fn worker_unavailable_response() -> Response {
3596 engine_error_response_with_retry_after(
3597 &worker::EngineError::overloaded("worker unavailable"),
3598 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3599 )
3600}
3601
3602fn stop_reason_to_finish(r: &str) -> &'static str {
3603 match r {
3604 "Eos" | "Callback" => "stop",
3605 "MaxNew" | "ContextFull" => "length",
3606 _ => "stop",
3607 }
3608}
3609
3610fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3614 match v {
3615 serde_json::Value::Null => Ok(String::new()),
3616 serde_json::Value::String(s) => Ok(s.clone()),
3617 serde_json::Value::Array(parts) => {
3618 let mut out = String::new();
3619 for p in parts {
3620 match p.get("type").and_then(|t| t.as_str()) {
3621 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3622 Some(t) => out.push_str(t),
3623 None => return Err("content part has no text field".into()),
3624 },
3625 Some(other) => {
3626 return Err(format!(
3627 "unsupported content part type {other:?} (text only)"
3628 ));
3629 }
3630 }
3631 }
3632 Ok(out)
3633 }
3634 _ => Err("content must be a string, null, or an array of text parts".into()),
3635 }
3636}
3637
3638pub(crate) static VISION_PLACEMENT_SERVING: std::sync::atomic::AtomicBool =
3656 std::sync::atomic::AtomicBool::new(true);
3657
3658fn vision_placement_serving() -> bool {
3659 VISION_PLACEMENT_SERVING.load(std::sync::atomic::Ordering::Acquire)
3660}
3661
3662fn vision_media_admissible(placement: bool, kind: &str) -> Result<(), String> {
3668 if placement {
3669 Ok(())
3670 } else {
3671 Err(format!(
3672 "{kind} input is not enabled on this deployment (vision overlay placement \
3673 inadmissible at boot: see the worker's IMAGE INPUT DISABLED line)"
3674 ))
3675 }
3676}
3677
3678fn vision_placement_admits(kind: &str) -> Result<(), String> {
3679 vision_media_admissible(vision_placement_serving(), kind)
3680}
3681
3682fn vision_enabled() -> bool {
3686 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3687 *ON.get_or_init(|| {
3688 std::env::var("MEMRA_VISION_DIR").is_ok()
3689 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3690 })
3691}
3692
3693fn gemma_vision_enabled() -> bool {
3698 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3699 *ON.get_or_init(|| {
3700 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3701 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3702 })
3703}
3704
3705pub(crate) static GLM5_VISION_SERVING: std::sync::atomic::AtomicBool =
3715 std::sync::atomic::AtomicBool::new(false);
3716
3717fn glm5_vision_enabled() -> bool {
3720 GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)
3721}
3722
3723fn step_vision_enabled() -> bool {
3731 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3732 *ON.get_or_init(|| {
3733 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3734 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3735 })
3736}
3737
3738const VISION_MAX_IMAGES: usize = 8;
3740
3741pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3747 std::sync::atomic::AtomicUsize::new(0);
3748pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3752 tokio::sync::Semaphore::const_new(1);
3753
3754#[allow(clippy::result_large_err)]
3757pub(crate) fn try_vision_preprocess(
3758 required: bool,
3759) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3760 try_vision_preprocess_with(required, &VISION_PREPROCESS_SEMAPHORE)
3761}
3762
3763#[allow(clippy::result_large_err)]
3764fn try_vision_preprocess_with(
3765 required: bool,
3766 semaphore: &'static tokio::sync::Semaphore,
3767) -> Result<Option<tokio::sync::SemaphorePermit<'static>>, Response> {
3768 if !required {
3769 return Ok(None);
3770 }
3771 match semaphore.try_acquire() {
3772 Ok(permit) => Ok(Some(permit)),
3773 Err(tokio::sync::TryAcquireError::NoPermits) => Err(retry_contract_response(
3774 error_response_coded(
3775 StatusCode::TOO_MANY_REQUESTS,
3776 "vision preprocessing is busy",
3777 "rate_limit_error",
3778 Some("messages"),
3779 Some("vision_preprocess_busy"),
3780 ),
3781 Some(BODY_ADMISSION_RETRY_AFTER_S),
3782 )),
3783 Err(tokio::sync::TryAcquireError::Closed) => Err(error_response_coded(
3784 StatusCode::SERVICE_UNAVAILABLE,
3785 "vision preprocessing is unavailable",
3786 "server_error",
3787 Some("messages"),
3788 Some("vision_preprocess_unavailable"),
3789 )),
3790 }
3791}
3792
3793pub(crate) struct VisionMemoryPermit {
3794 bytes: usize,
3795}
3796
3797#[derive(Debug)]
3798pub(crate) enum VisionMemoryError {
3799 Request(String),
3800 Capacity(String),
3801}
3802
3803impl Drop for VisionMemoryPermit {
3804 fn drop(&mut self) {
3805 if self.bytes != 0 {
3806 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3807 }
3808 }
3809}
3810
3811fn try_reserve_vision_memory(
3812 bytes: usize,
3813) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3814 if bytes == 0 {
3815 return Ok(None);
3816 }
3817 if bytes > MAX_VISION_PATCH_BYTES {
3818 return Err(VisionMemoryError::Request(format!(
3819 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3820 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3821 )));
3822 }
3823 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3824 loop {
3825 let Some(next) = in_use.checked_add(bytes) else {
3826 return Err(VisionMemoryError::Capacity(
3827 "vision patch memory reservation overflowed".into(),
3828 ));
3829 };
3830 if next > MAX_VISION_PATCH_BYTES {
3831 return Err(VisionMemoryError::Capacity(format!(
3832 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3833 in_use / (1024 * 1024),
3834 bytes / (1024 * 1024)
3835 )));
3836 }
3837 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3838 in_use,
3839 next,
3840 std::sync::atomic::Ordering::AcqRel,
3841 std::sync::atomic::Ordering::Acquire,
3842 ) {
3843 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3844 Err(actual) => in_use = actual,
3845 }
3846 }
3847}
3848
3849pub(crate) fn vision_memory_error_response(
3850 error: VisionMemoryError,
3851 param: Option<&str>,
3852) -> Response {
3853 match error {
3854 VisionMemoryError::Request(message) => bad_request(&message, param),
3855 VisionMemoryError::Capacity(message) => retry_contract_response(
3856 error_response_coded(
3857 StatusCode::SERVICE_UNAVAILABLE,
3858 &message,
3859 "server_error",
3860 None,
3861 Some("vision_memory_busy"),
3862 ),
3863 Some(RETRY_AFTER_S_OVERLOADED),
3864 ),
3865 }
3866}
3867
3868enum PendingVisionUnit {
3874 Still {
3875 bytes: Vec<u8>,
3876 gh: usize,
3877 gw: usize,
3878 },
3879 Video {
3880 bytes: Vec<u8>,
3881 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3882 video: usize,
3883 },
3884}
3885
3886struct PendingGemmaImage {
3888 bytes: Vec<u8>,
3889 gw: usize,
3890 gh: usize,
3891}
3892
3893struct PendingGlm5Image {
3896 bytes: Vec<u8>,
3897 gh: usize,
3898 gw: usize,
3899}
3900
3901struct PendingStepImage {
3904 bytes: Vec<u8>,
3905 plan: memra_engine::vision_step::StepImagePlan,
3906}
3907
3908fn content_to_text_vision_step(
3918 v: &serde_json::Value,
3919 step_images: &mut Vec<PendingStepImage>,
3920) -> Result<String, String> {
3921 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
3922 let parts = match v {
3923 serde_json::Value::Array(parts) => parts,
3924 _ => return content_to_text(v),
3925 };
3926 let mut out = String::new();
3927 let mut needs_sep = false;
3928 for p in parts {
3929 match p.get("type").and_then(|t| t.as_str()) {
3930 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3931 Some(t) => {
3932 if needs_sep {
3933 out.push(' ');
3934 }
3935 out.push_str(t);
3936 needs_sep = true;
3937 }
3938 None => return Err("content part has no text field".into()),
3939 },
3940 Some("image_url") => {
3941 vision_placement_admits("image")?;
3942 let url = p
3943 .get("image_url")
3944 .and_then(|u| {
3945 if u.is_string() {
3946 u.as_str()
3947 } else {
3948 u.get("url").and_then(|x| x.as_str())
3949 }
3950 })
3951 .ok_or("image_url part has no url")?;
3952 if !url.starts_with("data:") {
3953 return Err(
3954 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3955 );
3956 }
3957 if step_images.len() >= VISION_MAX_IMAGES {
3958 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3959 }
3960 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3964 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3965 let plan = memra_engine::vision_step::step_plan_image(&bytes)
3966 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3967 for i in 0..plan.n_tiles {
3968 out.push_str("<patch_start>");
3969 for _ in 0..SV_TILE_ROWS {
3970 out.push_str("<im_patch>");
3971 }
3972 out.push_str("<patch_end>");
3973 if plan.newline_mask[i] {
3974 out.push_str("<patch_newline>");
3975 }
3976 }
3977 out.push_str("<im_start>");
3978 for _ in 0..SV_MAIN_ROWS {
3979 out.push_str("<im_patch>");
3980 }
3981 out.push_str("<im_end>");
3982 step_images.push(PendingStepImage { bytes, plan });
3983 needs_sep = false;
3984 }
3985 Some("video_url") => {
3986 return Err("step37 has no video input (image-only processor)".into());
3987 }
3988 Some(other) => {
3989 return Err(format!("unsupported content part type {other:?}"));
3990 }
3991 }
3992 }
3993 Ok(out)
3994}
3995
3996fn content_to_text_vision(
4005 v: &serde_json::Value,
4006 images: &mut Vec<PendingVisionUnit>,
4007 gemma_images: &mut Vec<PendingGemmaImage>,
4008 glm5_images: &mut Vec<PendingGlm5Image>,
4009 step_images: &mut Vec<PendingStepImage>,
4010 next_video: &mut usize,
4011) -> Result<String, String> {
4012 if step_vision_enabled() {
4016 return content_to_text_vision_step(v, step_images);
4017 }
4018 let parts = match v {
4019 serde_json::Value::Array(parts) => parts,
4020 _ => return content_to_text(v),
4021 };
4022 let mut out = String::new();
4023 for p in parts {
4024 match p.get("type").and_then(|t| t.as_str()) {
4025 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
4026 Some(t) => out.push_str(t),
4027 None => return Err("content part has no text field".into()),
4028 },
4029 Some("image_url") if glm5_vision_enabled() => {
4030 let url = p
4031 .get("image_url")
4032 .and_then(|u| {
4033 if u.is_string() {
4034 u.as_str()
4035 } else {
4036 u.get("url").and_then(|x| x.as_str())
4037 }
4038 })
4039 .ok_or("image_url part has no url")?;
4040 if !url.starts_with("data:") {
4041 return Err(
4042 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4043 );
4044 }
4045 if glm5_images.len() >= VISION_MAX_IMAGES {
4046 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4047 }
4048 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4052 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4053 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes)
4054 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
4055 out.push_str("<|begin_of_image|>");
4059 for _ in 0..memra_engine::vision_glm5::n_merged_for_grid(gh, gw) {
4060 out.push_str("<|image|>");
4061 }
4062 out.push_str("<|end_of_image|>");
4063 glm5_images.push(PendingGlm5Image { bytes, gh, gw });
4064 }
4065 Some("video_url") if glm5_vision_enabled() => {
4066 return Err(
4067 "glm5 video input is not served (tensor census only; image input is the \
4068 supported surface)"
4069 .into(),
4070 );
4071 }
4072 Some("image_url") if gemma_vision_enabled() => {
4073 vision_placement_admits("image")?;
4074 let url = p
4075 .get("image_url")
4076 .and_then(|u| {
4077 if u.is_string() {
4078 u.as_str()
4079 } else {
4080 u.get("url").and_then(|x| x.as_str())
4081 }
4082 })
4083 .ok_or("image_url part has no url")?;
4084 if !url.starts_with("data:") {
4085 return Err(
4086 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4087 );
4088 }
4089 if gemma_images.len() >= VISION_MAX_IMAGES {
4090 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4091 }
4092 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
4096 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4097 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
4098 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
4099 out.push_str("<|image>");
4101 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
4102 out.push_str("<|image|>");
4103 }
4104 out.push_str("<image|>");
4105 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
4106 }
4107 Some("image_url") => {
4108 if !vision_enabled() {
4109 return Err("image input is not enabled on this deployment".into());
4110 }
4111 vision_placement_admits("image")?;
4112 let url = p
4113 .get("image_url")
4114 .and_then(|u| {
4115 if u.is_string() {
4116 u.as_str()
4117 } else {
4118 u.get("url").and_then(|x| x.as_str())
4119 }
4120 })
4121 .ok_or("image_url part has no url")?;
4122 if !url.starts_with("data:") {
4123 return Err(
4124 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4125 );
4126 }
4127 if images
4128 .iter()
4129 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
4130 .count()
4131 >= VISION_MAX_IMAGES
4132 {
4133 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
4134 }
4135 let bytes = memra_engine::vision_pre::decode_data_uri(url)
4140 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4141 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
4142 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
4143 out.push_str("<|vision_start|>");
4144 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
4145 out.push_str("<|image_pad|>");
4146 }
4147 out.push_str("<|vision_end|>");
4148 images.push(PendingVisionUnit::Still { bytes, gh, gw });
4149 }
4150 Some("video_url") if gemma_vision_enabled() => {
4151 return Err("gemma-4 has no video input (image-only projector)".into());
4152 }
4153 Some("video_url") => {
4154 if !vision_enabled() {
4155 return Err("video input is not enabled on this deployment".into());
4156 }
4157 vision_placement_admits("video")?;
4158 let url = p
4159 .get("video_url")
4160 .and_then(|u| {
4161 if u.is_string() {
4162 u.as_str()
4163 } else {
4164 u.get("url").and_then(|x| x.as_str())
4165 }
4166 })
4167 .ok_or("video_url part has no url")?;
4168 if !url.starts_with("data:") {
4169 return Err(
4170 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
4171 );
4172 }
4173 if *next_video >= 2 {
4174 return Err("too many videos (max 2)".into());
4175 }
4176 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
4179 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
4180 .map_err(|e| format!("video: {e}"))?;
4181 let vidx = *next_video;
4182 *next_video += 1;
4183 for group in &vid.groups {
4185 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
4186 out.push_str("<|vision_start|>");
4187 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
4188 out.push_str("<|video_pad|>");
4189 }
4190 out.push_str("<|vision_end|>");
4191 }
4192 images.push(PendingVisionUnit::Video {
4195 bytes,
4196 groups: vid.groups,
4197 video: vidx,
4198 });
4199 }
4200 Some(other) => {
4201 return Err(format!("unsupported content part type {other:?}"));
4202 }
4203 }
4204 }
4205 Ok(out)
4206}
4207
4208fn pyjson(v: &serde_json::Value, out: &mut String) {
4212 match v {
4213 serde_json::Value::Object(m) => {
4214 out.push('{');
4215 for (i, (k, val)) in m.iter().enumerate() {
4216 if i > 0 {
4217 out.push_str(", ");
4218 }
4219 out.push_str(&serde_json::Value::String(k.clone()).to_string());
4220 out.push_str(": ");
4221 pyjson(val, out);
4222 }
4223 out.push('}');
4224 }
4225 serde_json::Value::Array(a) => {
4226 out.push('[');
4227 for (i, val) in a.iter().enumerate() {
4228 if i > 0 {
4229 out.push_str(", ");
4230 }
4231 pyjson(val, out);
4232 }
4233 out.push(']');
4234 }
4235 scalar => out.push_str(&scalar.to_string()),
4236 }
4237}
4238
4239fn pyjson_str(v: &serde_json::Value) -> String {
4240 let mut s = String::new();
4241 pyjson(v, &mut s);
4242 s
4243}
4244
4245#[allow(clippy::too_many_arguments)] fn sampler_config(
4253 temperature: f32,
4254 top_k: usize,
4255 top_p: f32,
4256 min_p: f32,
4257 frequency_penalty: f32,
4258 presence_penalty: f32,
4259 repetition_penalty: f32,
4260 seed: Option<u64>,
4261) -> SamplerConfig {
4262 let penalties_on =
4263 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
4264 SamplerConfig {
4265 temperature,
4266 top_k,
4267 top_p,
4268 min_p,
4269 penalty_last_n: if penalties_on {
4270 memra_engine::spec::PEN_WINDOW_MAX
4271 } else {
4272 0
4273 },
4274 penalty_repeat: repetition_penalty,
4275 penalty_freq: frequency_penalty,
4276 penalty_present: presence_penalty,
4277 seed: seed.unwrap_or_else(fresh_seed),
4280 }
4281}
4282
4283fn fresh_seed() -> u64 {
4288 use std::sync::atomic::{AtomicU64, Ordering};
4289 static COUNTER: AtomicU64 = AtomicU64::new(0);
4290 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
4291 let nanos = std::time::SystemTime::now()
4292 .duration_since(std::time::UNIX_EPOCH)
4293 .map(|d| d.as_nanos() as u64)
4294 .unwrap_or(0);
4295 let mut z = nanos
4296 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
4297 .wrapping_add(0x9E3779B97F4A7C15);
4298 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
4299 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
4300 z ^= z >> 31;
4301 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
4304}
4305
4306fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
4311 for (param, present, why) in fields {
4312 if *present {
4313 return Err((format!("{param} is not supported{why}"), param.to_string()));
4314 }
4315 }
4316 Ok(())
4317}
4318
4319#[derive(PartialEq)]
4320enum ToolChoice {
4321 Auto,
4322 None,
4323}
4324
4325fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
4326 match v {
4327 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
4328 Some(serde_json::Value::String(s)) => match s.as_str() {
4329 "auto" => Ok(ToolChoice::Auto),
4330 "none" => Ok(ToolChoice::None),
4331 "required" => Err("tool_choice \"required\" is not supported (no constrained \
4332 decoding); use \"auto\""
4333 .into()),
4334 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
4335 },
4336 Some(serde_json::Value::Object(_)) => {
4337 Err("named-function tool_choice is not supported; use \"auto\"".into())
4338 }
4339 Some(other) => Err(format!("bad tool_choice: {other}")),
4340 }
4341}
4342
4343fn parse_think(
4393 reasoning_effort: &Option<String>,
4394 reasoning: &Option<serde_json::Value>,
4395 vllm_switch: Option<bool>,
4396 suppress_switch: Option<bool>,
4397 default_effort: Option<&str>,
4398 max_tier: bool,
4399) -> Result<(ThinkMode, Option<String>, bool), String> {
4400 let mut effort = reasoning_effort.clone();
4401 let ReasoningObject {
4402 mut enabled,
4403 effort: object_effort,
4404 exclude,
4405 } = parse_reasoning_object(reasoning)?;
4406 if let Some(e) = object_effort {
4407 effort = Some(e);
4408 }
4409 match (enabled, vllm_switch) {
4414 (Some(a), Some(b)) if a != b => {
4415 return Err(format!(
4416 "contradictory reasoning switches: reasoning.enabled={a} and \
4417 enable_thinking={b} — send one"
4418 ));
4419 }
4420 (None, Some(b)) => enabled = Some(b),
4421 _ => {}
4422 }
4423 let suppress = match (exclude, suppress_switch) {
4435 (Some(true), _) | (_, Some(false)) => Some(false),
4436 _ => None,
4437 };
4438 match (enabled, suppress) {
4439 (Some(true), Some(false)) => {
4440 return Err(
4441 "contradictory reasoning switches: reasoning is enabled but \
4442 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
4443 on this server not delivering reasoning means not generating it, so send one"
4444 .into(),
4445 );
4446 }
4447 (None, Some(b)) => enabled = Some(b),
4448 _ => {}
4449 }
4450 let client_explicit = effort.is_some() || enabled.is_some();
4454 if effort.is_none() && enabled.is_none() {
4459 effort = default_effort.map(str::to_string);
4460 }
4461 let effort_arm = match effort.as_deref() {
4466 None => None,
4467 Some(raw) => {
4468 let level = canonical_effort_for(raw, max_tier).ok_or_else(|| {
4469 format!(
4470 "bad reasoning_effort {raw:?} \
4471 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
4472 highest level this model's template distinguishes)"
4473 )
4474 })?;
4475 Some(match level {
4476 "none" | "minimal" => (ThinkMode::NoThink, "low"),
4477 "low" => (ThinkMode::Think, "low"),
4478 "medium" => (ThinkMode::Think, "medium"),
4479 "max" => (ThinkMode::Think, "max"),
4480 _ => (ThinkMode::Think, "high"),
4481 })
4482 }
4483 };
4484 let (think, level) = match (enabled, effort_arm) {
4485 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
4488 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
4489 (None, Some((think, level))) => (think, Some(level.to_string())),
4490 (None, None) => (ThinkMode::Default, None),
4491 };
4492 Ok((think, level, client_explicit))
4493}
4494
4495struct ReasoningObject {
4497 enabled: Option<bool>,
4498 effort: Option<String>,
4499 exclude: Option<bool>,
4500}
4501
4502fn parse_reasoning_object(
4521 reasoning: &Option<serde_json::Value>,
4522) -> Result<ReasoningObject, String> {
4523 let mut out = ReasoningObject {
4524 enabled: None,
4525 effort: None,
4526 exclude: None,
4527 };
4528 let Some(v) = reasoning else { return Ok(out) };
4529 let obj = match v {
4530 serde_json::Value::Null => return Ok(out),
4531 serde_json::Value::Object(obj) => obj,
4532 _ => return Err("reasoning must be an object".into()),
4533 };
4534 for (key, value) in obj {
4535 match key.as_str() {
4543 "enabled" => {
4544 if !value.is_null() {
4545 out.enabled = Some(
4546 value
4547 .as_bool()
4548 .ok_or("reasoning.enabled must be true or false")?,
4549 );
4550 }
4551 }
4552 "exclude" => {
4553 if !value.is_null() {
4554 out.exclude = Some(
4555 value
4556 .as_bool()
4557 .ok_or("reasoning.exclude must be true or false")?,
4558 );
4559 }
4560 }
4561 "effort" => {
4562 if !value.is_null() {
4563 out.effort = Some(
4564 value
4565 .as_str()
4566 .ok_or("reasoning.effort must be a string")?
4567 .to_string(),
4568 );
4569 }
4570 }
4571 "max_tokens" => {
4572 return Err(
4573 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4574 are output tokens here, and max_tokens is the ONE output budget covering \
4575 reasoning and content together — there is no separate reasoning budget to \
4576 spend against, so honouring this field is impossible rather than merely \
4577 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4578 reasoning.enabled:false) to spend less of it on reasoning"
4579 .into(),
4580 );
4581 }
4582 other => {
4583 return Err(format!(
4584 "reasoning.{other} is not a field this server implements (it would change \
4585 nothing about the request); the supported keys are enabled, effort and \
4586 exclude"
4587 ));
4588 }
4589 }
4590 }
4591 Ok(out)
4592}
4593
4594fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4615 let Some(v) = kwargs else { return Ok(None) };
4616 let obj = match v {
4617 serde_json::Value::Null => return Ok(None),
4618 serde_json::Value::Object(obj) => obj,
4619 _ => return Err("chat_template_kwargs must be an object".into()),
4620 };
4621 let mut switch = None;
4622 for (key, value) in obj {
4623 match key.as_str() {
4624 "enable_thinking" => {
4625 switch = Some(
4626 value
4627 .as_bool()
4628 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4629 );
4630 }
4631 "preserve_thinking" => {
4632 let preserve = value
4633 .as_bool()
4634 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4635 if !preserve {
4636 return Err(
4637 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4638 server: the renderer implements the vendor DEFAULT (replay every prior \
4639 assistant turn's <think> block, empty when no reasoning was sent) but \
4640 not the strip arm — serving replay bytes under a strip request would \
4641 misdescribe the prompt. Omit the flag or send true"
4642 .into(),
4643 );
4644 }
4645 }
4647 other => {
4648 return Err(format!(
4649 "chat_template_kwargs.{other} is not supported by this server's \
4650 template renderer (it would change nothing about the prompt); the only \
4651 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4652 refuses in both directions — see its own message)"
4653 ));
4654 }
4655 }
4656 }
4657 Ok(switch)
4658}
4659
4660fn resolve_vllm_think_switch(
4664 enable_thinking: Option<bool>,
4665 kwargs: &Option<serde_json::Value>,
4666) -> Result<Option<bool>, String> {
4667 let from_kwargs = parse_template_kwargs(kwargs)?;
4668 match (enable_thinking, from_kwargs) {
4669 (Some(a), Some(b)) if a != b => Err(format!(
4670 "contradictory reasoning switches: enable_thinking={a} and \
4671 chat_template_kwargs.enable_thinking={b} — send one"
4672 )),
4673 (Some(a), _) => Ok(Some(a)),
4674 (None, b) => Ok(b),
4675 }
4676}
4677
4678pub(crate) fn canonical_effort_for(value: &str, max_tier: bool) -> Option<&'static str> {
4698 match value {
4699 "none" => Some("none"),
4700 "minimal" => Some("minimal"),
4701 "low" => Some("low"),
4702 "medium" => Some("medium"),
4703 "high" => Some("high"),
4704 "xhigh" | "max" | "ultra" => Some(if max_tier { "max" } else { "high" }),
4710 _ => None,
4711 }
4712}
4713
4714pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4717 canonical_effort_for(value, false)
4718}
4719
4720fn json_to_val(v: &serde_json::Value) -> chat::Val {
4723 match v {
4724 serde_json::Value::Null => chat::Val::Null,
4725 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4726 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4727 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4728 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4729 serde_json::Value::Object(o) => chat::Val::Obj(
4732 o.iter()
4733 .map(|(k, val)| (k.clone(), json_to_val(val)))
4734 .collect(),
4735 ),
4736 }
4737}
4738
4739#[allow(clippy::type_complexity)]
4743fn prepare_tools(
4744 tools: &[serde_json::Value],
4745) -> Result<
4746 (
4747 Vec<String>,
4748 Vec<chat::Val>,
4749 HashMap<String, HashMap<String, String>>,
4750 ),
4751 String,
4752> {
4753 let mut tools_json = Vec::with_capacity(tools.len());
4754 let mut tools_struct = Vec::with_capacity(tools.len());
4755 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4756 for t in tools {
4757 let f = t
4758 .get("function")
4759 .ok_or("each tool needs a function object")?;
4760 let name = f
4761 .get("name")
4762 .and_then(|n| n.as_str())
4763 .ok_or("each tool needs function.name")?;
4764 let mut params: HashMap<String, String> = HashMap::new();
4765 if let Some(props) = f
4766 .get("parameters")
4767 .and_then(|p| p.get("properties"))
4768 .and_then(|p| p.as_object())
4769 {
4770 for (p, def) in props {
4771 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4772 params.insert(p.clone(), ty.to_string());
4773 }
4774 }
4775 }
4776 schemas.insert(name.to_string(), params);
4777 tools_json.push(pyjson_str(t));
4778 tools_struct.push(json_to_val(f));
4780 }
4781 Ok((tools_json, tools_struct, schemas))
4782}
4783
4784fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4789 let parsed: serde_json::Value = match &tc.function.arguments {
4790 serde_json::Value::Null => json!({}),
4791 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4792 serde_json::Value::String(s) => serde_json::from_str(s)
4793 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4794 v @ serde_json::Value::Object(_) => v.clone(),
4795 _ => return Err("tool_calls arguments must be a JSON object".into()),
4796 };
4797 let obj = parsed
4798 .as_object()
4799 .ok_or("tool_calls arguments must decode to a JSON object")?;
4800 let params = obj
4801 .iter()
4802 .map(|(k, v)| {
4803 let rendered = match v {
4804 serde_json::Value::String(s) => s.clone(),
4805 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4806 scalar => scalar.to_string(),
4807 };
4808 (k.clone(), rendered)
4809 })
4810 .collect();
4811 let args = obj
4814 .iter()
4815 .map(|(k, v)| (k.clone(), json_to_val(v)))
4816 .collect();
4817 Ok(TmplToolCall {
4818 name: tc.function.name.clone(),
4819 params,
4820 args,
4821 id: tc.id.clone(),
4822 })
4823}
4824
4825fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4827 json!({ "id": c.id, "type": "function",
4828 "function": { "name": c.name, "arguments": c.arguments } })
4829}
4830
4831async fn serve_bounded_http_with_limits<F>(
4835 listener: tokio::net::TcpListener,
4836 app: Router,
4837 shutdown: F,
4838 header_read_timeout: std::time::Duration,
4839 max_connections: usize,
4840 connection_max_lifetime: std::time::Duration,
4841) -> std::io::Result<()>
4842where
4843 F: std::future::Future<Output = ()> + Send,
4844{
4845 let connections = Arc::new(tokio::sync::Semaphore::new(max_connections));
4846 let (connection_shutdown, _) = tokio::sync::watch::channel(false);
4847 let mut connection_tasks = tokio::task::JoinSet::new();
4848 let mut shutdown = Box::pin(shutdown);
4849
4850 loop {
4851 tokio::select! {
4852 _ = &mut shutdown => break,
4853 joined = connection_tasks.join_next(), if !connection_tasks.is_empty() => {
4854 if let Some(Err(error)) = joined {
4855 eprintln!("[server] connection task failed: {error}");
4856 }
4857 }
4858 accepted = listener.accept() => {
4859 let (stream, _) = match accepted {
4860 Ok(connection) => connection,
4861 Err(error) => {
4862 eprintln!("[server] accept failed: {error}");
4863 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4864 continue;
4865 }
4866 };
4867 let permit = match connections.clone().try_acquire_owned() {
4868 Ok(permit) => permit,
4869 Err(_) => {
4870 drop(stream);
4871 continue;
4872 }
4873 };
4874 let service = app.clone().map_request(
4875 |request: hyper::Request<hyper::body::Incoming>| request.map(Body::new),
4876 );
4877 let service = hyper_util::service::TowerToHyperService::new(service);
4878 let io = hyper_util::rt::TokioIo::new(stream);
4879 let mut builder = hyper_util::server::conn::auto::Builder::new(
4880 hyper_util::rt::TokioExecutor::new(),
4881 );
4882 builder
4883 .http1()
4884 .timer(hyper_util::rt::TokioTimer::new())
4885 .header_read_timeout(header_read_timeout)
4886 .max_headers(64);
4887 builder
4888 .http2()
4889 .timer(hyper_util::rt::TokioTimer::new())
4890 .max_concurrent_streams(MAX_HTTP2_STREAMS_PER_CONNECTION)
4891 .keep_alive_interval(Some(std::time::Duration::from_secs(30)))
4892 .keep_alive_timeout(std::time::Duration::from_secs(10));
4893 let mut connection = Box::pin(builder
4894 .serve_connection_with_upgrades(io, service)
4895 .into_owned());
4896 let mut shutdown_rx = connection_shutdown.subscribe();
4897 connection_tasks.spawn(async move {
4898 let _permit = permit;
4899 tokio::select! {
4900 result = connection.as_mut() => {
4901 let _ = result;
4902 }
4903 _ = tokio::time::sleep(connection_max_lifetime) => {
4904 connection.as_mut().graceful_shutdown();
4909 let _ = connection.await;
4910 }
4911 _ = shutdown_rx.changed() => {
4912 connection.as_mut().graceful_shutdown();
4913 let _ = connection.await;
4914 }
4915 }
4916 });
4917 }
4918 }
4919 }
4920 drop(listener);
4921 let _ = connection_shutdown.send(true);
4922 let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
4923 while connection_tasks.join_next().await.is_some() {}
4924 })
4925 .await;
4926 if drained.is_err() {
4927 connection_tasks.abort_all();
4928 eprintln!("[server] WARN: HTTP connections exceeded the 5s graceful close deadline");
4929 }
4930 Ok(())
4931}
4932
4933async fn serve_bounded_http<F>(
4934 listener: tokio::net::TcpListener,
4935 app: Router,
4936 shutdown: F,
4937) -> std::io::Result<()>
4938where
4939 F: std::future::Future<Output = ()> + Send,
4940{
4941 serve_bounded_http_with_limits(
4942 listener,
4943 app,
4944 shutdown,
4945 HTTP1_HEADER_READ_TIMEOUT,
4946 MAX_HTTP_CONNECTIONS,
4947 HTTP_CONNECTION_MAX_LIFETIME,
4948 )
4949 .await
4950}
4951
4952#[tokio::main]
4953pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4954 serve_with(ServerWiring::stock()).await
4955}
4956
4957enum MeteringWiring {
4959 Stock,
4963 Custom(metering::MeteringFactory),
4968}
4969
4970pub struct ServerWiring {
4974 metering: MeteringWiring,
4975 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
4978 claimed_env: Vec<&'static str>,
4982}
4983
4984impl ServerWiring {
4985 pub fn stock() -> Self {
4987 ServerWiring {
4988 metering: MeteringWiring::Stock,
4989 on_ready: None,
4990 claimed_env: Vec::new(),
4991 }
4992 }
4993
4994 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
4997 ServerWiring {
4998 metering: MeteringWiring::Custom(factory),
4999 on_ready: None,
5000 claimed_env: Vec::new(),
5001 }
5002 }
5003
5004 pub fn claiming(mut self, var: &'static str) -> Self {
5008 self.claimed_env.push(var);
5009 self
5010 }
5011
5012 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
5013 self.on_ready = Some(Box::new(hook));
5014 self
5015 }
5016}
5017
5018pub struct RuntimeHandles {
5021 pub trim: TrimHandle,
5022 pub purge: PurgeHandle,
5025 pub kv_handoff: HostHandoffHandle,
5031 pub shutdown: tokio::sync::watch::Receiver<bool>,
5036}
5037
5038#[derive(Clone)]
5041pub struct TrimHandle {
5042 cmd_tx: Sender<Cmd>,
5043}
5044
5045impl TrimHandle {
5046 pub async fn trim(&self) -> Result<serde_json::Value, String> {
5048 let (tx, rx) = tokio::sync::oneshot::channel();
5049 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
5050 return Err("worker is down".into());
5051 }
5052 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5053 Ok(Ok(report)) => Ok(json!(report)),
5054 _ => Err("worker did not answer the trim within 30s".into()),
5055 }
5056 }
5057}
5058
5059#[derive(Clone)]
5068pub struct PurgeHandle {
5069 cmd_tx: Sender<Cmd>,
5070}
5071
5072impl PurgeHandle {
5073 pub async fn purge_tenant(&self, tenant: &str) -> Result<serde_json::Value, String> {
5075 let (tx, rx) = tokio::sync::oneshot::channel();
5076 let cmd = Cmd::PurgeTenantHost {
5077 tenant: tenant.to_string(),
5078 tx,
5079 };
5080 if self.cmd_tx.send(cmd).is_err() {
5081 return Err("worker is down".into());
5082 }
5083 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5084 Ok(Ok(report)) => Ok(json!(report)),
5085 _ => Err("worker did not answer the purge within 30s".into()),
5086 }
5087 }
5088}
5089
5090#[derive(Clone)]
5099pub struct HostHandoffHandle {
5100 cmd_tx: Sender<Cmd>,
5101}
5102
5103impl HostHandoffHandle {
5104 pub async fn export(&self, force: bool) -> Result<serde_json::Value, String> {
5107 let (tx, rx) = tokio::sync::oneshot::channel();
5108 if self
5109 .cmd_tx
5110 .send(Cmd::ExportHostHandoff { force, tx })
5111 .is_err()
5112 {
5113 return Err("worker is down".into());
5114 }
5115 match tokio::time::timeout(std::time::Duration::from_secs(900), rx).await {
5116 Ok(Ok(Ok(report))) => Ok(json!(report)),
5117 Ok(Ok(Err(refused))) => Err(refused),
5118 _ => Err("worker did not answer the export within 900s".into()),
5119 }
5120 }
5121
5122 pub async fn import(&self) -> Result<serde_json::Value, String> {
5125 let (tx, rx) = tokio::sync::oneshot::channel();
5126 if self.cmd_tx.send(Cmd::ImportHostHandoff { tx }).is_err() {
5127 return Err("worker is down".into());
5128 }
5129 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
5130 Ok(Ok(Ok(start))) => Ok(json!(start)),
5131 Ok(Ok(Err(refused))) => Err(refused),
5132 _ => Err("worker did not answer the import within 30s".into()),
5133 }
5134 }
5135}
5136
5137pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
5138 let args: Vec<String> = std::env::args().skip(1).collect();
5141 if args.iter().any(|a| a == "--version" || a == "-V") {
5146 println!("memra-server {}", env!("CARGO_PKG_VERSION"));
5147 println!("system_fingerprint {SYSTEM_FINGERPRINT}");
5148 println!("build_id_src {BUILD_ID_SRC}");
5149 println!("git_sha {BUILD_GIT_SHA}");
5150 if !BUILD_ID_NOTE.is_empty() {
5151 println!("degraded {BUILD_ID_NOTE}");
5152 }
5153 return Ok(());
5154 }
5155 if let Some(code) = auth::run_cli(&args) {
5156 std::process::exit(code);
5157 }
5158 eprintln!("{}", build_identity_line());
5162 if BUILD_ID_SRC != build_id::BUILD_ID_SRC_TREE {
5163 eprintln!(
5164 "[server] WARNING: build identity is DEGRADED: {BUILD_ID_NOTE}. \
5165 system_fingerprint {SYSTEM_FINGERPRINT} carries a version-only id, so it does \
5166 NOT identify the source this binary was compiled from and published \
5167 performance pins cannot be verified against it (darklanes \
5168 tools/check-claim-builds.mjs --live). Rebuild where the workspace source tree \
5169 is readable."
5170 );
5171 }
5172 auth::init_from_env();
5175 let api_auth = match ApiAuth::from_env() {
5176 Ok(auth) => auth,
5177 Err(err) => {
5178 eprintln!("[server] FATAL: {err}");
5179 std::process::exit(1);
5180 }
5181 };
5182 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
5183 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
5184 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
5185 Ok(resolved) => resolved,
5186 Err(err) => {
5187 eprintln!("[server] FATAL: {err}");
5188 std::process::exit(1);
5189 }
5190 };
5191 if let Err(message) = validate_bind_security(&addr, api_auth.configured(), allow_open_bind) {
5196 eprintln!("[server] FATAL: {message}");
5197 std::process::exit(1);
5198 }
5199 if !bind_loopback && !api_auth.configured() {
5200 eprintln!(
5201 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
5202 metrics remain bearer-protected"
5203 );
5204 }
5205 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
5206 Ok(token) if token.is_empty() => {
5207 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
5208 std::process::exit(1);
5209 }
5210 Ok(token) => Some(token),
5211 Err(std::env::VarError::NotPresent) => None,
5212 Err(std::env::VarError::NotUnicode(_)) => {
5213 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
5214 std::process::exit(1);
5215 }
5216 };
5217 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
5218
5219 let models = parse_models_config();
5220 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
5221 Ok(loaded) => loaded,
5222 Err(err) => {
5223 eprintln!("[server] FATAL: {err}");
5224 std::process::exit(1);
5225 }
5226 };
5227 let metering_obj: Option<Arc<dyn metering::Metering>> = {
5233 let factory = match wiring.metering {
5234 MeteringWiring::Stock => None,
5235 MeteringWiring::Custom(factory) => Some(factory),
5236 };
5237 for deployment_only in [
5238 "MEMRA_REQUEST_LEDGER",
5239 "MEMRA_TENANT_BUDGETS",
5240 "MEMRA_ADMIN_ADDR",
5241 "MEMRA_ADMIN_TOKEN_FILE",
5242 "MEMRA_CAPTURE_DIR",
5243 ] {
5244 if std::env::var_os(deployment_only).is_some()
5245 && !wiring.claimed_env.contains(&deployment_only)
5246 {
5247 eprintln!(
5248 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
5249 build ships no accounting/admin/capture. Wire a Metering implementation \
5250 through ServerWiring and claim the vars it consumes."
5251 );
5252 std::process::exit(1);
5253 }
5254 }
5255 match factory {
5256 None => None,
5257 Some(factory) => {
5258 let model_ids: Vec<String> =
5259 models.iter().map(|(name, _, _)| name.clone()).collect();
5260 match factory(&metering::MeteringInit { models: &model_ids }) {
5261 Ok(metering_obj) => metering_obj,
5262 Err(err) => {
5263 eprintln!("[server] FATAL: metering wiring: {err}");
5264 std::process::exit(1);
5265 }
5266 }
5267 }
5268 }
5269 };
5270 let budget_tokenizers = if metering_obj
5271 .as_ref()
5272 .is_some_and(|manager| manager.enforces_limits())
5273 {
5274 match load_budget_tokenizers(&models) {
5275 Ok(tokenizers) => Some(tokenizers),
5276 Err(err) => {
5277 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
5278 std::process::exit(1);
5279 }
5280 }
5281 } else {
5282 None
5283 };
5284 eprintln!("[server] starting; models config = {models:?}");
5285
5286 let health_state = health::WorkerHealth::new();
5291 health::spawn_gpu_watch(health_state.clone());
5295 health::spawn_sd_watchdog(health_state.clone());
5296
5297 let (cmd_tx, model_names, caps, metrics, worker_thread) =
5299 match worker::spawn(models, health_state.clone()) {
5300 Ok(v) => v,
5301 Err(err) => {
5302 eprintln!("[server] FATAL: worker init failed: {err}");
5303 health_state.mark_dead(format!("worker init failed: {err}"));
5304 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
5305 std::process::exit(1);
5306 }
5307 };
5308 eprintln!("[server] worker ready; serving models: {model_names:?}");
5309
5310 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
5317 if let Some(on_ready) = wiring.on_ready {
5318 on_ready(RuntimeHandles {
5319 trim: TrimHandle {
5320 cmd_tx: cmd_tx.clone(),
5321 },
5322 purge: PurgeHandle {
5323 cmd_tx: cmd_tx.clone(),
5324 },
5325 kv_handoff: HostHandoffHandle {
5326 cmd_tx: cmd_tx.clone(),
5327 },
5328 shutdown: drain_shutdown_rx.clone(),
5329 });
5330 }
5331
5332 let bg_handle = darklane::spawn_from_env(health_state.clone());
5335 let bg_state = bg_handle.as_ref().map(|h| {
5336 let mode = darklane::BgConfig::from_env()
5337 .map(|c| c.yield_mode.as_str())
5338 .unwrap_or("stop");
5339 (h.state.clone(), mode)
5340 });
5341
5342 let state = AppState {
5343 cmd_tx,
5344 models: model_names,
5345 caps,
5346 openrouter_metadata: Arc::new(openrouter_metadata),
5347 provider_metadata: Arc::new(provider_metadata),
5348 metering: metering_obj,
5349 budget_tokenizers,
5350 api_auth,
5351 metrics_auth,
5352 metrics,
5353 inflight: Arc::new(Default::default()),
5354 tenant_inflight: Arc::new(Default::default()),
5355 health: health_state.clone(),
5356 bg: bg_state,
5357 };
5358 let inflight_handle = state.inflight.clone();
5359 let drain_metering = state.metering.clone();
5362 worker::register_http_inflight(state.inflight.clone());
5367 let app = Router::new()
5368 .route("/health", get(health_live))
5373 .route("/livez", get(health_live))
5374 .route("/readyz", get(health_ready))
5375 .route("/models", get(list_models))
5376 .route("/v1/models", get(list_models_v1))
5377 .route("/v1/auth/check", get(auth_check))
5378 .route("/v1/completions", post(completions_admitted))
5379 .route("/v1/embeddings", post(embed_api::embeddings_admitted))
5380 .route("/v1/rerank", post(embed_api::rerank_admitted))
5381 .route("/v1/chat/completions", post(chat_completions_admitted))
5382 .route("/v1/messages", post(anthropic::messages_admitted))
5386 .route("/v1/responses", post(responses_api::responses_admitted))
5387 .route("/metrics", get(get_metrics))
5388 .route("/yield/metrics", get(yield_metrics))
5389 .with_state(state.clone());
5390 let app = apply_body_limit(app);
5393 let app = app.layer(middleware::from_fn_with_state(
5396 state,
5397 authenticate_inference_before_body,
5398 ));
5399 let app = if ttft::enabled() {
5400 app.layer(middleware::from_fn(ttft_request_start))
5401 } else {
5402 app
5403 };
5404
5405 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
5406 eprintln!("[server] listening on http://{bind_addr}");
5407 drop(drain_shutdown_rx);
5408 health::sd_notify("READY=1\nSTATUS=serving");
5412 let inflight = inflight_handle;
5419 let signal_admin_shutdown = drain_shutdown_tx.clone();
5420 let serve_result = serve_bounded_http(listener, app, async move {
5421 let mut sigterm =
5422 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
5423 Ok(s) => s,
5424 Err(err) => {
5425 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
5426 std::future::pending::<()>().await;
5427 unreachable!()
5428 }
5429 };
5430 sigterm.recv().await;
5431 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
5432 let _ = signal_admin_shutdown.send(true);
5433 health::sd_notify(&format!(
5437 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
5438 (drain_deadline_s() + 5) * 1_000_000
5439 ));
5440 let n: usize = inflight
5441 .iter()
5442 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5443 .sum();
5444 eprintln!(
5445 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
5446 drain_deadline_s()
5447 );
5448 let deadline = std::time::Duration::from_secs(drain_deadline_s());
5449 let t0 = std::time::Instant::now();
5450 loop {
5451 let n: usize = inflight
5452 .iter()
5453 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
5454 .sum();
5455 if n == 0 {
5456 eprintln!(
5457 "[server] drain complete in {:.1}s; exiting",
5458 t0.elapsed().as_secs_f64()
5459 );
5460 break;
5461 }
5462 if t0.elapsed() >= deadline {
5463 eprintln!(
5464 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
5465 drain_deadline_s()
5466 );
5467 if let Some(metering) = drain_metering.as_ref() {
5474 metering.drain_kill();
5475 }
5476 break;
5477 }
5478 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
5479 }
5480 })
5481 .await;
5482 let _ = drain_shutdown_tx.send(true);
5485 serve_result?;
5486 if let Some(h) = bg_handle {
5490 h.shutdown();
5491 }
5492 worker_thread.join().map_err(|_| {
5501 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
5502 })?;
5503 eprintln!("[server] GPU worker shutdown complete");
5504 Ok(())
5505}
5506
5507fn validate_model_path(path: &str) -> Result<(), String> {
5513 let p = std::path::Path::new(path);
5514 if !p.exists() {
5515 return Err(format!("model path {path:?} does not exist"));
5516 }
5517 if p.is_file() {
5518 return Ok(()); }
5520 if p.join("manifest.json").exists() {
5521 return Ok(()); }
5523 let has_st =
5524 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
5525 if !has_st {
5526 return Err(format!(
5527 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
5528 model.safetensors.index.json + config.json (HF safetensors dir), or \
5529 manifest.json (memra repack dir)"
5530 ));
5531 }
5532 if !p.join("config.json").exists() {
5533 return Err(format!(
5534 "model dir {path:?} has safetensors weights but no config.json"
5535 ));
5536 }
5537 Ok(())
5538}
5539
5540fn parse_models_config() -> Vec<(String, String, Option<String>)> {
5548 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
5549 let mut out = Vec::new();
5550 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
5551 if let Some((name, path)) = entry.split_once('=') {
5552 let (mpath, dpath) = match path.trim().split_once('+') {
5555 Some((m, d)) => (m.trim(), Some(d.trim())),
5556 None => (path.trim(), None),
5557 };
5558 let resolve = |p: &str| {
5559 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
5560 eprintln!("[server] FATAL: model {name:?}: {err}");
5561 std::process::exit(1);
5562 })
5563 };
5564 let mpath = resolve(mpath);
5565 if let Err(err) = validate_model_path(&mpath) {
5566 eprintln!("[server] FATAL: model {name:?}: {err}");
5567 std::process::exit(1);
5568 }
5569 let dpath = dpath.map(|d| {
5579 let d = resolve(d);
5580 let p = std::path::Path::new(&d);
5581 if !p.exists() {
5582 eprintln!(
5583 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
5584 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
5585 rather than serving plain decode under a config that asked \
5586 for speculative decoding."
5587 );
5588 std::process::exit(1);
5589 }
5590 if !p.is_file() {
5591 eprintln!(
5592 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
5593 file — a '+draft' attach must be a NextN/MTP GGUF file."
5594 );
5595 std::process::exit(1);
5596 }
5597 d
5598 });
5599 out.push((name.trim().to_string(), mpath, dpath));
5600 } else {
5601 eprintln!(
5602 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
5603 );
5604 }
5605 }
5606 if !out.is_empty() {
5607 return out;
5608 }
5609 }
5610 vec![
5612 (
5613 "main".into(),
5614 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
5615 None,
5616 ),
5617 (
5618 "judge".into(),
5619 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
5620 None,
5621 ),
5622 ]
5623}
5624
5625fn load_budget_tokenizers(
5626 models: &[(String, String, Option<String>)],
5627) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
5628 let mut tokenizers = HashMap::new();
5629 for (alias, path, _) in models {
5630 let path = std::path::Path::new(path);
5631 let tokenizer = if path.is_dir() {
5632 let tokenizer_dir = if path.join("manifest.json").exists() {
5633 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
5634 format!("model {alias:?}: open repack tokenizer source: {err}")
5635 })?;
5636 repack
5637 .source_dir()
5638 .filter(|source| source.join("tokenizer.json").exists())
5639 .unwrap_or(path)
5640 .to_path_buf()
5641 } else {
5642 path.to_path_buf()
5643 };
5644 Tokenizer::from_hf_dir(&tokenizer_dir)
5645 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5646 } else {
5647 let gguf = memra_gguf::GgufFile::open(path)
5648 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
5649 Tokenizer::from_gguf(&gguf)
5650 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
5651 };
5652 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
5653 }
5654 Ok(Arc::new(tokenizers))
5655}
5656
5657fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5659 let s = st.health.snapshot();
5660 let mut v = json!({
5661 "status": status,
5662 "models": *st.models,
5663 "worker": {
5664 "phase": health::phase_name(s.phase),
5665 "beat_age_ms": s.beat_age_ms,
5666 "tick_max_ms": s.tick_max_ms,
5667 "stall_threshold_ms": s.stall_threshold_ms,
5668 "generation": s.generation,
5669 "xid_warnings": s.xid_warns,
5670 },
5671 });
5672 if let Some(d) = detail {
5673 v["detail"] = json!(d);
5674 }
5675 v
5676}
5677
5678fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
5682 let mut v = health_payload(st, status, detail);
5683 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
5684 v
5685}
5686
5687async fn auth_check() -> impl IntoResponse {
5691 StatusCode::NO_CONTENT
5692}
5693
5694async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
5710 if draining() {
5711 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
5714 }
5715 match st.health.live() {
5716 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
5717 Err(why) => retry_contract_response(
5718 (
5719 StatusCode::SERVICE_UNAVAILABLE,
5720 Json(health_payload(&st, "unhealthy", Some(&why))),
5721 )
5722 .into_response(),
5723 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
5724 ),
5725 }
5726}
5727
5728async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
5739 let is_draining = draining();
5740 match st.health.ready(is_draining) {
5741 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
5742 Err(why) => retry_contract_response(
5743 (
5744 StatusCode::SERVICE_UNAVAILABLE,
5745 Json(readiness_payload(&st, "not_ready", Some(&why))),
5746 )
5747 .into_response(),
5748 Some(if is_draining {
5749 drain_deadline_s()
5750 } else {
5751 worker::WORKER_RESPAWN_BACKOFF_BASE_S
5752 }),
5753 ),
5754 }
5755}
5756
5757#[derive(Clone, Copy)]
5758struct DualPpMetricsSnapshot {
5759 stage_ns: [u64; 4],
5760 stage_samples: [usize; 4],
5761 dropped_timing_samples: usize,
5762 overlaps: usize,
5763 slot_pairs: usize,
5764 slot_uses: [usize; 2],
5765 slot_collisions: usize,
5766}
5767
5768impl DualPpMetricsSnapshot {
5769 fn current() -> Self {
5770 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
5771 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
5772 Self {
5773 stage_ns,
5774 stage_samples,
5775 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
5776 overlaps: memra_engine::pp::dual_pp_overlaps(),
5777 slot_pairs,
5778 slot_uses,
5779 slot_collisions,
5780 }
5781 }
5782
5783 fn populated(self) -> bool {
5784 self.stage_samples.iter().any(|&n| n > 0)
5785 || self.dropped_timing_samples > 0
5786 || self.slot_pairs > 0
5787 || self.slot_collisions > 0
5788 }
5789}
5790
5791fn insert_dual_pp_metrics(
5792 body: &mut serde_json::Value,
5793 metrics_scope: &MetricsScope,
5794 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
5795) {
5796 if !metrics_scope.operator() {
5799 return;
5800 }
5801 let snapshot = snapshot();
5802 if !snapshot.populated() {
5803 return;
5804 }
5805 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
5806 .iter()
5807 .enumerate()
5808 .map(|(i, name)| {
5809 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
5810 (
5811 name.to_string(),
5812 json!({
5813 "samples": snapshot.stage_samples[i],
5814 "total_ms": total_ms,
5815 "mean_ms": if snapshot.stage_samples[i] > 0 {
5816 total_ms / snapshot.stage_samples[i] as f64
5817 } else { 0.0 },
5818 }),
5819 )
5820 })
5821 .collect();
5822 body["dual_pp"] = json!({
5823 "overlaps": snapshot.overlaps,
5824 "slot_pairs": snapshot.slot_pairs,
5825 "slot_uses": snapshot.slot_uses,
5826 "slot_collisions": snapshot.slot_collisions,
5827 "cuda_event_spans": timings,
5828 "dropped_timing_samples": snapshot.dropped_timing_samples,
5829 });
5830}
5831
5832#[derive(Clone, Copy)]
5833struct PpWaveMetricsSnapshot {
5834 ticks: usize,
5835 cells: usize,
5836 overlaps: usize,
5837}
5838
5839impl PpWaveMetricsSnapshot {
5840 fn current() -> Self {
5841 let (ticks, cells, overlaps) = memra_engine::pp::pp_wave_snapshot();
5842 Self {
5843 ticks,
5844 cells,
5845 overlaps,
5846 }
5847 }
5848}
5849
5850fn insert_pp_wave_metrics(
5851 body: &mut serde_json::Value,
5852 metrics_scope: &MetricsScope,
5853 snapshot: impl FnOnce() -> PpWaveMetricsSnapshot,
5854) {
5855 if !metrics_scope.operator() {
5856 return;
5857 }
5858 let snapshot = snapshot();
5859 if snapshot.ticks == 0 && snapshot.cells == 0 {
5860 return;
5861 }
5862 body["pp_wave"] = json!({
5863 "ticks": snapshot.ticks,
5864 "cells": snapshot.cells,
5865 "overlaps": snapshot.overlaps,
5866 });
5867}
5868
5869fn insert_spec_acceptance_metrics(
5870 body: &mut serde_json::Value,
5871 metrics_scope: &MetricsScope,
5872 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
5873) {
5874 if !metrics_scope.operator() {
5877 return;
5878 }
5879 let snapshot = snapshot();
5880 if snapshot.is_empty() {
5881 return;
5882 }
5883
5884 let mut tau = serde_json::Map::new();
5885 let mut by_position = serde_json::Map::new();
5886 for (model, telemetry) in snapshot {
5887 if telemetry.rounds == 0 {
5888 continue;
5889 }
5890 let n_pos = telemetry
5891 .pos_drafted
5892 .iter()
5893 .rposition(|&n| n > 0)
5894 .map_or(0, |position| position + 1);
5895 tau.insert(model.clone(), json!(telemetry.tau()));
5896 by_position.insert(
5897 model,
5898 json!({
5899 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
5900 "rounds": telemetry.rounds,
5901 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
5902 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
5903 "accept_rate": (0..n_pos).map(|position| {
5904 let offered = telemetry.pos_drafted[position];
5905 if offered > 0 {
5906 telemetry.pos_accepted[position] as f64 / offered as f64
5907 } else {
5908 0.0
5909 }
5910 }).collect::<Vec<f64>>(),
5911 }),
5912 );
5913 }
5914 if !tau.is_empty() {
5915 body["spec_tau"] = serde_json::Value::Object(tau);
5916 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
5917 }
5918}
5919
5920fn insert_peer_probe_metrics(
5921 body: &mut serde_json::Value,
5922 metrics_scope: &MetricsScope,
5923 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
5924) {
5925 if !metrics_scope.operator() {
5928 return;
5929 }
5930 let snapshot = snapshot();
5931 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
5932 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
5933 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
5934 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
5935 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
5936 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
5937 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
5938}
5939
5940async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5942 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5943 Ok(scope) => scope,
5944 Err(response) => return response,
5945 };
5946 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5947 let mut body = if metrics_scope.process_wide() {
5951 json!({
5952 "admitted": m.admitted,
5953 "completed": m.completed,
5954 "tokens_out": m.tokens_out,
5955 "step_p50_ms": m.step_p50_ms,
5956 "step_p99_ms": m.step_p99_ms,
5957 "prompt_tokens_in": m.prompt_tokens_in,
5959 "cached_tokens_in": m.cached_tokens_in,
5960 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5963 "admission_session_defers": m.admission_session_defers,
5966 "admission_vram_defers": m.admission_vram_defers,
5967 "step_oom_parks": m.step_oom_parks,
5968 "continuation_pool_hits": m.continuation_pool_hits,
5969 "continuation_pool_evictions": m.continuation_pool_evictions,
5970 "plain_affinity_rewinds": m.plain_affinity_rewinds,
5971 "served_dspark": m.served_dspark,
5972 "served_spec": m.served_spec,
5973 "served_plain": m.served_plain,
5974 "spec_pool_hits": m.spec_pool_hits,
5975 "spec_pool_misses": m.spec_pool_misses,
5976 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
5977 "spec_pool_evictions": m.spec_pool_evictions,
5978 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
5981 })
5982 } else {
5983 json!({})
5984 };
5985 if metrics_scope.operator() {
5989 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
5990 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
5991 body["budget_source_reload_consecutive"] =
5992 json!(budget_health.source_reload_consecutive);
5993 body["budget_source_available"] = json!(budget_health.source_available);
5994 }
5995 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
5997 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
5998 } else {
5999 0.0
6000 });
6001 body["prefix_cache_hits"] = json!(m.prefix_hits);
6002 body["prefix_cache_misses"] = json!(m.prefix_misses);
6003 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
6004 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
6005 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
6006 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
6007 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
6008 body["prefix_host_entries"] = json!(m.prefix_host_entries);
6012 body["prefix_host_bytes"] = json!(m.prefix_host_bytes);
6013 body["prefix_host_demotions"] = json!(m.prefix_host_demotions);
6014 body["prefix_host_promotions"] = json!(m.prefix_host_promotions);
6015 body["prefix_host_demote_ms"] = json!(m.prefix_host_demote_ms);
6016 body["prefix_host_promote_ms"] = json!(m.prefix_host_promote_ms);
6017 body["prefix_host_rejected_allocs"] = json!(m.prefix_host_rejected_allocs);
6018 body["prefix_host_purges"] = json!(m.prefix_host_purges);
6019 body["prefix_host_purged_entries"] = json!(m.prefix_host_purged_entries);
6020 body["prefix_host_purged_bytes"] = json!(m.prefix_host_purged_bytes);
6021 body["prefix_host_tenant_rejects"] = json!(m.prefix_host_tenant_rejects);
6022 body["prefix_host_pause_demotes"] = json!(m.prefix_host_pause_demotes);
6026 body["prefix_host_pause_cancels"] = json!(m.prefix_host_pause_cancels);
6027 body["prefix_host_handoff_exports"] = json!(m.prefix_host_handoff_exports);
6028 body["prefix_host_handoff_imported_entries"] =
6029 json!(m.prefix_host_handoff_imported_entries);
6030 body["prefix_host_handoff_imported_bytes"] = json!(m.prefix_host_handoff_imported_bytes);
6031 body["prefix_host_handoff_skips"] = json!(m.prefix_host_handoff_skips);
6032 body["kv_flex_borrowed_bytes"] = json!(m.kv_flex_borrowed_bytes);
6037 body["kv_flex_sheds"] = json!(m.kv_flex_sheds);
6038 body["kv_flex_shed_ms"] = json!(m.kv_flex_shed_ms);
6039 body["lcp_histogram"] = json!({
6042 "edges": worker::LCP_HIST_EDGES.to_vec(),
6043 "counts": m.lcp_hist.to_vec(),
6044 });
6045 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
6049 body["prefix_cache_entries"] = json!(m.prefix_entries);
6050 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
6051 body["active_sessions"] = json!(m.active_sessions);
6052 body["queued_requests"] = json!(m.queued_requests);
6053 body["admission_inflight"] = json!(m.admission_inflight);
6057 body["admission_booked_bytes"] = json!(m.admission_booked_bytes);
6058 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
6059 body["spec_pool_entries"] = json!(m.spec_pool_entries);
6060 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
6061 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
6062 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
6063 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
6064 if !m.constraint_compiler_fail_closed.is_empty() {
6065 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
6066 m.constraint_compiler_fail_closed
6067 .iter()
6068 .map(|(model, gauge)| {
6069 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
6070 (model.clone(), json!(value))
6071 })
6072 .collect(),
6073 );
6074 }
6075 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
6076 }
6077 if !m.ns_tokens.is_empty() {
6082 let tenants: serde_json::Map<String, serde_json::Value> = m
6083 .ns_tokens
6084 .iter()
6085 .filter(|(ns, _)| metrics_scope.includes(ns))
6086 .map(|(ns, [p, c])| {
6087 (
6088 ns.clone(),
6089 json!({
6090 "prompt_tokens_in": p,
6091 "cached_tokens_in": c,
6092 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
6093 }),
6094 )
6095 })
6096 .collect();
6097 if !tenants.is_empty() {
6098 body["tenants"] = serde_json::Value::Object(tenants);
6099 }
6100 }
6101 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
6102 .adsd_suspect_total
6103 .iter()
6104 .filter(|(tenant, _)| metrics_scope.includes(tenant))
6105 .map(|(tenant, total)| (tenant.clone(), json!(total)))
6106 .collect();
6107 if !adsd_suspect_total.is_empty() {
6108 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
6109 }
6110 if metrics_scope.operator()
6112 && let Some((bg, mode)) = &st.bg
6113 {
6114 body["bg"] = bg.to_json(mode);
6115 }
6116 if metrics_scope.operator() {
6123 let spec: serde_json::Map<String, serde_json::Value> = m
6124 .spec
6125 .iter()
6126 .map(|(model, t)| {
6127 let n_pos = t
6128 .pos_drafted
6129 .iter()
6130 .rposition(|&d| d > 0)
6131 .map_or(0, |p| p + 1);
6132 (
6133 model.clone(),
6134 json!({
6135 "rounds": t.rounds,
6136 "drafted": t.drafted,
6137 "accepted": t.accepted,
6138 "acceptance_rate": if t.drafted > 0 {
6139 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
6140 "tokens_per_round": if t.rounds > 0 {
6141 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
6142 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
6143 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
6144 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
6145 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
6146 .collect::<Vec<f64>>(),
6147 }),
6148 )
6149 })
6150 .collect();
6151 if !spec.is_empty() {
6152 body["spec"] = serde_json::Value::Object(spec);
6153 }
6154 }
6155 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
6156 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
6157 insert_pp_wave_metrics(&mut body, &metrics_scope, PpWaveMetricsSnapshot::current);
6158 insert_peer_probe_metrics(
6159 &mut body,
6160 &metrics_scope,
6161 memra_engine::pp::peer_probe_metrics,
6162 );
6163 Json(body).into_response()
6164}
6165
6166#[derive(Debug, Default, Deserialize)]
6167struct ModelsQuery {
6168 #[serde(default)]
6169 schema: Option<String>,
6170}
6171
6172fn models_openai_body(models: &[String]) -> serde_json::Value {
6173 let data: Vec<_> = models
6174 .iter()
6175 .map(|m| json!({ "id": m, "object": "model" }))
6176 .collect();
6177 json!({ "object": "list", "data": data })
6178}
6179
6180fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
6185 match metadata.and_then(|m| m.surface.as_deref()) {
6186 Some("embedding") => "embedding",
6187 Some("rerank") => "rerank",
6188 _ => "chat",
6189 }
6190}
6191
6192fn openrouter_supported_parameters(
6193 caps: Option<&ModelCaps>,
6194 max_output_length: Option<u64>,
6195 is_chat: bool,
6196) -> serde_json::Value {
6197 let mut parameters = serde_json::Map::new();
6198 if !is_chat {
6206 return serde_json::Value::Object(parameters);
6207 }
6208 for name in [
6209 "temperature",
6210 "top_p",
6211 "min_p",
6212 "frequency_penalty",
6213 "presence_penalty",
6214 "repetition_penalty",
6215 "stop",
6216 ] {
6217 parameters.insert(name.into(), json!({ "type": "unknown" }));
6218 }
6219 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
6220 parameters.insert(
6221 "seed".into(),
6222 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
6223 );
6224 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
6225 if let Some(max) = max_output_length {
6226 max_tokens["max"] = json!(max);
6227 }
6228 parameters.insert("max_tokens".into(), max_tokens);
6229 if is_chat
6242 && caps.is_some_and(|c| {
6243 !c.dsv4 && !(c.qwen_think && !c.think_switch && c.think_close.is_empty())
6244 })
6245 {
6246 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
6247 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
6248 }
6249 if is_chat && caps.is_some_and(|c| c.tools_branch) {
6250 parameters.insert("tools".into(), json!({ "type": "boolean" }));
6251 parameters.insert(
6252 "tool_choice".into(),
6253 json!({ "type": "enum", "values": ["auto", "none"] }),
6254 );
6255 }
6256 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
6257 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
6258 }
6259 serde_json::Value::Object(parameters)
6260}
6261
6262fn published_context_length(
6278 caps: Option<&ModelCaps>,
6279 metadata: Option<&OpenRouterModelMetadata>,
6280) -> Option<u64> {
6281 let trained = caps
6282 .map(|c| c.context_length as u64)
6283 .filter(|&value| value > 0)?;
6284 let envelope = metadata.and_then(|m| {
6285 let prompt = m.max_prompt_length?;
6286 let output = m.max_output_length?;
6287 prompt.checked_add(output)
6288 });
6289 Some(envelope.map_or(trained, |envelope| trained.min(envelope)))
6290}
6291
6292fn model_entry_openrouter(
6293 name: &str,
6294 caps: Option<&ModelCaps>,
6295 metadata: Option<&OpenRouterModelMetadata>,
6296) -> serde_json::Value {
6297 let empty = OpenRouterModelMetadata::default();
6298 let metadata = metadata.unwrap_or(&empty);
6299 let context_length =
6300 published_context_length(caps, Some(metadata)).filter(|&v| v <= JSON_SAFE_INTEGER_MAX);
6301 let tokenizer = caps
6302 .map(|c| c.tokenizer.as_str())
6303 .filter(|tokenizer| !tokenizer.is_empty());
6304
6305 let mut input = serde_json::Map::new();
6306 input.insert("type".into(), json!("text"));
6307 let mut supported_inputs = serde_json::Map::new();
6308 if let Some(value) = context_length {
6309 supported_inputs.insert(
6310 "max_context_length".into(),
6311 json!({ "value": value, "unit": "token" }),
6312 );
6313 }
6314 if let Some(value) = metadata.max_prompt_length {
6315 supported_inputs.insert(
6316 "max_prompt_length".into(),
6317 json!({ "value": value, "unit": "token" }),
6318 );
6319 }
6320 if !supported_inputs.is_empty() {
6321 input.insert(
6322 "supported_inputs".into(),
6323 serde_json::Value::Object(supported_inputs),
6324 );
6325 }
6326 let mut input_pricing = Vec::new();
6327 for (kind, cost) in [
6328 ("prompt", metadata.pricing.prompt.as_deref()),
6329 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
6330 ("cache_write", metadata.pricing.cache_write.as_deref()),
6331 ] {
6332 if let Some(cost) = cost {
6333 input_pricing.push(json!({
6334 "type": kind,
6335 "unit": "token",
6336 "cost_usd": cost,
6337 }));
6338 }
6339 }
6340 if !input_pricing.is_empty() {
6341 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
6342 }
6343 let mut input_capacity = Vec::new();
6344 for (kind, value) in [
6345 ("prompt", metadata.capacity.prompt_tpm),
6346 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
6347 ] {
6348 if let Some(value) = value {
6349 input_capacity.push(json!({
6350 "type": kind,
6351 "unit": "token",
6352 "per": "minute",
6353 "value": value,
6354 }));
6355 }
6356 }
6357 if !input_capacity.is_empty() {
6358 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
6359 }
6360
6361 let or_surface = declared_surface(Some(metadata));
6362 let or_is_chat = or_surface == "chat";
6363 let mut output = serde_json::Map::new();
6364 output.insert(
6372 "type".into(),
6373 json!(match or_surface {
6374 "embedding" => "embeddings",
6375 "rerank" => "rerank",
6376 _ => "text",
6377 }),
6378 );
6379 output.insert(
6380 "supported_parameters".into(),
6381 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
6382 );
6383 if or_is_chat {
6387 output.insert("streaming".into(), json!(true));
6388 }
6389 if let Some(value) = metadata.max_output_length
6392 && or_is_chat
6393 {
6394 output.insert(
6395 "max_length".into(),
6396 json!({ "value": value, "unit": "token" }),
6397 );
6398 }
6399 let mut output_pricing = Vec::new();
6400 for (kind, cost) in [
6401 ("completion", metadata.pricing.completion.as_deref()),
6402 (
6403 "internal_reasoning",
6404 metadata.pricing.internal_reasoning.as_deref(),
6405 ),
6406 ] {
6407 if let Some(cost) = cost {
6408 output_pricing.push(json!({
6409 "type": kind,
6410 "unit": "token",
6411 "cost_usd": cost,
6412 }));
6413 }
6414 }
6415 if !output_pricing.is_empty() {
6416 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
6417 }
6418 let mut output_capacity = Vec::new();
6419 if let Some(value) = metadata.capacity.completion_tpm {
6420 output_capacity.push(json!({
6421 "type": "completion",
6422 "unit": "token",
6423 "per": "minute",
6424 "value": value,
6425 }));
6426 }
6427 if let Some(value) = metadata.capacity.concurrency {
6428 output_capacity.push(json!({
6429 "type": "concurrency",
6430 "unit": "request",
6431 "value": value,
6432 }));
6433 }
6434 if !output_capacity.is_empty() {
6435 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
6436 }
6437
6438 let mut entry = serde_json::Map::new();
6439 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
6440 entry.insert("id".into(), json!(name));
6441 entry.insert("name".into(), json!(name));
6442 if let Some(value) = metadata.hugging_face_id.as_deref() {
6443 entry.insert("hugging_face_id".into(), json!(value));
6444 }
6445 if let Some(value) = metadata.created {
6446 entry.insert("created".into(), json!(value));
6447 }
6448 if let Some(value) = metadata.quantization.as_deref() {
6449 entry.insert("quantization".into(), json!(value));
6450 }
6451 if let Some(value) = tokenizer {
6452 entry.insert("tokenizer".into(), json!(value));
6453 }
6454 if let Some(value) = metadata.description.as_deref() {
6455 entry.insert("description".into(), json!(value));
6456 }
6457 let mut input_modalities = vec![serde_json::Value::Object(input)];
6458 for m in &metadata.input_modalities {
6459 let mut extra = serde_json::Map::new();
6460 extra.insert("type".into(), json!(m));
6461 if let Some(cost) = metadata.pricing.prompt.as_deref() {
6462 extra.insert(
6464 "pricing".into(),
6465 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
6466 );
6467 }
6468 input_modalities.push(serde_json::Value::Object(extra));
6469 }
6470 entry.insert(
6471 "input_modalities".into(),
6472 serde_json::Value::Array(input_modalities),
6473 );
6474 entry.insert(
6475 "output_modalities".into(),
6476 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
6477 );
6478 if let Some(cost) = metadata.pricing.request.as_deref() {
6479 entry.insert(
6480 "pricing".into(),
6481 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
6482 );
6483 }
6484 if let Some(value) = metadata.capacity.request_rpm {
6485 entry.insert(
6486 "capacity".into(),
6487 json!([{
6488 "type": "request",
6489 "unit": "request",
6490 "per": "minute",
6491 "value": value,
6492 }]),
6493 );
6494 }
6495 if let Some(value) = metadata.is_ready {
6496 entry.insert("is_ready".into(), json!(value));
6497 }
6498 if let Some(value) = metadata.is_free {
6499 entry.insert("is_free".into(), json!(value));
6500 }
6501 if let Some(value) = metadata.discount_to_user {
6502 entry.insert("discount_to_user".into(), json!(value));
6503 }
6504 if let Some(value) = metadata.openrouter_slug.as_deref() {
6505 entry.insert("openrouter".into(), json!({ "slug": value }));
6506 }
6507 if !metadata.datacenters.is_empty() {
6508 entry.insert("datacenters".into(), json!(metadata.datacenters));
6509 }
6510 let mut compliance = serde_json::Map::new();
6511 if let Some(value) = metadata.zdr {
6512 compliance.insert("zdr".into(), json!(value));
6513 }
6514 if let Some(value) = metadata.hipaa {
6515 compliance.insert("hipaa".into(), json!(value));
6516 }
6517 if !compliance.is_empty() {
6518 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
6519 }
6520 serde_json::Value::Object(entry)
6521}
6522
6523fn models_openrouter_body(st: &AppState) -> serde_json::Value {
6524 let data: Vec<_> = st
6525 .models
6526 .iter()
6527 .map(|model| {
6528 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
6529 })
6530 .collect();
6531 json!({ "data": data })
6532}
6533
6534fn model_entry_openmodels(
6535 name: &str,
6536 caps: Option<&ModelCaps>,
6537 metadata: Option<&OpenRouterModelMetadata>,
6538) -> Result<serde_json::Value, String> {
6539 let metadata = metadata.ok_or_else(|| {
6540 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
6541 })?;
6542 let context_length = published_context_length(caps, Some(metadata))
6543 .filter(|&value| value <= JSON_SAFE_INTEGER_MAX)
6544 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
6545 let created = metadata
6546 .created
6547 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
6548 let max_output_length = metadata
6549 .max_output_length
6550 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
6551 let prompt = metadata
6552 .pricing
6553 .prompt
6554 .as_deref()
6555 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
6556 let completion =
6557 metadata.pricing.completion.as_deref().ok_or_else(|| {
6558 format!("OpenModels feed requires pricing.completion for model {name:?}")
6559 })?;
6560 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
6561 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
6562 })?;
6563 let is_ready = metadata
6564 .is_ready
6565 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
6566 let is_free = metadata
6567 .is_free
6568 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
6569 let discount_to_user = metadata
6570 .discount_to_user
6571 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
6572
6573 let mut pricing = serde_json::Map::new();
6574 pricing.insert("prompt".into(), json!(prompt));
6575 pricing.insert("completion".into(), json!(completion));
6576 pricing.insert("input_cache_read".into(), json!(input_cache_read));
6577 if let Some(value) = metadata.pricing.request.as_deref() {
6578 pricing.insert("request".into(), json!(value));
6579 }
6580
6581 let om_surface = declared_surface(Some(metadata));
6582 let om_is_chat = om_surface == "chat";
6583 let mut supported_features = Vec::new();
6584 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
6585 supported_features.push("tool_calling");
6586 }
6587 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
6588 supported_features.push("reasoning");
6589 }
6590
6591 let mut entry = serde_json::Map::new();
6592 entry.insert("id".into(), json!(name));
6593 entry.insert("name".into(), json!(name));
6594 entry.insert("created".into(), json!(created));
6595 entry.insert("input_modalities".into(), json!(["text"]));
6596 entry.insert(
6597 "output_modalities".into(),
6598 json!(match om_surface {
6599 "embedding" => ["embeddings"],
6600 "rerank" => ["rerank"],
6601 _ => ["text"],
6602 }),
6603 );
6604 entry.insert("context_length".into(), json!(context_length));
6605 entry.insert("max_output_length".into(), json!(max_output_length));
6606 entry.insert("currency".into(), json!("USD"));
6609 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
6610 entry.insert("supported_features".into(), json!(supported_features));
6611 entry.insert("is_ready".into(), json!(is_ready));
6612 entry.insert("is_free".into(), json!(is_free));
6613 entry.insert("discount_to_user".into(), json!(discount_to_user));
6614 Ok(serde_json::Value::Object(entry))
6615}
6616
6617fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
6618 let data: Result<Vec<_>, _> = st
6619 .models
6620 .iter()
6621 .map(|model| {
6622 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
6623 })
6624 .collect();
6625 Ok(json!({ "data": data? }))
6626}
6627
6628async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
6629 match query.schema.as_deref() {
6630 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
6631 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
6632 Some("openmodels") => match models_openmodels_body(&st) {
6633 Ok(body) => Json(body).into_response(),
6634 Err(error) => bad_request(&error, Some("schema")),
6635 },
6636 Some(schema) => bad_request(
6637 &format!(
6638 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
6639 ),
6640 Some("schema"),
6641 ),
6642 }
6643}
6644
6645fn model_entry_v1(
6653 name: &str,
6654 caps: Option<&ModelCaps>,
6655 metadata: Option<&OpenRouterModelMetadata>,
6656) -> serde_json::Value {
6657 let ctx = published_context_length(caps, metadata);
6658 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
6662 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
6665 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
6666 Some(p) => json!(p),
6667 None => serde_json::Value::Null,
6668 };
6669 let owned_by = metadata
6670 .and_then(|m| m.owned_by.as_deref())
6671 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
6672 let mut input_modalities = vec!["text"];
6673 if let Some(meta) = metadata {
6674 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
6675 }
6676 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
6677 let reliability = metadata.and_then(|m| m.reliability.as_ref());
6678 let surface = declared_surface(metadata);
6684 let (model_type, endpoints, output_modalities) = match surface {
6685 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
6689 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
6690 _ => ("chat", vec!["chat/completions"], vec!["text"]),
6691 };
6692 let is_chat = surface == "chat";
6693 json!({
6694 "id": name,
6695 "name": name,
6696 "object": "model",
6697 "owned_by": owned_by,
6698 "type": model_type,
6699 "context_length": ctx,
6700 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
6703 "endpoints": endpoints,
6704 "input_modalities": input_modalities,
6705 "output_modalities": output_modalities,
6706 "capabilities": {
6707 "streaming": is_chat,
6710 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
6711 "structured_output": is_chat
6719 && !is_dsv4
6720 && !caps.is_some_and(|c| c.qwen_think && !c.think_switch && c.think_close.is_empty()),
6721 "reasoning": is_chat && thinking,
6722 "prompt_caching": is_chat && !is_dsv4,
6723 },
6724 "pricing": {
6725 "currency": "USD",
6726 "unit": "per_1m_tokens",
6727 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
6728 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
6729 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
6730 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
6731 "minimum_request": metadata
6733 .and_then(|m| m.pricing.request.as_deref())
6734 .unwrap_or("0"),
6735 },
6736 "lifecycle": {
6737 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
6738 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
6739 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
6740 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
6741 },
6742 "reliability": {
6743 "first_token_timeout_seconds":
6744 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
6745 "completion_timeout_seconds":
6746 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
6747 "stream_idle_timeout_seconds":
6748 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
6749 "capacity_scope":
6750 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
6751 },
6752 })
6753}
6754
6755async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
6758 let data: Vec<_> = st
6759 .models
6760 .iter()
6761 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
6762 .collect();
6763 let mut body = json!({
6764 "object": "list",
6765 "contract_version": "2.0",
6766 "data": data,
6767 });
6768 if let Some(provider) = st.provider_metadata.as_ref() {
6773 body["provider"] = json!({
6774 "id": provider.id,
6775 "status_url": provider.status_url,
6776 "support_contact": provider.support_contact,
6777 "incident_contact": provider.incident_contact,
6778 "regions": provider.regions,
6779 "request_id_header": "x-request-id",
6780 "error_contract": {
6781 "rate_limit_status": 429,
6782 "overload_status": 503,
6783 "retry_after_header": "Retry-After",
6784 "account_quota_error_codes": ["insufficient_balance"],
6785 },
6786 });
6787 }
6788 Json(body)
6789}
6790
6791async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
6794 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
6795 Ok(scope) => scope,
6796 Err(response) => return response,
6797 };
6798 if !metrics_scope.process_wide() {
6799 return error_response(
6800 StatusCode::FORBIDDEN,
6801 "completion api keys do not authorize process-wide yield metrics; configure \
6802 MEMRA_METRICS_TOKEN",
6803 "authentication_error",
6804 None,
6805 );
6806 }
6807 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
6808 let lane = |i: usize| {
6809 json!({
6810 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
6811 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
6812 })
6813 };
6814 let mut body = json!({
6815 "lanes": {
6816 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
6817 },
6818 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
6819 });
6820 if metrics_scope.operator() {
6821 body["batch_size_last"] = json!(m.batch_size_last);
6822 }
6823 Json(body).into_response()
6824}
6825
6826async fn peek_admission(
6839 mut rx: worker::EventReceiver,
6840) -> Result<worker::EventReceiver, (Response, &'static str)> {
6841 match rx.recv().await {
6842 Some(Event::Error(e)) => {
6847 let error_code = engine_error_code(e.class);
6848 Err((engine_error_response(&e), error_code))
6849 }
6850 first => {
6851 let (tx2, rx2) = worker::event_channel();
6852 if let Some(ev) = first {
6853 let _ = tx2.send(ev);
6854 }
6855 tokio::spawn(forward_events(rx, tx2));
6856 Ok(rx2)
6857 }
6858 }
6859}
6860
6861async fn forward_events(mut rx: worker::EventReceiver, tx2: worker::EventSender) {
6872 loop {
6873 tokio::select! {
6874 biased;
6875 () = tx2.closed() => break,
6876 ev = rx.recv() => match ev {
6877 Some(ev) => {
6878 if tx2.send(ev).is_err() {
6879 break;
6880 }
6881 }
6882 None => break,
6883 },
6884 }
6885 }
6886}
6887
6888async fn peek_first_token(
6903 mut rx: worker::EventReceiver,
6904 deadline: RequestDeadline,
6905) -> Result<worker::EventReceiver, ()> {
6906 let mut buffered: Vec<Event> = Vec::new();
6907 loop {
6908 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
6909 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
6912 let first_delivery = matches!(
6913 ev,
6914 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
6915 );
6916 buffered.push(ev);
6917 if first_delivery {
6918 break;
6919 }
6920 }
6921 }
6922 }
6923 let (tx2, rx2) = worker::event_channel();
6924 for ev in buffered {
6925 let _ = tx2.send(ev);
6926 }
6927 tokio::spawn(forward_events(rx, tx2));
6928 Ok(rx2)
6929}
6930
6931#[cfg(test)]
6933fn build_request(
6937 req: &CompletionReq,
6938 tx: worker::EventSender,
6939 lane: lanes::Lane,
6940 affinity: Option<String>,
6941) -> Request {
6942 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
6943}
6944
6945fn build_request_with_trace(
6946 req: &CompletionReq,
6947 tx: worker::EventSender,
6948 lane: lanes::Lane,
6949 affinity: Option<String>,
6950 ttft: Option<Arc<ttft::Trace>>,
6951 sampling_defaults: &SamplingDefaults,
6952) -> Request {
6953 let params = GenParams {
6954 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6955 max_ctx: req.max_ctx,
6956 eos: Vec::new(), };
6958 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
6963 Request {
6964 model: req.model.clone(),
6965 prompt_ids: req.prompt_ids.clone(),
6966 prompt_text: req.prompt.clone(),
6967 chat: req.chat,
6968 chat_turns: Vec::new(),
6969 tools_json: Vec::new(),
6970 tools_struct: Vec::new(),
6971 think: ThinkMode::Default,
6972 reasoning_effort: None, params,
6974 sampler_cfg,
6975 stop_strings: req.stop.clone().into_vec(),
6976 trace_id: req.trace_id.clone(),
6977 request_id: String::new(),
6980 admit_predict_logged: false,
6981 max_prompt_tokens: None,
6982 cache_ns: cache_namespace(&req.cache_salt),
6983 affinity,
6984 lane,
6985 grammar: None, prepared_constraint: None,
6987 constraint_ready: None,
6988 oom_retries: 0, spec_k_replay: None,
6990 prepared_prompt: None,
6991 capture: None, images: Vec::new(), gemma_images: Vec::new(),
6994 glm5_images: Vec::new(),
6995 step_images: Vec::new(),
6996 vision_memory: None,
6997 wire_deadline: None, ttft,
6999 tx,
7000 }
7001}
7002
7003struct ChatPlan {
7006 request: Request,
7007 parser: Option<ToolStreamParser>,
7010 pending_images: Vec<PendingVisionUnit>,
7013 pending_gemma: Vec<PendingGemmaImage>,
7014 pending_glm5: Vec<PendingGlm5Image>,
7015 pending_step: Vec<PendingStepImage>,
7016 vision_memory: Option<VisionMemoryPermit>,
7020}
7021
7022pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
7023 req.messages.iter().any(|message| {
7024 message.content.as_array().is_some_and(|parts| {
7025 parts.iter().any(|part| {
7026 matches!(
7027 part.get("type").and_then(serde_json::Value::as_str),
7028 Some("image_url" | "video_url")
7029 )
7030 })
7031 })
7032 })
7033}
7034
7035fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
7036 let mut total = 0usize;
7037 let mut add = |bytes: usize| {
7038 total = total.checked_add(bytes).ok_or_else(|| {
7039 "vision patch memory reservation overflowed while planning".to_string()
7040 })?;
7041 Ok::<(), String>(())
7042 };
7043 for unit in &plan.pending_images {
7044 let bytes = match unit {
7045 PendingVisionUnit::Still { gh, gw, .. } => gh
7046 .checked_mul(*gw)
7047 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7048 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7049 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
7050 PendingVisionUnit::Video { groups, .. } => {
7051 groups.iter().try_fold(0usize, |total, group| {
7052 let bytes = group
7053 .gh
7054 .checked_mul(group.gw)
7055 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
7056 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7057 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7058 total.checked_add(bytes).ok_or_else(|| {
7059 "vision patch memory reservation overflowed while planning".to_string()
7060 })
7061 })?
7062 }
7063 };
7064 add(bytes)?;
7065 }
7066 for unit in &plan.pending_gemma {
7067 let bytes = unit
7068 .gw
7069 .checked_mul(unit.gh)
7070 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
7071 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7072 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7073 add(bytes)?;
7074 }
7075 for unit in &plan.pending_glm5 {
7076 let bytes = unit
7077 .gh
7078 .checked_mul(unit.gw)
7079 .and_then(|n| n.checked_mul(memra_engine::vision_glm5::G5V_PATCH_IN))
7080 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7081 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7082 add(bytes)?;
7083 }
7084 for unit in &plan.pending_step {
7085 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
7086 let patches = unit
7088 .plan
7089 .n_tiles
7090 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
7091 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
7092 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7093 let bytes = patches
7094 .checked_mul(SV_PATCH_IN)
7095 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
7096 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
7097 add(bytes)?;
7098 }
7099 Ok(total)
7100}
7101
7102pub(crate) fn reserve_vision_memory(
7103 plan: &ChatPlan,
7104) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
7105 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
7106 try_reserve_vision_memory(bytes)
7107}
7108
7109#[cfg(test)]
7110fn build_chat_request(
7111 req: ChatCompletionReq,
7112 caps: Option<&ModelCaps>,
7113 tx: worker::EventSender,
7114 lane: lanes::Lane,
7115 affinity: Option<String>,
7116) -> Result<ChatPlan, String> {
7117 let defaults = ModelSamplingDefaults::resolve(None, caps);
7120 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
7121}
7122
7123#[allow(clippy::too_many_arguments)]
7137fn build_chat_request_with_trace(
7138 req: ChatCompletionReq,
7139 caps: Option<&ModelCaps>,
7140 tx: worker::EventSender,
7141 lane: lanes::Lane,
7142 affinity: Option<String>,
7143 ttft: Option<Arc<ttft::Trace>>,
7144 default_effort: Option<&str>,
7145 sampling_defaults: &ModelSamplingDefaults,
7146) -> Result<ChatPlan, String> {
7147 req.stop.validate()?;
7148 let client_sampling: ClientSampling = (&req).into();
7151 let tool_choice = parse_tool_choice(&req.tool_choice)?;
7152 if let Some(c) = caps
7157 && !c.chat_ok
7158 {
7159 return Err(format!(
7160 "model {:?} has no chat template (checkpoint carries neither \
7161 tokenizer_config.json chat_template nor chat_template.jinja) — \
7162 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
7163 req.model
7164 ));
7165 }
7166 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
7167 let (mut think, effort_level, think_client_explicit) = parse_think(
7168 &req.reasoning_effort,
7169 &req.reasoning,
7170 vllm_switch,
7171 req.include_reasoning,
7172 default_effort,
7173 caps.is_some_and(|c| c.dsv4 || c.glm5),
7177 )?;
7178 let level_template = caps
7182 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort || c.glm5)
7183 .unwrap_or(false);
7184 if think_client_explicit
7205 && think == ThinkMode::NoThink
7206 && let Some(c) = caps
7207 && c.qwen_think
7208 && !c.think_switch
7209 && !c.dsv4
7210 {
7211 return Err(format!(
7212 "model {:?} cannot disable reasoning: its chat template opens a think \
7213 tail unconditionally and carries no enable_thinking switch, so \
7214 reasoning_effort/enable_thinking cannot turn it off on this model",
7215 req.model
7216 ));
7217 }
7218 let reasoning_effort = if level_template { effort_level } else { None };
7251 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
7254 if grammar.is_some()
7268 && let Some(c) = caps
7269 && c.qwen_think
7270 && think != ThinkMode::NoThink
7271 {
7272 if c.think_switch {
7273 think = ThinkMode::NoThink;
7274 } else if c.think_close.is_empty() {
7275 return Err(
7276 "response_format requires the model's think channel to close \
7277 before the grammar can engage, but this chat template has \
7278 neither an enable_thinking switch nor a recognizable \
7279 think-close token sequence"
7280 .into(),
7281 );
7282 }
7283 }
7287
7288 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
7295
7296 let (tools_json, tools_struct, schemas) =
7299 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
7300 prepare_tools(&req.tools)?
7301 } else {
7302 (Vec::new(), Vec::new(), HashMap::new())
7303 };
7304
7305 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
7306 let mut images: Vec<PendingVisionUnit> = Vec::new();
7307 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
7308 let mut glm5_images: Vec<PendingGlm5Image> = Vec::new();
7309 let mut step_images: Vec<PendingStepImage> = Vec::new();
7310 let mut next_video = 0usize;
7311 for msg in &req.messages {
7312 let content = content_to_text_vision(
7313 &msg.content,
7314 &mut images,
7315 &mut gemma_images,
7316 &mut glm5_images,
7317 &mut step_images,
7318 &mut next_video,
7319 )
7320 .map_err(|e| format!("{} message: {e}", msg.role))?;
7321 let tool_calls = msg
7322 .tool_calls
7323 .iter()
7324 .map(render_req_tool_call)
7325 .collect::<Result<Vec<_>, _>>()?;
7326 if !tool_calls.is_empty() && msg.role != "assistant" {
7327 return Err("tool_calls are only valid on assistant messages".into());
7328 }
7329 let role = if msg.role == "developer" {
7332 "system".to_string()
7333 } else {
7334 msg.role.clone()
7335 };
7336 turns.push(TmplTurn {
7337 role,
7338 content,
7339 tool_calls,
7340 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
7342 tool_call_id: msg.tool_call_id.clone(),
7343 tool_name: msg.name.clone(),
7344 tool_responses: Vec::new(),
7345 task: None,
7349 tools: Vec::new(),
7350 });
7351 }
7352
7353 let has_tool_features = !tools_json.is_empty()
7356 || turns
7357 .iter()
7358 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
7359 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
7360 return Err(format!(
7361 "model {:?} chat template has no tools branch",
7362 req.model
7363 ));
7364 }
7365
7366 let think_open = caps
7369 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
7370 .unwrap_or(false);
7371 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
7389 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
7395 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
7396 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
7397 let glm5 = caps.map(|c| c.glm5).unwrap_or(false);
7403 let is_hy3 = caps.map(|c| c.hy3).unwrap_or(false);
7406 let hy3_think_open = is_hy3 && think == ThinkMode::Think;
7407 let hy3_tools = is_hy3 && !tools_json.is_empty();
7408 let parser = if glm5 {
7409 Some(ToolStreamParser::glm5(think_open, schemas))
7410 } else if is_hy3 && (hy3_tools || hy3_think_open) {
7411 Some(ToolStreamParser::hy3(schemas, hy3_think_open))
7412 } else if is_dsv4 && (dsv4_tools || dsv4_think_open) {
7413 Some(ToolStreamParser::dsv4(dsv4_think_open))
7414 } else if gemma_tools {
7415 Some(ToolStreamParser::gemma_tools())
7416 } else if !tools_json.is_empty() {
7417 Some(ToolStreamParser::new(schemas, think_open))
7418 } else if think_open {
7419 Some(ToolStreamParser::reasoning_only())
7420 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
7421 Some(ToolStreamParser::gemma_thought())
7429 } else {
7430 None
7431 };
7432
7433 Ok(ChatPlan {
7434 request: Request {
7435 model: req.model,
7436 prompt_ids: Vec::new(),
7437 prompt_text: String::new(),
7438 chat: false,
7439 chat_turns: turns,
7440 tools_json,
7441 tools_struct,
7442 think,
7443 reasoning_effort,
7444 params: GenParams {
7445 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
7446 max_ctx: req.max_ctx,
7447 eos: Vec::new(),
7448 },
7449 sampler_cfg,
7450 stop_strings: {
7451 let mut stops = req.stop.into_vec();
7456 if gemma_tools {
7457 stops.push("<tool_call|>".to_string());
7458 }
7459 if dsv4_tools {
7464 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
7465 }
7466 if hy3_tools {
7469 stops.push("</tool_calls:opensource>".to_string());
7470 }
7471 stops
7472 },
7473 trace_id: None,
7474 request_id: String::new(),
7477 admit_predict_logged: false,
7478 max_prompt_tokens: None,
7479 cache_ns: cache_namespace(&req.cache_salt),
7480 affinity,
7481 lane,
7482 grammar,
7483 prepared_constraint: None,
7484 constraint_ready: None,
7485 oom_retries: 0, spec_k_replay: None,
7487 prepared_prompt: None,
7488 images: Vec::new(),
7493 gemma_images: Vec::new(),
7494 glm5_images: Vec::new(),
7495 step_images: Vec::new(),
7496 capture: None, vision_memory: None,
7498 wire_deadline: None, ttft,
7500 tx,
7501 },
7502 parser,
7503 pending_images: images,
7504 pending_gemma: gemma_images,
7505 pending_glm5: glm5_images,
7506 pending_step: step_images,
7507 vision_memory: None,
7508 })
7509}
7510
7511fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
7517 for (i, unit) in plan.pending_images.drain(..).enumerate() {
7518 match unit {
7519 PendingVisionUnit::Still { bytes, gh, gw } => {
7520 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
7521 .map_err(|e| format!("image {}: {e}", i + 1))?;
7522 if (prep.gh, prep.gw) != (gh, gw) {
7523 return Err(format!(
7524 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
7525 i + 1,
7526 prep.gh,
7527 prep.gw
7528 ));
7529 }
7530 plan.request
7531 .images
7532 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
7533 }
7534 PendingVisionUnit::Video {
7535 bytes,
7536 groups,
7537 video,
7538 } => {
7539 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
7540 .map_err(|e| format!("video {}: {e}", i + 1))?;
7541 if prepared.groups.len() != groups.len() {
7542 return Err(format!(
7543 "video {}: decoded {} groups differ from its header-planned {} groups",
7544 i + 1,
7545 prepared.groups.len(),
7546 groups.len()
7547 ));
7548 }
7549 for ((group, prep), timestamp) in
7550 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
7551 {
7552 if (prep.gh, prep.gw) != (group.gh, group.gw) {
7553 return Err(format!(
7554 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
7555 i + 1,
7556 prep.gh,
7557 prep.gw,
7558 group.gh,
7559 group.gw
7560 ));
7561 }
7562 if (timestamp - group.timestamp).abs() > 0.001 {
7563 return Err(format!(
7564 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
7565 i + 1,
7566 group.timestamp
7567 ));
7568 }
7569 plan.request
7570 .images
7571 .push(memra_engine::vision_pre::VisionUnit {
7572 prep,
7573 video: Some(video),
7574 });
7575 }
7576 }
7577 }
7578 }
7579 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
7580 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
7581 .map_err(|e| format!("image {}: {e}", i + 1))?;
7582 if (gw, gh) != (unit.gw, unit.gh) {
7583 return Err(format!(
7584 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
7585 i + 1,
7586 unit.gw,
7587 unit.gh
7588 ));
7589 }
7590 plan.request
7591 .gemma_images
7592 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
7593 }
7594 for (i, unit) in plan.pending_glm5.drain(..).enumerate() {
7595 let (patches, gh, gw) = memra_engine::vision_glm5::glm5_prep_image(&unit.bytes)
7596 .map_err(|e| format!("image {}: {e}", i + 1))?;
7597 if (gh, gw) != (unit.gh, unit.gw) {
7598 return Err(format!(
7599 "image {}: decoded grid {gh}x{gw} differs from its header-planned grid {}x{} — refusing (placeholder runs already rendered)",
7600 i + 1,
7601 unit.gh,
7602 unit.gw
7603 ));
7604 }
7605 plan.request
7606 .glm5_images
7607 .push(memra_engine::vision_glm5::Glm5VisionUnit { patches, gh, gw });
7608 }
7609 for (i, unit) in plan.pending_step.drain(..).enumerate() {
7610 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
7611 .map_err(|e| format!("image {}: {e}", i + 1))?;
7612 if prepped.tiles.len() != unit.plan.n_tiles
7613 || prepped.newline_mask != unit.plan.newline_mask
7614 {
7615 return Err(format!(
7616 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
7617 ({} tiles) — refusing (pad runs already rendered)",
7618 i + 1,
7619 prepped.tiles.len(),
7620 unit.plan.n_tiles
7621 ));
7622 }
7623 plan.request.step_images.push(prepped);
7624 }
7625 Ok(())
7626}
7627
7628fn bearer_token(headers: &HeaderMap) -> Option<&str> {
7636 headers
7637 .get("authorization")
7638 .and_then(|value| value.to_str().ok())
7639 .and_then(|value| value.strip_prefix("Bearer "))
7640}
7641
7642fn authentication_error(why: auth::AuthDenied) -> Response {
7643 match why {
7644 auth::AuthDenied::Unknown => error_response(
7645 StatusCode::UNAUTHORIZED,
7646 "invalid api key",
7647 "authentication_error",
7648 None,
7649 ),
7650 auth::AuthDenied::Disabled => error_response(
7651 StatusCode::FORBIDDEN,
7652 "api key is disabled",
7653 "authentication_error",
7654 None,
7655 ),
7656 }
7657}
7658
7659#[allow(clippy::result_large_err)] fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
7661 auth::authenticate_with(
7662 api_auth.keyring,
7663 api_auth.single_key.as_deref(),
7664 bearer_token(headers),
7665 )
7666 .map_err(authentication_error)
7667}
7668
7669#[derive(Debug, Clone, PartialEq, Eq)]
7670enum MetricsScope {
7671 All,
7672 CompletionDomain,
7673 Tenant(String),
7674}
7675
7676impl MetricsScope {
7677 fn operator(&self) -> bool {
7678 matches!(self, MetricsScope::All)
7679 }
7680
7681 fn process_wide(&self) -> bool {
7682 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
7683 }
7684
7685 fn includes(&self, tenant_row: &str) -> bool {
7686 match self {
7687 MetricsScope::All | MetricsScope::CompletionDomain => true,
7688 MetricsScope::Tenant(tenant) => tenant == tenant_row,
7689 }
7690 }
7691}
7692
7693#[allow(clippy::result_large_err)] fn authorize_metrics(
7695 api_auth: &ApiAuth,
7696 metrics_auth: &MetricsAuth,
7697 headers: &HeaderMap,
7698) -> Result<MetricsScope, Response> {
7699 if !metrics_auth.required {
7700 return Ok(MetricsScope::All);
7701 }
7702 let Some(candidate) = bearer_token(headers) else {
7703 return Err(authentication_error(auth::AuthDenied::Unknown));
7704 };
7705 if let Some(token) = metrics_auth.token.as_deref() {
7706 if auth::constant_time_secret_eq(token, candidate) {
7707 return Ok(MetricsScope::All);
7708 }
7709 if api_auth.configured() {
7710 return match auth::authenticate_with(
7711 api_auth.keyring,
7712 api_auth.single_key.as_deref(),
7713 Some(candidate),
7714 ) {
7715 Ok(_) => Err(error_response(
7716 StatusCode::FORBIDDEN,
7717 "completion api keys do not authorize metrics while \
7718 MEMRA_METRICS_TOKEN is configured",
7719 "authentication_error",
7720 None,
7721 )),
7722 Err(why) => Err(authentication_error(why)),
7723 };
7724 }
7725 return Err(authentication_error(auth::AuthDenied::Unknown));
7726 }
7727 if api_auth.configured() {
7728 let tenant = authenticate(api_auth, headers)?;
7729 return Ok(if api_auth.keyring.is_some() {
7730 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
7731 } else {
7732 MetricsScope::CompletionDomain
7736 });
7737 }
7738 Err(authentication_error(auth::AuthDenied::Unknown))
7739}
7740
7741#[allow(clippy::result_large_err)] fn lane_for_tenant(
7748 headers: &axum::http::HeaderMap,
7749 tenant: &auth::TenantCtx,
7750) -> Result<lanes::Lane, Response> {
7751 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
7752 None => None,
7753 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
7758 error_response_coded(
7759 StatusCode::BAD_REQUEST,
7760 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
7761 "invalid_request_error",
7762 Some("x-lane"),
7763 Some("invalid_lane"),
7764 )
7765 })?),
7766 };
7767 match tenant.lane_class {
7768 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
7769 auth::LaneClass::Batch => match requested {
7770 None => Ok(lanes::Lane::Harvest),
7771 Some(lanes::Lane::Interactive) => Err(error_response(
7772 StatusCode::FORBIDDEN,
7773 "this api key is batch-class: x-lane interactive is not permitted \
7774 (use judge or harvest)",
7775 "authentication_error",
7776 Some("x-lane"),
7777 )),
7778 Some(l) => Ok(l),
7779 },
7780 }
7781}
7782
7783fn tenant_namespace(
7787 tenant: &auth::TenantCtx,
7788 cache_salt: &Option<String>,
7789) -> Result<String, &'static str> {
7790 let keyring_configured = auth::global().is_some();
7791 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
7792 if keyring_configured {
7793 Ok(auth::scope_namespace(&tenant.tenant, &raw))
7794 } else {
7795 Ok(raw)
7796 }
7797}
7798
7799fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
7804 eprintln!(
7805 "[meter] admit id={} tenant={} lane={} model={:?}",
7806 env.id,
7807 tenant.tenant,
7808 lane.as_str(),
7809 model
7810 );
7811}
7812
7813fn apply_model_request_limits(
7814 request: &mut Request,
7815 metadata: Option<&OpenRouterModelMetadata>,
7816 caps: Option<&ModelCaps>,
7817) -> Result<(), (String, &'static str)> {
7818 let Some(metadata) = metadata else {
7819 return Ok(());
7820 };
7821 let max_prompt = metadata
7822 .max_prompt_length
7823 .map(usize::try_from)
7824 .transpose()
7825 .map_err(|_| {
7826 (
7827 "configured model prompt limit does not fit this platform".into(),
7828 "model",
7829 )
7830 })?;
7831 let max_output = metadata
7832 .max_output_length
7833 .map(usize::try_from)
7834 .transpose()
7835 .map_err(|_| {
7836 (
7837 "configured model output limit does not fit this platform".into(),
7838 "model",
7839 )
7840 })?;
7841
7842 request.max_prompt_tokens = max_prompt;
7843 if let Some(max_output) = max_output {
7844 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
7845 request.params.max_new = metadata
7846 .default_output_length
7847 .map(usize::try_from)
7848 .transpose()
7849 .map_err(|_| {
7850 (
7851 "configured default output length does not fit this platform".into(),
7852 "model",
7853 )
7854 })?
7855 .unwrap_or(max_output);
7856 } else if request.params.max_new > max_output {
7857 return Err((
7858 format!(
7859 "max_tokens {} exceeds configured model maximum {max_output}",
7860 request.params.max_new
7861 ),
7862 "max_tokens",
7863 ));
7864 }
7865 }
7866
7867 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
7871 (max_prompt, max_output, request.params.max_ctx)
7872 {
7873 let operational_ctx = max_prompt
7874 .checked_add(max_output)
7875 .and_then(|value| value.checked_add(8))
7876 .ok_or_else(|| {
7877 (
7878 "configured model context envelope overflowed".into(),
7879 "model",
7880 )
7881 })?;
7882 let operational_ctx = caps
7883 .map(|caps| caps.context_length)
7884 .filter(|&context| context > 0)
7885 .map_or(operational_ctx, |context| operational_ctx.min(context));
7886 if requested_ctx > operational_ctx {
7887 return Err((
7888 format!(
7889 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
7890 ),
7891 "max_ctx",
7892 ));
7893 }
7894 }
7895 Ok(())
7896}
7897
7898fn effective_max_tokens(request: &worker::Request) -> Option<u64> {
7902 (request.params.max_new != worker::MAX_NEW_CTX_BOUNDED).then_some(request.params.max_new as u64)
7903}
7904
7905#[allow(clippy::too_many_arguments)]
7906fn start_request_receipt(
7907 st: &AppState,
7908 env: &Envelope,
7909 tenant: &auth::TenantCtx,
7910 model: &str,
7911 route: &'static str,
7912 lane: lanes::Lane,
7913 stream: bool,
7914 max_tokens: Option<u64>,
7915 reserved_ctx: Option<u64>,
7916 budget_permit: Option<metering::Permit>,
7917) -> Option<Box<dyn metering::Receipt>> {
7918 st.metering.as_ref().map(|accounting| {
7919 accounting.open(
7920 &metering::RequestMeta {
7921 request_id: &env.id,
7922 tenant: &tenant.tenant,
7923 principal: tenant.key_prefix.as_deref(),
7924 model,
7925 route,
7926 lane: lane.as_str(),
7927 stream,
7928 max_tokens,
7929 reserved_ctx,
7930 },
7931 budget_permit,
7932 )
7933 })
7934}
7935
7936fn arm_capture(
7942 mut receipt: Option<Box<dyn metering::Receipt>>,
7943 prompt: impl FnOnce() -> serde_json::Value,
7944) -> Option<Box<dyn metering::Receipt>> {
7945 if let Some(receipt) = receipt.as_mut()
7946 && receipt.wants_capture()
7947 {
7948 receipt.arm_capture(prompt());
7949 }
7950 receipt
7951}
7952
7953fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
7957 serde_json::Value::Array(
7958 messages
7959 .iter()
7960 .map(|message| {
7961 let mut row = json!({ "role": message.role, "content": message.content });
7962 if !message.tool_calls.is_empty() {
7963 row["tool_calls"] = serde_json::Value::Array(
7964 message
7965 .tool_calls
7966 .iter()
7967 .map(|call| {
7968 json!({
7969 "id": call.id,
7970 "function": {
7971 "name": call.function.name,
7972 "arguments": call.function.arguments,
7973 },
7974 })
7975 })
7976 .collect(),
7977 );
7978 }
7979 row
7980 })
7981 .collect(),
7982 )
7983}
7984
7985enum BudgetRejection {
7986 Invalid(String),
7987 Insufficient,
7988 Unenrolled,
7989 PrincipalCapped,
7992 Unavailable(String),
7993}
7994
7995impl BudgetRejection {
7996 fn into_response(self) -> (Response, &'static str) {
7997 match self {
7998 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
7999 Self::Insufficient => (
8000 error_response_coded(
8001 StatusCode::PAYMENT_REQUIRED,
8002 "tenant prepaid balance is insufficient for this request",
8003 "insufficient_balance",
8004 None,
8005 Some("insufficient_balance"),
8006 ),
8007 "insufficient_balance",
8008 ),
8009 Self::Unenrolled => (
8010 error_response_coded(
8011 StatusCode::PAYMENT_REQUIRED,
8012 "tenant is not enrolled for prepaid billing",
8013 "tenant_not_enrolled",
8014 None,
8015 Some("tenant_not_enrolled"),
8016 ),
8017 "tenant_not_enrolled",
8018 ),
8019 Self::PrincipalCapped => (
8020 error_response_coded(
8021 StatusCode::PAYMENT_REQUIRED,
8022 "this API key's spend cap is reached; raise or clear the key's cap to continue",
8023 "key_spend_cap_reached",
8024 None,
8025 Some("key_spend_cap_reached"),
8026 ),
8027 "key_spend_cap_reached",
8028 ),
8029 Self::Unavailable(err) => {
8030 eprintln!("[budget] ERROR: admission unavailable: {err}");
8031 (
8032 error_response_coded(
8033 StatusCode::SERVICE_UNAVAILABLE,
8034 "tenant budget accounting is unavailable",
8035 "server_error",
8036 None,
8037 Some("tenant_budget_unavailable"),
8038 ),
8039 "tenant_budget_unavailable",
8040 )
8041 }
8042 }
8043 }
8044}
8045
8046fn prepare_budget_prompt(
8047 request: &mut Request,
8048 tokenizer: Option<&Tokenizer>,
8049) -> Result<usize, String> {
8050 if let Some(error) = worker::prompt_source_limit_error(request) {
8051 return Err(error);
8052 }
8053 if request.prepared_prompt.is_none() {
8054 if let Some(trace) = request.ttft.as_ref() {
8055 trace.mark_tokenize_start();
8056 }
8057 let prompt = if !request.prompt_ids.is_empty() {
8058 request.prompt_ids.clone()
8059 } else if !request.chat_turns.is_empty() {
8060 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8061 let plain = worker::plain_chat_render_path(
8068 &request.tools_json,
8069 &request.think,
8070 request.reasoning_effort.as_deref(),
8071 &request.chat_turns,
8072 tokenizer.has_qwen_effort_ladder(),
8073 );
8074 let rendered = if plain {
8075 let messages: Vec<_> = request
8076 .chat_turns
8077 .iter()
8078 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
8079 .collect();
8080 tokenizer.apply_chat_template(&messages, true)
8081 } else {
8082 tokenizer
8083 .apply_chat_template_tools_ex(
8084 &request.chat_turns,
8085 true,
8086 &request.tools_json,
8087 &request.tools_struct,
8088 request.think,
8089 request.reasoning_effort.as_deref(),
8090 )
8091 .map_err(|err| format!("chat template: {err}"))?
8092 };
8093 tokenizer.encode(&rendered, true)
8094 } else if request.chat {
8095 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8096 let rendered =
8097 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
8098 tokenizer.encode(&rendered, true)
8099 } else {
8100 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
8101 tokenizer.encode(&request.prompt_text, true)
8102 };
8103 if prompt.is_empty() {
8104 return Err("empty prompt after tokenization".into());
8105 }
8106 if let Some(trace) = request.ttft.as_ref() {
8107 trace.mark_tokenize_end(prompt.len());
8108 }
8109 request.prepared_prompt = Some(prompt);
8110 }
8111 let prompt_tokens = request
8112 .prepared_prompt
8113 .as_ref()
8114 .expect("budget prompt was prepared")
8115 .len();
8116 if let Some(limit) = request.max_prompt_tokens
8117 && prompt_tokens > limit
8118 {
8119 return Err(format!(
8120 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
8121 ));
8122 }
8123 Ok(prompt_tokens)
8124}
8125
8126fn budget_completion_bound(
8127 request: &Request,
8128 prompt_tokens: usize,
8129 caps: Option<&ModelCaps>,
8130) -> Result<usize, String> {
8131 let max_new = request.params.max_new;
8132 let requested_ctx = match (request.params.max_ctx, max_new) {
8133 (Some(cap), _) => cap,
8134 (None, worker::MAX_NEW_CTX_BOUNDED) => {
8135 let server_ctx = std::env::var("MEMRA_CTX")
8136 .ok()
8137 .and_then(|value| value.parse().ok())
8138 .unwrap_or(8192usize);
8139 let mut cap = server_ctx;
8140 if prompt_tokens.saturating_add(16) > cap {
8141 cap = prompt_tokens.saturating_add(server_ctx);
8142 }
8143 cap
8144 }
8145 (None, max_new) => prompt_tokens
8146 .checked_add(max_new)
8147 .and_then(|value| value.checked_add(8))
8148 .ok_or_else(|| "request context bound overflowed".to_string())?,
8149 };
8150 let ctx_cap = caps
8151 .map(|caps| caps.context_length)
8152 .filter(|&context| context > 0)
8153 .map_or(requested_ctx, |context| requested_ctx.min(context));
8154 if prompt_tokens >= ctx_cap {
8155 return Err(format!(
8156 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
8157 ));
8158 }
8159 Ok(max_new.min(ctx_cap - prompt_tokens))
8160}
8161
8162struct BudgetAdmission {
8167 permit: Option<metering::Permit>,
8168 reserved_ctx: Option<u64>,
8169}
8170
8171impl std::fmt::Debug for BudgetAdmission {
8173 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
8174 f.debug_struct("BudgetAdmission")
8175 .field("permit", &self.permit.is_some())
8176 .field("reserved_ctx", &self.reserved_ctx)
8177 .finish()
8178 }
8179}
8180
8181fn admit_tenant_budget(
8182 st: &AppState,
8183 tenant: &auth::TenantCtx,
8184 request: &mut Request,
8185) -> Result<BudgetAdmission, BudgetRejection> {
8186 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
8187 return Ok(BudgetAdmission {
8188 permit: None,
8189 reserved_ctx: None,
8190 });
8191 };
8192 match accounting.is_limited(&tenant.tenant) {
8193 Ok(false) => return Err(BudgetRejection::Unenrolled),
8194 Ok(true) => {}
8195 Err(metering::AdmitError::Unavailable(err)) => {
8196 return Err(BudgetRejection::Unavailable(err));
8197 }
8198 Err(other) => {
8199 return Err(BudgetRejection::Unavailable(format!(
8200 "unexpected budget enrollment result: {other:?}"
8201 )));
8202 }
8203 }
8204 let tokenizer = st
8205 .budget_tokenizers
8206 .as_ref()
8207 .and_then(|tokenizers| tokenizers.get(&request.model))
8208 .map(Arc::as_ref);
8209 if request.prompt_ids.is_empty() && tokenizer.is_none() {
8210 return Err(BudgetRejection::Unavailable(format!(
8211 "no reservation tokenizer for model {:?}",
8212 request.model
8213 )));
8214 }
8215 let prompt_tokens =
8216 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
8217 let completion_tokens =
8218 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
8219 .map_err(BudgetRejection::Invalid)?;
8220 let prompt_tokens = u64::try_from(prompt_tokens)
8221 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
8222 let completion_tokens = u64::try_from(completion_tokens)
8223 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
8224 match accounting.reserve(
8225 &tenant.tenant,
8226 tenant.key_prefix.as_deref(),
8227 &request.model,
8228 prompt_tokens,
8229 completion_tokens,
8230 ) {
8231 Ok(permit) => Ok(BudgetAdmission {
8232 permit,
8233 reserved_ctx: Some(prompt_tokens.saturating_add(completion_tokens)),
8234 }),
8235 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
8236 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
8237 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
8241 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
8242 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
8243 }
8244}
8245
8246fn request_ledger_error_response() -> Response {
8247 error_response_coded(
8248 StatusCode::INTERNAL_SERVER_ERROR,
8249 "request completion could not be committed to the billing ledger",
8250 "server_error",
8251 None,
8252 Some("request_ledger_unavailable"),
8253 )
8254}
8255
8256fn request_ledger_error_body() -> serde_json::Value {
8257 error_body(
8258 "request completion could not be committed to the billing ledger",
8259 "server_error",
8260 None,
8261 Some("request_ledger_unavailable"),
8262 )
8263}
8264
8265fn ledger_rejected(
8266 mut receipt: Option<Box<dyn metering::Receipt>>,
8267 response: Response,
8268 error_code: &str,
8269 request_id: &str,
8270) -> Response {
8271 let status = response.status().as_u16();
8272 if let Some(receipt) = receipt.as_mut()
8273 && let Err(err) = receipt.reject(status, error_code)
8274 {
8275 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
8276 return with_request_id(request_id, request_ledger_error_response());
8277 }
8278 with_request_id(request_id, response)
8279}
8280
8281fn ledger_unbilled(
8286 mut receipt: Option<Box<dyn metering::Receipt>>,
8287 response: Response,
8288 outcome: &'static str,
8289 error_code: &str,
8290 request_id: &str,
8291) -> Response {
8292 let status = response.status().as_u16();
8293 if let Some(receipt) = receipt.as_mut()
8294 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
8295 {
8296 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
8297 return with_request_id(request_id, request_ledger_error_response());
8298 }
8299 with_request_id(request_id, response)
8300}
8301
8302fn engine_error_code(class: worker::ErrClass) -> &'static str {
8303 use worker::ErrClass as C;
8304 match class {
8305 C::InvalidRequest => "invalid_request",
8306 C::ContextLength => "context_length_exceeded",
8307 C::ModelNotFound => "model_not_found",
8308 C::RateLimit => "rate_limit_exceeded",
8309 C::Overloaded => "overloaded",
8310 C::Engine => "engine_error",
8311 }
8312}
8313
8314fn model_not_found_response(models: &[String], requested: &str) -> Response {
8334 error_response_coded(
8335 StatusCode::BAD_REQUEST,
8336 &format!("unknown model {requested:?}; loaded: {models:?}"),
8337 "invalid_request_error",
8338 Some("model"),
8339 Some("model_not_found"),
8340 )
8341}
8342
8343fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
8352 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
8353 return Ok(());
8354 };
8355 if let Some((pos, &id)) = ids
8356 .iter()
8357 .enumerate()
8358 .find(|&(_, &id)| id as usize >= n_vocab)
8359 {
8360 return Err(format!(
8361 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
8362 ));
8363 }
8364 Ok(())
8365}
8366
8367#[cfg(test)]
8368mod prompt_ids_tests {
8369 use super::*;
8370
8371 #[test]
8372 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
8373 let caps = ModelCaps {
8374 n_vocab: 8,
8375 ..Default::default()
8376 };
8377 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
8379 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
8380 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
8382 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
8383 assert!(err.contains("vocab size 8"), "{err}");
8384 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
8385 assert!(err.contains("4294967295"), "{err}");
8386 let unknown = ModelCaps::default();
8388 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
8389 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
8390 }
8391}
8392
8393fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
8394 if models.iter().any(|m| m == requested) {
8395 return Some(requested.to_string());
8396 }
8397 if requested.is_empty() || requested.contains('/') {
8398 return None;
8399 }
8400 let mut matches = models.iter().filter(|m| {
8401 m.rsplit('/')
8402 .next()
8403 .is_some_and(|suffix| suffix == requested)
8404 });
8405 match (matches.next(), matches.next()) {
8406 (Some(only), None) => Some(only.clone()),
8407 _ => None,
8408 }
8409}
8410
8411async fn completions_admitted(
8412 state: State<AppState>,
8413 headers: axum::http::HeaderMap,
8414 trace: Option<Extension<TtftRequestTrace>>,
8415 AdmittedJson(req, admission): AdmittedJson<CompletionReq>,
8416) -> Response {
8417 completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
8418}
8419
8420#[cfg(test)]
8421async fn completions(
8422 State(st): State<AppState>,
8423 headers: axum::http::HeaderMap,
8424 trace: Option<Extension<TtftRequestTrace>>,
8425 request: Json<CompletionReq>,
8426) -> Response {
8427 completions_with_admission(State(st), headers, trace, request, None).await
8428}
8429
8430async fn completions_with_admission(
8431 State(st): State<AppState>,
8432 headers: axum::http::HeaderMap,
8433 trace: Option<Extension<TtftRequestTrace>>,
8434 Json(mut req): Json<CompletionReq>,
8435 mut body_admission: Option<BodyAdmissionLease>,
8436) -> Response {
8437 let env = Envelope::new(false);
8438 if let Err(msg) = req.stop.validate() {
8439 return with_request_id(&env.id, bad_request(&msg, Some("stop")));
8440 }
8441 if let Err(msg) = validate_client_identifier(req.trace_id.as_deref(), "trace_id") {
8442 return with_request_id(&env.id, bad_request(&msg, Some("trace_id")));
8443 }
8444 match canonical_model_id(&st.models, &req.model) {
8445 Some(canonical) => req.model = canonical,
8446 None => {
8447 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
8448 }
8449 }
8450 let ttft = trace.and_then(|Extension(trace)| trace.0);
8453 if let Some(trace) = ttft.as_ref() {
8454 trace.mark_parsed();
8455 trace.bind_request(&env.id, &req.model);
8456 }
8457 let tenant = match authenticate(&st.api_auth, &headers) {
8458 Ok(t) => t,
8459 Err(resp) => return with_request_id(&env.id, resp),
8460 };
8461 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
8462 Ok(ns) => ns,
8463 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
8464 };
8465 if let Err((msg, param)) = reject_unsupported(&[
8467 (
8468 "logit_bias",
8469 req.logit_bias.is_some(),
8470 " (device-side sampling has no bias hook yet)",
8471 ),
8472 ("logprobs", req.logprobs.is_some(), ""),
8473 (
8474 "n",
8475 req.n.is_some_and(|n| n != 1),
8476 " for n != 1 (single choice only)",
8477 ),
8478 (
8479 "best_of",
8480 req.best_of.is_some_and(|n| n != 1),
8481 " (single choice only)",
8482 ),
8483 ]) {
8484 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8485 }
8486 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
8489 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
8490 }
8491 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
8494 Ok(ms) => RequestDeadline::starting_now(ms),
8495 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8496 };
8497 let lane = match lane_for_tenant(&headers, &tenant) {
8498 Ok(l) => l,
8499 Err(resp) => return resp,
8500 };
8501 let (tx, rx) = worker::event_channel();
8502 let model = req.model.clone();
8503 let stream = req.stream;
8504 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
8505 Ok(affinity) => affinity,
8506 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
8507 };
8508 let mut request = build_request_with_trace(
8509 &req,
8510 tx,
8511 lane,
8512 affinity,
8513 ttft.clone(),
8514 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
8518 );
8519 request.cache_ns = cache_ns;
8520 request.request_id = env.id.clone();
8521 request.wire_deadline = Some(deadline.at.into_std());
8524 if let Err((message, param)) = apply_model_request_limits(
8525 &mut request,
8526 st.openrouter_metadata.get(&model),
8527 st.caps.get(&model),
8528 ) {
8529 return with_request_id(&env.id, bad_request(&message, Some(param)));
8530 }
8531 if let Err(msg) = nonstream_deadline_gate(
8536 &request,
8537 req.stream,
8538 deadline,
8539 req.max_tokens.is_some(),
8540 st.budget_tokenizers
8541 .as_ref()
8542 .and_then(|t| t.get(&req.model))
8543 .map(Arc::as_ref),
8544 ) {
8545 return with_request_id(
8546 &env.id,
8547 error_response_coded(
8548 StatusCode::BAD_REQUEST,
8549 &msg,
8550 "invalid_request_error",
8551 Some("max_tokens"),
8552 Some("nonstream_deadline_infeasible"),
8553 ),
8554 );
8555 }
8556 if draining() {
8559 let receipt = start_request_receipt(
8560 &st,
8561 &env,
8562 &tenant,
8563 &req.model,
8564 "/v1/completions",
8565 lane,
8566 req.stream,
8567 effective_max_tokens(&request),
8568 None,
8569 None,
8570 );
8571 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
8572 }
8573 let budget = match admit_tenant_budget(&st, &tenant, &mut request) {
8574 Ok(budget) => budget,
8575 Err(rejection) => {
8576 let (response, error_code) = rejection.into_response();
8577 let receipt = start_request_receipt(
8578 &st,
8579 &env,
8580 &tenant,
8581 &req.model,
8582 "/v1/completions",
8583 lane,
8584 req.stream,
8585 effective_max_tokens(&request),
8586 None,
8587 None,
8588 );
8589 return ledger_rejected(receipt, response, error_code, &env.id);
8590 }
8591 };
8592 let receipt = start_request_receipt(
8593 &st,
8594 &env,
8595 &tenant,
8596 &req.model,
8597 "/v1/completions",
8598 lane,
8599 req.stream,
8600 effective_max_tokens(&request),
8601 budget.reserved_ctx,
8602 budget.permit,
8603 );
8604 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
8605 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
8608 Ok(slot) => slot,
8609 Err(resp) => {
8610 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
8611 }
8612 };
8613 if let Some(admission) = body_admission.as_mut() {
8614 admission.release();
8615 }
8616 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
8619 Ok(guard) => guard,
8620 Err((resp, outcome)) => {
8621 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
8622 }
8623 };
8624 meter_admit(&env, &tenant, &model, lane);
8625 let stop_strings = request.stop_strings.clone();
8626
8627 if let Some(trace) = ttft.as_ref() {
8632 trace.mark_submitted();
8633 }
8634 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
8635 drop(pending_admit);
8636 return ledger_rejected(
8637 receipt,
8638 rl.attach(worker_unavailable_response()),
8639 "worker_unavailable",
8640 &env.id,
8641 );
8642 }
8643 pending_admit.commit();
8644 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
8648 Ok(Ok(rx)) => rx,
8649 Ok(Err((resp, error_code))) => {
8650 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
8651 }
8652 Err(_) => {
8653 return ledger_unbilled(
8654 receipt,
8655 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8656 "deadline_exceeded",
8657 "deadline_exceeded",
8658 &env.id,
8659 );
8660 }
8661 };
8662
8663 let resp = if stream {
8664 let rx = match peek_first_token(rx, deadline).await {
8668 Ok(rx) => rx,
8669 Err(()) => {
8670 return ledger_unbilled(
8671 receipt,
8672 rl.attach(deadline_exceeded_response(deadline.ms, true)),
8673 "deadline_exceeded",
8674 "deadline_exceeded",
8675 &env.id,
8676 );
8677 }
8678 };
8679 sse_response_with_receipt(
8680 rx,
8681 model,
8682 false,
8683 None,
8684 env.clone(),
8685 stop_strings,
8686 Some(guard),
8687 receipt,
8688 )
8689 .into_response()
8690 } else {
8691 let mut receipt = receipt;
8697 let resp = blocking_response_with_receipt(
8698 rx,
8699 model,
8700 false,
8701 stop_strings,
8702 None,
8703 env.clone(),
8704 &mut receipt,
8705 Some(deadline),
8706 )
8707 .await;
8708 drop(guard); resp.into_response()
8710 };
8711 rl.attach(with_request_id(&env.id, resp))
8712}
8713
8714async fn chat_completions_admitted(
8715 state: State<AppState>,
8716 headers: axum::http::HeaderMap,
8717 trace: Option<Extension<TtftRequestTrace>>,
8718 AdmittedJson(req, admission): AdmittedJson<ChatCompletionReq>,
8719) -> Response {
8720 chat_completions_with_admission(state, headers, trace, Json(req), Some(admission)).await
8721}
8722
8723#[cfg(test)]
8724async fn chat_completions(
8725 State(st): State<AppState>,
8726 headers: axum::http::HeaderMap,
8727 trace: Option<Extension<TtftRequestTrace>>,
8728 request: Json<ChatCompletionReq>,
8729) -> Response {
8730 chat_completions_with_admission(State(st), headers, trace, request, None).await
8731}
8732
8733async fn chat_completions_with_admission(
8734 State(st): State<AppState>,
8735 headers: axum::http::HeaderMap,
8736 trace: Option<Extension<TtftRequestTrace>>,
8737 Json(mut req): Json<ChatCompletionReq>,
8738 mut body_admission: Option<BodyAdmissionLease>,
8739) -> Response {
8740 let env = Envelope::new(true);
8741 match canonical_model_id(&st.models, &req.model) {
8746 Some(canonical) => req.model = canonical,
8747 None => {
8748 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
8749 }
8750 }
8751 let ttft = trace.and_then(|Extension(trace)| trace.0);
8752 if let Some(trace) = ttft.as_ref() {
8753 trace.mark_parsed();
8754 trace.bind_request(&env.id, &req.model);
8755 }
8756 let tenant = match authenticate(&st.api_auth, &headers) {
8757 Ok(t) => t,
8758 Err(resp) => return with_request_id(&env.id, resp),
8759 };
8760 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
8761 Ok(ns) => ns,
8762 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
8763 };
8764 if req.messages.is_empty()
8765 || req.messages.iter().any(|message| {
8766 !matches!(
8767 message.role.as_str(),
8768 "system" | "developer" | "user" | "assistant" | "tool"
8769 )
8770 })
8771 {
8772 return with_request_id(
8773 &env.id,
8774 bad_request(
8775 "messages must use system/developer/user/assistant/tool roles",
8776 Some("messages"),
8777 ),
8778 );
8779 }
8780 if let Err((msg, param)) = reject_unsupported(&[
8785 (
8786 "logit_bias",
8787 req.logit_bias.is_some(),
8788 " (device-side sampling has no bias hook yet)",
8789 ),
8790 (
8791 "logprobs",
8792 req.logprobs
8793 .as_ref()
8794 .is_some_and(|v| v.as_bool() != Some(false)),
8795 "",
8796 ),
8797 ("top_logprobs", req.top_logprobs.is_some(), ""),
8798 (
8799 "n",
8800 req.n.is_some_and(|n| n != 1),
8801 " for n != 1 (single choice only)",
8802 ),
8803 ]) {
8804 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8805 }
8806 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
8809 Ok(ms) => RequestDeadline::starting_now(ms),
8810 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8811 };
8812 let lane = match lane_for_tenant(&headers, &tenant) {
8813 Ok(l) => l,
8814 Err(resp) => return resp,
8815 };
8816 let model = req.model.clone();
8817 let stream = req.stream;
8818 let capture_prompt = st
8821 .metering
8822 .as_ref()
8823 .filter(|m| m.captures(&tenant.tenant))
8824 .map(|_| capture_chat_messages(&req.messages));
8825 let declared_max_tokens = req.max_tokens.is_some();
8829 let vision_preprocess_permit = match try_vision_preprocess(request_has_vision(&req)) {
8833 Ok(permit) => permit,
8834 Err(response) => return with_request_id(&env.id, response),
8835 };
8836 let (tx, rx) = worker::event_channel();
8837 let affinity = match affinity_key(&req.session_id, &req.user, &headers) {
8838 Ok(affinity) => affinity,
8839 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("session_id"))),
8840 };
8841 let mut plan = match build_chat_request_with_trace(
8842 req,
8843 st.caps.get(&model),
8844 tx,
8845 lane,
8846 affinity,
8847 ttft.clone(),
8848 st.openrouter_metadata
8849 .get(&model)
8850 .and_then(|m| m.default_reasoning_effort.as_deref()),
8851 &st.sampling_defaults(&model),
8852 ) {
8853 Ok(plan) => plan,
8854 Err(err) => {
8855 return with_request_id(&env.id, bad_request(&err, None));
8856 }
8857 };
8858 plan.request.cache_ns = cache_ns;
8859 plan.request.request_id = env.id.clone();
8860 plan.request.wire_deadline = Some(deadline.at.into_std());
8861 if let Err((message, param)) = apply_model_request_limits(
8862 &mut plan.request,
8863 st.openrouter_metadata.get(&model),
8864 st.caps.get(&model),
8865 ) {
8866 return with_request_id(&env.id, bad_request(&message, Some(param)));
8867 }
8868 if let Err(msg) = nonstream_deadline_gate(
8871 &plan.request,
8872 stream,
8873 deadline,
8874 declared_max_tokens,
8875 st.budget_tokenizers
8876 .as_ref()
8877 .and_then(|t| t.get(&model))
8878 .map(Arc::as_ref),
8879 ) {
8880 return with_request_id(
8881 &env.id,
8882 error_response_coded(
8883 StatusCode::BAD_REQUEST,
8884 &msg,
8885 "invalid_request_error",
8886 Some("max_tokens"),
8887 Some("nonstream_deadline_infeasible"),
8888 ),
8889 );
8890 }
8891 plan.vision_memory = match reserve_vision_memory(&plan) {
8892 Ok(permit) => permit,
8893 Err(err) => {
8894 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
8895 }
8896 };
8897 if draining() {
8900 let receipt = start_request_receipt(
8901 &st,
8902 &env,
8903 &tenant,
8904 &model,
8905 "/v1/chat/completions",
8906 lane,
8907 stream,
8908 effective_max_tokens(&plan.request),
8909 None,
8910 None,
8911 );
8912 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
8913 }
8914 let budget = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
8915 Ok(budget) => budget,
8916 Err(rejection) => {
8917 let (response, error_code) = rejection.into_response();
8918 let receipt = start_request_receipt(
8919 &st,
8920 &env,
8921 &tenant,
8922 &model,
8923 "/v1/chat/completions",
8924 lane,
8925 stream,
8926 effective_max_tokens(&plan.request),
8927 None,
8928 None,
8929 );
8930 return ledger_rejected(receipt, response, error_code, &env.id);
8931 }
8932 };
8933 let receipt = start_request_receipt(
8934 &st,
8935 &env,
8936 &tenant,
8937 &model,
8938 "/v1/chat/completions",
8939 lane,
8940 stream,
8941 effective_max_tokens(&plan.request),
8942 budget.reserved_ctx,
8943 budget.permit,
8944 );
8945 let receipt = if let Some(prompt) = capture_prompt {
8946 arm_capture(receipt, move || prompt)
8947 } else {
8948 receipt
8949 };
8950 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
8954 Ok(slot) => slot,
8955 Err(resp) => {
8956 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
8957 }
8958 };
8959 if let Some(admission) = body_admission.as_mut() {
8960 admission.release();
8961 }
8962 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
8965 Ok(guard) => guard,
8966 Err((resp, outcome)) => {
8967 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
8968 }
8969 };
8970 if let Err(err) = decode_pending_vision(&mut plan) {
8975 return ledger_rejected(
8976 receipt,
8977 rl.attach(bad_request(&err, Some("messages"))),
8978 "invalid_request_error",
8979 &env.id,
8980 );
8981 }
8982 plan.request.vision_memory = plan.vision_memory.take();
8983 drop(vision_preprocess_permit);
8984 let constraint_ready = if plan.request.grammar.is_some() {
8985 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
8986 plan.request.constraint_ready = Some(ready_tx);
8987 Some(ready_rx)
8988 } else {
8989 None
8990 };
8991 meter_admit(&env, &tenant, &model, lane);
8992 let stop_strings = plan.request.stop_strings.clone();
8993 if let Some(trace) = ttft.as_ref() {
8995 trace.mark_submitted();
8996 }
8997 if st
8998 .cmd_tx
8999 .send(Cmd::Generate(Box::new(plan.request)))
9000 .is_err()
9001 {
9002 drop(pending_admit);
9003 return ledger_rejected(
9004 receipt,
9005 rl.attach(worker_unavailable_response()),
9006 "worker_unavailable",
9007 &env.id,
9008 );
9009 }
9010 pending_admit.commit();
9011 if let Some(ready) = constraint_ready {
9017 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
9018 match tokio::time::timeout(bound, ready).await {
9019 Ok(Ok(Ok(()))) => {}
9020 Ok(Ok(Err(err))) => {
9021 return ledger_rejected(
9022 receipt,
9023 rl.attach(engine_error_response(&err)),
9024 engine_error_code(err.class),
9025 &env.id,
9026 );
9027 }
9028 Ok(Err(_)) => {
9029 return ledger_rejected(
9030 receipt,
9031 rl.attach(worker_unavailable_response()),
9032 "worker_unavailable",
9033 &env.id,
9034 );
9035 }
9036 Err(_) if deadline.remaining().is_zero() => {
9037 return ledger_unbilled(
9038 receipt,
9039 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
9040 "deadline_exceeded",
9041 "deadline_exceeded",
9042 &env.id,
9043 );
9044 }
9045 Err(_) => {
9046 return ledger_rejected(
9047 receipt,
9048 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
9049 "constraint_compile_timeout",
9050 &env.id,
9051 );
9052 }
9053 }
9054 }
9055 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
9057 Ok(Ok(rx)) => rx,
9058 Ok(Err((resp, error_code))) => {
9059 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
9060 }
9061 Err(_) => {
9062 return ledger_unbilled(
9063 receipt,
9064 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
9065 "deadline_exceeded",
9066 "deadline_exceeded",
9067 &env.id,
9068 );
9069 }
9070 };
9071 let resp = if stream {
9072 let rx = match peek_first_token(rx, deadline).await {
9074 Ok(rx) => rx,
9075 Err(()) => {
9076 return ledger_unbilled(
9077 receipt,
9078 rl.attach(deadline_exceeded_response(deadline.ms, true)),
9079 "deadline_exceeded",
9080 "deadline_exceeded",
9081 &env.id,
9082 );
9083 }
9084 };
9085 sse_response_with_receipt(
9086 rx,
9087 model,
9088 true,
9089 plan.parser,
9090 env.clone(),
9091 stop_strings,
9092 Some(guard),
9093 receipt,
9094 )
9095 .into_response()
9096 } else {
9097 let mut receipt = receipt;
9100 let resp = blocking_response_with_receipt(
9101 rx,
9102 model,
9103 true,
9104 stop_strings,
9105 plan.parser,
9106 env.clone(),
9107 &mut receipt,
9108 Some(deadline),
9109 )
9110 .await;
9111 drop(guard); resp.into_response()
9113 };
9114 rl.attach(with_request_id(&env.id, resp))
9115}
9116
9117#[cfg(test)]
9126fn sse_response(
9127 rx: worker::EventReceiver,
9128 model: String,
9129 chat: bool,
9130 parser: Option<ToolStreamParser>,
9131 env: Envelope,
9132 stop_strings: Vec<String>,
9133 guard: Option<InflightGuard>,
9134) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9135 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
9136}
9137
9138#[allow(clippy::too_many_arguments)] fn sse_response_with_receipt(
9140 mut rx: worker::EventReceiver,
9141 model: String,
9142 chat: bool,
9143 mut parser: Option<ToolStreamParser>,
9144 env: Envelope,
9145 stop_strings: Vec<String>,
9146 guard: Option<InflightGuard>,
9147 mut receipt: Option<Box<dyn metering::Receipt>>,
9148) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
9149 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
9153 .then(|| StopScrubber::new(stop_strings));
9154 let stream = async_stream::stream! {
9155 let _guard = guard;
9158 let mut call_index: usize = 0;
9159 let mut role_sent = false;
9162 macro_rules! chat_chunk {
9163 ($delta:expr, $finish:expr) => {{
9164 let mut delta = $delta;
9165 if chat && !role_sent {
9166 role_sent = true;
9167 delta["role"] = json!("assistant");
9168 }
9169 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
9170 "choices": [{ "index": 0, "delta": delta,
9171 "finish_reason": $finish }] }))
9172 .to_string()
9173 }};
9174 }
9175 macro_rules! piece_chunks {
9177 ($piece:expr) => {{
9178 let mut payloads: Vec<String> = Vec::new();
9179 match $piece {
9180 Piece::Content(text) => {
9181 let text = match scrub.as_mut() {
9182 Some(sc) => sc.push(&text),
9183 None => text,
9184 };
9185 if !text.is_empty() {
9186 payloads.push(chat_chunk!(json!({ "content": text }),
9187 serde_json::Value::Null));
9188 }
9189 }
9190 Piece::Reasoning(text) => payloads.push(
9194 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
9195 Piece::Call(call) => {
9196 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9197 "index": call_index, "id": call.id, "type": "function",
9198 "function": { "name": call.name, "arguments": "" } }] }),
9199 serde_json::Value::Null));
9200 payloads.push(chat_chunk!(json!({ "tool_calls": [{
9201 "index": call_index,
9202 "function": { "arguments": call.arguments } }] }),
9203 serde_json::Value::Null));
9204 call_index += 1;
9205 }
9206 }
9207 payloads
9208 }};
9209 }
9210 let mut terminal = false;
9214 while let Some(ev) = rx.recv().await {
9215 match ev {
9216 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9218 if let Some(receipt) = receipt.as_mut()
9219 && let Err(err) = receipt.record_prompt_usage(
9220 n_prompt as u64,
9221 n_cached as u64,
9222 )
9223 {
9224 eprintln!(
9225 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9226 env.id
9227 );
9228 let _ = receipt.reject(500, "request_ledger_unavailable");
9231 let payload = request_ledger_error_body().to_string();
9232 if chat || openai_compat() {
9233 yield Ok(SseEvent::default().data(payload));
9234 yield Ok(SseEvent::default().data("[DONE]"));
9235 } else {
9236 yield Ok(SseEvent::default().event("error").data(payload));
9237 }
9238 terminal = true;
9239 break;
9240 }
9241 }
9242 Event::Token { id, text } => {
9243 if let Some(receipt) = receipt.as_mut()
9244 && let Err(err) = receipt.record_completion_token()
9245 {
9246 eprintln!(
9247 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9248 env.id
9249 );
9250 let _ = receipt.reject(500, "request_ledger_unavailable");
9251 let payload = request_ledger_error_body().to_string();
9252 if chat || openai_compat() {
9253 yield Ok(SseEvent::default().data(payload));
9254 yield Ok(SseEvent::default().data("[DONE]"));
9255 } else {
9256 yield Ok(SseEvent::default().event("error").data(payload));
9257 }
9258 terminal = true;
9259 break;
9260 }
9261 if let Some(receipt) = receipt.as_mut() {
9264 receipt.capture_completion_delta(&text);
9265 }
9266 if let Some(p) = parser.as_mut() {
9267 for piece in p.push(&text) {
9268 for payload in piece_chunks!(piece) {
9269 yield Ok(SseEvent::default().data(payload));
9270 }
9271 }
9272 continue;
9273 }
9274 let text = match scrub.as_mut() {
9275 Some(sc) => sc.push(&text),
9276 None => text,
9277 };
9278 if text.is_empty() && scrub.is_some() {
9279 continue; }
9281 let payload = if chat {
9282 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
9283 } else if openai_compat() {
9284 env.stamp(json!({ "object": "text_completion", "model": model,
9285 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
9286 .to_string()
9287 } else {
9288 json!({ "model": model, "id": id, "text": text }).to_string()
9289 };
9290 yield Ok(SseEvent::default().data(payload));
9291 }
9292 Event::TokenSnapshot(_) => {}
9296 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
9297 let mut finish = stop_reason_to_finish(&stop_reason);
9298 if let Some(p) = parser.as_mut() {
9299 for piece in p.finish() {
9300 for payload in piece_chunks!(piece) {
9301 yield Ok(SseEvent::default().data(payload));
9302 }
9303 }
9304 if p.n_calls() > 0 { finish = "tool_calls"; }
9305 }
9306 if let Some(sc) = scrub.as_mut() {
9308 let tail = sc.finish();
9309 if !tail.is_empty() {
9310 let payload = if chat {
9311 chat_chunk!(json!({ "content": tail }),
9312 serde_json::Value::Null)
9313 } else {
9314 env.stamp(json!({ "object": "text_completion",
9315 "model": model,
9316 "choices": [{ "index": 0, "text": tail,
9317 "finish_reason": null }] })).to_string()
9318 };
9319 yield Ok(SseEvent::default().data(payload));
9320 }
9321 }
9322 if let Some(receipt) = receipt.as_mut()
9323 && let Err(err) = receipt.complete(
9324 metering::UsageCounts {
9325 prompt_tokens: n_prompt as u64,
9326 cached_prompt_tokens: n_cached as u64,
9327 completion_tokens: n_tokens as u64,
9328 },
9329 elapsed_s,
9330 )
9331 {
9332 eprintln!(
9333 "[ledger] ERROR: request {} completion receipt failed: {err}",
9334 env.id
9335 );
9336 let _ = receipt.reject(500, "request_ledger_unavailable");
9340 let payload = request_ledger_error_body().to_string();
9341 if chat || openai_compat() {
9342 yield Ok(SseEvent::default().data(payload));
9343 yield Ok(SseEvent::default().data("[DONE]"));
9344 } else {
9345 yield Ok(SseEvent::default().event("error").data(payload));
9346 }
9347 terminal = true;
9348 break;
9349 }
9350 if chat || openai_compat() {
9351 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
9352 let fin = if chat {
9353 let mut v = env.stamp(json!({
9354 "object": "chat.completion.chunk", "model": model,
9355 "choices": [{ "index": 0, "delta": {},
9356 "finish_reason": finish }],
9357 "usage": usage }));
9358 if !role_sent {
9360 v["choices"][0]["delta"]["role"] = json!("assistant");
9361 }
9362 v
9363 } else {
9364 env.stamp(json!({ "object": "text_completion", "model": model,
9365 "choices": [{ "index": 0, "text": "",
9366 "finish_reason": finish }],
9367 "usage": usage }))
9368 }.to_string();
9369 yield Ok(SseEvent::default().data(fin));
9370 yield Ok(SseEvent::default().data("[DONE]"));
9371 } else {
9372 let payload = json!({
9373 "stop_reason": stop_reason, "n_tokens": n_tokens,
9374 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
9375 "elapsed_s": elapsed_s
9376 }).to_string();
9377 yield Ok(SseEvent::default().event("done").data(payload));
9378 }
9379 terminal = true;
9380 break;
9381 }
9382 Event::Error(err) => {
9383 let ledger_error = if let Some(receipt) = receipt.as_mut() {
9393 receipt
9394 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
9395 .err()
9396 } else {
9397 None
9398 };
9399 if let Some(ref ledger_error) = ledger_error {
9400 eprintln!(
9401 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
9402 env.id
9403 );
9404 }
9405 let payload = if ledger_error.is_some() {
9406 request_ledger_error_body().to_string()
9407 } else {
9408 engine_error_body(&err).to_string()
9409 };
9410 if chat || openai_compat() {
9411 yield Ok(SseEvent::default().data(payload));
9414 yield Ok(SseEvent::default().data("[DONE]"));
9415 } else {
9416 yield Ok(SseEvent::default().event("error").data(payload));
9419 }
9420 terminal = true;
9421 break;
9422 }
9423 }
9424 }
9425 if !terminal {
9426 let e = worker::EngineError::overloaded(
9432 "worker closed the stream without completing (worker restart in progress)",
9433 );
9434 if let Some(receipt) = receipt.as_mut()
9435 && let Err(ledger_err) = receipt.reject(
9436 class_http(e.class).0.as_u16(),
9437 engine_error_code(e.class),
9438 )
9439 {
9440 eprintln!(
9441 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9442 env.id
9443 );
9444 }
9445 let payload = engine_error_body(&e).to_string();
9446 if chat || openai_compat() {
9447 yield Ok(SseEvent::default().data(payload));
9448 yield Ok(SseEvent::default().data("[DONE]"));
9449 } else {
9450 yield Ok(SseEvent::default().event("error").data(payload));
9451 }
9452 }
9453 };
9454 Sse::new(stream).keep_alive(
9455 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
9458 )
9459}
9460
9461fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
9463 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
9464 text.truncate(offset);
9465 }
9466}
9467
9468fn partial_stop_suffix(s: &str, tag: &str) -> usize {
9471 let mut best = 0;
9472 for (k, _) in tag.char_indices().skip(1) {
9473 if k <= s.len() && s.ends_with(&tag[..k]) {
9474 best = k;
9475 }
9476 }
9477 best
9478}
9479
9480struct StopScrubber {
9486 stops: Vec<String>,
9487 buf: String,
9488 done: bool,
9489}
9490
9491impl StopScrubber {
9492 fn new(stops: Vec<String>) -> Self {
9493 Self {
9494 stops,
9495 buf: String::new(),
9496 done: false,
9497 }
9498 }
9499
9500 fn push(&mut self, text: &str) -> String {
9502 if self.done {
9503 return String::new();
9504 }
9505 self.buf.push_str(text);
9506 if let Some(i) = self
9507 .stops
9508 .iter()
9509 .filter_map(|s| self.buf.find(s.as_str()))
9510 .min()
9511 {
9512 self.done = true;
9513 let out = self.buf[..i].to_string();
9514 self.buf.clear();
9515 return out;
9516 }
9517 let keep = self
9518 .stops
9519 .iter()
9520 .map(|s| partial_stop_suffix(&self.buf, s))
9521 .max()
9522 .unwrap_or(0);
9523 let emit_to = self.buf.len() - keep;
9524 let out = self.buf[..emit_to].to_string();
9525 self.buf.drain(..emit_to);
9526 out
9527 }
9528
9529 fn finish(&mut self) -> String {
9531 if self.done {
9532 self.buf.clear();
9533 return String::new();
9534 }
9535 std::mem::take(&mut self.buf)
9536 }
9537}
9538
9539#[cfg(test)]
9540async fn blocking_response(
9541 rx: worker::EventReceiver,
9542 model: String,
9543 chat: bool,
9544 stop_strings: Vec<String>,
9545 parser: Option<ToolStreamParser>,
9546 env: Envelope,
9547) -> Response {
9548 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
9549 .await
9550}
9551
9552struct BlockingPayload<'a> {
9556 env: &'a Envelope,
9557 model: String,
9558 chat: bool,
9559 finish: &'static str,
9560 text: String,
9561 reasoning: String,
9562 calls: Vec<ParsedToolCall>,
9563 tokens: Vec<u32>,
9564 stop_reason: String,
9565 n_prompt: usize,
9566 n_tokens: usize,
9567 n_cached: usize,
9568 elapsed_s: f64,
9569 spec: Option<worker::SpecUsage>,
9570 deadline_error: Option<serde_json::Value>,
9576}
9577
9578fn blocking_payload(p: BlockingPayload<'_>) -> Response {
9579 let BlockingPayload {
9580 env,
9581 model,
9582 chat,
9583 finish,
9584 text,
9585 reasoning,
9586 calls,
9587 tokens,
9588 stop_reason,
9589 n_prompt,
9590 n_tokens,
9591 n_cached,
9592 elapsed_s,
9593 spec,
9594 deadline_error,
9595 } = p;
9596 if chat {
9597 let content = if !calls.is_empty() && text.is_empty() {
9599 serde_json::Value::Null
9600 } else {
9601 serde_json::Value::String(text)
9602 };
9603 let mut message = json!({ "role": "assistant", "content": content });
9604 if !reasoning.is_empty() {
9607 message["reasoning"] = json!(reasoning);
9608 message["reasoning_details"] = json!([{
9609 "type": "reasoning.text", "text": reasoning }]);
9610 }
9611 if !calls.is_empty() {
9612 message["tool_calls"] =
9613 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
9614 }
9615 let mut body = json!({
9616 "object": "chat.completion", "model": model,
9617 "choices": [{ "index": 0,
9618 "message": message,
9619 "finish_reason": finish }],
9620 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
9621 });
9622 if let Some(err) = deadline_error {
9623 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
9624 body["error"] = err;
9625 }
9626 return Json(env.stamp(body)).into_response();
9627 }
9628 if openai_compat() {
9629 let mut body = json!({
9630 "object": "text_completion", "model": model,
9631 "choices": [{ "index": 0, "text": text,
9632 "finish_reason": finish }],
9633 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
9634 });
9635 if let Some(err) = deadline_error {
9636 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
9637 body["error"] = err;
9638 }
9639 return Json(env.stamp(body)).into_response();
9640 }
9641 Json(CompletionResp {
9642 model,
9643 text,
9644 tokens,
9645 stop_reason,
9646 error: deadline_error,
9647 n_tokens,
9648 prompt_tokens: n_prompt,
9649 cached_tokens: n_cached,
9650 elapsed_s,
9651 })
9652 .into_response()
9653}
9654
9655#[allow(clippy::too_many_arguments)] async fn blocking_response_with_receipt(
9681 mut rx: worker::EventReceiver,
9682 model: String,
9683 chat: bool,
9684 stop_strings: Vec<String>,
9685 mut parser: Option<ToolStreamParser>,
9686 env: Envelope,
9687 receipt: &mut Option<Box<dyn metering::Receipt>>,
9688 deadline: Option<RequestDeadline>,
9689) -> Response {
9690 let mut text = String::new();
9691 let mut reasoning = String::new();
9692 let mut tokens: Vec<u32> = Vec::new();
9693 let mut calls: Vec<ParsedToolCall> = Vec::new();
9694 let consume = |pieces: Vec<Piece>,
9695 text: &mut String,
9696 reasoning: &mut String,
9697 calls: &mut Vec<ParsedToolCall>| {
9698 for piece in pieces {
9699 match piece {
9700 Piece::Content(t) => text.push_str(&t),
9701 Piece::Reasoning(t) => reasoning.push_str(&t),
9702 Piece::Call(c) => calls.push(c),
9703 }
9704 }
9705 };
9706 let started = std::time::Instant::now();
9708 let mut seen_prompt: usize = 0;
9709 let mut seen_cached: usize = 0;
9710 let mut seen_tokens: usize = 0;
9711 loop {
9712 let ev = match deadline {
9713 Some(d) => tokio::select! {
9714 biased;
9715 ev = rx.recv() => ev,
9716 () = tokio::time::sleep_until(d.at) => {
9717 drop(rx);
9720 if seen_tokens == 0 {
9721 if let Some(receipt) = receipt.as_mut()
9725 && let Err(err) = receipt.settle_unbilled(
9726 "deadline_exceeded",
9727 StatusCode::REQUEST_TIMEOUT.as_u16(),
9728 "deadline_exceeded",
9729 )
9730 {
9731 eprintln!(
9732 "[ledger] ERROR: request {} deadline receipt failed: {err}",
9733 env.id
9734 );
9735 return request_ledger_error_response();
9736 }
9737 return deadline_exceeded_response(d.ms, false);
9738 }
9739 if let Some(p) = parser.as_mut() {
9740 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
9741 }
9742 truncate_at_stop(&mut text, &stop_strings);
9743 let elapsed_s = started.elapsed().as_secs_f64();
9744 if let Some(receipt) = receipt.as_mut()
9747 && let Err(err) = receipt.complete_deadline_partial(
9748 metering::UsageCounts {
9749 prompt_tokens: seen_prompt as u64,
9750 cached_prompt_tokens: seen_cached as u64,
9751 completion_tokens: seen_tokens as u64,
9752 },
9753 elapsed_s,
9754 )
9755 {
9756 eprintln!(
9757 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
9758 env.id
9759 );
9760 let _ = receipt.reject(500, "request_ledger_unavailable");
9761 return request_ledger_error_response();
9762 }
9763 eprintln!(
9764 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
9765 {} ms deadline (prompt {}); non-streaming caller advised to stream",
9766 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
9767 );
9768 let err_obj = json!({
9769 "message": format!(
9770 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
9771 the {} tokens produced before the cut are delivered above and are \
9772 billed. Set \"stream\": true for work this long — a stream's \
9773 deadline bounds only the time to first token — or lower max_tokens.",
9774 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
9775 ),
9776 "code": "deadline_exceeded",
9777 "metadata": { "error_type": "timeout", "provider_name": "memra" }
9778 });
9779 return blocking_payload(BlockingPayload {
9780 env: &env,
9781 model,
9782 chat,
9783 finish: "error",
9784 text,
9785 reasoning,
9786 calls,
9787 tokens,
9788 stop_reason: "Deadline".to_string(),
9789 n_prompt: seen_prompt,
9790 n_tokens: seen_tokens,
9791 n_cached: seen_cached,
9792 elapsed_s,
9793 spec: None,
9794 deadline_error: Some(err_obj),
9795 });
9796 }
9797 },
9798 None => rx.recv().await,
9799 };
9800 let Some(ev) = ev else { break };
9801 match ev {
9802 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9804 if let Some(receipt) = receipt.as_mut()
9805 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
9806 {
9807 eprintln!(
9808 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9809 env.id
9810 );
9811 let _ = receipt.reject(500, "request_ledger_unavailable");
9814 return request_ledger_error_response();
9815 }
9816 seen_prompt = n_prompt;
9817 seen_cached = n_cached;
9818 }
9819 Event::Token { id, text: delta } => {
9820 if let Some(receipt) = receipt.as_mut()
9821 && let Err(err) = receipt.record_completion_token()
9822 {
9823 eprintln!(
9824 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9825 env.id
9826 );
9827 let _ = receipt.reject(500, "request_ledger_unavailable");
9828 return request_ledger_error_response();
9829 }
9830 if let Some(receipt) = receipt.as_mut() {
9832 receipt.capture_completion_delta(&delta);
9833 }
9834 tokens.push(id);
9835 seen_tokens += 1;
9836 match parser.as_mut() {
9837 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
9838 None => text.push_str(&delta),
9839 }
9840 }
9841 Event::TokenSnapshot(ids) => tokens = ids,
9842 Event::Done {
9843 stop_reason,
9844 n_tokens,
9845 n_prompt,
9846 n_cached,
9847 elapsed_s,
9848 spec,
9849 } => {
9850 if let Some(p) = parser.as_mut() {
9851 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
9852 }
9853 truncate_at_stop(&mut text, &stop_strings);
9854 let finish = if calls.is_empty() {
9855 stop_reason_to_finish(&stop_reason)
9856 } else {
9857 "tool_calls"
9858 };
9859 if let Some(receipt) = receipt.as_mut()
9860 && let Err(err) = receipt.complete(
9861 metering::UsageCounts {
9862 prompt_tokens: n_prompt as u64,
9863 cached_prompt_tokens: n_cached as u64,
9864 completion_tokens: n_tokens as u64,
9865 },
9866 elapsed_s,
9867 )
9868 {
9869 eprintln!(
9870 "[ledger] ERROR: request {} completion receipt failed: {err}",
9871 env.id
9872 );
9873 let _ = receipt.reject(500, "request_ledger_unavailable");
9876 return request_ledger_error_response();
9877 }
9878 return blocking_payload(BlockingPayload {
9879 env: &env,
9880 model,
9881 chat,
9882 finish,
9883 text,
9884 reasoning,
9885 calls,
9886 tokens,
9887 stop_reason,
9888 n_prompt,
9889 n_tokens,
9890 n_cached,
9891 elapsed_s,
9892 spec,
9893 deadline_error: None,
9894 });
9895 }
9896 Event::Error(err) => {
9897 if let Some(receipt) = receipt.as_mut()
9901 && let Err(ledger_err) = receipt.reject(
9902 class_http(err.class).0.as_u16(),
9903 engine_error_code(err.class),
9904 )
9905 {
9906 eprintln!(
9907 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
9908 env.id
9909 );
9910 return request_ledger_error_response();
9911 }
9912 return engine_error_response(&err);
9913 }
9914 }
9915 }
9916 let e = worker::EngineError::overloaded(
9921 "worker closed the stream without completing (worker restart in progress)",
9922 );
9923 if let Some(receipt) = receipt.as_mut()
9924 && let Err(ledger_err) =
9925 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
9926 {
9927 eprintln!(
9928 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9929 env.id
9930 );
9931 return request_ledger_error_response();
9932 }
9933 engine_error_response(&e)
9934}
9935
9936#[cfg(test)]
9937mod tests {
9938 use super::*;
9939
9940 #[derive(Debug, Clone, PartialEq)]
9948 enum MeterEvent {
9949 Reserve {
9950 tenant: String,
9951 principal: Option<String>,
9952 model: String,
9953 },
9954 Open {
9955 request_id: String,
9956 tenant: String,
9957 model: String,
9958 route: &'static str,
9959 stream: bool,
9960 with_permit: bool,
9961 },
9962 PromptUsage {
9963 prompt: u64,
9964 cached: u64,
9965 },
9966 Token,
9967 CapturePrompt(serde_json::Value),
9968 CaptureDelta(String),
9969 Complete {
9970 prompt: u64,
9971 cached: u64,
9972 completion: u64,
9973 },
9974 DeadlinePartial {
9975 prompt: u64,
9976 cached: u64,
9977 completion: u64,
9978 },
9979 Reject {
9980 status: u16,
9981 code: String,
9982 },
9983 Unbilled {
9984 outcome: &'static str,
9985 status: u16,
9986 code: String,
9987 },
9988 Dropped {
9991 prompt: u64,
9992 cached: u64,
9993 completion: u64,
9994 },
9995 }
9996
9997 enum ReserveScript {
10000 Admit { with_permit: bool },
10001 Insufficient,
10002 Blocked,
10003 PrincipalCapped,
10004 }
10005
10006 struct MockMetering {
10007 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10008 limits: bool,
10009 limited: bool,
10010 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
10011 captures: bool,
10012 }
10013
10014 impl MockMetering {
10015 fn admit_all() -> Arc<Self> {
10016 Arc::new(MockMetering {
10017 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10018 limits: false,
10019 limited: true,
10020 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10021 captures: false,
10022 })
10023 }
10024
10025 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
10026 Arc::new(MockMetering {
10027 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10028 limits: true,
10029 limited: true,
10030 reserve_script: std::sync::Mutex::new(script.into()),
10031 captures: false,
10032 })
10033 }
10034
10035 fn capturing() -> Arc<Self> {
10036 Arc::new(MockMetering {
10037 events: Arc::new(std::sync::Mutex::new(Vec::new())),
10038 limits: false,
10039 limited: true,
10040 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
10041 captures: true,
10042 })
10043 }
10044
10045 fn events(&self) -> Vec<MeterEvent> {
10046 self.events.lock().unwrap().clone()
10047 }
10048 }
10049
10050 impl metering::Metering for MockMetering {
10051 fn enforces_limits(&self) -> bool {
10052 self.limits
10053 }
10054
10055 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
10056 Ok(self.limited)
10057 }
10058
10059 fn reserve(
10060 &self,
10061 tenant: &str,
10062 principal: Option<&str>,
10063 model: &str,
10064 _prompt_tokens: u64,
10065 _completion_bound: u64,
10066 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
10067 self.events.lock().unwrap().push(MeterEvent::Reserve {
10068 tenant: tenant.into(),
10069 principal: principal.map(str::to_owned),
10070 model: model.into(),
10071 });
10072 match self.reserve_script.lock().unwrap().pop_front() {
10073 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
10074 Some(ReserveScript::Admit { with_permit: true }) => {
10075 Ok(Some(Box::new(()) as metering::Permit))
10076 }
10077 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
10078 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
10079 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
10080 }
10081 }
10082
10083 fn open(
10084 &self,
10085 meta: &metering::RequestMeta<'_>,
10086 permit: Option<metering::Permit>,
10087 ) -> Box<dyn metering::Receipt> {
10088 self.events.lock().unwrap().push(MeterEvent::Open {
10089 request_id: meta.request_id.into(),
10090 tenant: meta.tenant.into(),
10091 model: meta.model.into(),
10092 route: meta.route,
10093 stream: meta.stream,
10094 with_permit: permit.is_some(),
10095 });
10096 Box::new(MockReceipt {
10097 events: self.events.clone(),
10098 wants_capture: self.captures,
10099 prompt: 0,
10100 cached: 0,
10101 completion: 0,
10102 finalized: false,
10103 })
10104 }
10105
10106 fn captures(&self, _tenant: &str) -> bool {
10107 self.captures
10108 }
10109
10110 fn limits_health(&self) -> Option<metering::LimitsHealth> {
10111 self.limits.then_some(metering::LimitsHealth {
10112 source_reload_failed: 0,
10113 source_reload_consecutive: 0,
10114 source_available: true,
10115 })
10116 }
10117 }
10118
10119 struct MockReceipt {
10120 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
10121 wants_capture: bool,
10122 prompt: u64,
10123 cached: u64,
10124 completion: u64,
10125 finalized: bool,
10126 }
10127
10128 impl metering::Receipt for MockReceipt {
10129 fn wants_capture(&self) -> bool {
10130 self.wants_capture
10131 }
10132
10133 fn arm_capture(&mut self, prompt: serde_json::Value) {
10134 self.events
10135 .lock()
10136 .unwrap()
10137 .push(MeterEvent::CapturePrompt(prompt));
10138 }
10139
10140 fn capture_completion_delta(&mut self, text: &str) {
10141 if self.wants_capture {
10142 self.events
10143 .lock()
10144 .unwrap()
10145 .push(MeterEvent::CaptureDelta(text.into()));
10146 }
10147 }
10148
10149 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
10150 self.prompt = prompt;
10151 self.cached = cached;
10152 self.events
10153 .lock()
10154 .unwrap()
10155 .push(MeterEvent::PromptUsage { prompt, cached });
10156 Ok(())
10157 }
10158
10159 fn record_completion_token(&mut self) -> Result<(), String> {
10160 self.completion += 1;
10161 self.events.lock().unwrap().push(MeterEvent::Token);
10162 Ok(())
10163 }
10164
10165 fn complete(
10166 &mut self,
10167 usage: metering::UsageCounts,
10168 _worker_elapsed_s: f64,
10169 ) -> Result<(), String> {
10170 self.finalized = true;
10171 self.events.lock().unwrap().push(MeterEvent::Complete {
10172 prompt: usage.prompt_tokens,
10173 cached: usage.cached_prompt_tokens,
10174 completion: usage.completion_tokens,
10175 });
10176 Ok(())
10177 }
10178
10179 fn complete_deadline_partial(
10180 &mut self,
10181 usage: metering::UsageCounts,
10182 _worker_elapsed_s: f64,
10183 ) -> Result<(), String> {
10184 self.finalized = true;
10185 self.events
10186 .lock()
10187 .unwrap()
10188 .push(MeterEvent::DeadlinePartial {
10189 prompt: usage.prompt_tokens,
10190 cached: usage.cached_prompt_tokens,
10191 completion: usage.completion_tokens,
10192 });
10193 Ok(())
10194 }
10195
10196 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
10197 self.finalized = true;
10198 self.events.lock().unwrap().push(MeterEvent::Reject {
10199 status,
10200 code: error_code.into(),
10201 });
10202 Ok(())
10203 }
10204
10205 fn settle_unbilled(
10206 &mut self,
10207 outcome: &'static str,
10208 status: u16,
10209 error_code: &str,
10210 ) -> Result<(), String> {
10211 self.finalized = true;
10212 self.events.lock().unwrap().push(MeterEvent::Unbilled {
10213 outcome,
10214 status,
10215 code: error_code.into(),
10216 });
10217 Ok(())
10218 }
10219 }
10220
10221 impl Drop for MockReceipt {
10222 fn drop(&mut self) {
10223 if !self.finalized {
10224 self.events.lock().unwrap().push(MeterEvent::Dropped {
10225 prompt: self.prompt,
10226 cached: self.cached,
10227 completion: self.completion,
10228 });
10229 }
10230 }
10231 }
10232
10233 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
10238
10239 fn gate_env_lock() -> std::sync::MutexGuard<'static, ()> {
10246 let guard = GATE_ENV_LOCK.lock().unwrap_or_else(|poisoned| {
10247 GATE_ENV_LOCK.clear_poison();
10250 poisoned.into_inner()
10251 });
10252 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10253 guard
10254 }
10255
10256 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
10259 let req: CompletionReq = serde_json::from_value(json!({
10260 "model": "qwen/qwen3.8-27b",
10261 "prompt_ids": vec![7u32; prompt_ids],
10262 }))
10263 .unwrap();
10264 let (tx, _rx) = worker::event_channel();
10265 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
10266 request.params.max_new = max_new;
10267 request
10268 }
10269
10270 #[test]
10277 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
10278 let prompt = 30_278u64;
10279 let deadline_ms = TIMEOUT_MS_DEFAULT;
10280 let margin = |max_new: u64| {
10281 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
10282 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
10283 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
10284 };
10285 for allowed in [64u64, 2048, 4096, 5120, 6144] {
10286 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
10287 }
10288 for refused in [8192u64, 16384, 262_144] {
10289 assert!(
10290 !margin(refused),
10291 "{refused} measured as a 408 and must be refused"
10292 );
10293 }
10294 }
10295
10296 #[test]
10297 fn the_gate_names_a_max_tokens_that_actually_fits() {
10298 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
10301 assert!(
10302 fits > 0 && fits < 7_800,
10303 "advice {fits} must fit the measured ceiling"
10304 );
10305 assert_eq!(
10307 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
10308 None
10309 );
10310 }
10311
10312 #[test]
10313 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
10314 let req = gate_request(262_144, 30_000);
10315 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
10316 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
10318 assert!(
10319 err.contains("stream"),
10320 "message must name the streaming alternative: {err}"
10321 );
10322 assert!(
10323 err.contains("max_tokens"),
10324 "message must name the knob: {err}"
10325 );
10326 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
10328 let _l = gate_env_lock(); for off in ["0", "off", "false"] {
10335 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
10336 assert!(
10337 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
10338 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
10339 );
10340 }
10341 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
10342 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
10343 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
10344 assert!(
10345 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
10346 "unset means ON (the documented default)"
10347 );
10348 }
10349
10350 #[test]
10357 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
10358 let strip = |src: &str| -> String {
10365 src.lines()
10366 .map(|line| match line.find("//") {
10367 Some(i) => line[..i].to_string(),
10368 None => line.to_string(),
10369 })
10370 .collect::<Vec<_>>()
10371 .join("\n")
10372 };
10373 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
10375 let start = src
10376 .find(signature)
10377 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
10378 let rest = &src[start + signature.len()..];
10379 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
10380 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
10381 &rest[..end]
10382 }
10383 let main_src = strip(include_str!("lib.rs"));
10384 let surfaces_src = strip(include_str!("surfaces.rs"));
10385 for (surface, src, signature) in [
10386 (
10387 "/v1/completions",
10388 &main_src,
10389 "async fn completions_with_admission(",
10390 ),
10391 (
10392 "/v1/chat/completions",
10393 &main_src,
10394 "async fn chat_completions_with_admission(",
10395 ),
10396 (
10397 "/v1/messages + /v1/responses (shared admission)",
10398 &surfaces_src,
10399 "pub(crate) async fn admit_translated(",
10400 ),
10401 ] {
10402 let handler = body(src, signature);
10403 assert!(
10404 handler.contains("nonstream_deadline_gate("),
10405 "{surface} must CALL the feasibility gate inside {signature}"
10406 );
10407 let limits = handler
10410 .find("apply_model_request_limits(")
10411 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
10412 let gate = handler.find("nonstream_deadline_gate(").unwrap();
10413 assert!(
10414 limits < gate,
10415 "{surface}: the gate must run after apply_model_request_limits"
10416 );
10417 }
10418 }
10419
10420 #[test]
10424 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
10425 let err = json!({"code": "deadline_exceeded",
10426 "metadata": {"error_type": "timeout"}});
10427 let cut = CompletionResp {
10428 model: "m".into(),
10429 text: "partial".into(),
10430 tokens: vec![1, 2],
10431 stop_reason: "Deadline".into(),
10432 error: Some(err.clone()),
10433 n_tokens: 2,
10434 prompt_tokens: 9,
10435 cached_tokens: 0,
10436 elapsed_s: 1.0,
10437 };
10438 let v = serde_json::to_value(&cut).unwrap();
10439 assert_eq!(v["stop_reason"], "Deadline");
10440 assert_eq!(v["error"]["code"], "deadline_exceeded");
10441 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
10442 let whole = CompletionResp {
10444 error: None,
10445 stop_reason: "Eos".into(),
10446 ..cut
10447 };
10448 let v = serde_json::to_value(&whole).unwrap();
10449 assert!(
10450 v.get("error").is_none(),
10451 "a complete response must not grow an error key: {v}"
10452 );
10453 }
10454
10455 #[test]
10456 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
10457 let _l = gate_env_lock();
10458 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
10462 assert!(
10463 nonstream_deadline_gate(
10464 &req,
10465 false,
10466 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10467 false,
10468 None,
10469 )
10470 .is_ok(),
10471 "an omitted max_tokens is never gated — context is its only limit"
10472 );
10473 let resolved = gate_request(32_768, 30_000);
10478 assert!(
10479 nonstream_deadline_gate(
10480 &resolved,
10481 false,
10482 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10483 false,
10484 None,
10485 )
10486 .is_ok(),
10487 "a resolved-but-undeclared cap is not the caller's number to be refused over"
10488 );
10489 assert!(
10491 nonstream_deadline_gate(
10492 &resolved,
10493 false,
10494 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
10495 true,
10496 None,
10497 )
10498 .is_err()
10499 );
10500 }
10501
10502 #[test]
10503 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
10504 let req = gate_request(64, 1234);
10505 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
10506 let mut text = gate_request(64, 0);
10507 text.prompt_ids.clear();
10508 text.prompt_text = "x".repeat(6_000);
10509 assert_eq!(
10510 prompt_tokens_estimate(&text, None),
10511 1_000,
10512 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
10513 that would have succeeded"
10514 );
10515 }
10516
10517 #[test]
10518 fn vision_memory_reservation_is_bounded_and_released() {
10519 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
10520 let Err(capacity) = try_reserve_vision_memory(1) else {
10521 panic!("a full process vision budget admitted another request");
10522 };
10523 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
10524 let response = vision_memory_error_response(capacity, Some("messages"));
10525 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
10526 assert_eq!(response.headers()["retry-after"], "5");
10527 assert_eq!(response.headers()["retry-after-ms"], "5000");
10528 drop(permit);
10529 assert!(try_reserve_vision_memory(1).is_ok());
10530 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
10531 panic!("an over-limit vision request was admitted");
10532 };
10533 assert!(matches!(request, VisionMemoryError::Request(_)));
10534 let response = vision_memory_error_response(request, Some("messages"));
10535 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
10536 assert_eq!(response.headers()["x-should-retry"], "false");
10537 let _ = try_reserve_vision_memory(1);
10538 }
10539
10540 #[test]
10541 fn header_auth_gate_covers_only_inference_dialects() {
10542 for path in [
10543 "/v1/auth/check",
10544 "/v1/completions",
10545 "/v1/chat/completions",
10546 "/v1/messages",
10547 "/v1/responses",
10548 "/v1/embeddings",
10549 "/v1/rerank",
10550 ] {
10551 assert!(protected_inference_path(path), "{path}");
10552 }
10553 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
10554 assert!(!protected_inference_path(path), "{path}");
10555 }
10556 }
10557 #[tokio::test]
10564 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
10565 use crate::metering::Metering as _;
10566 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
10567
10568 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
10569 let (tx, rx) = worker::event_channel();
10570 tx.send(Event::PromptUsage {
10571 n_prompt: 7,
10572 n_cached: 0,
10573 })
10574 .unwrap();
10575 tx.send(Event::Token {
10576 id: 1,
10577 text: "Hel".into(),
10578 })
10579 .unwrap();
10580 tx.send(Event::Token {
10581 id: 2,
10582 text: "lo".into(),
10583 })
10584 .unwrap();
10585 tx.send(Event::Done {
10586 stop_reason: "eos".into(),
10587 n_tokens: 2,
10588 n_prompt: 7,
10589 n_cached: 0,
10590 elapsed_s: 0.05,
10591 spec: None,
10592 })
10593 .unwrap();
10594 drop(tx);
10595 let mut receipt = receipt;
10596 blocking_response_with_receipt(
10597 rx,
10598 "m".into(),
10599 true,
10600 Vec::new(),
10601 None,
10602 Envelope::new(true),
10603 &mut receipt,
10604 None,
10605 )
10606 .await
10607 };
10608
10609 let plain = MockMetering::admit_all();
10611 let receipt = plain.open(
10612 &metering::RequestMeta {
10613 request_id: "cap-unmarked",
10614 tenant: "unmarked",
10615 principal: None,
10616 model: "m",
10617 route: "/v1/chat/completions",
10618 lane: "interactive",
10619 stream: false,
10620 max_tokens: None,
10621 reserved_ctx: None,
10622 },
10623 None,
10624 );
10625 let response = drive(Some(receipt)).await;
10626 assert_eq!(response.status(), StatusCode::OK);
10627 assert!(
10628 !plain.events().iter().any(|e| matches!(
10629 e,
10630 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
10631 )),
10632 "an unarmed receipt must see no capture traffic: {:?}",
10633 plain.events()
10634 );
10635
10636 let capturing = MockMetering::capturing();
10639 let mut receipt = capturing.open(
10640 &metering::RequestMeta {
10641 request_id: "cap-marked",
10642 tenant: "marked",
10643 principal: None,
10644 model: "m",
10645 route: "/v1/chat/completions",
10646 lane: "interactive",
10647 stream: false,
10648 max_tokens: None,
10649 reserved_ctx: None,
10650 },
10651 None,
10652 );
10653 assert!(receipt.wants_capture());
10654 receipt.arm_capture(prompt.clone());
10655 let response = drive(Some(receipt)).await;
10656 assert_eq!(response.status(), StatusCode::OK);
10657 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
10658 .await
10659 .unwrap();
10660 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
10661 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
10662
10663 let events = capturing.events();
10664 assert!(
10665 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
10666 "prompt must arm byte-exact: {events:?}"
10667 );
10668 let completion: String = events
10669 .iter()
10670 .filter_map(|e| match e {
10671 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
10672 _ => None,
10673 })
10674 .collect();
10675 assert_eq!(
10676 completion, "Hello",
10677 "the deltas must reassemble the served completion byte-exact: {events:?}"
10678 );
10679 assert!(
10680 events.contains(&MeterEvent::Complete {
10681 prompt: 7,
10682 cached: 0,
10683 completion: 2,
10684 }),
10685 "worker-truth usage settles alongside the capture: {events:?}"
10686 );
10687 }
10688
10689 fn tool_caps() -> ModelCaps {
10690 ModelCaps {
10691 tools_branch: true,
10692 qwen_think: true,
10693 think_switch: true,
10694 chat_ok: true,
10695 ..Default::default()
10696 }
10697 }
10698
10699 fn ladder_caps() -> ModelCaps {
10704 ModelCaps {
10705 qwen_effort: true,
10706 ..tool_caps()
10707 }
10708 }
10709
10710 fn gemma_tool_caps() -> ModelCaps {
10711 ModelCaps {
10712 tools_branch: true,
10713 gemma_think: true,
10714 chat_ok: true,
10715 instruct_type: Some("gemma".into()),
10716 ..Default::default()
10717 }
10718 }
10719
10720 fn hy3_tool_caps() -> ModelCaps {
10721 ModelCaps {
10722 tools_branch: true,
10723 hy3: true,
10724 chat_ok: true,
10725 effort_levels: true,
10726 instruct_type: Some("hy3".into()),
10727 ..Default::default()
10728 }
10729 }
10730
10731 fn gemma_template(kind: &str) -> String {
10732 let file = match kind {
10733 "qat" => "qat-trunk-template.jinja",
10734 _ => "official-tooluse-template.jinja",
10735 };
10736 let path = format!(
10737 "{}/../../research/gemma4-tools-20260817/{file}",
10738 env!("CARGO_MANIFEST_DIR")
10739 );
10740 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
10741 }
10742
10743 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
10748 let tools_arr = request
10749 .get("tools")
10750 .and_then(|t| t.as_array())
10751 .cloned()
10752 .unwrap_or_default();
10753 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
10754 (Vec::new(), Vec::new(), HashMap::new())
10755 } else {
10756 prepare_tools(&tools_arr).unwrap()
10757 };
10758 let effort = request
10759 .get("reasoning_effort")
10760 .and_then(|v| v.as_str())
10761 .map(String::from);
10762 let (think, _lvl, _explicit) =
10763 parse_think(&effort, &None, None, None, None, false).unwrap();
10764
10765 let mut turns: Vec<TmplTurn> = Vec::new();
10766 for msg in request["messages"].as_array().unwrap() {
10767 let role = msg["role"].as_str().unwrap();
10768 let role = if role == "developer" { "system" } else { role };
10769 let content =
10770 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
10771 let tool_calls = msg
10772 .get("tool_calls")
10773 .and_then(|a| a.as_array())
10774 .map(|a| {
10775 a.iter()
10776 .map(|tc| {
10777 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
10778 render_req_tool_call(&rtc).unwrap()
10779 })
10780 .collect()
10781 })
10782 .unwrap_or_default();
10783 let tool_responses = msg
10784 .get("tool_responses")
10785 .and_then(|a| a.as_array())
10786 .map(|a| {
10787 a.iter()
10788 .map(|tr| {
10789 (
10790 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
10791 json_to_val(&tr["response"]),
10792 )
10793 })
10794 .collect()
10795 })
10796 .unwrap_or_default();
10797 turns.push(TmplTurn {
10798 role: role.to_string(),
10799 content,
10800 tool_calls,
10801 reasoning: msg
10802 .get("reasoning")
10803 .and_then(|r| r.as_str())
10804 .map(String::from)
10805 .filter(|s| !s.is_empty()),
10806 tool_call_id: msg
10807 .get("tool_call_id")
10808 .and_then(|s| s.as_str())
10809 .map(String::from),
10810 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
10811 tool_responses,
10812 task: None,
10813 tools: Vec::new(),
10814 });
10815 }
10816 chat::apply_chat_template_tools_ex(
10817 Some(template),
10818 &turns,
10819 true,
10820 &tools_json,
10821 &tools_struct,
10822 think,
10823 None,
10824 None,
10825 )
10826 .unwrap()
10827 }
10828
10829 #[test]
10833 fn gemma4_tools_fixtures_match_the_official_jinja() {
10834 let dir = format!(
10835 "{}/../../research/gemma4-tools-20260817/fixtures",
10836 env!("CARGO_MANIFEST_DIR")
10837 );
10838 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10839 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10840 .map(|e| e.unwrap().path())
10841 .filter(|p| p.is_dir())
10842 .collect();
10843 entries.sort();
10844 assert!(
10845 entries.len() >= 14,
10846 "expected >=14 fixtures, found {}",
10847 entries.len()
10848 );
10849 let (mut official, mut qat) = (0u32, 0u32);
10850 for d in entries {
10851 let input: serde_json::Value =
10852 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10853 .unwrap();
10854 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10855 let kind = input
10856 .get("template")
10857 .and_then(|t| t.as_str())
10858 .unwrap_or("official");
10859 match kind {
10860 "qat" => qat += 1,
10861 _ => official += 1,
10862 }
10863 let tmpl = gemma_template(kind);
10864 let got = render_fixture(&input["request"], &tmpl);
10865 assert_eq!(
10866 got, expected,
10867 "fixture {:?} diverged from the jinja oracle",
10868 d
10869 );
10870 }
10871 assert!(
10872 official >= 12 && qat >= 2,
10873 "coverage: {official} official, {qat} qat"
10874 );
10875 }
10876
10877 #[test]
10883 fn gemma4_tools_flow_through_build_chat_request() {
10884 let tmpl = gemma_template("official");
10885 for name in [
10886 "01-system-tools-basic",
10887 "04-single-call-cycle",
10888 "07-multi-cycle-agentic",
10889 ] {
10890 let path = format!(
10891 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
10892 env!("CARGO_MANIFEST_DIR")
10893 );
10894 let input: serde_json::Value =
10895 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
10896 let expected_path = format!(
10897 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
10898 env!("CARGO_MANIFEST_DIR")
10899 );
10900 let expected = std::fs::read_to_string(&expected_path).unwrap();
10901 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
10902 let (tx, _rx) = worker::event_channel();
10903 let plan = build_chat_request(
10904 req,
10905 Some(&gemma_tool_caps()),
10906 tx,
10907 lanes::Lane::Interactive,
10908 None,
10909 )
10910 .unwrap();
10911 let got = chat::apply_chat_template_tools_ex(
10912 Some(&tmpl),
10913 &plan.request.chat_turns,
10914 true,
10915 &plan.request.tools_json,
10916 &plan.request.tools_struct,
10917 plan.request.think,
10918 plan.request.reasoning_effort.as_deref(),
10919 None,
10920 )
10921 .unwrap();
10922 assert_eq!(got, expected, "pipeline render diverged for {name}");
10923 }
10924 }
10925
10926 fn glm5_template() -> String {
10935 let path = format!(
10936 "{}/../../research/glm53-flash-bringup-20260827/chat_template.jinja",
10937 env!("CARGO_MANIFEST_DIR")
10938 );
10939 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
10940 }
10941
10942 fn glm5_caps() -> ModelCaps {
10947 ModelCaps {
10948 tools_branch: true,
10949 qwen_think: true,
10950 think_switch: false,
10951 chat_ok: true,
10952 context_length: 1_048_576,
10953 tokenizer: "glm4".into(),
10954 instruct_type: Some("glm".into()),
10955 effort_levels: true,
10956 glm5: true,
10957 ..Default::default()
10958 }
10959 }
10960
10961 fn glm5_render(body: serde_json::Value) -> Result<String, String> {
10963 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
10964 let (tx, _rx) = worker::event_channel();
10965 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)?;
10966 chat::apply_chat_template_tools_ex(
10967 Some(&glm5_template()),
10968 &plan.request.chat_turns,
10969 true,
10970 &plan.request.tools_json,
10971 &plan.request.tools_struct,
10972 plan.request.think,
10973 plan.request.reasoning_effort.as_deref(),
10974 None,
10975 )
10976 }
10977
10978 #[test]
10983 fn glm5_fixtures_match_the_vendor_jinja() {
10984 let dir = format!(
10985 "{}/../../research/glm53-flash-bringup-20260827/surface-fixtures",
10986 env!("CARGO_MANIFEST_DIR")
10987 );
10988 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10989 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10990 .map(|e| e.unwrap().path())
10991 .filter(|p| p.is_dir())
10992 .collect();
10993 entries.sort();
10994 assert!(
10995 entries.len() >= 22,
10996 "expected >=22 fixtures, found {}",
10997 entries.len()
10998 );
10999 for d in entries {
11000 let input: serde_json::Value =
11001 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11002 .unwrap();
11003 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11004 let got = glm5_render(input["request"].clone())
11005 .unwrap_or_else(|e| panic!("fixture {d:?} refused: {e}"));
11006 assert_eq!(
11007 got, expected,
11008 "fixture {d:?} diverged from the jinja oracle"
11009 );
11010 }
11011 }
11012
11013 #[test]
11019 fn glm5_never_renders_chatml() {
11020 let tmpl = glm5_template();
11021 assert!(tmpl.contains("<think>") && tmpl.contains("add_generation_prompt"));
11023 assert!(tmpl.contains("<tools>"));
11024 assert!(chat::template_is_glm5(&tmpl));
11025 for body in [
11026 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11027 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11028 "tools": [{"type": "function", "function": {"name": "f",
11029 "parameters": {"type": "object", "properties": {}}}}]}),
11030 ] {
11031 let got = glm5_render(body).unwrap();
11032 assert!(
11033 !got.contains("<|im_start|>") && !got.contains("<|im_end|>"),
11034 "glm5 rendered ChatML frames: {got:?}"
11035 );
11036 assert!(
11037 got.starts_with("[gMASK]<sop><|system|>Reasoning Effort: "),
11038 "{got:?}"
11039 );
11040 assert!(got.ends_with("<|assistant|><think>"), "{got:?}");
11041 }
11042 }
11043
11044 #[test]
11048 fn glm5_reasoning_effort_renders_and_keeps_its_max_tier() {
11049 for (sent, line) in [
11050 (None, "Max"),
11051 (Some("low"), "Low"),
11052 (Some("medium"), "Low"),
11055 (Some("high"), "High"),
11056 (Some("xhigh"), "Max"),
11057 (Some("max"), "Max"),
11058 (Some("ultra"), "Max"),
11059 ] {
11060 let mut body = json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]});
11061 if let Some(v) = sent {
11062 body["reasoning_effort"] = json!(v);
11063 }
11064 let got = glm5_render(body).unwrap();
11065 assert!(
11066 got.starts_with(&format!("[gMASK]<sop><|system|>Reasoning Effort: {line}<|")),
11067 "reasoning_effort {sent:?} should render {line:?}: {got:?}"
11068 );
11069 }
11070 let sampler_of = |v: &str| {
11073 let req: ChatCompletionReq = serde_json::from_value(
11074 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11077 "reasoning_effort": v, "seed": 7}),
11078 )
11079 .unwrap();
11080 let (tx, _rx) = worker::event_channel();
11081 let plan =
11082 build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11083 .unwrap();
11084 format!("{:?}", plan.request.sampler_cfg)
11085 };
11086 assert_eq!(sampler_of("low"), sampler_of("max"));
11087 assert_eq!(canonical_effort_for("max", true), Some("max"));
11089 assert_eq!(canonical_effort_for("xhigh", true), Some("max"));
11090 assert_eq!(canonical_effort_for("max", false), Some("high"));
11091 }
11092
11093 #[test]
11097 fn glm5_refuses_what_its_template_cannot_honour() {
11098 for (value, needle) in [
11099 ("none", "cannot disable reasoning"),
11100 ("minimal", "cannot disable reasoning"),
11101 ("bogus", "bad reasoning_effort"),
11102 ] {
11103 let err = glm5_render(json!({"model": "m",
11104 "messages": [{"role": "user", "content": "hi"}],
11105 "reasoning_effort": value}))
11106 .err()
11107 .unwrap_or_else(|| panic!("reasoning_effort {value:?} must be refused"));
11108 assert!(err.contains(needle), "{value}: {err}");
11109 }
11110 }
11111
11112 #[test]
11116 fn one_glm5_request_renders_identical_bytes_on_all_three_surfaces() {
11117 let chat = json!({
11124 "model": "m",
11125 "reasoning_effort": "high",
11126 "messages": [
11127 {"role": "user", "content": "Weather in Paris and Rome?"},
11128 {"role": "assistant", "content": null,
11129 "tool_calls": [
11130 {"id": "c1", "type": "function",
11131 "function": {"name": "get_weather",
11132 "arguments": "{\"city\": \"Paris\"}"}},
11133 {"id": "c2", "type": "function",
11134 "function": {"name": "get_weather",
11135 "arguments": "{\"city\": \"Rome\"}"}}]},
11136 {"role": "tool", "tool_call_id": "c2", "content": "rome:27"},
11137 {"role": "tool", "tool_call_id": "c1", "content": "paris:21"}
11138 ],
11139 "tools": [{"type": "function", "function": {
11140 "name": "get_weather", "description": "Get the current weather for a city",
11141 "parameters": {"type": "object",
11142 "properties": {"city": {"type": "string"}},
11143 "required": ["city"]}}}]
11144 });
11145 let responses = responses_api::translate(&json!({
11146 "model": "m",
11147 "reasoning": {"effort": "high"},
11148 "input": [
11149 {"type": "message", "role": "user",
11150 "content": [{"type": "input_text", "text": "Weather in Paris and Rome?"}]},
11151 {"type": "function_call", "call_id": "c1", "name": "get_weather",
11152 "arguments": "{\"city\": \"Paris\"}"},
11153 {"type": "function_call", "call_id": "c2", "name": "get_weather",
11154 "arguments": "{\"city\": \"Rome\"}"},
11155 {"type": "function_call_output", "call_id": "c2", "output": "rome:27"},
11156 {"type": "function_call_output", "call_id": "c1", "output": "paris:21"}
11157 ],
11158 "tools": [{"type": "function", "name": "get_weather",
11159 "description": "Get the current weather for a city",
11160 "parameters": {"type": "object",
11161 "properties": {"city": {"type": "string"}},
11162 "required": ["city"]}}]
11163 }))
11164 .expect("/v1/responses translate");
11165 let messages = anthropic::translate(&json!({
11166 "model": "m",
11167 "max_tokens": 256,
11168 "output_config": {"effort": "high"},
11169 "messages": [
11170 {"role": "user", "content": "Weather in Paris and Rome?"},
11171 {"role": "assistant", "content": [
11172 {"type": "tool_use", "id": "c1", "name": "get_weather",
11173 "input": {"city": "Paris"}},
11174 {"type": "tool_use", "id": "c2", "name": "get_weather",
11175 "input": {"city": "Rome"}}]},
11176 {"role": "user", "content": [
11177 {"type": "tool_result", "tool_use_id": "c2", "content": "rome:27"},
11178 {"type": "tool_result", "tool_use_id": "c1", "content": "paris:21"}]}
11179 ],
11180 "tools": [{"name": "get_weather",
11181 "description": "Get the current weather for a city",
11182 "input_schema": {"type": "object",
11183 "properties": {"city": {"type": "string"}},
11184 "required": ["city"]}}]
11185 }))
11186 .expect("/v1/messages translate");
11187 let want = glm5_render(chat).expect("chat");
11188 assert!(
11190 want.contains(
11191 "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value>\
11192 </tool_call><tool_call>get_weather<arg_key>city</arg_key>\
11193 <arg_value>Rome</arg_value></tool_call>"
11194 ),
11195 "{want:?}"
11196 );
11197 assert!(
11201 want.contains(
11202 "<|observation|><tool_response>paris:21</tool_response>\
11203 <tool_response>rome:27</tool_response>"
11204 ),
11205 "{want:?}"
11206 );
11207 assert!(
11208 want.contains("<|system|>Reasoning Effort: High"),
11209 "{want:?}"
11210 );
11211 for (surface, body) in [
11212 ("/v1/responses", responses),
11213 ("/v1/messages", messages.clone()),
11214 ] {
11215 let got = glm5_render(body).unwrap_or_else(|e| panic!("{surface}: {e}"));
11216 assert_eq!(
11217 got, want,
11218 "{surface} rendered DIFFERENT glm5 prompt bytes than /v1/chat/completions"
11219 );
11220 }
11221 let mut idless = messages;
11227 for m in idless["messages"].as_array_mut().unwrap() {
11228 if m["role"] == "tool" {
11229 m.as_object_mut().unwrap().remove("tool_call_id");
11230 }
11231 }
11232 let got = glm5_render(idless).expect("id-less render");
11233 assert_ne!(
11234 got, want,
11235 "dropping tool_call_id must change the rendered order — this test cannot detect \
11236 a surface that loses ids otherwise"
11237 );
11238 assert!(
11239 got.contains(
11240 "<|observation|><tool_response>rome:27</tool_response>\
11241 <tool_response>paris:21</tool_response>"
11242 ),
11243 "{got:?}"
11244 );
11245 }
11246
11247 #[test]
11250 fn glm5_chat_arms_the_native_tool_parser() {
11251 let req: ChatCompletionReq = serde_json::from_value(json!({
11252 "model": "m", "messages": [{"role": "user", "content": "weather?"}],
11253 "tools": [{"type": "function", "function": {"name": "get_weather",
11254 "parameters": {"type": "object",
11255 "properties": {"city": {"type": "string"}}}}}]}))
11256 .unwrap();
11257 let (tx, _rx) = worker::event_channel();
11258 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11259 .unwrap();
11260 let mut parser = plan.parser.expect("glm5 tools request must carry a parser");
11261 let pieces = parser.push(
11262 "reasoning here</think><tool_call>get_weather<arg_key>city</arg_key>\
11263 <arg_value>Paris</arg_value></tool_call>",
11264 );
11265 let calls: Vec<_> = pieces
11266 .iter()
11267 .filter_map(|p| match p {
11268 toolcall::Piece::Call(c) => Some((c.name.as_str(), c.arguments.as_str())),
11269 _ => None,
11270 })
11271 .collect();
11272 assert_eq!(
11273 calls,
11274 vec![("get_weather", r#"{"city":"Paris"}"#)],
11275 "{pieces:?}"
11276 );
11277 assert!(
11278 pieces
11279 .iter()
11280 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "reasoning here")),
11281 "{pieces:?}"
11282 );
11283 assert!(
11285 !pieces
11286 .iter()
11287 .any(|p| matches!(p, toolcall::Piece::Content(_))),
11288 "{pieces:?}"
11289 );
11290 let req: ChatCompletionReq = serde_json::from_value(
11294 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11295 )
11296 .unwrap();
11297 let (tx, _rx) = worker::event_channel();
11298 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
11299 .unwrap();
11300 let mut parser = plan
11301 .parser
11302 .expect("glm5 non-tools request must still split reasoning");
11303 let pieces = parser.push("weighing it</think>The answer.");
11304 assert!(
11305 pieces
11306 .iter()
11307 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "weighing it")),
11308 "{pieces:?}"
11309 );
11310 assert!(
11311 pieces
11312 .iter()
11313 .any(|p| matches!(p, toolcall::Piece::Content(c) if c == "The answer.")),
11314 "{pieces:?}"
11315 );
11316 }
11317
11318 #[test]
11325 fn glm5_plain_fast_path_never_drops_replayed_reasoning() {
11326 let with_reasoning = vec![
11327 chat::Turn {
11328 role: "user".into(),
11329 content: "a".into(),
11330 ..Default::default()
11331 },
11332 chat::Turn {
11333 role: "assistant".into(),
11334 content: "A".into(),
11335 reasoning: Some("I considered a.".into()),
11336 ..Default::default()
11337 },
11338 chat::Turn {
11339 role: "user".into(),
11340 content: "b".into(),
11341 ..Default::default()
11342 },
11343 ];
11344 assert!(!worker::plain_chat_render_path(
11346 &[],
11347 &chat::ThinkMode::Default,
11348 None,
11349 &with_reasoning,
11350 false,
11351 ));
11352 let plain_turns: Vec<chat::Turn> = with_reasoning
11355 .iter()
11356 .cloned()
11357 .map(|mut t| {
11358 t.reasoning = None;
11359 t
11360 })
11361 .collect();
11362 assert!(worker::plain_chat_render_path(
11363 &[],
11364 &chat::ThinkMode::Default,
11365 None,
11366 &plain_turns,
11367 false,
11368 ));
11369 let tmpl = glm5_template();
11372 let via_tools = chat::apply_chat_template_tools_ex(
11373 Some(&tmpl),
11374 &with_reasoning,
11375 true,
11376 &[],
11377 &[],
11378 chat::ThinkMode::Default,
11379 None,
11380 None,
11381 )
11382 .unwrap();
11383 let msgs: Vec<(&str, &str)> = with_reasoning
11384 .iter()
11385 .map(|t| (t.role.as_str(), t.content.as_str()))
11386 .collect();
11387 let via_plain = chat::apply_chat_template_str(Some(&tmpl), &msgs, true);
11388 assert!(
11389 via_tools.contains("<think>I considered a.</think>"),
11390 "{via_tools:?}"
11391 );
11392 assert_ne!(via_tools, via_plain);
11393 let plain_msgs: Vec<(&str, &str)> = plain_turns
11396 .iter()
11397 .map(|t| (t.role.as_str(), t.content.as_str()))
11398 .collect();
11399 assert_eq!(
11400 chat::apply_chat_template_tools_ex(
11401 Some(&tmpl),
11402 &plain_turns,
11403 true,
11404 &[],
11405 &[],
11406 chat::ThinkMode::Default,
11407 None,
11408 None,
11409 )
11410 .unwrap(),
11411 chat::apply_chat_template_str(Some(&tmpl), &plain_msgs, true)
11412 );
11413 }
11414
11415 #[test]
11419 fn glm5_model_row_does_not_claim_structured_output() {
11420 let caps = glm5_caps();
11421 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), None);
11422 assert_eq!(row["capabilities"]["structured_output"], json!(false));
11423 assert_eq!(row["capabilities"]["tools"], json!(true));
11424 assert_eq!(row["capabilities"]["reasoning"], json!(true));
11425 let err = glm5_render(json!({"model": "m",
11427 "messages": [{"role": "user", "content": "hi"}],
11428 "response_format": {"type": "json_object"}}))
11429 .expect_err("response_format must be refused on a switchless think template");
11430 assert!(
11434 err.contains("neither an enable_thinking switch nor a recognizable"),
11435 "{err}"
11436 );
11437 let switchable = model_entry_v1("q", Some(&tool_caps()), None);
11439 assert_eq!(switchable["capabilities"]["structured_output"], json!(true));
11440 let glm_params = openrouter_supported_parameters(Some(&caps), None, true);
11443 assert!(
11444 glm_params.get("structured_outputs").is_none(),
11445 "{glm_params}"
11446 );
11447 let step_like = ModelCaps {
11453 chat_ok: true,
11454 qwen_think: true,
11455 think_switch: false,
11456 think_close: vec![128799],
11457 ..caps.clone()
11458 };
11459 let step_row = model_entry_v1("stepfun/step-3.7-flash", Some(&step_like), None);
11460 assert_eq!(step_row["capabilities"]["structured_output"], json!(true));
11461 let step_params = openrouter_supported_parameters(Some(&step_like), None, true);
11462 assert!(
11463 step_params.get("structured_outputs").is_some(),
11464 "{step_params}"
11465 );
11466 assert!(glm_params.get("json_mode").is_none(), "{glm_params}");
11467 assert!(glm_params.get("tools").is_some(), "{glm_params}");
11468 let qwen_params = openrouter_supported_parameters(Some(&tool_caps()), None, true);
11469 assert!(
11470 qwen_params.get("structured_outputs").is_some(),
11471 "{qwen_params}"
11472 );
11473 assert!(qwen_params.get("json_mode").is_some(), "{qwen_params}");
11474 }
11475
11476 #[test]
11484 fn catalog_context_claim_is_capped_by_the_deployment_envelope() {
11485 let caps = glm5_caps();
11486 assert_eq!(caps.context_length, 1_048_576);
11487 let metadata = OpenRouterModelMetadata {
11488 max_prompt_length: Some(126_976),
11489 max_output_length: Some(4_096),
11490 ..Default::default()
11491 };
11492 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
11494 assert_eq!(row["context_length"], json!(131_072));
11495 let or_row = model_entry_openrouter("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
11496 assert_eq!(
11497 or_row["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
11498 json!(131_072)
11499 );
11500 assert_eq!(
11501 published_context_length(Some(&caps), Some(&metadata)),
11502 Some(131_072)
11503 );
11504 assert_eq!(published_context_length(Some(&caps), None), Some(1_048_576));
11506 let half = OpenRouterModelMetadata {
11507 max_output_length: Some(4_096),
11508 ..Default::default()
11509 };
11510 assert_eq!(
11511 published_context_length(Some(&caps), Some(&half)),
11512 Some(1_048_576)
11513 );
11514 let wide = OpenRouterModelMetadata {
11516 max_prompt_length: Some(2_000_000),
11517 max_output_length: Some(2_000_000),
11518 ..Default::default()
11519 };
11520 assert_eq!(
11521 published_context_length(Some(&caps), Some(&wide)),
11522 Some(1_048_576)
11523 );
11524 }
11525
11526 fn dsv4_sentinel() -> String {
11534 let path = format!(
11535 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
11536 env!("CARGO_MANIFEST_DIR")
11537 );
11538 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
11539 }
11540
11541 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
11546 let role = msg["role"].as_str().unwrap().to_string();
11547 let content =
11548 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
11549 let reasoning = msg
11550 .get("reasoning")
11551 .or_else(|| msg.get("reasoning_content"))
11552 .and_then(|r| r.as_str())
11553 .map(String::from)
11554 .filter(|s| !s.is_empty());
11555 let tool_calls = msg
11556 .get("tool_calls")
11557 .and_then(|a| a.as_array())
11558 .map(|a| {
11559 a.iter()
11560 .map(|tc| {
11561 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
11562 render_req_tool_call(&rtc).unwrap()
11563 })
11564 .collect()
11565 })
11566 .unwrap_or_default();
11567 let tools = msg
11568 .get("tools")
11569 .and_then(|a| a.as_array())
11570 .map(|a| {
11571 a.iter()
11572 .filter_map(|t| t.get("function").map(json_to_val))
11573 .collect()
11574 })
11575 .unwrap_or_default();
11576 TmplTurn {
11577 role,
11578 content,
11579 tool_calls,
11580 reasoning,
11581 tool_call_id: msg
11582 .get("tool_call_id")
11583 .and_then(|s| s.as_str())
11584 .map(String::from),
11585 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
11586 tool_responses: Vec::new(),
11587 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
11588 tools,
11589 }
11590 }
11591
11592 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
11593 v.and_then(|t| t.as_array())
11594 .map(|a| {
11595 a.iter()
11596 .filter_map(|t| t.get("function").map(json_to_val))
11597 .collect()
11598 })
11599 .unwrap_or_default()
11600 }
11601
11602 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
11606 let dir = format!(
11607 "{}/../../research/dsv4-template-20260818/{subdir}",
11608 env!("CARGO_MANIFEST_DIR")
11609 );
11610 let tmpl = dsv4_sentinel();
11611 let mut entries: Vec<_> = std::fs::read_dir(&dir)
11612 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
11613 .map(|e| e.unwrap().path())
11614 .filter(|p| p.is_dir())
11615 .collect();
11616 entries.sort();
11617 assert!(
11618 entries.len() >= min_fixtures,
11619 "expected >={min_fixtures} fixtures, found {}",
11620 entries.len()
11621 );
11622 for d in &entries {
11623 let input: serde_json::Value =
11624 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
11625 .unwrap();
11626 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
11627 let turns: Vec<TmplTurn> = input["turns"]
11628 .as_array()
11629 .unwrap()
11630 .iter()
11631 .map(dsv4_turn)
11632 .collect();
11633 let think = match input["think"].as_str().unwrap() {
11634 "chat" => ThinkMode::NoThink,
11635 _ => ThinkMode::Think,
11636 };
11637 let effort = input
11638 .get("reasoning_effort")
11639 .and_then(|v| v.as_str())
11640 .map(String::from);
11641 let req_tools = dsv4_req_tools(input.get("req_tools"));
11642 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
11643 let got = chat::apply_chat_template_tools_ex(
11644 Some(&tmpl),
11645 &turns,
11646 agp,
11647 &[],
11648 &req_tools,
11649 think,
11650 effort.as_deref(),
11651 Some(encoding),
11652 )
11653 .unwrap();
11654 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
11655 }
11656 }
11657
11658 #[test]
11659 fn dsv4_template_fixtures_match_the_oracle() {
11660 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
11661 }
11662
11663 #[test]
11669 fn dsv4_0731_fixtures_match_the_oracle() {
11670 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
11671 }
11672
11673 #[test]
11674 fn dsv4_artifact_fixtures_are_byte_identical() {
11675 let base = format!(
11679 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
11680 env!("CARGO_MANIFEST_DIR")
11681 );
11682 let tmpl = dsv4_sentinel();
11683 for (n, think) in [
11684 (1u32, ThinkMode::Think),
11685 (2, ThinkMode::Think),
11686 (3, ThinkMode::Think),
11687 (4, ThinkMode::NoThink),
11688 ] {
11689 let td: serde_json::Value = serde_json::from_str(
11690 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
11691 )
11692 .unwrap();
11693 let (messages, tools) = if td.is_object() {
11694 (td["messages"].clone(), td.get("tools").cloned())
11695 } else {
11696 (td.clone(), None)
11697 };
11698 let mut turns: Vec<TmplTurn> = Vec::new();
11699 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
11700 let mut t = dsv4_turn(msg);
11701 if i == 0
11702 && let Some(tl) = &tools
11703 {
11704 t.tools = tl
11705 .as_array()
11706 .unwrap()
11707 .iter()
11708 .filter_map(|x| x.get("function").map(json_to_val))
11709 .collect();
11710 }
11711 turns.push(t);
11712 }
11713 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
11714 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
11718 let got = chat::apply_chat_template_tools_ex(
11719 Some(&tmpl),
11720 &turns,
11721 true,
11722 &[],
11723 &[],
11724 think,
11725 None,
11726 Some(encoding),
11727 )
11728 .unwrap();
11729 assert_eq!(
11730 got, expected,
11731 "artifact fixture {n} diverged from the oracle under {encoding:?}"
11732 );
11733 }
11734 }
11735 }
11736
11737 #[test]
11738 fn dsv4_default_thinkmode_renders_thinking() {
11739 let tmpl = dsv4_sentinel();
11742 let turns = vec![TmplTurn {
11743 role: "user".into(),
11744 content: "Hi".into(),
11745 ..Default::default()
11746 }];
11747 let dflt = chat::apply_chat_template_tools_ex(
11748 Some(&tmpl),
11749 &turns,
11750 true,
11751 &[],
11752 &[],
11753 ThinkMode::Default,
11754 None,
11755 None,
11756 )
11757 .unwrap();
11758 let think = chat::apply_chat_template_tools_ex(
11759 Some(&tmpl),
11760 &turns,
11761 true,
11762 &[],
11763 &[],
11764 ThinkMode::Think,
11765 None,
11766 None,
11767 )
11768 .unwrap();
11769 assert_eq!(dflt, think);
11770 assert!(
11771 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
11772 "{dflt:?}"
11773 );
11774 let chat_mode = chat::apply_chat_template_tools_ex(
11775 Some(&tmpl),
11776 &turns,
11777 true,
11778 &[],
11779 &[],
11780 ThinkMode::NoThink,
11781 None,
11782 None,
11783 )
11784 .unwrap();
11785 assert!(
11786 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
11787 "{chat_mode:?}"
11788 );
11789 }
11790
11791 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
11796 let base = format!(
11797 "{}/../../research/dsv4-template-20260818",
11798 env!("CARGO_MANIFEST_DIR")
11799 );
11800 let refdir = std::path::Path::new(&base).join("ref");
11801 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
11802 .expect("load dsv4 tokenizer from ref dir");
11803 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
11804 let banked: serde_json::Value = serde_json::from_str(
11805 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
11806 .unwrap(),
11807 )
11808 .unwrap();
11809 let obj = banked.as_object().unwrap();
11810 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
11811 for (name, ids_v) in obj {
11812 let rendered =
11813 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
11814 let want: Vec<u32> = ids_v
11815 .as_array()
11816 .unwrap()
11817 .iter()
11818 .map(|v| v.as_u64().unwrap() as u32)
11819 .collect();
11820 let got = tok.encode(&rendered, true);
11821 assert_eq!(got, want, "tokenization diverged for {name}");
11822 }
11823 }
11824
11825 #[test]
11826 fn dsv4_tokenization_crosscheck_matches_official_ids() {
11827 dsv4_run_tokenization_crosscheck("fixtures");
11828 }
11829
11830 #[test]
11834 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
11835 dsv4_run_tokenization_crosscheck("fixtures-0731");
11836 }
11837
11838 #[test]
11839 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
11840 let base = format!(
11850 "{}/../../research/dsv4-template-20260818",
11851 env!("CARGO_MANIFEST_DIR")
11852 );
11853 let refdir = std::path::Path::new(&base).join("ref");
11854 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
11855 .expect("load dsv4 tokenizer from ref dir");
11856 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
11857 let tmpl = dsv4_sentinel();
11858 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
11859 {"type": "function", "function": {
11860 "name": "get_data",
11861 "description": "Fetch a blob",
11862 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
11863 "required": ["key"]}
11864 }}
11865 ])));
11866
11867 let cases: Vec<(&str, String)> = vec![
11868 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
11870 ("space-131k", " ".repeat(131_072)),
11871 ("digit-131k", "7".repeat(131_072)),
11872 (
11873 "mixed-runs",
11874 format!(
11875 "{}{}{}{}",
11876 "Z".repeat(65_536),
11877 " ".repeat(65_536),
11878 "7".repeat(65_536),
11879 "\n".repeat(65_536)
11880 ),
11881 ),
11882 ("cjk-64k", "中".repeat(65_536)),
11883 ("accented-letter-64k", "é".repeat(65_536)),
11884 ];
11885 for (name, blob) in &cases {
11886 let msgs = serde_json::json!([
11887 {"role": "system", "content": "You are a tool-using assistant."},
11888 {"role": "user", "content": "Fetch the blob."},
11889 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
11890 "tool_calls": [{"id": "call_001", "type": "function",
11891 "function": {"name": "get_data",
11892 "arguments": "{\"key\": \"blob\"}"}}]},
11893 {"role": "tool", "tool_call_id": "call_001", "content": blob}
11894 ]);
11895 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
11896 let rendered = chat::apply_chat_template_tools_ex(
11897 Some(&tmpl),
11898 &turns,
11899 true,
11900 &[],
11901 &req_tools,
11902 ThinkMode::Think,
11903 None,
11904 None,
11905 )
11906 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
11907 assert!(
11908 rendered.contains(blob.as_str()),
11909 "{name}: tool result missing from render"
11910 );
11911 let t0 = std::time::Instant::now();
11912 let ids = tok.encode(&rendered, true);
11913 let encode_dt = t0.elapsed();
11914 assert!(!ids.is_empty(), "{name}: empty encode");
11915 let back = tok.decode(&ids);
11916 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
11917 assert!(
11921 encode_dt < std::time::Duration::from_secs(60),
11922 "{name}: encode took {encode_dt:?}"
11923 );
11924 if *name == "ascii-letter-131k"
11928 && let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR")
11929 {
11930 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
11931 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
11932 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
11933 }
11934 }
11935 }
11936
11937 #[test]
11938 fn models_v1_entry_advertises_thinking_support() {
11939 let step_caps = ModelCaps {
11942 effort_levels: true,
11943 ..tool_caps()
11944 };
11945 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
11946 assert_eq!(entry["capabilities"]["reasoning"], true);
11947 assert_eq!(entry["capabilities"]["tools"], true);
11948
11949 let plain = ModelCaps {
11951 chat_ok: true,
11952 ..Default::default()
11953 };
11954 let entry = model_entry_v1("plain", Some(&plain), None);
11955 assert_eq!(entry["capabilities"]["reasoning"], false);
11956 assert_eq!(entry["capabilities"]["tools"], false);
11957 let entry = model_entry_v1("unknown", None, None);
11959 assert_eq!(entry["capabilities"]["reasoning"], false);
11960 assert_eq!(entry["capabilities"]["streaming"], true);
11961 }
11962
11963 #[test]
11964 fn chat_request_preserves_turns_and_openai_stop_forms() {
11965 let payload = serde_json::json!({
11966 "model": "plain_quant",
11967 "messages": [
11968 {"role": "system", "content": "rules"},
11969 {"role": "developer", "content": "dev rules"},
11970 {"role": "user", "content": "task"},
11971 {"role": "assistant", "content": "work"}
11972 ],
11973 "max_tokens": 64,
11974 "temperature": 0.0,
11975 "stop": "<stop>"
11976 });
11977 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11978 let (tx, _rx) = worker::event_channel();
11979 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
11980 let request = plan.request;
11981 assert!(
11982 plan.parser.is_none(),
11983 "no tools -> no parser (isolation contract)"
11984 );
11985 assert!(request.tools_json.is_empty());
11986 assert_eq!(request.think, ThinkMode::Default);
11987 assert_eq!(request.model, "plain_quant");
11988 assert_eq!(request.params.max_new, 64);
11989 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11991 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
11992 }))
11993 .unwrap();
11994 let (tx, _rx) = worker::event_channel();
11995 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
11996 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
11997 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11999 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12000 "max_completion_tokens": 7
12001 }))
12002 .unwrap();
12003 let (tx, _rx) = worker::event_channel();
12004 assert_eq!(
12005 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12006 .unwrap()
12007 .request
12008 .params
12009 .max_new,
12010 7
12011 );
12012 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12014 "model": "plain_quant", "prompt": "task"
12015 }))
12016 .unwrap();
12017 let (tx, _rx) = worker::event_channel();
12018 assert_eq!(
12019 build_request(&req, tx, lanes::Lane::Interactive, None)
12020 .params
12021 .max_new,
12022 worker::MAX_NEW_CTX_BOUNDED
12023 );
12024 let turns: Vec<(String, String)> = request
12025 .chat_turns
12026 .iter()
12027 .map(|t| (t.role.clone(), t.content.clone()))
12028 .collect();
12029 assert_eq!(
12030 turns,
12031 vec![
12032 ("system".into(), "rules".into()),
12033 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
12035 ("assistant".into(), "work".into()),
12036 ]
12037 );
12038 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
12039 assert_eq!(request.stop_strings, vec!["<stop>"]);
12040
12041 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12042 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12043 "stop": ["a", "b"]
12044 }))
12045 .unwrap();
12046 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
12047
12048 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12052 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12053 "stop": ["", "real", ""]
12054 }))
12055 .unwrap();
12056 assert_eq!(req.stop.into_vec(), vec!["real"]);
12057 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12058 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12059 "stop": ""
12060 }))
12061 .unwrap();
12062 assert!(req.stop.into_vec().is_empty());
12063
12064 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12065 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
12066 "stop": null
12067 }))
12068 .unwrap();
12069 assert!(req.stop.into_vec().is_empty());
12070 }
12071
12072 #[test]
12073 fn stop_sequence_limits_bound_count_individual_and_aggregate_work() {
12074 let at_limit = StopSequences::Many(vec!["x".repeat(256); MAX_STOP_SEQUENCES]);
12075 assert!(at_limit.validate().is_ok());
12076 assert!(
12077 StopSequences::Many(vec![String::new(); MAX_STOP_SEQUENCES + 1])
12078 .validate()
12079 .unwrap_err()
12080 .contains("at most")
12081 );
12082 assert!(
12083 StopSequences::One("x".repeat(MAX_STOP_SEQUENCE_BYTES + 1))
12084 .validate()
12085 .unwrap_err()
12086 .contains("each stop")
12087 );
12088 assert!(
12089 StopSequences::Many(vec!["x".repeat(300); MAX_STOP_SEQUENCES])
12090 .validate()
12091 .unwrap_err()
12092 .contains("total at most")
12093 );
12094 }
12095
12096 #[tokio::test]
12097 async fn chat_response_has_openai_message_shape() {
12098 let (tx, rx) = worker::event_channel();
12099 tx.send(Event::Token {
12100 id: 1,
12101 text: "hello".into(),
12102 })
12103 .unwrap();
12104 tx.send(Event::Done {
12105 stop_reason: "Eos".into(),
12106 n_tokens: 1,
12107 n_prompt: 42,
12108 n_cached: 30,
12109 elapsed_s: 0.5,
12110 spec: None,
12111 })
12112 .unwrap();
12113 drop(tx);
12114 let response = blocking_response(
12115 rx,
12116 "plain_quant".into(),
12117 true,
12118 Vec::new(),
12119 None,
12120 Envelope::new(true),
12121 )
12122 .await;
12123 assert_eq!(response.status(), StatusCode::OK);
12124 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12125 .await
12126 .unwrap();
12127 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12128 assert_eq!(payload["object"], "chat.completion");
12129 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
12131 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
12132 let fingerprint = payload["system_fingerprint"].as_str().unwrap();
12136 assert!(
12137 build_id::fingerprint_is_well_formed(fingerprint),
12138 "system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
12139 );
12140 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
12141 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
12142 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
12143 assert_eq!(payload["usage"]["prompt_tokens"], 42);
12145 assert_eq!(payload["usage"]["completion_tokens"], 1);
12146 assert_eq!(payload["usage"]["total_tokens"], 43);
12147 assert_eq!(
12148 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12149 30
12150 );
12151 assert!(payload["usage"].get("spec").is_none());
12154 }
12155
12156 #[tokio::test]
12157 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
12158 let (tx, rx) = worker::event_channel();
12159 tx.send(Event::Token {
12161 id: 4,
12162 text: "hello".into(),
12163 })
12164 .unwrap();
12165 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
12166 tx.send(Event::Done {
12167 stop_reason: "MaxNew".into(),
12168 n_tokens: 4,
12169 n_prompt: 2,
12170 n_cached: 0,
12171 elapsed_s: 0.5,
12172 spec: None,
12173 })
12174 .unwrap();
12175 drop(tx);
12176
12177 let response = blocking_response(
12178 rx,
12179 "plain_quant".into(),
12180 false,
12181 Vec::new(),
12182 None,
12183 Envelope::new(false),
12184 )
12185 .await;
12186 assert_eq!(response.status(), StatusCode::OK);
12187 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12188 .await
12189 .unwrap();
12190 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12191 assert_eq!(payload["text"], "hello");
12192 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
12193 assert_eq!(payload["n_tokens"], 4);
12194 }
12195
12196 #[tokio::test]
12199 async fn chat_usage_carries_spec_acceptance_summary() {
12200 let (tx, rx) = worker::event_channel();
12201 tx.send(Event::Token {
12202 id: 1,
12203 text: "hello".into(),
12204 })
12205 .unwrap();
12206 tx.send(Event::Done {
12207 stop_reason: "Eos".into(),
12208 n_tokens: 1,
12209 n_prompt: 42,
12210 n_cached: 0,
12211 elapsed_s: 0.5,
12212 spec: Some(worker::SpecUsage {
12213 rounds: 10,
12214 drafted: 30,
12215 accepted: 21,
12216 }),
12217 })
12218 .unwrap();
12219 drop(tx);
12220 let response = blocking_response(
12221 rx,
12222 "plain_quant".into(),
12223 true,
12224 Vec::new(),
12225 None,
12226 Envelope::new(true),
12227 )
12228 .await;
12229 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12230 .await
12231 .unwrap();
12232 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12233 let sp = &payload["usage"]["spec"];
12234 assert_eq!(sp["rounds"], 10);
12235 assert_eq!(sp["drafted"], 30);
12236 assert_eq!(sp["accepted"], 21);
12237 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
12238 assert_eq!(payload["usage"]["total_tokens"], 43);
12240 }
12241
12242 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
12243 let mut payload = serde_json::json!({
12244 "model": "m",
12245 "messages": [{"role": "user", "content": "Weather in Paris?"}],
12246 "tools": [{"type": "function", "function": {
12247 "name": "get_weather",
12248 "description": "Get current weather",
12249 "parameters": {"type": "object",
12250 "properties": {"city": {"type": "string"},
12251 "days": {"type": "integer"}},
12252 "required": ["city"]}}}],
12253 });
12254 if let Some(obj) = extra.as_object() {
12255 for (k, v) in obj {
12256 payload[k] = v.clone();
12257 }
12258 }
12259 serde_json::from_value(payload).unwrap()
12260 }
12261
12262 #[test]
12266 fn glm5_vision_decode_is_deferred_and_grid_pinned() {
12267 let (tx, _rx) = worker::event_channel();
12268 let req: ChatCompletionReq = serde_json::from_value(json!({
12269 "model": "m", "messages": [{"role": "user", "content": "hi"}],
12270 }))
12271 .unwrap();
12272 let mut plan = build_chat_request(
12273 req,
12274 Some(&ModelCaps {
12275 chat_ok: true,
12276 ..Default::default()
12277 }),
12278 tx,
12279 lanes::Lane::Interactive,
12280 None,
12281 )
12282 .unwrap();
12283 let bmp = |w: u32, h: u32| -> Vec<u8> {
12286 let row = (w * 3).div_ceil(4) * 4;
12287 let size = 54 + row * h;
12288 let mut b = vec![0x42u8, 0x4d];
12289 b.extend_from_slice(&size.to_le_bytes());
12290 b.extend_from_slice(&[0; 4]);
12291 b.extend_from_slice(&54u32.to_le_bytes());
12292 b.extend_from_slice(&40u32.to_le_bytes());
12293 b.extend_from_slice(&w.to_le_bytes());
12294 b.extend_from_slice(&h.to_le_bytes());
12295 b.extend_from_slice(&1u16.to_le_bytes());
12296 b.extend_from_slice(&24u16.to_le_bytes());
12297 b.extend_from_slice(&[0u8; 24]);
12298 b.extend(std::iter::repeat_n(0x7fu8, (row * h) as usize));
12299 b
12300 };
12301 let bytes = bmp(112, 112);
12302 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes).unwrap();
12303 assert_eq!((gh, gw), (8, 8), "identity resize grid");
12304 assert_eq!(memra_engine::vision_glm5::n_merged_for_grid(gh, gw), 16);
12305 plan.pending_glm5.push(PendingGlm5Image {
12306 bytes: bytes.clone(),
12307 gh,
12308 gw,
12309 });
12310 decode_pending_vision(&mut plan).unwrap();
12311 assert_eq!(plan.request.glm5_images.len(), 1);
12312 let unit = &plan.request.glm5_images[0];
12313 assert_eq!((unit.gh, unit.gw), (gh, gw));
12314 assert_eq!(
12315 unit.patches.len(),
12316 gh * gw * memra_engine::vision_glm5::G5V_PATCH_IN
12317 );
12318 plan.request.glm5_images.clear();
12320 plan.pending_glm5.push(PendingGlm5Image {
12321 bytes,
12322 gh: gh + 2,
12323 gw,
12324 });
12325 let err = decode_pending_vision(&mut plan).unwrap_err();
12326 assert!(err.contains("header-planned"), "got: {err}");
12327 }
12328
12329 #[test]
12330 fn vision_decode_is_deferred_and_grid_pinned() {
12331 let (tx, _rx) = worker::event_channel();
12336 let req: ChatCompletionReq = serde_json::from_value(json!({
12337 "model": "m", "messages": [{"role": "user", "content": "hi"}],
12338 }))
12339 .unwrap();
12340 let mut plan = build_chat_request(
12341 req,
12342 Some(&ModelCaps {
12343 chat_ok: true,
12344 ..Default::default()
12345 }),
12346 tx,
12347 lanes::Lane::Interactive,
12348 None,
12349 )
12350 .unwrap();
12351 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
12355 let mut b = Vec::new();
12356 b.extend_from_slice(b"BM");
12357 b.extend_from_slice(&54u32.to_le_bytes());
12358 b.extend_from_slice(&0u32.to_le_bytes());
12359 b.extend_from_slice(&54u32.to_le_bytes());
12360 b.extend_from_slice(&40u32.to_le_bytes());
12361 b.extend_from_slice(&w.to_le_bytes());
12362 b.extend_from_slice(&h.to_le_bytes());
12363 b.extend_from_slice(&1u16.to_le_bytes());
12364 b.extend_from_slice(&24u16.to_le_bytes());
12365 b.extend_from_slice(&[0u8; 24]);
12366 if with_pixels {
12367 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
12368 }
12369 b
12370 };
12371 let bytes = bmp(64, 64, true);
12372 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
12373 plan.pending_images.push(PendingVisionUnit::Still {
12374 bytes: bytes.clone(),
12375 gh,
12376 gw,
12377 });
12378 decode_pending_vision(&mut plan).unwrap();
12379 assert_eq!(plan.request.images.len(), 1);
12380 assert_eq!(
12381 (
12382 plan.request.images[0].prep.gh,
12383 plan.request.images[0].prep.gw
12384 ),
12385 (gh, gw),
12386 "decoded grid must equal the header-planned grid the pad run was rendered from"
12387 );
12388 plan.request.images.clear();
12390 plan.pending_images.push(PendingVisionUnit::Still {
12391 bytes,
12392 gh: gh + 2,
12393 gw,
12394 });
12395 let err = decode_pending_vision(&mut plan).unwrap_err();
12396 assert!(err.contains("header-planned"), "got: {err}");
12397 let bomb = bmp(16_000, 16_000, false);
12400 plan.pending_images.clear();
12401 plan.pending_images.push(PendingVisionUnit::Still {
12402 bytes: bomb,
12403 gh: 2,
12404 gw: 2,
12405 });
12406 let err = decode_pending_vision(&mut plan).unwrap_err();
12407 assert!(err.contains("exceeds the decode budget"), "got: {err}");
12408 }
12409
12410 #[test]
12411 fn tools_request_renders_client_key_order_and_arms_parser() {
12412 let (tx, _rx) = worker::event_channel();
12413 let plan = build_chat_request(
12414 weather_request(json!({})),
12415 Some(&tool_caps()),
12416 tx,
12417 lanes::Lane::Interactive,
12418 None,
12419 )
12420 .unwrap();
12421 assert!(plan.parser.is_some());
12422 assert_eq!(plan.request.tools_json.len(), 1);
12423 assert_eq!(
12425 plan.request.tools_json[0],
12426 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
12427 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
12428 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
12429 \"integer\"}}, \"required\": [\"city\"]}}}"
12430 );
12431 }
12432
12433 #[test]
12434 fn hy3_tools_and_reasoning_flow_through_the_real_chat_plan() {
12435 let (tx, _rx) = worker::event_channel();
12436 let plan = build_chat_request(
12437 weather_request(json!({"reasoning_effort": "high"})),
12438 Some(&hy3_tool_caps()),
12439 tx,
12440 lanes::Lane::Interactive,
12441 None,
12442 )
12443 .unwrap();
12444 assert_eq!(plan.request.think, ThinkMode::Think);
12445 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12446 assert!(
12447 plan.request
12448 .stop_strings
12449 .iter()
12450 .any(|stop| stop == "</tool_calls:opensource>")
12451 );
12452 let rendered = chat::apply_chat_template_tools_ex(
12453 Some("... hy_User ... <tools> ..."),
12454 &plan.request.chat_turns,
12455 true,
12456 &plan.request.tools_json,
12457 &plan.request.tools_struct,
12458 plan.request.think,
12459 plan.request.reasoning_effort.as_deref(),
12460 None,
12461 )
12462 .unwrap();
12463 assert!(rendered.contains("<tool_calls:opensource>"));
12464 assert!(rendered.ends_with("<think:opensource>"));
12465
12466 let mut parser = plan.parser.expect("HY3 tools arm its native parser");
12467 let pieces = parser.push(concat!(
12468 "Need weather.</think:opensource>",
12469 "<tool_calls:opensource><tool_call:opensource>get_weather",
12470 "<tool_sep:opensource>\n<arg_key:opensource>city</arg_key:opensource>\n",
12471 "<arg_value:opensource>Paris</arg_value:opensource>\n",
12472 "</tool_call:opensource></tool_calls:opensource>",
12473 ));
12474 assert!(pieces.contains(&Piece::Reasoning("Need weather.".into())));
12475 assert!(pieces.iter().any(|piece| matches!(piece, Piece::Call(call)
12476 if call.name == "get_weather" && call.arguments == r#"{"city":"Paris"}"#)));
12477 }
12478
12479 #[test]
12480 fn tool_choice_none_strips_tools_and_parser() {
12481 let (tx, _rx) = worker::event_channel();
12482 let plan = build_chat_request(
12483 weather_request(json!({"tool_choice": "none"})),
12484 Some(&tool_caps()),
12485 tx,
12486 lanes::Lane::Interactive,
12487 None,
12488 )
12489 .unwrap();
12490 let mut p = plan
12493 .parser
12494 .expect("think-open chat arms the reasoning splitter");
12495 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
12496 assert_eq!(
12497 pieces,
12498 vec![
12499 Piece::Reasoning("x".into()),
12500 Piece::Content("<tool_call> stays prose".into()),
12501 ]
12502 );
12503 assert!(plan.request.tools_json.is_empty());
12504 let (tx, _rx) = worker::event_channel();
12506 assert!(
12507 build_chat_request(
12508 weather_request(json!({"tool_choice": "required"})),
12509 Some(&tool_caps()),
12510 tx,
12511 lanes::Lane::Interactive,
12512 None
12513 )
12514 .is_err()
12515 );
12516 let (tx, _rx) = worker::event_channel();
12517 assert!(
12518 build_chat_request(
12519 weather_request(json!({"tool_choice":
12520 {"type": "function", "function": {"name": "get_weather"}}})),
12521 Some(&tool_caps()),
12522 tx,
12523 lanes::Lane::Interactive,
12524 None
12525 )
12526 .is_err()
12527 );
12528 }
12529
12530 #[test]
12531 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
12532 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
12533 let _ = std::fs::remove_dir_all(&root);
12534
12535 let st = root.join("st_single");
12537 std::fs::create_dir_all(&st).unwrap();
12538 std::fs::write(st.join("config.json"), "{}").unwrap();
12539 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
12540 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
12541
12542 let sh = root.join("st_sharded");
12544 std::fs::create_dir_all(&sh).unwrap();
12545 std::fs::write(sh.join("config.json"), "{}").unwrap();
12546 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
12547 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
12548
12549 let rp = root.join("repack");
12551 std::fs::create_dir_all(&rp).unwrap();
12552 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
12553 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
12554
12555 let bogus = root.join("bogus");
12557 std::fs::create_dir_all(&bogus).unwrap();
12558 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
12559 assert!(
12560 err.contains("model.safetensors"),
12561 "error should say what is missing: {err}"
12562 );
12563 assert!(
12564 err.contains("manifest.json"),
12565 "error should mention the repack form: {err}"
12566 );
12567
12568 let nc = root.join("no_config");
12570 std::fs::create_dir_all(&nc).unwrap();
12571 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
12572 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
12573 assert!(
12574 err.contains("config.json"),
12575 "error should name config.json: {err}"
12576 );
12577
12578 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
12580 assert!(err.contains("does not exist"), "{err}");
12581
12582 let f = root.join("model.gguf");
12584 std::fs::write(&f, b"g").unwrap();
12585 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
12586
12587 let _ = std::fs::remove_dir_all(&root);
12588 }
12589
12590 #[test]
12591 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
12592 let caps = ModelCaps {
12595 tools_branch: false,
12596 qwen_think: false,
12597 think_switch: false,
12598 chat_ok: false,
12599 ..Default::default()
12600 };
12601 let payload = serde_json::json!({
12602 "model": "st_model",
12603 "messages": [{"role": "user", "content": "hello"}],
12604 });
12605 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12606 let (tx, _rx) = worker::event_channel();
12607 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
12608 Err(e) => e,
12609 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
12610 };
12611 assert!(
12612 err.contains("no chat template"),
12613 "message should name the cause: {err}"
12614 );
12615 assert!(
12616 err.contains("/v1/completions"),
12617 "message should point at the raw-prompt escape hatch: {err}"
12618 );
12619 }
12620
12621 #[test]
12622 fn tools_on_model_without_tools_branch_is_rejected() {
12623 let (tx, _rx) = worker::event_channel();
12624 let caps = ModelCaps {
12625 chat_ok: true,
12626 ..Default::default()
12627 };
12628 assert!(
12629 build_chat_request(
12630 weather_request(json!({})),
12631 Some(&caps),
12632 tx,
12633 lanes::Lane::Interactive,
12634 None
12635 )
12636 .is_err()
12637 );
12638 let (tx, _rx) = worker::event_channel();
12639 assert!(
12640 build_chat_request(
12641 weather_request(json!({})),
12642 None,
12643 tx,
12644 lanes::Lane::Interactive,
12645 None
12646 )
12647 .is_err()
12648 );
12649 }
12650
12651 #[test]
12652 fn reasoning_effort_maps_to_think_switch() {
12653 for (extra, want) in [
12659 (json!({}), ThinkMode::Default),
12660 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
12661 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
12662 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
12663 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
12664 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
12665 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
12666 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
12667 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
12668 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
12672 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
12673 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
12674 (
12678 json!({"reasoning": {"enabled": true, "effort": "none"}}),
12679 ThinkMode::Think,
12680 ),
12681 (
12682 json!({"reasoning": {"enabled": false, "effort": "high"}}),
12683 ThinkMode::NoThink,
12684 ),
12685 ] {
12686 let (tx, _rx) = worker::event_channel();
12687 let plan = build_chat_request(
12688 weather_request(extra.clone()),
12689 Some(&ladder_caps()),
12694 tx,
12695 lanes::Lane::Interactive,
12696 None,
12697 )
12698 .unwrap();
12699 assert_eq!(plan.request.think, want, "extra={extra}");
12700 }
12701 for extra in [
12705 json!({"reasoning_effort": "extreme"}),
12706 json!({"reasoning": {"effort": "banana"}}),
12707 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
12708 json!({"reasoning": {"enabled": true, "effort": ""}}),
12709 ] {
12710 let (tx, _rx) = worker::event_channel();
12711 assert!(
12712 build_chat_request(
12713 weather_request(extra.clone()),
12714 Some(&tool_caps()),
12715 tx,
12716 lanes::Lane::Interactive,
12717 None
12718 )
12719 .is_err(),
12720 "extra={extra} must be rejected by the one allowlist"
12721 );
12722 }
12723 for (raw, want) in [
12726 ("none", Some("none")),
12727 ("minimal", Some("minimal")),
12728 ("low", Some("low")),
12729 ("medium", Some("medium")),
12730 ("high", Some("high")),
12731 ("xhigh", Some("high")),
12732 ("max", Some("high")),
12733 ("ultra", Some("high")),
12734 ("banana", None),
12735 ("", None),
12736 ("HIGH", None),
12737 ] {
12738 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
12739 }
12740 for (raw, want) in [
12743 ("none", Some("none")),
12744 ("minimal", Some("minimal")),
12745 ("low", Some("low")),
12746 ("medium", Some("medium")),
12747 ("high", Some("high")),
12748 ("xhigh", Some("max")),
12749 ("max", Some("max")),
12750 ("ultra", Some("max")),
12751 ("banana", None),
12752 ("", None),
12753 ("MAX", None),
12754 ] {
12755 assert_eq!(
12756 canonical_effort_for(raw, true),
12757 want,
12758 "canonical_effort_for({raw:?}, dsv4)"
12759 );
12760 }
12761 }
12762
12763 #[test]
12764 fn dsv4_reasoning_effort_max_survives_canonicalization() {
12765 let dsv4_caps = ModelCaps {
12771 chat_ok: true,
12772 dsv4: true,
12773 ..Default::default()
12774 };
12775 let build = |caps: &ModelCaps, effort: &str| {
12776 let (tx, _rx) = worker::event_channel();
12777 let req: ChatCompletionReq = serde_json::from_value(json!({
12778 "model": "m",
12779 "messages": [{"role": "user", "content": "hi"}],
12780 "reasoning_effort": effort,
12781 }))
12782 .unwrap();
12783 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
12784 };
12785 for raw in ["max", "xhigh", "ultra"] {
12786 let plan = build(&dsv4_caps, raw).unwrap();
12787 assert_eq!(
12788 plan.request.reasoning_effort.as_deref(),
12789 Some("max"),
12790 "dsv4 {raw:?} must reach the renderer as the max rung"
12791 );
12792 assert_eq!(plan.request.think, chat::ThinkMode::Think);
12793 }
12794 let plan = build(&dsv4_caps, "high").unwrap();
12796 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12797 let step_caps = ModelCaps {
12799 chat_ok: true,
12800 effort_levels: true,
12801 ..Default::default()
12802 };
12803 let plan = build(&step_caps, "max").unwrap();
12804 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
12805 }
12806
12807 #[test]
12808 fn default_reasoning_effort_flips_only_the_unset_request() {
12809 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
12816 let (tx, _rx) = worker::event_channel();
12817 build_chat_request_with_trace(
12818 weather_request(extra),
12819 Some(&ladder_caps()),
12820 tx,
12821 lanes::Lane::Interactive,
12822 None,
12823 None,
12824 default_effort,
12825 &ModelSamplingDefaults::default(),
12826 )
12827 .unwrap()
12828 };
12829 for (extra, want) in [
12830 (json!({}), ThinkMode::Think),
12832 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
12836 (json!({"include_reasoning": false}), ThinkMode::NoThink),
12837 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
12839 (json!({"include_reasoning": true}), ThinkMode::Think),
12840 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
12842 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
12843 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
12844 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
12846 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
12847 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
12848 ] {
12849 let plan = build(extra.clone(), Some("high"));
12850 assert_eq!(plan.request.think, want, "extra={extra}");
12851 }
12852 assert_eq!(
12854 build(json!({}), Some("none")).request.think,
12855 ThinkMode::NoThink
12856 );
12857 assert_eq!(
12858 build(json!({"reasoning_effort": "high"}), Some("none"))
12859 .request
12860 .think,
12861 ThinkMode::Think
12862 );
12863 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
12867 }
12868
12869 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
12875 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
12876 {%- else %}'<think>\\n'{%- endif %}";
12877
12878 #[test]
12879 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
12880 let build = |extra: serde_json::Value| {
12886 let (tx, _rx) = worker::event_channel();
12887 build_chat_request(
12888 weather_request(extra),
12889 Some(&tool_caps()),
12890 tx,
12891 lanes::Lane::Interactive,
12892 None,
12893 )
12894 };
12895 for (extra, want) in [
12896 (json!({"enable_thinking": false}), ThinkMode::NoThink),
12897 (json!({"enable_thinking": true}), ThinkMode::Think),
12898 (
12899 json!({"chat_template_kwargs": {"enable_thinking": false}}),
12900 ThinkMode::NoThink,
12901 ),
12902 (
12903 json!({"chat_template_kwargs": {"enable_thinking": true}}),
12904 ThinkMode::Think,
12905 ),
12906 (
12909 json!({"enable_thinking": false, "reasoning_effort": "high"}),
12910 ThinkMode::NoThink,
12911 ),
12912 (
12914 json!({"enable_thinking": false,
12915 "chat_template_kwargs": {"enable_thinking": false}}),
12916 ThinkMode::NoThink,
12917 ),
12918 ] {
12919 let plan = build(extra.clone()).unwrap_or_else(|e| {
12920 panic!("{extra} must be accepted and honored, got 400: {e}");
12921 });
12922 assert_eq!(
12923 plan.request.think, want,
12924 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
12925 );
12926 }
12927 let render = |extra: serde_json::Value| -> String {
12930 let plan = build(extra).unwrap();
12931 chat::apply_chat_template_tools_ex(
12932 Some(SWITCHED_QWEN_TMPL),
12933 &plan.request.chat_turns,
12934 true,
12935 &plan.request.tools_json,
12936 &plan.request.tools_struct,
12937 plan.request.think,
12938 plan.request.reasoning_effort.as_deref(),
12939 None,
12940 )
12941 .unwrap()
12942 };
12943 let off = render(json!({"enable_thinking": false}));
12944 assert!(
12945 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
12946 "enable_thinking:false must render the CLOSED think pair: {off:?}"
12947 );
12948 let on = render(json!({}));
12949 assert!(
12950 on.ends_with("<|im_start|>assistant\n<think>\n"),
12951 "an unset request must still render the template's OPEN think tail: {on:?}"
12952 );
12953 assert_eq!(
12954 off,
12955 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
12956 "both vLLM spellings must render byte-identically"
12957 );
12958 assert_eq!(
12959 off,
12960 render(json!({"reasoning_effort": "none"})),
12961 "the vLLM spelling must render byte-identically to the OpenAI spelling"
12962 );
12963 }
12964
12965 #[test]
12966 fn unknown_chat_template_kwarg_refuses_by_name() {
12967 let build = |extra: serde_json::Value| {
12970 let (tx, _rx) = worker::event_channel();
12971 build_chat_request(
12972 weather_request(extra),
12973 Some(&tool_caps()),
12974 tx,
12975 lanes::Lane::Interactive,
12976 None,
12977 )
12978 };
12979 let refusal = |extra: serde_json::Value, why: &str| -> String {
12980 build(extra).err().unwrap_or_else(|| panic!("{why}"))
12981 };
12982 let err = refusal(
12983 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
12984 "an unimplementable template kwarg must not be accepted",
12985 );
12986 assert!(
12987 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
12988 "the refusal must name the offending key AND the supported one: {err}"
12989 );
12990 let err = refusal(
12991 json!({"chat_template_kwargs": "enable_thinking=false"}),
12992 "a non-object chat_template_kwargs must not be accepted",
12993 );
12994 assert!(
12995 err.contains("must be an object"),
12996 "refusal must say what shape is expected: {err}"
12997 );
12998 let err = refusal(
12999 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
13000 "a stringly-typed switch must not be accepted",
13001 );
13002 assert!(
13003 err.contains("true or false"),
13004 "refusal must name the expected type: {err}"
13005 );
13006 let plan = build(json!({"chat_template_kwargs": null}))
13008 .expect("null chat_template_kwargs is the unset case");
13009 assert_eq!(plan.request.think, ThinkMode::Default);
13010 }
13011
13012 const Q38_TMPL: &str =
13027 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
13028
13029 fn render_with(
13032 tmpl: &str,
13033 caps: &ModelCaps,
13034 extra: serde_json::Value,
13035 default_effort: Option<&str>,
13036 ) -> Result<String, String> {
13037 let mut payload = serde_json::json!({
13038 "model": "m",
13039 "messages": [{"role": "user", "content": "hi"}],
13040 });
13041 if let Some(obj) = extra.as_object() {
13042 for (k, v) in obj {
13043 payload[k] = v.clone();
13044 }
13045 }
13046 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13047 let (tx, _rx) = worker::event_channel();
13048 let plan = build_chat_request_with_trace(
13049 req,
13050 Some(caps),
13051 tx,
13052 lanes::Lane::Interactive,
13053 None,
13054 None,
13055 default_effort,
13056 &ModelSamplingDefaults::default(),
13057 )?;
13058 Ok(chat::apply_chat_template_tools_ex(
13059 Some(tmpl),
13060 &plan.request.chat_turns,
13061 true,
13062 &plan.request.tools_json,
13063 &plan.request.tools_struct,
13064 plan.request.think,
13065 plan.request.reasoning_effort.as_deref(),
13066 None,
13067 )
13068 .unwrap())
13069 }
13070
13071 #[test]
13072 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
13073 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
13080 let xhigh = "Reasoning effort is set to xhigh.";
13081 let low = "Reasoning effort is set to low.";
13082 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
13084 assert!(
13085 r(json!({"reasoning_effort": "high"}))
13086 .unwrap()
13087 .contains(xhigh)
13088 );
13089 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
13092 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
13093 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
13095 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
13096 assert_ne!(low_p, high_p);
13097 assert_ne!(low_p, medium);
13098 assert_ne!(high_p, medium);
13099 for alias in ["xhigh", "max", "ultra"] {
13102 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
13103 }
13104 assert_eq!(r(json!({})).unwrap(), high_p);
13107 assert_eq!(
13110 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
13111 medium
13112 );
13113 let off = r(json!({"reasoning_effort": "none"})).unwrap();
13116 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
13117 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
13118 }
13119
13120 #[test]
13121 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
13122 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
13135focused, moving directly to the conclusion without unnecessary elaboration.";
13136 let expected = format!(
13137 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
13138 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
13139 );
13140 let after_fix = render_with(
13142 Q38_TMPL,
13143 &ladder_caps(),
13144 json!({"reasoning_effort": "low"}),
13145 None,
13146 )
13147 .unwrap();
13148 assert_eq!(
13149 after_fix, expected,
13150 "the shipped prompt for reasoning_effort:\"low\""
13151 );
13152 const ORNITH_TMPL: &str = include_str!(
13155 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
13156 );
13157 let on_deployed_binary = render_with(
13158 ORNITH_TMPL,
13159 &tool_caps(),
13160 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
13161 {"role": "user", "content": "hi"}]}),
13162 None,
13163 )
13164 .unwrap();
13165 assert_eq!(
13166 on_deployed_binary, expected,
13167 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
13168 level, or its reasoning-volume numbers do not describe the shipped prompt"
13169 );
13170 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
13173 assert!(
13174 !ladderless_unset.contains("Reasoning effort is set to"),
13175 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
13176 );
13177 assert_eq!(
13178 ladderless_unset,
13179 render_with(
13180 Q38_TMPL,
13181 &ladder_caps(),
13182 json!({"reasoning_effort": "medium"}),
13183 None
13184 )
13185 .unwrap(),
13186 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
13187 );
13188 }
13189
13190 #[test]
13191 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
13192 let off = render_with(
13199 Q38_TMPL,
13200 &ladder_caps(),
13201 json!({"reasoning_effort": "none"}),
13202 None,
13203 )
13204 .unwrap();
13205 for extra in [
13206 json!({"include_reasoning": false}),
13207 json!({"reasoning": {"exclude": true}}),
13208 ] {
13209 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
13210 assert!(
13211 got.ends_with("<think>\n\n</think>\n\n"),
13212 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
13213 );
13214 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
13215 }
13216 for extra in [
13221 json!({"enable_thinking": true, "include_reasoning": false}),
13222 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
13223 json!({"reasoning": {"enabled": true, "exclude": true}}),
13224 ] {
13225 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
13226 .err()
13227 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
13228 assert!(e.contains("contradictory"), "{extra}: {e}");
13229 assert!(
13230 e.contains("include_reasoning") || e.contains("exclude"),
13231 "{extra}: the refusal must name the suppression field the caller sent: {e}"
13232 );
13233 }
13234 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
13237 for extra in [
13238 json!({"include_reasoning": true}),
13239 json!({"reasoning": {"exclude": false}}),
13240 ] {
13241 assert_eq!(
13242 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
13243 dflt,
13244 "{extra} must not perturb the model's default"
13245 );
13246 }
13247 let switchless = ModelCaps {
13251 think_switch: false,
13252 ..tool_caps()
13253 };
13254 let err = render_with(
13255 Q38_TMPL,
13256 &switchless,
13257 json!({"include_reasoning": false}),
13258 None,
13259 )
13260 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
13261 assert!(err.contains("cannot disable reasoning"), "{err}");
13262 }
13263
13264 #[test]
13265 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
13266 let build = |extra: serde_json::Value| {
13267 let (tx, _rx) = worker::event_channel();
13268 build_chat_request(
13269 weather_request(extra),
13270 Some(&ladder_caps()),
13271 tx,
13272 lanes::Lane::Interactive,
13273 None,
13274 )
13275 };
13276 let err = |extra: serde_json::Value, why: &str| -> String {
13277 build(extra).err().unwrap_or_else(|| panic!("{why}"))
13278 };
13279 let e = err(
13283 json!({"reasoning": {"max_tokens": 1024}}),
13284 "reasoning.max_tokens must not be accepted-and-ignored",
13285 );
13286 assert!(e.contains("reasoning.max_tokens"), "{e}");
13287 assert!(e.contains("ONE output budget"), "{e}");
13288 for extra in [
13293 json!({"reasoning": {"max_tokens": null}}),
13294 json!({"reasoning": {"banana": null}}),
13295 ] {
13296 let e = err(
13297 extra.clone(),
13298 "a null-valued unhonourable key must still refuse",
13299 );
13300 assert!(
13301 e.contains("max_tokens") || e.contains("banana"),
13302 "{extra}: {e}"
13303 );
13304 }
13305 let e = err(
13307 json!({"reasoning": {"budget": 5}}),
13308 "an unknown reasoning key must not be accepted",
13309 );
13310 assert!(
13311 e.contains("reasoning.budget") && e.contains("enabled"),
13312 "{e}"
13313 );
13314 for (extra, want) in [
13318 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
13319 (json!({"reasoning": {"exclude": 1}}), "true or false"),
13320 (json!({"reasoning": {"effort": 3}}), "must be a string"),
13321 ] {
13322 let e = err(
13323 extra.clone(),
13324 "a wrong-typed reasoning key must not be ignored",
13325 );
13326 assert!(e.contains(want), "{extra}: {e}");
13327 }
13328 for extra in [
13333 json!({"reasoning": {"enabled": true}}),
13334 json!({"reasoning": {"effort": "low"}}),
13335 json!({"reasoning": {"exclude": false}}),
13336 json!({"reasoning": null}),
13337 json!({"reasoning": {"effort": null}}),
13338 json!({"reasoning": {"enabled": null, "exclude": null}}),
13339 ] {
13340 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
13341 }
13342 }
13343
13344 #[test]
13345 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
13346 const ORNITH_TMPL: &str = include_str!(
13354 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
13355 );
13356 let explicit_on = render_with(
13360 ORNITH_TMPL,
13361 &tool_caps(),
13362 json!({"reasoning": {"enabled": true}}),
13363 None,
13364 )
13365 .unwrap();
13366 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
13367 for extra in [
13368 json!({"reasoning_effort": "low"}),
13369 json!({"reasoning_effort": "medium"}),
13370 json!({"reasoning_effort": "high"}),
13371 json!({"reasoning_effort": "xhigh"}),
13373 json!({"reasoning": {"effort": "xhigh"}}),
13374 ] {
13375 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
13376 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
13377 assert_eq!(
13378 got, explicit_on,
13379 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
13380 documented translation, not a decorative accept"
13381 );
13382 }
13383 for extra in [
13385 json!({}),
13386 json!({"reasoning_effort": "none"}),
13387 json!({"reasoning_effort": "minimal"}),
13388 json!({"enable_thinking": false}),
13389 ] {
13390 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
13391 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
13392 }
13393 let minimal = render_with(
13396 ORNITH_TMPL,
13397 &tool_caps(),
13398 json!({"reasoning_effort": "minimal"}),
13399 None,
13400 )
13401 .unwrap();
13402 assert!(
13403 minimal.ends_with("<think>\n\n</think>\n\n"),
13404 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
13405 );
13406 let ladder_low = render_with(
13409 Q38_TMPL,
13410 &ladder_caps(),
13411 json!({"reasoning_effort": "low"}),
13412 None,
13413 )
13414 .unwrap();
13415 assert!(
13416 ladder_low.contains("Reasoning effort is set to low."),
13417 "{ladder_low:?}"
13418 );
13419 assert_ne!(
13420 ladder_low,
13421 render_with(
13422 Q38_TMPL,
13423 &ladder_caps(),
13424 json!({"reasoning_effort": "high"}),
13425 None
13426 )
13427 .unwrap(),
13428 "the ladder model's rungs stay distinct prompts"
13429 );
13430 }
13431
13432 #[test]
13433 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
13434 let render_chat = |body: serde_json::Value| -> Result<String, String> {
13445 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13446 let (tx, _rx) = worker::event_channel();
13447 let plan = build_chat_request(
13448 req,
13449 Some(&ladder_caps()),
13450 tx,
13451 lanes::Lane::Interactive,
13452 None,
13453 )?;
13454 Ok(chat::apply_chat_template_tools_ex(
13455 Some(Q38_TMPL),
13456 &plan.request.chat_turns,
13457 true,
13458 &plan.request.tools_json,
13459 &plan.request.tools_struct,
13460 plan.request.think,
13461 plan.request.reasoning_effort.as_deref(),
13462 None,
13463 )
13464 .unwrap())
13465 };
13466 for (intent, chat_body, responses_body, messages_body) in [
13471 (
13472 "reasoning OFF",
13473 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13474 "reasoning_effort": "none"}),
13475 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
13476 json!({"model": "m", "max_tokens": 16,
13477 "messages": [{"role": "user", "content": "hi"}],
13478 "thinking": {"type": "disabled"}}),
13479 ),
13480 (
13481 "reasoning ON at the top rung",
13482 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13483 "reasoning_effort": "xhigh"}),
13484 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
13485 json!({"model": "m", "max_tokens": 16,
13486 "messages": [{"role": "user", "content": "hi"}],
13487 "output_config": {"effort": "xhigh"}}),
13488 ),
13489 (
13490 "reasoning ON at the bottom rung",
13491 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
13492 "reasoning_effort": "low"}),
13493 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
13494 json!({"model": "m", "max_tokens": 16,
13495 "messages": [{"role": "user", "content": "hi"}],
13496 "output_config": {"effort": "low"}}),
13497 ),
13498 (
13499 "the model's own default",
13500 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
13501 json!({"model": "m", "input": "hi"}),
13502 json!({"model": "m", "max_tokens": 16,
13503 "messages": [{"role": "user", "content": "hi"}]}),
13504 ),
13505 ] {
13506 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
13507 let via_responses = responses_api::translate(&responses_body)
13508 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
13509 let via_messages = anthropic::translate(&messages_body)
13510 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
13511 for (surface, translated) in [
13512 ("/v1/responses", via_responses),
13513 ("/v1/messages", via_messages),
13514 ] {
13515 let got = render_chat(translated)
13516 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
13517 assert_eq!(
13518 got, chat,
13519 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
13520 /v1/chat/completions — the parameter is honoured on one format and not \
13521 the other"
13522 );
13523 }
13524 }
13525 let switchless = ModelCaps {
13528 think_switch: false,
13529 ..ladder_caps()
13530 };
13531 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
13532 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13533 let (tx, _rx) = worker::event_channel();
13534 let plan =
13535 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
13536 Ok(format!("{:?}", plan.request.think))
13537 };
13538 for (surface, body) in [
13539 (
13540 "/v1/responses",
13541 responses_api::translate(&json!({
13542 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
13543 .unwrap(),
13544 ),
13545 (
13546 "/v1/messages",
13547 anthropic::translate(&json!({
13548 "model": "m", "max_tokens": 16,
13549 "messages": [{"role": "user", "content": "hi"}],
13550 "thinking": {"type": "disabled"}}))
13551 .unwrap(),
13552 ),
13553 ] {
13554 let err = render_switchless(body)
13555 .err()
13556 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
13557 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
13558 }
13559 }
13560
13561 #[test]
13562 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
13563 let build = |extra: serde_json::Value| {
13571 let (tx, _rx) = worker::event_channel();
13572 build_chat_request(
13573 weather_request(extra),
13574 Some(&ladder_caps()),
13575 tx,
13576 lanes::Lane::Interactive,
13577 None,
13578 )
13579 };
13580 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
13581 .expect("preserve_thinking:true is the vendor default the renderer implements");
13582 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
13583 .err()
13584 .expect("preserve_thinking:false (the strip arm) must refuse");
13585 assert!(e.contains("preserve_thinking"), "{e}");
13586 assert!(e.contains("strip"), "{e}");
13587 assert_eq!(
13590 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
13591 .unwrap()
13592 .request
13593 .think,
13594 ThinkMode::NoThink
13595 );
13596 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
13598 .err()
13599 .expect("a stringly-typed preserve_thinking must not be accepted");
13600 assert!(e.contains("true or false"), "{e}");
13601 }
13602
13603 #[test]
13604 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
13605 let dsv4_caps = ModelCaps {
13611 qwen_think: true,
13612 think_switch: false,
13613 dsv4: true,
13614 ..tool_caps()
13615 };
13616 for extra in [
13617 json!({"reasoning_effort": "none"}),
13618 json!({"reasoning": {"enabled": false}}),
13619 json!({"enable_thinking": false}),
13620 json!({"include_reasoning": false}),
13621 ] {
13622 let (tx, _rx) = worker::event_channel();
13623 let plan = build_chat_request(
13624 weather_request(extra.clone()),
13625 Some(&dsv4_caps),
13626 tx,
13627 lanes::Lane::Interactive,
13628 None,
13629 )
13630 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
13631 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
13632 }
13633 }
13634
13635 #[test]
13636 fn contradictory_think_switches_refuse_instead_of_picking_one() {
13637 let build = |extra: serde_json::Value| {
13640 let (tx, _rx) = worker::event_channel();
13641 build_chat_request(
13642 weather_request(extra),
13643 Some(&tool_caps()),
13644 tx,
13645 lanes::Lane::Interactive,
13646 None,
13647 )
13648 };
13649 for extra in [
13650 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
13651 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
13652 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
13653 ] {
13654 match build(extra.clone()) {
13655 Err(err) => assert!(
13656 err.contains("contradictory"),
13657 "the refusal must say the switches contradict: {err}"
13658 ),
13659 Ok(plan) => panic!(
13660 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
13661 plan.request.think
13662 ),
13663 }
13664 }
13665 for extra in [
13667 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
13668 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
13669 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
13670 ] {
13671 build(extra.clone())
13672 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
13673 }
13674 }
13675
13676 #[test]
13677 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
13678 let switchless = ModelCaps {
13683 tools_branch: true,
13684 qwen_think: true,
13685 think_switch: false,
13686 chat_ok: true,
13687 ..Default::default()
13688 };
13689 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
13690 let (tx, _rx) = worker::event_channel();
13691 build_chat_request_with_trace(
13692 weather_request(extra),
13693 Some(caps),
13694 tx,
13695 lanes::Lane::Interactive,
13696 None,
13697 None,
13698 default_effort,
13699 &ModelSamplingDefaults::default(),
13700 )
13701 };
13702 for extra in [
13703 json!({"reasoning_effort": "none"}),
13704 json!({"reasoning_effort": "minimal"}),
13705 json!({"reasoning": {"enabled": false}}),
13706 json!({"enable_thinking": false}),
13707 json!({"chat_template_kwargs": {"enable_thinking": false}}),
13708 ] {
13709 let err = build(extra.clone(), &switchless, None)
13710 .err()
13711 .unwrap_or_else(|| {
13712 panic!(
13713 "{extra} on a switchless think template must not be accepted-and-ignored"
13714 )
13715 });
13716 assert!(
13717 err.contains("cannot disable reasoning"),
13718 "the refusal must say the model cannot disable reasoning: {err}"
13719 );
13720 }
13721 for (extra, default_effort) in [
13725 (json!({}), None),
13726 (json!({"reasoning_effort": "high"}), None),
13729 (json!({"reasoning": {"enabled": true}}), None),
13730 (json!({"enable_thinking": true}), None),
13731 (json!({}), Some("none")),
13732 (json!({}), Some("minimal")),
13733 (json!({}), Some("high")),
13734 ] {
13735 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
13736 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
13737 });
13738 }
13739 assert_eq!(
13742 build(json!({"enable_thinking": false}), &tool_caps(), None)
13743 .unwrap()
13744 .request
13745 .think,
13746 ThinkMode::NoThink
13747 );
13748 }
13749
13750 #[test]
13751 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
13752 let gemma_caps = ModelCaps {
13759 tools_branch: true,
13760 chat_ok: true,
13761 gemma_think: true,
13762 instruct_type: Some("gemma".into()),
13763 ..Default::default()
13764 };
13765 let render =
13766 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
13767 let mut payload = serde_json::json!({
13768 "model": "google/gemma-4-31b-it",
13769 "messages": [{"role": "user", "content": "Weather in Paris?"}],
13770 });
13771 if let Some(obj) = extra.as_object() {
13772 for (k, v) in obj {
13773 payload[k] = v.clone();
13774 }
13775 }
13776 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13777 let (tx, _rx) = worker::event_channel();
13778 let plan = build_chat_request_with_trace(
13779 req,
13780 Some(&gemma_caps),
13781 tx,
13782 lanes::Lane::Interactive,
13783 None,
13784 None,
13785 default_effort,
13786 &ModelSamplingDefaults::default(),
13787 )
13788 .unwrap();
13789 chat::apply_chat_template_tools_ex(
13790 Some(tmpl),
13791 &plan.request.chat_turns,
13792 true,
13793 &plan.request.tools_json,
13794 &plan.request.tools_struct,
13795 plan.request.think,
13796 plan.request.reasoning_effort.as_deref(),
13797 None, )
13799 .unwrap()
13800 };
13801 let official = gemma_template("official");
13802 let unset_with_knob = render(&official, json!({}), Some("high"));
13803 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
13804 assert_eq!(
13805 unset_with_knob, explicit_on,
13806 "knob render must be byte-identical to the explicit think-on render"
13807 );
13808 assert!(
13809 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
13810 "think-on injects the <|think|> system token: {unset_with_knob:?}"
13811 );
13812 assert!(
13813 unset_with_knob.ends_with("<|turn>model\n"),
13814 "think-on generation turn is OPEN: {unset_with_knob:?}"
13815 );
13816 let explicit_off_with_knob =
13820 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
13821 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
13822 assert_eq!(explicit_off_with_knob, explicit_off);
13823 assert!(
13824 !explicit_off_with_knob.contains("<|think|>")
13825 && explicit_off_with_knob.ends_with("<|turn>model\n"),
13826 "explicit off keeps the official template's thinking-off bytes: \
13827 {explicit_off_with_knob:?}"
13828 );
13829 let unset_no_knob = render(&official, json!({}), None);
13831 assert_eq!(
13832 unset_no_knob, explicit_off,
13833 "knobless unset stays the template's own thinking-off default"
13834 );
13835 assert_ne!(unset_no_knob, unset_with_knob);
13836 let qat = gemma_template("qat");
13839 assert!(
13840 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
13841 "QAT knobless unset keeps the closed-channel default"
13842 );
13843 assert_eq!(
13844 render(&qat, json!({}), Some("high")),
13845 render(&qat, json!({"reasoning_effort": "high"}), None),
13846 "QAT knob render must equal the explicit think-on render"
13847 );
13848 }
13849
13850 #[test]
13851 fn default_reasoning_effort_is_validated_at_metadata_load() {
13852 let parsed = OpenRouterMetadataFile::from_toml(
13854 r#"
13855[models.g]
13856default_reasoning_effort = "high"
13857"#,
13858 )
13859 .unwrap();
13860 assert_eq!(
13861 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
13862 Some("high")
13863 );
13864 let err = OpenRouterMetadataFile::from_toml(
13865 r#"
13866[models.g]
13867default_reasoning_effort = "always"
13868"#,
13869 )
13870 .unwrap_err();
13871 assert!(err.contains("default_reasoning_effort"), "{err}");
13872 }
13873
13874 #[test]
13875 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
13876 let effort_caps = ModelCaps {
13888 effort_levels: true,
13889 think_switch: false,
13890 ..tool_caps()
13891 };
13892 for (extra, want) in [
13893 (json!({}), None),
13894 (json!({"reasoning_effort": "low"}), Some("low")),
13895 (json!({"reasoning_effort": "medium"}), Some("medium")),
13896 (json!({"reasoning_effort": "high"}), Some("high")),
13897 (json!({"reasoning": {"effort": "high"}}), Some("high")),
13898 (json!({"reasoning_effort": "xhigh"}), Some("high")),
13900 (json!({"reasoning": {"effort": "max"}}), Some("high")),
13901 ] {
13902 let (tx, _rx) = worker::event_channel();
13903 let plan = build_chat_request(
13904 weather_request(extra.clone()),
13905 Some(&effort_caps),
13906 tx,
13907 lanes::Lane::Interactive,
13908 None,
13909 )
13910 .unwrap();
13911 assert_eq!(
13912 plan.request.reasoning_effort.as_deref(),
13913 want,
13914 "extra={extra}"
13915 );
13916 }
13917 for extra in [
13924 json!({"reasoning_effort": "none"}),
13925 json!({"reasoning_effort": "minimal"}),
13926 json!({"reasoning": {"enabled": false}}),
13927 json!({"enable_thinking": false}),
13928 json!({"include_reasoning": false}),
13929 ] {
13930 let (tx, _rx) = worker::event_channel();
13931 let err = build_chat_request(
13932 weather_request(extra.clone()),
13933 Some(&effort_caps),
13934 tx,
13935 lanes::Lane::Interactive,
13936 None,
13937 )
13938 .err()
13939 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
13940 assert!(
13941 err.contains("cannot disable reasoning"),
13942 "extra={extra}: {err}"
13943 );
13944 }
13945 for extra in [
13952 json!({"reasoning_effort": "high"}),
13953 json!({"reasoning": {"effort": "low"}}),
13954 ] {
13955 let (tx, _rx) = worker::event_channel();
13956 let plan = build_chat_request(
13957 weather_request(extra.clone()),
13958 Some(&tool_caps()),
13959 tx,
13960 lanes::Lane::Interactive,
13961 None,
13962 )
13963 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
13964 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
13965 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
13966 }
13967 let (tx, _rx) = worker::event_channel();
13969 let plan = build_chat_request(
13970 weather_request(json!({})),
13971 Some(&tool_caps()),
13972 tx,
13973 lanes::Lane::Interactive,
13974 None,
13975 )
13976 .unwrap();
13977 assert_eq!(plan.request.reasoning_effort, None);
13978 }
13979
13980 #[test]
13981 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
13982 let payload = serde_json::json!({
13983 "model": "m",
13984 "messages": [
13985 {"role": "user", "content": "Weather in Paris?"},
13986 {"role": "assistant", "content": null, "tool_calls": [
13987 {"id": "call_x", "type": "function", "function": {
13988 "name": "get_weather",
13989 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
13990 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
13991 ],
13992 });
13993 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13994 let (tx, _rx) = worker::event_channel();
13995 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
13996 .unwrap();
13997 let turns = &plan.request.chat_turns;
13998 assert_eq!(turns[1].tool_calls.len(), 1);
13999 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
14000 assert_eq!(
14001 turns[1].tool_calls[0].params,
14002 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
14003 );
14004 assert_eq!(turns[2].role, "tool");
14005 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
14006 let mut p = plan
14009 .parser
14010 .expect("think-open chat arms the reasoning splitter");
14011 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
14012 assert_eq!(
14013 pieces,
14014 vec![
14015 Piece::Reasoning("thought".into()),
14016 Piece::Content("answer <tool_call> is prose here".into()),
14017 ]
14018 );
14019 }
14020
14021 #[tokio::test]
14022 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
14023 let (tx, rx) = worker::event_channel();
14024 tx.send(Event::Token {
14025 id: 1,
14026 text: "plan</think>\n\n".into(),
14027 })
14028 .unwrap();
14029 tx.send(Event::Token {
14030 id: 2,
14031 text: "<tool_call>\n<function=get_weather>\n\
14032<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
14033 .into(),
14034 })
14035 .unwrap();
14036 tx.send(Event::Done {
14037 stop_reason: "Eos".into(),
14038 n_tokens: 2,
14039 n_prompt: 40,
14040 n_cached: 0,
14041 elapsed_s: 0.5,
14042 spec: None,
14043 })
14044 .unwrap();
14045 drop(tx);
14046 let parser = ToolStreamParser::new(HashMap::new(), true);
14047 let response = blocking_response(
14048 rx,
14049 "m".into(),
14050 true,
14051 Vec::new(),
14052 Some(parser),
14053 Envelope::new(true),
14054 )
14055 .await;
14056 assert_eq!(response.status(), StatusCode::OK);
14057 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
14058 .await
14059 .unwrap();
14060 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14061 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
14062 assert_eq!(
14065 payload["choices"][0]["message"]["content"],
14066 serde_json::Value::Null
14067 );
14068 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
14069 assert_eq!(
14070 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
14071 "plan"
14072 );
14073 let call = &payload["choices"][0]["message"]["tool_calls"][0];
14074 assert_eq!(call["type"], "function");
14075 assert_eq!(call["function"]["name"], "get_weather");
14076 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
14077 assert_eq!(payload["usage"]["prompt_tokens"], 40);
14080 assert_eq!(payload["usage"]["completion_tokens"], 2);
14081 assert_eq!(payload["usage"]["total_tokens"], 42);
14082 assert_eq!(
14083 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
14084 0
14085 );
14086 }
14087
14088 #[test]
14089 fn cache_salt_plumbs_to_the_worker_namespace() {
14090 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14092 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
14093 }))
14094 .unwrap();
14095 let (tx, _rx) = worker::event_channel();
14096 assert_eq!(
14097 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14098 "tenant-a"
14099 );
14100
14101 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14102 "model": "m", "messages": [{"role": "user", "content": "task"}],
14103 "cache_salt": "tenant-b"
14104 }))
14105 .unwrap();
14106 let (tx, _rx) = worker::event_channel();
14107 assert_eq!(
14108 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14109 .unwrap()
14110 .request
14111 .cache_ns,
14112 "tenant-b"
14113 );
14114
14115 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14117 "model": "m", "prompt": "task"
14118 }))
14119 .unwrap();
14120 let (tx, _rx) = worker::event_channel();
14121 assert_eq!(
14122 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
14123 ""
14124 );
14125 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14126 "model": "m", "messages": [{"role": "user", "content": "task"}]
14127 }))
14128 .unwrap();
14129 let (tx, _rx) = worker::event_channel();
14130 assert_eq!(
14131 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14132 .unwrap()
14133 .request
14134 .cache_ns,
14135 ""
14136 );
14137 }
14138
14139 #[test]
14140 fn cache_salt_validation_rejects_oversized_value() {
14141 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
14142 assert_eq!(
14143 validate_cache_namespace(&salt, false),
14144 Err("cache_salt must be at most 64 bytes")
14145 );
14146 }
14147
14148 #[test]
14149 fn cache_salt_validation_rejects_reserved_open_namespace() {
14150 let salt = Some("t:acme\u{1f}private".to_string());
14151 assert_eq!(
14152 validate_cache_namespace(&salt, false),
14153 Err("cache_salt must not use the reserved t: prefix without a keyring")
14154 );
14155 }
14156
14157 #[test]
14158 fn cache_salt_validation_accepts_normal_value() {
14159 let raw = "tenant-A_7.c2VjcmV0LXNjb3Bl+/=";
14160 let salt = Some(raw.to_string());
14161 assert_eq!(validate_cache_namespace(&salt, false).unwrap(), raw);
14162 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
14163 let max_raw = "a".repeat(CACHE_SALT_MAX_BYTES);
14164 let max = Some(max_raw.clone());
14165 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max_raw);
14166 }
14167
14168 #[test]
14169 fn cache_salt_validation_rejects_unsupported_characters() {
14170 let salt = Some("tenant salt".to_string());
14171 assert_eq!(
14172 validate_cache_namespace(&salt, false),
14173 Err("cache_salt contains unsupported characters")
14174 );
14175 }
14176
14177 #[test]
14178 fn affinity_key_honors_both_client_conventions_in_priority_order() {
14179 use axum::http::HeaderMap;
14180 let hdr = |v: &str| {
14181 let mut h = HeaderMap::new();
14182 h.insert("x-session-id", v.parse().unwrap());
14183 h
14184 };
14185 let empty = HeaderMap::new();
14186 let s = |v: &str| Some(v.to_string());
14187 assert_eq!(
14189 affinity_key(&s("explicit"), &None, &empty).unwrap(),
14190 s("explicit")
14191 );
14192 assert_eq!(
14193 affinity_key(&None, &s("openai-user"), &empty).unwrap(),
14194 s("openai-user")
14195 );
14196 assert_eq!(
14197 affinity_key(&None, &None, &hdr("hdr-id")).unwrap(),
14198 s("hdr-id")
14199 );
14200 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")).unwrap(), s("a"));
14203 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")).unwrap(), s("b"));
14204 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")).unwrap(), None);
14207 assert_eq!(affinity_key(&s(""), &s("real"), &empty).unwrap(), s("real"));
14208 assert_eq!(
14210 affinity_key(&s(" padded "), &None, &empty).unwrap(),
14211 s("padded")
14212 );
14213 assert_eq!(affinity_key(&None, &None, &empty).unwrap(), None);
14215 assert!(
14216 affinity_key(
14217 &s(&"x".repeat(MAX_CLIENT_IDENTIFIER_BYTES + 1)),
14218 &None,
14219 &empty,
14220 )
14221 .unwrap_err()
14222 .contains("at most")
14223 );
14224 assert!(
14225 affinity_key(&s("forged\nlog"), &None, &empty)
14226 .unwrap_err()
14227 .contains("control")
14228 );
14229 }
14230
14231 #[test]
14232 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
14233 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14234 "model": "m", "prompt": "task", "session_id": "conv-1"
14235 }))
14236 .unwrap();
14237 let (tx, _rx) = worker::event_channel();
14238 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14239 assert_eq!(
14240 build_request(&req, tx, lanes::Lane::Interactive, key)
14241 .affinity
14242 .as_deref(),
14243 Some("conv-1")
14244 );
14245 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14247 "model": "m", "messages": [{"role": "user", "content": "task"}],
14248 "user": "conv-2"
14249 }))
14250 .unwrap();
14251 let (tx, _rx) = worker::event_channel();
14252 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new()).unwrap();
14253 assert_eq!(
14254 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
14255 .unwrap()
14256 .request
14257 .affinity
14258 .as_deref(),
14259 Some("conv-2")
14260 );
14261 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14263 "model": "m", "prompt": "task"
14264 }))
14265 .unwrap();
14266 let (tx, _rx) = worker::event_channel();
14267 assert!(
14268 build_request(&req, tx, lanes::Lane::Interactive, None)
14269 .affinity
14270 .is_none()
14271 );
14272 }
14273
14274 async fn sse_data_lines(resp: Response) -> Vec<String> {
14276 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14277 .await
14278 .unwrap();
14279 String::from_utf8(bytes.to_vec())
14280 .unwrap()
14281 .lines()
14282 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
14283 .collect()
14284 }
14285
14286 #[tokio::test]
14287 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
14288 let feed = |think: bool| {
14294 let (tx, rx) = worker::event_channel();
14295 let body = if think {
14296 "a plan</think>\n\nanswer"
14297 } else {
14298 "answer"
14299 };
14300 tx.send(Event::Token {
14301 id: 1,
14302 text: body.into(),
14303 })
14304 .unwrap();
14305 tx.send(Event::Done {
14306 stop_reason: "Eos".into(),
14307 n_tokens: 3,
14308 n_prompt: 10,
14309 n_cached: 0,
14310 elapsed_s: 0.1,
14311 spec: None,
14312 })
14313 .unwrap();
14314 drop(tx);
14315 rx
14316 };
14317 let resp = blocking_response(
14319 feed(true),
14320 "m".into(),
14321 true,
14322 Vec::new(),
14323 Some(ToolStreamParser::reasoning_only()),
14324 Envelope::new(true),
14325 )
14326 .await;
14327 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14328 .await
14329 .unwrap();
14330 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14331 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
14332 assert_eq!(
14333 v["choices"][0]["message"]["reasoning_details"][0]["text"],
14334 "a plan"
14335 );
14336 assert_eq!(v["choices"][0]["message"]["content"], "answer");
14337 let resp = blocking_response(
14340 feed(false),
14341 "m".into(),
14342 true,
14343 Vec::new(),
14344 None,
14345 Envelope::new(true),
14346 )
14347 .await;
14348 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14349 .await
14350 .unwrap();
14351 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14352 assert!(
14353 v["choices"][0]["message"].get("reasoning").is_none(),
14354 "a reasoning-off response must carry no reasoning field: {v}"
14355 );
14356 assert_eq!(v["choices"][0]["message"]["content"], "answer");
14357 let resp = sse_response(
14359 feed(true),
14360 "m".into(),
14361 true,
14362 Some(ToolStreamParser::reasoning_only()),
14363 Envelope::new(true),
14364 Vec::new(),
14365 None,
14366 )
14367 .into_response();
14368 let lines = sse_data_lines(resp).await;
14369 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14370 .iter()
14371 .map(|l| serde_json::from_str(l).unwrap())
14372 .collect();
14373 let reasoning: String = chunks
14374 .iter()
14375 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
14376 .collect();
14377 assert_eq!(
14378 reasoning, "a plan",
14379 "think text must stream as delta.reasoning"
14380 );
14381 let content: String = chunks
14382 .iter()
14383 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
14384 .collect();
14385 assert_eq!(content, "answer", "content must exclude the think segment");
14386 let resp = sse_response(
14388 feed(false),
14389 "m".into(),
14390 true,
14391 None,
14392 Envelope::new(true),
14393 Vec::new(),
14394 None,
14395 )
14396 .into_response();
14397 let lines = sse_data_lines(resp).await;
14398 for l in &lines[..lines.len() - 1] {
14399 let c: serde_json::Value = serde_json::from_str(l).unwrap();
14400 assert!(
14401 c["choices"][0]["delta"].get("reasoning").is_none(),
14402 "a reasoning-off stream must carry no reasoning deltas: {c}"
14403 );
14404 }
14405 }
14406
14407 #[tokio::test]
14408 async fn stream_chunks_carry_envelope_and_first_delta_role() {
14409 let (tx, rx) = worker::event_channel();
14410 tx.send(Event::Token {
14411 id: 1,
14412 text: "he".into(),
14413 })
14414 .unwrap();
14415 tx.send(Event::Token {
14416 id: 2,
14417 text: "llo".into(),
14418 })
14419 .unwrap();
14420 tx.send(Event::Done {
14421 stop_reason: "Eos".into(),
14422 n_tokens: 2,
14423 n_prompt: 10,
14424 n_cached: 0,
14425 elapsed_s: 0.1,
14426 spec: None,
14427 })
14428 .unwrap();
14429 drop(tx);
14430 let resp = sse_response(
14431 rx,
14432 "m".into(),
14433 true,
14434 None,
14435 Envelope::new(true),
14436 Vec::new(),
14437 None,
14438 )
14439 .into_response();
14440 let lines = sse_data_lines(resp).await;
14441 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
14442 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14443 .iter()
14444 .map(|l| serde_json::from_str(l).unwrap())
14445 .collect();
14446 let id = chunks[0]["id"].as_str().unwrap().to_string();
14448 assert!(id.starts_with("chatcmpl-"));
14449 for c in &chunks {
14450 assert_eq!(c["id"], id.as_str());
14451 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
14452 let fingerprint = c["system_fingerprint"].as_str().unwrap();
14453 assert!(
14454 build_id::fingerprint_is_well_formed(fingerprint),
14455 "chunk system_fingerprint {fingerprint:?} is not memra-<version>-<12 hex>"
14456 );
14457 assert_eq!(c["object"], "chat.completion.chunk");
14458 }
14459 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
14461 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
14462 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
14463 let fin = chunks.last().unwrap();
14465 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
14466 assert_eq!(fin["usage"]["prompt_tokens"], 10);
14467 }
14468
14469 #[tokio::test]
14470 async fn stream_token_events_equal_usage_on_every_finish_path() {
14471 for (stop_reason, expected_finish) in [
14472 ("Eos", "stop"),
14473 ("Callback", "stop"),
14474 ("MaxNew", "length"),
14475 ("ContextFull", "length"),
14476 ] {
14477 let (tx, rx) = worker::event_channel();
14478 tx.send(Event::Token {
14481 id: 248_046,
14482 text: String::new(),
14483 })
14484 .unwrap();
14485 tx.send(Event::Done {
14486 stop_reason: stop_reason.into(),
14487 n_tokens: 1,
14488 n_prompt: 8,
14489 n_cached: 8,
14490 elapsed_s: 0.1,
14491 spec: None,
14492 })
14493 .unwrap();
14494 drop(tx);
14495
14496 let resp = sse_response(
14497 rx,
14498 "m".into(),
14499 true,
14500 None,
14501 Envelope::new(true),
14502 Vec::new(),
14503 None,
14504 )
14505 .into_response();
14506 let lines = sse_data_lines(resp).await;
14507 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
14508 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
14509 .iter()
14510 .map(|line| serde_json::from_str(line).unwrap())
14511 .collect();
14512 let token_events = chunks
14513 .iter()
14514 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
14515 .count();
14516 let terminal = chunks.last().unwrap();
14517 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
14518 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
14519 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
14520 }
14521 }
14522
14523 #[tokio::test]
14524 async fn stream_excludes_stop_text_like_non_stream_does() {
14525 let (tx, rx) = worker::event_channel();
14529 tx.send(Event::Token {
14530 id: 1,
14531 text: "answer\nPro".into(),
14532 })
14533 .unwrap();
14534 tx.send(Event::Token {
14535 id: 2,
14536 text: "blem: leaked prompt".into(),
14537 })
14538 .unwrap();
14539 tx.send(Event::Done {
14540 stop_reason: "Callback".into(),
14541 n_tokens: 2,
14542 n_prompt: 8,
14543 n_cached: 0,
14544 elapsed_s: 0.1,
14545 spec: None,
14546 })
14547 .unwrap();
14548 drop(tx);
14549 let resp = sse_response(
14550 rx,
14551 "m".into(),
14552 true,
14553 None,
14554 Envelope::new(true),
14555 vec!["Problem:".into()],
14556 None,
14557 )
14558 .into_response();
14559 let lines = sse_data_lines(resp).await;
14560 let content: String = lines
14561 .iter()
14562 .filter(|l| *l != "[DONE]")
14563 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
14564 .filter_map(|c| {
14565 c["choices"][0]["delta"]["content"]
14566 .as_str()
14567 .map(str::to_string)
14568 })
14569 .collect();
14570 assert_eq!(content, "answer\n");
14571
14572 let (tx, rx) = worker::event_channel();
14574 tx.send(Event::Token {
14575 id: 1,
14576 text: "ends in Pro".into(),
14577 })
14578 .unwrap();
14579 tx.send(Event::Done {
14580 stop_reason: "Eos".into(),
14581 n_tokens: 1,
14582 n_prompt: 8,
14583 n_cached: 0,
14584 elapsed_s: 0.1,
14585 spec: None,
14586 })
14587 .unwrap();
14588 drop(tx);
14589 let resp = sse_response(
14590 rx,
14591 "m".into(),
14592 true,
14593 None,
14594 Envelope::new(true),
14595 vec!["Problem:".into()],
14596 None,
14597 )
14598 .into_response();
14599 let lines = sse_data_lines(resp).await;
14600 let content: String = lines
14601 .iter()
14602 .filter(|l| *l != "[DONE]")
14603 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
14604 .filter_map(|c| {
14605 c["choices"][0]["delta"]["content"]
14606 .as_str()
14607 .map(str::to_string)
14608 })
14609 .collect();
14610 assert_eq!(content, "ends in Pro");
14611 }
14612
14613 #[tokio::test]
14614 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
14615 let (tx, rx) = worker::event_channel();
14616 tx.send(Event::Error(worker::EngineError::engine("boom")))
14617 .unwrap();
14618 drop(tx);
14619 let resp = sse_response(
14620 rx,
14621 "m".into(),
14622 true,
14623 None,
14624 Envelope::new(true),
14625 Vec::new(),
14626 None,
14627 )
14628 .into_response();
14629 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14630 .await
14631 .unwrap();
14632 let body = String::from_utf8(bytes.to_vec()).unwrap();
14633 assert!(
14635 !body.contains("event: error"),
14636 "named SSE event leaked: {body}"
14637 );
14638 let lines: Vec<&str> = body
14639 .lines()
14640 .filter_map(|l| l.strip_prefix("data: "))
14641 .collect();
14642 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
14643 assert_eq!(err["error"]["message"], "boom");
14644 assert_eq!(err["error"]["type"], "server_error");
14645 assert_eq!(err["error"]["code"], "engine_error");
14646 assert_eq!(lines.last(), Some(&"[DONE]"));
14647 }
14648
14649 #[test]
14650 fn ttft_sse_marker_ignores_keepalive_comments() {
14651 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
14652 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
14653 assert!(is_sse_data_frame(
14654 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
14655 ));
14656 }
14657
14658 #[tokio::test]
14659 async fn error_bodies_use_the_openai_object_shape() {
14660 let (tx, rx) = worker::event_channel();
14661 tx.send(Event::Error(worker::EngineError::model_not_found(
14662 "unknown model \"x\"",
14663 )))
14664 .unwrap();
14665 drop(tx);
14666 let response =
14667 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
14668 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
14669 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
14670 .await
14671 .unwrap();
14672 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14673 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
14675 assert_eq!(payload["error"]["type"], "invalid_request_error");
14676 assert_eq!(payload["error"]["param"], "model");
14677 assert_eq!(payload["error"]["code"], "model_not_found");
14678 }
14679
14680 fn retry_after(resp: &Response) -> Option<String> {
14687 resp.headers()
14688 .get(axum::http::header::RETRY_AFTER)
14689 .and_then(|v| v.to_str().ok())
14690 .map(str::to_string)
14691 }
14692
14693 async fn body_value(resp: Response) -> serde_json::Value {
14696 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14697 .await
14698 .expect("body");
14699 serde_json::from_slice(&bytes).expect("json body")
14700 }
14701
14702 async fn chat_completion_admitted(st: &AppState, req: serde_json::Value) -> Response {
14718 let mut last_shed = serde_json::Value::Null;
14719 for _ in 0..50 {
14720 let resp = chat_completions(
14721 State(st.clone()),
14722 HeaderMap::new(),
14723 None,
14724 Json(serde_json::from_value(req.clone()).unwrap()),
14725 )
14726 .await;
14727 if resp.status() != StatusCode::TOO_MANY_REQUESTS {
14728 return resp;
14729 }
14730 let body = body_value(resp).await;
14731 let code = body["error"]["code"].as_str().unwrap_or_default();
14732 assert!(
14733 code.starts_with("shed_"),
14734 "only a contention shed may be retried; any other 429 is a finding: {body}"
14735 );
14736 last_shed = body;
14737 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
14738 }
14739 panic!(
14742 "still shed after 50 attempts — either load the retry budget cannot absorb \
14743 or a shed that no longer clears; last refusal: {last_shed}"
14744 );
14745 }
14746
14747 #[test]
14748 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
14749 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
14751 assert_eq!(
14752 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
14753 TIMEOUT_MS_DEFAULT
14754 );
14755 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
14759 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
14760 }
14761 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
14763 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
14764 assert!(err.contains("timeout_ms"), "{err}");
14765 assert!(
14766 err.contains(&TIMEOUT_MS_MIN.to_string())
14767 && err.contains(&TIMEOUT_MS_MAX.to_string()),
14768 "the message must state the range: {err}"
14769 );
14770 assert!(
14771 err.contains("stream"),
14772 "the message must point at streaming for longer work: {err}"
14773 );
14774 }
14775 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
14777 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
14778 assert!(
14779 err.contains("timeout_ms") && err.contains("stream"),
14780 "{err}"
14781 );
14782 }
14783 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
14785 }
14786
14787 #[tokio::test]
14790 #[allow(clippy::await_holding_lock)] async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
14792 let _l = drain_lock();
14793 let st = fake_worker_state();
14794
14795 let comp = completions(
14796 State(st.clone()),
14797 HeaderMap::new(),
14798 None,
14799 Json(
14800 serde_json::from_value(json!({
14801 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
14802 .unwrap(),
14803 ),
14804 )
14805 .await;
14806 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
14807 let chat = chat_completions(
14808 State(st.clone()),
14809 HeaderMap::new(),
14810 None,
14811 Json(
14812 serde_json::from_value(json!({
14813 "model": "m", "messages": [{"role": "user", "content": "t"}],
14814 "timeout_ms": 90_001}))
14815 .unwrap(),
14816 ),
14817 )
14818 .await;
14819 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
14820 let resp_api = responses_api::responses(
14821 State(st.clone()),
14822 HeaderMap::new(),
14823 None,
14824 axum::body::Bytes::from(
14825 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
14826 ),
14827 )
14828 .await;
14829 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
14830 let msgs = anthropic::messages(
14831 State(st.clone()),
14832 HeaderMap::new(),
14833 None,
14834 axum::body::Bytes::from(
14835 json!({"model": "m", "max_tokens": 16,
14836 "messages": [{"role": "user", "content": "t"}],
14837 "timeout_ms": 90_001})
14838 .to_string(),
14839 ),
14840 )
14841 .await;
14842 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
14843
14844 for (surface, resp) in [
14846 ("/v1/completions", comp),
14847 ("/v1/chat/completions", chat),
14848 ("/v1/responses", resp_api),
14849 ] {
14850 let body = body_value(resp).await;
14851 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
14852 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
14853 let m = body["error"]["message"].as_str().unwrap();
14854 assert!(
14855 m.contains("90000") && m.contains("stream"),
14856 "{surface}: {m}"
14857 );
14858 }
14859 let body = body_value(msgs).await;
14861 assert_eq!(body["error"]["type"], "invalid_request_error");
14862 let m = body["error"]["message"].as_str().unwrap();
14863 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
14864 }
14865
14866 #[tokio::test]
14869 #[allow(clippy::await_holding_lock)] async fn a_non_integer_timeout_ms_is_a_named_400() {
14871 let _l = drain_lock();
14872 let st = fake_worker_state();
14873 let resp = chat_completions(
14874 State(st),
14875 HeaderMap::new(),
14876 None,
14877 Json(
14878 serde_json::from_value(json!({
14879 "model": "m", "messages": [{"role": "user", "content": "t"}],
14880 "timeout_ms": "30s"}))
14881 .unwrap(),
14882 ),
14883 )
14884 .await;
14885 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
14886 let body = body_value(resp).await;
14887 assert_eq!(body["error"]["param"], "timeout_ms");
14888 }
14889
14890 #[tokio::test]
14896 #[allow(clippy::await_holding_lock)] async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
14898 let _l = drain_lock();
14899 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
14904 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
14905 let worker_cancel = cancel_seen.clone();
14906 let health = health::WorkerHealth::new();
14907 let h = health.clone();
14908 std::thread::spawn(move || {
14909 h.mark_ready();
14910 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
14911 worker::release_pending_admit();
14912 worker::release_admission_reservation(req.lane);
14913 let _ = req.tx.send(Event::PromptUsage {
14914 n_prompt: 1,
14915 n_cached: 0,
14916 });
14917 let _ = req.tx.send(Event::Token {
14918 id: 1,
14919 text: "partial".into(),
14920 });
14921 for _ in 0..5_000 {
14924 if req.tx.is_closed() {
14925 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
14926 break;
14927 }
14928 std::thread::sleep(std::time::Duration::from_millis(1));
14929 }
14930 }
14931 });
14932 for _ in 0..2_000 {
14933 if health.live().is_ok() {
14934 break;
14935 }
14936 std::thread::sleep(std::time::Duration::from_millis(1));
14937 }
14938 let mut st = fake_worker_state();
14939 st.cmd_tx = cmd_tx;
14940 st.health = health;
14941 let mock = MockMetering::admit_all();
14942 st.metering = Some(mock.clone());
14943
14944 let resp = chat_completion_admitted(
14945 &st,
14946 json!({
14947 "model": "m", "messages": [{"role": "user", "content": "t"}],
14948 "timeout_ms": 1_000}),
14949 )
14950 .await;
14951
14952 assert_eq!(resp.status(), StatusCode::OK);
14957 let body = body_value(resp).await;
14958 assert!(
14959 body["choices"][0]["message"]["content"]
14960 .as_str()
14961 .unwrap()
14962 .contains("partial"),
14963 "the tokens generated before the cut must be delivered: {body}"
14964 );
14965 assert_eq!(body["choices"][0]["finish_reason"], "error");
14969 assert_eq!(
14970 body["choices"][0]["native_finish_reason"],
14971 "deadline_exceeded"
14972 );
14973 assert_eq!(body["error"]["code"], "deadline_exceeded");
14974 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
14975 let message = body["error"]["message"].as_str().unwrap();
14976 assert!(
14977 message.contains("1000") && message.contains("stream"),
14978 "the partial must name the deadline and the streaming alternative: {message}"
14979 );
14980 assert_eq!(body["usage"]["completion_tokens"], 1);
14981
14982 let mut cancelled = false;
14987 for _ in 0..500 {
14988 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
14989 cancelled = true;
14990 break;
14991 }
14992 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
14993 }
14994 assert!(
14995 cancelled,
14996 "the deadline must CANCEL generation (worker's event channel closed)"
14997 );
14998
14999 let events = mock.events();
15004 assert!(
15005 events.contains(&MeterEvent::DeadlinePartial {
15006 prompt: 1,
15007 cached: 0,
15008 completion: 1,
15009 }),
15010 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
15011 );
15012 assert!(
15013 !events
15014 .iter()
15015 .any(|e| matches!(e, MeterEvent::Complete { .. })),
15016 "a deadline cut must stay distinguishable from a full answer: {events:?}"
15017 );
15018 }
15019
15020 #[tokio::test]
15024 #[allow(clippy::await_holding_lock)] async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
15026 let _l = drain_lock();
15027 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15028 let health = health::WorkerHealth::new();
15029 let h = health.clone();
15030 std::thread::spawn(move || {
15031 h.mark_ready();
15032 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15035 worker::release_pending_admit();
15036 worker::release_admission_reservation(req.lane);
15037 let _ = req.tx.send(Event::PromptUsage {
15038 n_prompt: 1,
15039 n_cached: 0,
15040 });
15041 for _ in 0..5_000 {
15042 if req.tx.is_closed() {
15043 break;
15044 }
15045 std::thread::sleep(std::time::Duration::from_millis(1));
15046 }
15047 }
15048 });
15049 for _ in 0..2_000 {
15050 if health.live().is_ok() {
15051 break;
15052 }
15053 std::thread::sleep(std::time::Duration::from_millis(1));
15054 }
15055 let mut st = fake_worker_state();
15056 st.cmd_tx = cmd_tx;
15057 st.health = health;
15058 let mock = MockMetering::admit_all();
15059 st.metering = Some(mock.clone());
15060 let resp = chat_completion_admitted(
15061 &st,
15062 json!({
15063 "model": "m", "messages": [{"role": "user", "content": "t"}],
15064 "timeout_ms": 1_000}),
15065 )
15066 .await;
15067 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15068 assert!(resp.headers().get("x-should-retry").is_none());
15070 assert_eq!(retry_after(&resp), None);
15071 let body = body_value(resp).await;
15072 assert_eq!(body["error"]["code"], "deadline_exceeded");
15073 assert!(
15074 body["error"]["message"]
15075 .as_str()
15076 .unwrap()
15077 .contains("not billed"),
15078 "the zero-token 408 keeps the billing promise: {body}"
15079 );
15080 let events = mock.events();
15081 assert!(
15082 events.contains(&MeterEvent::Unbilled {
15083 outcome: "deadline_exceeded",
15084 status: 408,
15085 code: "deadline_exceeded".into(),
15086 }),
15087 "the named zero-debit census outcome, not the generic reject — every sibling \
15088 deadline path settles this one: {events:?}"
15089 );
15090 }
15091
15092 #[tokio::test]
15095 #[allow(clippy::await_holding_lock)] async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
15097 let _l = drain_lock();
15098 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15100 let health = health::WorkerHealth::new();
15101 let h = health.clone();
15102 std::thread::spawn(move || {
15103 h.mark_ready();
15104 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
15105 worker::release_pending_admit();
15106 worker::release_admission_reservation(req.lane);
15107 let _ = req.tx.send(Event::PromptUsage {
15108 n_prompt: 1,
15109 n_cached: 0,
15110 });
15111 while !req.tx.is_closed() {
15112 std::thread::sleep(std::time::Duration::from_millis(1));
15113 }
15114 }
15115 });
15116 for _ in 0..2_000 {
15117 if health.live().is_ok() {
15118 break;
15119 }
15120 std::thread::sleep(std::time::Duration::from_millis(1));
15121 }
15122 let mut st = fake_worker_state();
15123 st.cmd_tx = cmd_tx;
15124 st.health = health;
15125 let mock = MockMetering::admit_all();
15126 st.metering = Some(mock.clone());
15127
15128 let resp = chat_completion_admitted(
15129 &st,
15130 json!({
15131 "model": "m", "messages": [{"role": "user", "content": "t"}],
15132 "stream": true, "timeout_ms": 1_000}),
15133 )
15134 .await;
15135 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
15138 let body = body_value(resp).await;
15139 assert_eq!(body["error"]["code"], "deadline_exceeded");
15140 assert!(
15141 body["error"]["message"]
15142 .as_str()
15143 .unwrap()
15144 .contains("first token"),
15145 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
15146 );
15147 let events = mock.events();
15148 assert!(
15149 events.contains(&MeterEvent::Unbilled {
15150 outcome: "deadline_exceeded",
15151 status: 408,
15152 code: "deadline_exceeded".into(),
15153 }),
15154 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
15155 );
15156 }
15157
15158 #[tokio::test]
15162 #[allow(clippy::await_holding_lock)] async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
15164 let _l = drain_lock();
15165 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
15168 let mock = MockMetering::admit_all();
15169 st.metering = Some(mock.clone());
15170 let resp = chat_completion_admitted(
15171 &st,
15172 json!({
15173 "model": "m", "messages": [{"role": "user", "content": "t"}],
15174 "stream": true, "timeout_ms": 1_000}),
15175 )
15176 .await;
15177 assert_eq!(
15178 resp.status(),
15179 StatusCode::OK,
15180 "TTFT was met — 200 is correct"
15181 );
15182 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15183 .await
15184 .expect("the stream must run to completion past the deadline");
15185 let text = String::from_utf8(bytes.to_vec()).unwrap();
15186 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
15187 let events = mock.events();
15188 assert!(
15189 events
15190 .iter()
15191 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
15192 "a stream past its deadline after first token still settles as COMPLETE with \
15193 all four tokens: {events:?}"
15194 );
15195 }
15196
15197 fn admission_counters_guard() -> std::sync::MutexGuard<'static, ()> {
15204 static COUNTERS: std::sync::Mutex<()> = std::sync::Mutex::new(());
15205 COUNTERS
15206 .lock()
15207 .unwrap_or_else(|poisoned| poisoned.into_inner())
15208 }
15209
15210 struct CounterRestore<'a>(&'a std::sync::atomic::AtomicUsize, usize);
15216 impl Drop for CounterRestore<'_> {
15217 fn drop(&mut self) {
15218 self.0.store(self.1, std::sync::atomic::Ordering::Release);
15219 }
15220 }
15221
15222 #[allow(clippy::result_large_err)] fn reserve_interactive_through_contention(
15233 st: &AppState,
15234 rl: &RateLimit,
15235 deadline_ms: u64,
15236 ) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
15237 let reserve = || {
15238 reserve_pending_admit(
15239 st,
15240 lanes::Lane::Interactive,
15241 rl,
15242 RequestDeadline::starting_now(deadline_ms),
15243 )
15244 };
15245 let mut g = reserve();
15246 for _ in 0..50 {
15247 match &g {
15248 Ok(_) => break,
15249 Err((_, "shed_deadline")) => {
15250 std::thread::sleep(std::time::Duration::from_millis(10));
15251 g = reserve();
15252 }
15253 Err((_, outcome)) => panic!("unexpected refusal: {outcome}"),
15254 }
15255 }
15256 g
15257 }
15258
15259 #[test]
15262 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
15263 let _counters = admission_counters_guard();
15264 let st = fake_worker_state();
15265 let lane = lanes::Lane::Interactive;
15266 let cap = lane_cap(lane);
15267 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15268 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
15269 let _restore = CounterRestore(counter, prev);
15270 let rl = RateLimit {
15271 limit: cap,
15272 remaining: 0,
15273 reset_s: 1,
15274 };
15275 let (resp, outcome) = reserve_pending_admit(
15276 &st,
15277 lane,
15278 &rl,
15279 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15280 )
15281 .map(|_| ())
15282 .expect_err("a backlog at the bound must shed");
15283 assert_eq!(outcome, "shed_queue");
15284 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15285 assert!(
15286 retry_after(&resp).is_some(),
15287 "a shed must carry Retry-After so the router's spill can act on it"
15288 );
15289 let stamped = rl.attach(resp);
15291 for h in [
15292 "x-ratelimit-limit",
15293 "x-ratelimit-remaining",
15294 "x-ratelimit-reset",
15295 ] {
15296 assert!(stamped.headers().get(h).is_some(), "missing {h}");
15297 }
15298 }
15299
15300 #[test]
15304 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
15305 let _counters = admission_counters_guard();
15306 let st = fake_worker_state();
15307 let lane = lanes::Lane::Interactive;
15308 let cap = lane_cap(lane);
15309 {
15310 let mut m = st.metrics.lock().unwrap();
15311 m.completed = 10;
15312 m.tokens_out = 1_000;
15313 m.step_p50_ms = 10.0; }
15315 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15316 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15318 let rl = RateLimit {
15319 limit: cap,
15320 remaining: 0,
15321 reset_s: 1,
15322 };
15323 let admitted = reserve_pending_admit(
15325 &st,
15326 lane,
15327 &rl,
15328 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15329 );
15330 assert!(
15331 admitted.is_ok(),
15332 "a request whose deadline covers the estimate must be admitted"
15333 );
15334 drop(admitted); let (resp, outcome) = reserve_pending_admit(
15337 &st,
15338 lane,
15339 &rl,
15340 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15341 )
15342 .map(|_| ())
15343 .expect_err("a deadline shorter than the estimated wait must shed");
15344 assert_eq!(outcome, "shed_deadline");
15345 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15346 assert!(retry_after(&resp).is_some());
15347 }
15348
15349 #[test]
15352 fn deadline_shed_is_interactive_only_and_silent_with_free_slots() {
15353 let _counters = admission_counters_guard();
15354 let st = fake_worker_state();
15355 let cap = lane_cap(lanes::Lane::Interactive);
15356 {
15357 let mut m = st.metrics.lock().unwrap();
15358 m.completed = 10;
15359 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
15361 }
15362 let free = RateLimit {
15364 limit: cap,
15365 remaining: 1,
15366 reset_s: 0,
15367 };
15368 let g = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
15373 assert!(
15374 g.is_ok(),
15375 "free capacity must admit regardless of the estimate"
15376 );
15377 drop(g);
15378 let full = RateLimit {
15381 limit: cap,
15382 remaining: 0,
15383 reset_s: 5,
15384 };
15385 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
15386 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15387 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15389 let g = reserve_pending_admit(
15390 &st,
15391 lane,
15392 &full,
15393 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15394 );
15395 assert!(
15396 g.is_ok(),
15397 "{lane:?} must not be deadline-shed by the interactive gate"
15398 );
15399 drop(g);
15400 }
15401 }
15402
15403 #[test]
15410 fn a_saturated_queue_with_free_http_slots_queues_silently_without_a_ceiling() {
15411 let _counters = admission_counters_guard();
15412 let st = fake_worker_state();
15413 let lane = lanes::Lane::Interactive;
15414 let cap = lane_cap(lane);
15415 {
15416 let mut m = st.metrics.lock().unwrap();
15417 m.completed = 10;
15418 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
15421 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15422 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15424 let free = RateLimit {
15427 limit: cap,
15428 remaining: 1,
15429 reset_s: 0,
15430 };
15431 let g = reserve_pending_admit(
15432 &st,
15433 lane,
15434 &free,
15435 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15436 );
15437 assert!(
15438 g.is_ok(),
15439 "flag off: a ~20 s projected wait whose deadline can absorb it queues \
15440 silently (no 429) - the darklanes#5 defect shape, preserved by default"
15441 );
15442 drop(g);
15443 }
15444
15445 #[test]
15450 fn the_queue_wait_ceiling_sheds_with_429_retry_after_and_the_ratelimit_trio() {
15451 let _counters = admission_counters_guard();
15452 let st = fake_worker_state();
15453 let lane = lanes::Lane::Interactive;
15454 let cap = lane_cap(lane);
15455 {
15456 let mut m = st.metrics.lock().unwrap();
15457 m.completed = 10;
15458 m.tokens_out = 1_000; m.step_p50_ms = 100.0; }
15461 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15462 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15464 let free = RateLimit {
15465 limit: cap,
15466 remaining: 1,
15467 reset_s: 0,
15468 };
15469 let (resp, outcome) = reserve_pending_admit_with_ceiling(
15470 &st,
15471 lane,
15472 &free,
15473 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15474 5, )
15476 .map(|_| ())
15477 .expect_err("a projected wait past the ceiling must shed");
15478 assert_eq!(outcome, "shed_queue_wait");
15479 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
15480 assert_eq!(
15481 retry_after(&resp).as_deref(),
15482 Some("20"),
15483 "Retry-After must carry the estimate (~10 s/wave x 2 waves)"
15484 );
15485 assert_eq!(
15486 resp.headers()
15487 .get("retry-after-ms")
15488 .and_then(|v| v.to_str().ok()),
15489 Some("20000"),
15490 "the ms twin must match"
15491 );
15492 let stamped = free.attach(resp);
15493 for h in [
15494 "x-ratelimit-limit",
15495 "x-ratelimit-remaining",
15496 "x-ratelimit-reset",
15497 ] {
15498 assert!(stamped.headers().get(h).is_some(), "missing {h}");
15499 }
15500 }
15501
15502 #[test]
15506 fn the_queue_wait_ceiling_admits_under_it_and_never_touches_dark_lanes() {
15507 let _counters = admission_counters_guard();
15508 let st = fake_worker_state();
15509 let lane = lanes::Lane::Interactive;
15510 let cap = lane_cap(lane);
15511 {
15512 let mut m = st.metrics.lock().unwrap();
15513 m.completed = 10;
15514 m.tokens_out = 1_000;
15515 m.step_p50_ms = 100.0; }
15517 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15518 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel);
15519 let _restore = CounterRestore(counter, prev);
15520 let free = RateLimit {
15521 limit: cap,
15522 remaining: 1,
15523 reset_s: 0,
15524 };
15525 let g = reserve_pending_admit_with_ceiling(
15526 &st,
15527 lane,
15528 &free,
15529 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15530 60, );
15532 assert!(
15533 g.is_ok(),
15534 "an estimate under the ceiling must admit and queue as before"
15535 );
15536 drop(g);
15537 let full = RateLimit {
15539 limit: cap,
15540 remaining: 0,
15541 reset_s: 5,
15542 };
15543 for dark in [lanes::Lane::Judge, lanes::Lane::Harvest] {
15544 let counter = &worker::ADMISSION_RESERVATIONS[dark.idx()];
15545 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let _restore = CounterRestore(counter, prev);
15547 let g = reserve_pending_admit_with_ceiling(
15548 &st,
15549 dark,
15550 &full,
15551 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15552 1,
15553 );
15554 assert!(
15555 g.is_ok(),
15556 "{dark:?} must not be shed by the interactive queue-wait ceiling"
15557 );
15558 drop(g);
15559 }
15560 }
15561
15562 #[test]
15566 fn the_queue_wait_ceiling_leaves_the_existing_shed_arms_first_and_unchanged() {
15567 let _counters = admission_counters_guard();
15568 let st = fake_worker_state();
15569 let lane = lanes::Lane::Interactive;
15570 let cap = lane_cap(lane);
15571 {
15572 let mut m = st.metrics.lock().unwrap();
15573 m.completed = 10;
15574 m.tokens_out = 1_000;
15575 m.step_p50_ms = 100.0;
15576 }
15577 let rl = RateLimit {
15578 limit: cap,
15579 remaining: 0,
15580 reset_s: 1,
15581 };
15582 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
15583 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
15585 let _restore = CounterRestore(counter, prev);
15586 assert!(matches!(
15587 reserve_pending_admit_with_ceiling(
15588 &st,
15589 lane,
15590 &rl,
15591 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15592 1,
15593 ),
15594 Err((_, "shed_queue"))
15595 ));
15596 counter.store(cap, std::sync::atomic::Ordering::Release);
15598 assert!(matches!(
15599 reserve_pending_admit_with_ceiling(
15600 &st,
15601 lane,
15602 &rl,
15603 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
15604 1,
15605 ),
15606 Err((_, "shed_deadline"))
15607 ));
15608 }
15609
15610 #[test]
15615 fn the_queue_wait_ceiling_is_wired_through_the_production_wrapper() {
15616 let src = include_str!("lib.rs");
15617 let code: String = src
15618 .lines()
15619 .map(|l| l.split("//").next().unwrap_or(""))
15620 .collect::<Vec<_>>()
15621 .join("\n");
15622 let start = code
15623 .find("pub(crate) fn reserve_pending_admit(")
15624 .expect("the production wrapper exists");
15625 let rest = &code[start..];
15626 let end = rest.find("\nfn ").unwrap_or(rest.len());
15627 let wrapper = &rest[..end];
15628 assert!(
15629 wrapper.contains(
15630 "reserve_pending_admit_with_ceiling(st, lane, rl, deadline, queue_wait_ceiling_s())"
15631 ),
15632 "every production ingress must judge the ceiling the env read armed"
15633 );
15634 }
15635
15636 #[test]
15637 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
15638 let _counters = admission_counters_guard();
15639 let st = fake_worker_state();
15640 let cap = lane_cap(lanes::Lane::Interactive);
15641 let bound = max_queue_depth(cap);
15642 assert!(bound > 0, "the queue bound must admit at least one request");
15643 let rl = RateLimit {
15644 limit: cap,
15645 remaining: 0,
15646 reset_s: 1,
15647 };
15648 let _ = worker::PENDING_ADMITS.fetch_update(
15649 std::sync::atomic::Ordering::AcqRel,
15650 std::sync::atomic::Ordering::Acquire,
15651 |_| Some(0),
15652 );
15653 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
15654 let _restore = CounterRestore(counter, 0);
15655 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
15656 let guard = reserve_pending_admit(
15657 &st,
15658 lanes::Lane::Interactive,
15659 &rl,
15660 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15661 )
15662 .expect("the final queue slot should be reservable");
15663 assert_eq!(
15664 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
15665 1
15666 );
15667 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
15668 drop(guard);
15669 assert_eq!(
15670 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
15671 0
15672 );
15673 assert_eq!(
15674 counter.load(std::sync::atomic::Ordering::Acquire),
15675 bound - 1
15676 );
15677
15678 counter.store(bound, std::sync::atomic::Ordering::Release);
15679 let rejected = reserve_pending_admit(
15680 &st,
15681 lanes::Lane::Interactive,
15682 &rl,
15683 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15684 );
15685 assert!(matches!(rejected, Err((_, "shed_queue"))));
15686 }
15687
15688 #[test]
15689 fn admission_reservations_are_lane_scoped() {
15690 let _counters = admission_counters_guard();
15691 let st = fake_worker_state();
15692 let harvest = lanes::Lane::Harvest;
15693 let interactive = lanes::Lane::Interactive;
15694 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
15695 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
15696 let _restore = CounterRestore(harvest_counter, 0);
15697 harvest_counter.store(
15698 max_queue_depth(lane_cap(harvest)),
15699 std::sync::atomic::Ordering::Release,
15700 );
15701 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
15702 let free = RateLimit {
15703 limit: lane_cap(interactive),
15704 remaining: 1,
15705 reset_s: 0,
15706 };
15707 let guard = reserve_pending_admit(
15718 &st,
15719 interactive,
15720 &free,
15721 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
15722 )
15723 .expect("a full harvest queue must not consume interactive capacity");
15724 drop(guard);
15725 let tight = reserve_interactive_through_contention(&st, &free, TIMEOUT_MS_MIN);
15726 assert!(
15727 tight.is_ok(),
15728 "a full harvest queue must not deadline-shed a tight interactive request \
15729 (a backlog that outlasts the retry budget here is a cross-lane leak, not \
15730 contention)"
15731 );
15732 drop(tight);
15733 let harvest_rl = RateLimit {
15734 limit: lane_cap(harvest),
15735 remaining: 0,
15736 reset_s: 1,
15737 };
15738 assert!(matches!(
15739 reserve_pending_admit(
15740 &st,
15741 harvest,
15742 &harvest_rl,
15743 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
15744 ),
15745 Err((_, "shed_queue"))
15746 ));
15747 }
15748
15749 #[test]
15750 fn taxonomy_maps_every_class_to_its_status_and_code() {
15751 use worker::{EngineError as E, ErrClass as C};
15752 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
15753 (
15754 E::invalid_param("bad json", "response_format"),
15755 StatusCode::BAD_REQUEST,
15756 "invalid_request_error",
15757 "",
15758 ),
15759 (
15760 E::context_length("prompt (9000 tok) >= context cap (8192)"),
15761 StatusCode::BAD_REQUEST,
15762 "invalid_request_error",
15763 "context_length_exceeded",
15764 ),
15765 (
15766 E::model_not_found("unknown model \"nope\""),
15767 StatusCode::BAD_REQUEST,
15768 "invalid_request_error",
15769 "model_not_found",
15770 ),
15771 (
15772 E::rate_limit("lane judge is at capacity, retry"),
15773 StatusCode::TOO_MANY_REQUESTS,
15774 "rate_limit_error",
15775 "rate_limit_exceeded",
15776 ),
15777 (
15778 E::overloaded("no VRAM for a new session"),
15779 StatusCode::SERVICE_UNAVAILABLE,
15780 "server_error",
15781 "overloaded",
15782 ),
15783 (
15784 E::engine("graph step failed: launch error"),
15785 StatusCode::INTERNAL_SERVER_ERROR,
15786 "server_error",
15787 "engine_error",
15788 ),
15789 ];
15790 for (err, want_status, want_type, want_code) in cases {
15791 let (status, etype, code) = class_http(err.class);
15792 assert_eq!(status, want_status, "{:?}", err);
15793 assert_eq!(etype, want_type, "{:?}", err);
15794 if !want_code.is_empty() {
15795 assert_eq!(code, Some(want_code), "{:?}", err);
15796 }
15797 let body = engine_error_body(&err);
15799 assert_eq!(body["error"]["message"], err.message);
15800 assert_eq!(body["error"]["type"], want_type);
15801 }
15802 for c in [
15804 C::InvalidRequest,
15805 C::ContextLength,
15806 C::ModelNotFound,
15807 C::RateLimit,
15808 C::Overloaded,
15809 C::Engine,
15810 ] {
15811 let (s, t, _) = class_http(c);
15812 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
15813 assert!(!t.is_empty());
15814 }
15815 }
15816
15817 #[test]
15818 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
15819 let e = worker::EngineError::engine(
15824 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
15825 );
15826 let resp = engine_error_response(&e);
15827 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
15828 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
15829 }
15830
15831 #[test]
15832 fn retry_headers_follow_the_sdk_contract() {
15833 for e in [
15837 worker::EngineError::rate_limit("shed"),
15838 worker::EngineError::overloaded("no VRAM"),
15839 ] {
15840 let resp = engine_error_response(&e);
15841 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
15842 let secs: u64 = ra
15843 .parse()
15844 .expect("Retry-After must be integer delay-seconds");
15845 assert!(
15846 secs > 0 && secs <= 60,
15847 "Retry-After {secs}s outside the honored window"
15848 );
15849 let ms = resp
15850 .headers()
15851 .get("retry-after-ms")
15852 .unwrap()
15853 .to_str()
15854 .unwrap();
15855 assert_eq!(
15856 ms.parse::<u64>().unwrap(),
15857 secs * 1000,
15858 "the two headers disagree"
15859 );
15860 assert!(
15861 resp.headers().get("x-should-retry").is_none(),
15862 "a retryable class must not say x-should-retry: false"
15863 );
15864 }
15865 }
15866
15867 #[tokio::test]
15874 async fn admit_predict_reject_matches_shed_contract() {
15875 let shed = retry_contract_response(
15877 (
15878 StatusCode::TOO_MANY_REQUESTS,
15879 Json(error_body(
15880 "interactive queue is at its bound",
15881 "rate_limit_error",
15882 None,
15883 Some("shed_queue"),
15884 )),
15885 )
15886 .into_response(),
15887 Some(7),
15888 );
15889 let predict = engine_error_response(&worker::EngineError::rate_limit_after(
15892 "predicted KV-to-completion exceeds the box budget; retry",
15893 7,
15894 ));
15895 assert_eq!(shed.status(), predict.status());
15896 for header in ["retry-after", "retry-after-ms"] {
15897 assert_eq!(
15898 shed.headers().get(header),
15899 predict.headers().get(header),
15900 "header {header} must be byte-identical to the shed contract"
15901 );
15902 }
15903 let shed_body: serde_json::Value = serde_json::from_slice(
15904 &axum::body::to_bytes(shed.into_body(), usize::MAX)
15905 .await
15906 .unwrap(),
15907 )
15908 .unwrap();
15909 let predict_body: serde_json::Value = serde_json::from_slice(
15910 &axum::body::to_bytes(predict.into_body(), usize::MAX)
15911 .await
15912 .unwrap(),
15913 )
15914 .unwrap();
15915 assert_eq!(shed_body["error"]["type"], predict_body["error"]["type"]);
15916 assert_eq!(predict_body["error"]["type"], "rate_limit_error");
15917 let shed_keys: Vec<&String> = shed_body["error"].as_object().unwrap().keys().collect();
15918 let predict_keys: Vec<&String> =
15919 predict_body["error"].as_object().unwrap().keys().collect();
15920 assert_eq!(shed_keys, predict_keys, "same body schema, key for key");
15921 assert_eq!(predict_body["error"]["code"], "rate_limit_exceeded");
15922
15923 let clamped = engine_error_response(&worker::EngineError::rate_limit_after("m", 400));
15925 assert_eq!(retry_after(&clamped).as_deref(), Some("60"));
15926 let plain = engine_error_response(&worker::EngineError::rate_limit("m"));
15928 assert_eq!(retry_after(&plain).as_deref(), Some("2"));
15929 }
15930
15931 #[tokio::test]
15932 #[allow(clippy::await_holding_lock)] async fn command_send_failure_obeys_the_retry_contract() {
15934 let _l = drain_lock();
15935 let mut st = fake_worker_state();
15936 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15937 drop(cmd_rx);
15938 st.cmd_tx = cmd_tx;
15939
15940 let completion = completions(
15941 State(st.clone()),
15942 axum::http::HeaderMap::new(),
15943 None,
15944 Json(
15945 serde_json::from_value(serde_json::json!({
15946 "model": "m", "prompt": "test"
15947 }))
15948 .unwrap(),
15949 ),
15950 )
15951 .await;
15952 let chat = chat_completions(
15953 State(st),
15954 axum::http::HeaderMap::new(),
15955 None,
15956 Json(
15957 serde_json::from_value(serde_json::json!({
15958 "model": "m", "messages": [{"role": "user", "content": "test"}]
15959 }))
15960 .unwrap(),
15961 ),
15962 )
15963 .await;
15964
15965 for resp in [completion, chat] {
15966 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
15967 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
15968 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
15969 assert_ne!(
15970 resp.headers()
15971 .get("x-should-retry")
15972 .and_then(|v| v.to_str().ok()),
15973 Some("false")
15974 );
15975 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
15976 .await
15977 .unwrap();
15978 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15979 assert_eq!(payload["error"]["type"], "server_error");
15980 assert_eq!(payload["error"]["code"], "overloaded");
15981 }
15982 }
15983
15984 #[test]
15985 fn unfixable_client_errors_say_x_should_retry_false() {
15986 for e in [
15989 worker::EngineError::model_not_found("unknown model \"x\""),
15990 worker::EngineError::context_length("prompt too long"),
15991 worker::EngineError::invalid_param("bad", "messages"),
15992 ] {
15993 let resp = engine_error_response(&e);
15994 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
15995 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
15996 assert!(
15997 retry_after(&resp).is_none(),
15998 "a 400 must not promise a retry window"
15999 );
16000 }
16001 }
16002
16003 #[tokio::test]
16004 async fn a_closed_worker_channel_is_503_not_500() {
16005 let (tx, rx) = worker::event_channel();
16009 drop(tx);
16010 let resp =
16011 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
16012 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16013 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
16014 }
16015
16016 #[tokio::test]
16017 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
16018 let (tx, rx) = worker::event_channel();
16021 tx.send(Event::Error(worker::EngineError::rate_limit(
16022 "lane judge shed: interactive p99 over budget, retry",
16023 )))
16024 .unwrap();
16025 let (resp, error_code) = peek_admission(rx)
16026 .await
16027 .expect_err("a shed must not be forwarded into the stream");
16028 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16029 assert_eq!(error_code, "rate_limit_exceeded");
16030 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16031 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16032 .await
16033 .unwrap();
16034 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16035 assert!(
16036 payload["error"].is_object(),
16037 "bare-string error body: {payload}"
16038 );
16039 assert_eq!(payload["error"]["type"], "rate_limit_error");
16040 assert!(
16041 payload["error"]["message"]
16042 .as_str()
16043 .unwrap()
16044 .contains("shed")
16045 );
16046 }
16047
16048 #[tokio::test]
16049 async fn interactive_admission_error_is_a_preheader_429() {
16050 let (tx, rx) = worker::event_channel();
16053 tx.send(Event::Error(worker::EngineError::rate_limit(
16054 "KV capacity unavailable",
16055 )))
16056 .unwrap();
16057 let (resp, error_code) = peek_admission(rx)
16058 .await
16059 .expect_err("admission error must stay pre-header");
16060 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
16061 assert_eq!(error_code, "rate_limit_exceeded");
16062 }
16063
16064 #[tokio::test]
16065 async fn admission_peek_preserves_context_error_for_the_ledger() {
16066 let (tx, rx) = worker::event_channel();
16067 tx.send(Event::Error(worker::EngineError::context_length(
16068 "prompt exceeds configured model maximum",
16069 )))
16070 .unwrap();
16071 let (resp, error_code) = peek_admission(rx)
16072 .await
16073 .expect_err("context rejection must stay pre-header");
16074 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16075 assert_eq!(error_code, "context_length_exceeded");
16076 }
16077
16078 #[tokio::test]
16079 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
16080 let (tx, rx) = worker::event_channel();
16081 tx.send(Event::PromptUsage {
16082 n_prompt: 262_143,
16083 n_cached: 0,
16084 })
16085 .unwrap();
16086 let mut replay = peek_admission(rx).await.expect("successful admission");
16087 assert!(matches!(
16088 replay.recv().await,
16089 Some(Event::PromptUsage {
16090 n_prompt: 262_143,
16091 n_cached: 0
16092 }),
16093 ));
16094 }
16095
16096 #[test]
16097 fn penalties_plumb_from_http_to_sampler_config() {
16098 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
16101 "model": "m", "messages": [{"role": "user", "content": "task"}],
16102 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
16103 }))
16104 .unwrap();
16105 let (tx, _rx) = worker::event_channel();
16106 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16107 .unwrap()
16108 .request
16109 .sampler_cfg;
16110 assert_eq!(cfg.penalty_freq, 0.5);
16111 assert_eq!(cfg.penalty_present, 0.25);
16112 assert_eq!(cfg.penalty_repeat, 1.1);
16113 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16114
16115 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16116 "model": "m", "prompt": "task", "frequency_penalty": 1.5
16117 }))
16118 .unwrap();
16119 let (tx, _rx) = worker::event_channel();
16120 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16121 assert_eq!(cfg.penalty_freq, 1.5);
16122 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16123
16124 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16126 "model": "m", "prompt": "task"
16127 }))
16128 .unwrap();
16129 let (tx, _rx) = worker::event_channel();
16130 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16131 assert_eq!(cfg.penalty_last_n, 0);
16132 assert_eq!(cfg.penalty_repeat, 1.0);
16133 }
16134
16135 #[test]
16136 fn omitted_temperature_is_openai_default_not_greedy() {
16137 let chat_temp = |body: serde_json::Value| {
16153 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16154 let (tx, _rx) = worker::event_channel();
16155 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16156 .unwrap()
16157 .request
16158 .sampler_cfg
16159 .temperature
16160 };
16161 let comp_temp = |body: serde_json::Value| {
16162 let req: CompletionReq = serde_json::from_value(body).unwrap();
16163 let (tx, _rx) = worker::event_channel();
16164 build_request(&req, tx, lanes::Lane::Interactive, None)
16165 .sampler_cfg
16166 .temperature
16167 };
16168
16169 assert_eq!(
16171 chat_temp(serde_json::json!({
16172 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
16173 1.0,
16174 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
16175 );
16176 assert_eq!(
16177 comp_temp(serde_json::json!({
16178 "model": "m", "prompt": "t"})),
16179 1.0,
16180 "omitted completions temperature must be the OpenAI 1.0 default"
16181 );
16182
16183 assert_eq!(
16185 chat_temp(serde_json::json!({
16186 "model": "m", "messages": [{"role": "user", "content": "t"}],
16187 "temperature": 0.0})),
16188 0.0,
16189 "explicit temperature 0 must stay greedy"
16190 );
16191 assert_eq!(
16192 comp_temp(serde_json::json!({
16193 "model": "m", "prompt": "t", "temperature": 0})),
16194 0.0,
16195 "explicit temperature 0 must stay greedy"
16196 );
16197 assert!(
16199 memra_engine::sampler::Sampler::new(sampler_config(
16200 0.0,
16201 0,
16202 1.0,
16203 0.0,
16204 0.0,
16205 0.0,
16206 1.0,
16207 Some(0)
16208 ))
16209 .is_greedy()
16210 );
16211 assert!(
16212 !memra_engine::sampler::Sampler::new(sampler_config(
16213 1.0,
16214 0,
16215 1.0,
16216 0.0,
16217 0.0,
16218 0.0,
16219 1.0,
16220 Some(0)
16221 ))
16222 .is_greedy()
16223 );
16224
16225 assert_eq!(
16227 chat_temp(serde_json::json!({
16228 "model": "m", "messages": [{"role": "user", "content": "t"}],
16229 "temperature": 0.7})),
16230 0.7
16231 );
16232
16233 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16237 "model": "m", "prompt": "t"}))
16238 .unwrap();
16239 let (tx, _rx) = worker::event_channel();
16240 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
16241 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
16242 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
16243 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
16244 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
16245 assert!(
16250 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
16251 "the omitted-temperature default must ride sampled spec's pure-temp regime"
16252 );
16253 }
16254
16255 #[test]
16256 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
16257 let caps = ModelCaps {
16258 chat_temperature_default: Some(0.5),
16259 chat_top_p_default: Some(0.9),
16260 chat_ok: true,
16261 ..Default::default()
16262 };
16263 let cfg = |extra: serde_json::Value| {
16264 let mut body = serde_json::json!({
16265 "model": "step35",
16266 "messages": [{"role": "user", "content": "task"}]
16267 });
16268 body.as_object_mut()
16269 .unwrap()
16270 .extend(extra.as_object().unwrap().clone());
16271 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16272 let (tx, _rx) = worker::event_channel();
16273 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
16274 .unwrap()
16275 .request
16276 .sampler_cfg
16277 };
16278
16279 let omitted = cfg(serde_json::json!({}));
16280 assert_eq!(omitted.temperature, 0.5);
16281 assert_eq!(omitted.top_p, 0.9);
16282
16283 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
16284 assert_eq!(explicit_temp.temperature, 0.7);
16285 assert_eq!(
16286 explicit_temp.top_p, 0.9,
16287 "omitting top_p must retain StepFun's nucleus default"
16288 );
16289
16290 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
16291 assert_eq!(
16292 explicit.temperature, 0.0,
16293 "explicit greedy must remain authoritative"
16294 );
16295 assert_eq!(
16296 explicit.top_p, 1.0,
16297 "explicit untruncated sampling must remain authoritative"
16298 );
16299 }
16300
16301 fn qwen38_vendor_defaults() -> SamplingDefaults {
16305 SamplingDefaults {
16306 temperature: Some(1.0),
16307 top_p: Some(0.95),
16308 top_k: Some(20),
16309 min_p: Some(0.0),
16310 presence_penalty: Some(0.0),
16311 repetition_penalty: Some(1.0),
16312 frequency_penalty: None,
16313 }
16314 }
16315
16316 fn gemma4_vendor_defaults() -> SamplingDefaults {
16321 SamplingDefaults {
16322 temperature: Some(1.0),
16323 top_p: Some(0.95),
16324 top_k: Some(64),
16325 ..Default::default()
16326 }
16327 }
16328
16329 #[test]
16330 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
16331 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
16336 let chat = |extra: serde_json::Value| {
16337 let mut body = serde_json::json!({
16338 "model": "google/gemma-4-31b-it",
16339 "messages": [{"role": "user", "content": "task"}],
16340 "seed": 7
16342 });
16343 body.as_object_mut()
16344 .unwrap()
16345 .extend(extra.as_object().unwrap().clone());
16346 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16347 let (tx, _rx) = worker::event_channel();
16348 build_chat_request_with_trace(
16349 req,
16350 Some(&ModelCaps {
16351 chat_ok: true,
16352 ..Default::default()
16353 }),
16354 tx,
16355 lanes::Lane::Interactive,
16356 None,
16357 None,
16358 None,
16359 &d,
16360 )
16361 .unwrap()
16362 .request
16363 .sampler_cfg
16364 };
16365
16366 let omitted = chat(serde_json::json!({}));
16368 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
16369 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
16370 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
16371 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
16373 assert_eq!(omitted.penalty_repeat, 1.0);
16374 assert_eq!(omitted.penalty_freq, 0.0);
16375 assert_eq!(omitted.penalty_present, 0.0);
16376 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
16377 assert!(
16378 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
16379 "the vendor default must NOT be greedy — that is the whole point of the lane"
16380 );
16381
16382 let greedy = chat(serde_json::json!({"temperature": 0}));
16385 assert_eq!(
16386 greedy.temperature, 0.0,
16387 "explicit temperature 0 stays greedy"
16388 );
16389 assert!(
16390 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
16391 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
16392 spec/graph exactness arms"
16393 );
16394
16395 let one_field = chat(serde_json::json!({"top_k": 3}));
16397 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
16398 assert_eq!(
16399 one_field.temperature, 1.0,
16400 "omitting temperature still takes the vendor value"
16401 );
16402 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
16403
16404 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
16407 assert_eq!(
16408 disabled.top_k, 0,
16409 "an explicit top_k 0 means KEEP ALL, not 'unset'"
16410 );
16411 assert_eq!(
16412 disabled.top_p, 1.0,
16413 "an explicit top_p 1.0 means untruncated"
16414 );
16415
16416 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
16418 assert_eq!(penal.penalty_present, 1.5);
16419 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
16420 }
16421
16422 #[test]
16423 fn vendor_sampling_defaults_are_identical_on_every_surface() {
16424 let d = qwen38_vendor_defaults();
16436 let md = ModelSamplingDefaults::single(d);
16437 let comp = |extra: serde_json::Value| {
16438 let mut body = serde_json::json!({
16439 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
16440 body.as_object_mut()
16441 .unwrap()
16442 .extend(extra.as_object().unwrap().clone());
16443 let req: CompletionReq = serde_json::from_value(body).unwrap();
16444 let (tx, _rx) = worker::event_channel();
16445 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
16446 };
16447 let chat = |extra: serde_json::Value| {
16448 let mut body = serde_json::json!({
16449 "model": "qwen/qwen3.8-27b",
16450 "messages": [{"role": "user", "content": "task"}],
16451 "seed": 11 });
16452 body.as_object_mut()
16453 .unwrap()
16454 .extend(extra.as_object().unwrap().clone());
16455 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16456 let (tx, _rx) = worker::event_channel();
16457 build_chat_request_with_trace(
16458 req,
16459 Some(&ModelCaps {
16460 chat_ok: true,
16461 ..Default::default()
16462 }),
16463 tx,
16464 lanes::Lane::Interactive,
16465 None,
16466 None,
16467 None,
16468 &md,
16469 )
16470 .unwrap()
16471 .request
16472 .sampler_cfg
16473 };
16474
16475 for extra in [
16476 serde_json::json!({}),
16477 serde_json::json!({"temperature": 0}),
16478 serde_json::json!({"temperature": 0.0}),
16479 serde_json::json!({"temperature": 0.7}),
16480 serde_json::json!({"top_p": 1.0}),
16481 serde_json::json!({"top_k": 0}),
16482 serde_json::json!({"min_p": 0.05}),
16483 serde_json::json!({"repetition_penalty": 1.1}),
16484 serde_json::json!({"frequency_penalty": 0.5}),
16485 serde_json::json!({"presence_penalty": 1.5}),
16486 serde_json::json!({
16487 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
16488 "frequency_penalty": 0.1, "presence_penalty": 0.2,
16489 "repetition_penalty": 1.05 }),
16490 ] {
16491 let c = comp(extra.clone());
16492 let h = chat(extra.clone());
16493 assert_eq!(
16494 (
16495 c.temperature,
16496 c.top_p,
16497 c.top_k,
16498 c.min_p,
16499 c.penalty_repeat,
16500 c.penalty_freq,
16501 c.penalty_present,
16502 c.penalty_last_n,
16503 c.seed
16504 ),
16505 (
16506 h.temperature,
16507 h.top_p,
16508 h.top_k,
16509 h.min_p,
16510 h.penalty_repeat,
16511 h.penalty_freq,
16512 h.penalty_present,
16513 h.penalty_last_n,
16514 h.seed
16515 ),
16516 "/v1/completions and /v1/chat/completions disagree on {extra} — \
16517 standard-surface-law violation"
16518 );
16519 }
16520
16521 let omitted = comp(serde_json::json!({}));
16523 assert_eq!(
16524 omitted.temperature, 1.0,
16525 "qwen3.8 card thinking temperature"
16526 );
16527 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
16528 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
16529 assert!(
16531 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
16532 .is_greedy()
16533 );
16534 }
16535
16536 #[tokio::test]
16549 #[allow(clippy::await_holding_lock)] async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
16551 let _l = drain_lock();
16552 let step_caps = ModelCaps {
16553 chat_ok: true,
16554 chat_temperature_default: Some(0.5),
16555 chat_top_p_default: Some(0.9),
16556 ..Default::default()
16557 };
16558 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
16559 let st = fake_worker_state_full(
16560 1,
16561 std::time::Duration::ZERO,
16562 HashMap::from([("m".to_string(), step_caps)]),
16563 Some(cfg_tx),
16564 );
16565 let fields = |saw: &WorkerSaw| {
16570 let c = &saw.sampler_cfg;
16571 (
16572 c.temperature,
16573 c.top_p,
16574 c.top_k,
16575 c.min_p,
16576 c.penalty_repeat,
16577 c.penalty_freq,
16578 c.penalty_present,
16579 c.penalty_last_n,
16580 )
16581 };
16582 let worker_saw = |surface: &str| {
16583 cfg_rx
16584 .recv_timeout(std::time::Duration::from_secs(10))
16585 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
16586 };
16587
16588 let resp = completions(
16589 State(st.clone()),
16590 axum::http::HeaderMap::new(),
16591 None,
16592 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
16593 )
16594 .await;
16595 assert_eq!(
16596 resp.status(),
16597 StatusCode::OK,
16598 "/v1/completions rejected the omitted-sampling request"
16599 );
16600 let comp = worker_saw("/v1/completions");
16601
16602 let resp = chat_completions(
16603 State(st.clone()),
16604 axum::http::HeaderMap::new(),
16605 None,
16606 Json(
16607 serde_json::from_value(serde_json::json!({
16608 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
16609 .unwrap(),
16610 ),
16611 )
16612 .await;
16613 assert_eq!(
16614 resp.status(),
16615 StatusCode::OK,
16616 "/v1/chat/completions rejected the omitted-sampling request"
16617 );
16618 let chat = worker_saw("/v1/chat/completions");
16619
16620 let resp = anthropic::messages(
16621 State(st.clone()),
16622 axum::http::HeaderMap::new(),
16623 None,
16624 axum::body::Bytes::from(
16625 serde_json::json!({
16626 "model": "m", "max_tokens": 16,
16627 "messages": [{"role": "user", "content": "t"}]})
16628 .to_string(),
16629 ),
16630 )
16631 .await;
16632 assert_eq!(
16633 resp.status(),
16634 StatusCode::OK,
16635 "/v1/messages rejected the omitted-sampling request"
16636 );
16637 let msg = worker_saw("/v1/messages");
16638
16639 let resp = responses_api::responses(
16640 State(st.clone()),
16641 axum::http::HeaderMap::new(),
16642 None,
16643 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
16644 )
16645 .await;
16646 assert_eq!(
16647 resp.status(),
16648 StatusCode::OK,
16649 "/v1/responses rejected the omitted-sampling request"
16650 );
16651 let rsp = worker_saw("/v1/responses");
16652
16653 for (surface, cfg) in [
16654 ("/v1/completions", &comp),
16655 ("/v1/messages", &msg),
16656 ("/v1/responses", &rsp),
16657 ] {
16658 assert_eq!(
16659 fields(cfg),
16660 fields(&chat),
16661 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
16662 for the same omitted-sampling request — standard-surface-law violation \
16663 (hermes d991b51699218285)"
16664 );
16665 }
16666 assert_eq!(
16669 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
16670 (0.5, 0.9),
16671 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
16672 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
16673 );
16674 }
16675
16676 #[tokio::test]
16686 #[allow(clippy::await_holding_lock)] async fn same_effort_value_resolves_identically_on_every_surface() {
16688 let _l = drain_lock();
16689 let caps = ModelCaps {
16692 chat_ok: true,
16693 effort_levels: true,
16694 ..Default::default()
16695 };
16696 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
16697 let st = fake_worker_state_full(
16698 1,
16699 std::time::Duration::ZERO,
16700 HashMap::from([("m".to_string(), caps)]),
16701 Some(saw_tx),
16702 );
16703 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
16704 match surface {
16705 "/v1/chat/completions" => {
16706 chat_completions(
16707 State(st),
16708 axum::http::HeaderMap::new(),
16709 None,
16710 Json(
16711 serde_json::from_value(serde_json::json!({
16712 "model": "m", "max_tokens": 8,
16713 "reasoning_effort": effort,
16714 "messages": [{"role": "user", "content": "t"}]}))
16715 .unwrap(),
16716 ),
16717 )
16718 .await
16719 }
16720 "/v1/responses" => {
16721 responses_api::responses(
16722 State(st),
16723 axum::http::HeaderMap::new(),
16724 None,
16725 axum::body::Bytes::from(
16726 serde_json::json!({
16727 "model": "m", "max_output_tokens": 8, "input": "t",
16728 "reasoning": {"effort": effort}})
16729 .to_string(),
16730 ),
16731 )
16732 .await
16733 }
16734 "/v1/messages" => {
16735 anthropic::messages(
16736 State(st),
16737 axum::http::HeaderMap::new(),
16738 None,
16739 axum::body::Bytes::from(
16740 serde_json::json!({
16741 "model": "m", "max_tokens": 8,
16742 "messages": [{"role": "user", "content": "t"}],
16743 "output_config": {"effort": effort}})
16744 .to_string(),
16745 ),
16746 )
16747 .await
16748 }
16749 other => panic!("unknown surface {other}"),
16750 }
16751 };
16752 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
16753
16754 for (effort, want_think, want_level) in [
16757 ("none", ThinkMode::NoThink, Some("low")),
16758 ("minimal", ThinkMode::NoThink, Some("low")),
16759 ("low", ThinkMode::Think, Some("low")),
16760 ("medium", ThinkMode::Think, Some("medium")),
16761 ("high", ThinkMode::Think, Some("high")),
16762 ("xhigh", ThinkMode::Think, Some("high")),
16764 ] {
16765 for surface in SURFACES {
16766 let resp = send(st.clone(), surface, effort).await;
16767 assert_eq!(
16768 resp.status(),
16769 StatusCode::OK,
16770 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
16771 diverged again (issue #31)"
16772 );
16773 let saw = saw_rx
16774 .recv_timeout(std::time::Duration::from_secs(10))
16775 .unwrap_or_else(|_| {
16776 panic!("{surface}: effort {effort:?} request never reached the worker")
16777 });
16778 assert_eq!(
16779 (saw.think, saw.reasoning_effort.as_deref()),
16780 (want_think, want_level),
16781 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
16782 reasoning surface — the parameter was dropped or remapped before \
16783 parse_think (issue #31 regression)"
16784 );
16785 }
16786 }
16787
16788 for effort in ["bogus", "banana", ""] {
16791 for surface in SURFACES {
16792 let resp = send(st.clone(), surface, effort).await;
16793 assert_eq!(
16794 resp.status(),
16795 StatusCode::BAD_REQUEST,
16796 "{surface} accepted effort {effort:?} — silent-accept regression \
16797 (issue #31: the value never reached parse_think's allowlist)"
16798 );
16799 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
16801 .await
16802 .unwrap();
16803 let v: serde_json::Value = serde_json::from_slice(&body)
16804 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
16805 match surface {
16806 "/v1/messages" => {
16807 assert_eq!(v["type"], "error", "{surface} error envelope");
16808 assert_eq!(
16809 v["error"]["type"], "invalid_request_error",
16810 "{surface} error type"
16811 );
16812 }
16813 _ => {
16814 assert!(
16815 v["error"]["message"].is_string(),
16816 "{surface} OpenAI-shaped error body: {v}"
16817 );
16818 }
16819 }
16820 }
16821 }
16822
16823 let resp = anthropic::messages(
16827 State(st.clone()),
16828 axum::http::HeaderMap::new(),
16829 None,
16830 axum::body::Bytes::from(
16831 serde_json::json!({
16832 "model": "m", "max_tokens": 8,
16833 "messages": [{"role": "user", "content": "t"}],
16834 "thinking": {"type": "enabled"},
16835 "output_config": {"effort": "none"}})
16836 .to_string(),
16837 ),
16838 )
16839 .await;
16840 assert_eq!(resp.status(), StatusCode::OK);
16841 let saw = saw_rx
16842 .recv_timeout(std::time::Duration::from_secs(10))
16843 .expect("thinking+effort request never reached the worker");
16844 assert_eq!(
16845 saw.think,
16846 ThinkMode::Think,
16847 "thinking.type (the documented Anthropic lever) must win the switch over \
16848 output_config.effort"
16849 );
16850 let resp = anthropic::messages(
16851 State(st.clone()),
16852 axum::http::HeaderMap::new(),
16853 None,
16854 axum::body::Bytes::from(
16855 serde_json::json!({
16856 "model": "m", "max_tokens": 8,
16857 "messages": [{"role": "user", "content": "t"}],
16858 "thinking": {"type": "enabled"},
16859 "output_config": {"effort": "banana"}})
16860 .to_string(),
16861 ),
16862 )
16863 .await;
16864 assert_eq!(
16865 resp.status(),
16866 StatusCode::BAD_REQUEST,
16867 "an invalid effort must 400 even next to an explicit thinking.type — \
16868 precedence must not re-open the silent-accept hole"
16869 );
16870 }
16871
16872 #[test]
16873 fn vendor_sampling_defaults_are_boot_validated() {
16874 let parsed = OpenRouterMetadataFile::from_toml(
16877 r#"
16878[models.g]
16879default_temperature = 1.0
16880default_top_p = 0.95
16881default_top_k = 64
16882default_min_p = 0.0
16883default_presence_penalty = 0.0
16884default_frequency_penalty = 0.0
16885default_repetition_penalty = 1.0
16886"#,
16887 )
16888 .unwrap();
16889 let g = parsed.get("g").unwrap();
16890 assert_eq!(g.default_temperature, Some(1.0));
16891 assert_eq!(g.default_top_p, Some(0.95));
16892 assert_eq!(g.default_top_k, Some(64));
16893
16894 let err = OpenRouterMetadataFile::from_toml(
16898 r#"
16899[models.g]
16900default_temperature = 0.0
16901"#,
16902 )
16903 .unwrap_err();
16904 assert!(err.contains("default_temperature"), "{err}");
16905 assert!(
16906 err.contains("greedy"),
16907 "the refusal must say WHY a zero default is refused: {err}"
16908 );
16909
16910 for bad in [
16911 "default_temperature = 2.5",
16912 "default_temperature = -1.0",
16913 "default_top_p = 0.0",
16914 "default_top_p = 1.5",
16915 "default_min_p = 1.0",
16916 "default_min_p = -0.1",
16917 "default_presence_penalty = 3.0",
16918 "default_frequency_penalty = -2.5",
16919 "default_repetition_penalty = 0.0",
16920 ] {
16921 let err =
16922 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
16923 let key = bad.split(' ').next().unwrap();
16924 assert!(err.contains(key), "{bad} must be refused by name: {err}");
16925 }
16926
16927 let err = OpenRouterMetadataFile::from_toml(
16931 r#"
16932[models.g]
16933default_temperture = 1.0
16934"#,
16935 )
16936 .unwrap_err();
16937 assert!(
16938 err.contains("unknown field"),
16939 "an unknown key must be fatal, which is what makes binary-first ordering \
16940 mandatory: {err}"
16941 );
16942 }
16943
16944 #[test]
16945 fn non_thinking_sampling_arm_is_boot_validated() {
16946 let parsed = OpenRouterMetadataFile::from_toml(
16950 r#"
16951[models.q]
16952default_temperature = 1.0
16953default_top_p = 0.95
16954default_top_k = 20
16955
16956[models.q.non_thinking_sampling]
16957temperature = 0.7
16958top_p = 0.8
16959top_k = 20
16960presence_penalty = 1.5
16961"#,
16962 )
16963 .unwrap();
16964 let arm = parsed
16965 .get("q")
16966 .unwrap()
16967 .non_thinking_sampling
16968 .as_ref()
16969 .unwrap();
16970 assert_eq!(arm.temperature, Some(0.7));
16971 assert_eq!(arm.top_p, Some(0.8));
16972 assert_eq!(arm.top_k, Some(20));
16973 assert_eq!(arm.presence_penalty, Some(1.5));
16974 assert_eq!(
16975 arm.min_p, None,
16976 "undeclared arm fields stay undeclared, never invented"
16977 );
16978
16979 let err = OpenRouterMetadataFile::from_toml(
16983 r#"
16984[models.q]
16985[models.q.non_thinking_sampling]
16986temperature = 0.0
16987"#,
16988 )
16989 .unwrap_err();
16990 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
16991 assert!(err.contains("greedy"), "{err}");
16992
16993 let err = OpenRouterMetadataFile::from_toml(
16996 r#"
16997[models.q]
16998[models.q.non_thinking_sampling]
16999"#,
17000 )
17001 .unwrap_err();
17002 assert!(err.contains("non_thinking_sampling"), "{err}");
17003 assert!(err.contains("declare"), "{err}");
17004
17005 for bad in [
17007 "temperature = 2.5",
17008 "top_p = 0.0",
17009 "top_p = 1.5",
17010 "min_p = 1.0",
17011 "presence_penalty = 3.0",
17012 "frequency_penalty = -2.5",
17013 "repetition_penalty = 0.0",
17014 ] {
17015 let err = OpenRouterMetadataFile::from_toml(&format!(
17016 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
17017 ))
17018 .unwrap_err();
17019 let key = bad.split(' ').next().unwrap();
17020 assert!(
17021 err.contains(&format!("non_thinking_sampling.{key}")),
17022 "the refusal for {bad:?} must name the nested key: {err}"
17023 );
17024 }
17025
17026 let err = OpenRouterMetadataFile::from_toml(
17030 r#"
17031[models.q]
17032[models.q.non_thinking_sampling]
17033temperture = 0.7
17034"#,
17035 )
17036 .unwrap_err();
17037 assert!(err.contains("unknown field"), "{err}");
17038 }
17039
17040 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
17045 SamplingDefaults {
17046 temperature: Some(0.7),
17047 top_p: Some(0.8),
17048 top_k: Some(20),
17049 presence_penalty: Some(1.5),
17050 ..Default::default()
17051 }
17052 }
17053
17054 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
17055 ModelSamplingDefaults {
17056 thinking: qwen38_vendor_defaults(),
17057 non_thinking: Some(qwen38_non_thinking_defaults()),
17058 }
17059 }
17060
17061 fn qwen38_caps() -> ModelCaps {
17065 ModelCaps {
17066 chat_ok: true,
17067 qwen_think: true,
17068 think_switch: true,
17069 ..Default::default()
17070 }
17071 }
17072
17073 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
17076 (
17077 c.temperature,
17078 c.top_p,
17079 c.top_k,
17080 c.min_p,
17081 c.penalty_present,
17082 c.penalty_freq,
17083 c.penalty_repeat,
17084 c.penalty_last_n,
17085 c.seed,
17086 )
17087 }
17088
17089 fn build_with_arms(
17090 defaults: &ModelSamplingDefaults,
17091 caps: &ModelCaps,
17092 default_effort: Option<&str>,
17093 extra: serde_json::Value,
17094 ) -> Request {
17095 let mut body = serde_json::json!({
17096 "model": "m",
17097 "messages": [{"role": "user", "content": "task"}],
17098 "seed": 3
17100 });
17101 body.as_object_mut()
17102 .unwrap()
17103 .extend(extra.as_object().unwrap().clone());
17104 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17105 let (tx, _rx) = worker::event_channel();
17106 build_chat_request_with_trace(
17107 req,
17108 Some(caps),
17109 tx,
17110 lanes::Lane::Interactive,
17111 None,
17112 None,
17113 default_effort,
17114 defaults,
17115 )
17116 .unwrap()
17117 .request
17118 }
17119
17120 #[test]
17121 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
17122 let two_arm = qwen38_two_arm_defaults();
17127 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
17128 let caps = qwen38_caps();
17129
17130 let off_spellings = [
17132 serde_json::json!({"reasoning_effort": "none"}),
17133 serde_json::json!({"enable_thinking": false}),
17134 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
17135 serde_json::json!({"reasoning": {"enabled": false}}),
17136 ];
17137 for extra in &off_spellings {
17138 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
17139 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
17140 let c = &r.sampler_cfg;
17141 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
17142 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
17143 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
17144 assert_eq!(
17145 c.penalty_present, 1.5,
17146 "{extra}: non-thinking presence_penalty"
17147 );
17148 assert_eq!(
17149 c.penalty_last_n,
17150 memra_engine::spec::PEN_WINDOW_MAX,
17151 "{extra}: the arm's presence penalty uses the cross-path history window"
17152 );
17153 assert_eq!(
17154 c.min_p, 0.0,
17155 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
17156 );
17157
17158 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
17161 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
17162 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
17163 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
17164 assert_eq!(
17165 s.sampler_cfg.penalty_present, 0.0,
17166 "{extra}: single-arm model"
17167 );
17168 }
17169
17170 for extra in [
17173 serde_json::json!({}),
17174 serde_json::json!({"enable_thinking": true}),
17175 serde_json::json!({"reasoning_effort": "high"}),
17176 serde_json::json!({"reasoning": {"enabled": true}}),
17177 ] {
17178 for defaults in [&two_arm, &single_arm] {
17179 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
17180 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
17181 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
17182 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
17183 assert_eq!(
17184 c.penalty_present, 0.0,
17185 "{extra}: thinking arm has no presence"
17186 );
17187 }
17188 }
17189
17190 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
17193 assert_eq!(
17194 c.temperature, 0.7,
17195 "deployment-default off = non-thinking arm"
17196 );
17197 let c = build_with_arms(
17199 &two_arm,
17200 &caps,
17201 Some("none"),
17202 serde_json::json!({"enable_thinking": true}),
17203 )
17204 .sampler_cfg;
17205 assert_eq!(
17206 c.temperature, 1.0,
17207 "explicit ON beats the deployment default"
17208 );
17209
17210 let c = build_with_arms(
17212 &two_arm,
17213 &caps,
17214 None,
17215 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
17216 )
17217 .sampler_cfg;
17218 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
17219 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
17220 let c = build_with_arms(
17221 &two_arm,
17222 &caps,
17223 None,
17224 serde_json::json!({
17225 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
17226 )
17227 .sampler_cfg;
17228 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
17229 assert_eq!(
17230 c.penalty_present, 0.0,
17231 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
17232 is a value, not an absence"
17233 );
17234 assert_eq!(
17235 c.penalty_last_n, 0,
17236 "all penalties off => no history window"
17237 );
17238 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
17239
17240 let c = build_with_arms(
17243 &two_arm,
17244 &caps,
17245 None,
17246 serde_json::json!({"enable_thinking": false, "temperature": 0}),
17247 )
17248 .sampler_cfg;
17249 assert!(
17250 memra_engine::sampler::Sampler::new(c).is_greedy(),
17251 "explicit temperature 0 must stay greedy on the non-thinking arm"
17252 );
17253
17254 let c = build_with_arms(
17257 &single_arm,
17258 &caps,
17259 None,
17260 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
17261 )
17262 .sampler_cfg;
17263 assert_eq!(c.temperature, 0.55);
17264 assert_eq!(
17265 c.top_p, 0.95,
17266 "single-arm model: unset top_p takes its one arm"
17267 );
17268 }
17269
17270 #[test]
17271 fn sampling_arms_never_blend_field_by_field() {
17272 let parsed = OpenRouterMetadataFile::from_toml(
17277 r#"
17278[models.m]
17279default_temperature = 1.0
17280default_min_p = 0.05
17281
17282[models.m.non_thinking_sampling]
17283temperature = 0.6
17284"#,
17285 )
17286 .unwrap();
17287 let caps = ModelCaps {
17288 chat_temperature_default: Some(0.5),
17289 chat_top_p_default: Some(0.9),
17290 ..Default::default()
17291 };
17292 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
17293 let client = ClientSampling {
17294 seed: Some(1),
17295 ..Default::default()
17296 };
17297
17298 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
17299 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
17300 assert_eq!(
17301 off.min_p, 0.0,
17302 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
17303 );
17304 assert_eq!(
17305 off.top_p, 1.0,
17306 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
17307 );
17308
17309 for mode in [ThinkMode::Default, ThinkMode::Think] {
17311 let on = resolve_sampler_config(client, d.for_mode(mode));
17312 assert_eq!(on.temperature, 1.0);
17313 assert_eq!(on.min_p, 0.05);
17314 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
17315 }
17316 }
17317
17318 #[test]
17319 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
17320 let caps = qwen38_caps();
17329 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
17330 let two_arm = qwen38_two_arm_defaults();
17331
17332 let bodies = [
17333 serde_json::json!({}),
17334 serde_json::json!({"enable_thinking": true}),
17335 serde_json::json!({"reasoning_effort": "high"}),
17336 serde_json::json!({"reasoning_effort": "none"}),
17337 serde_json::json!({"enable_thinking": false}),
17338 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
17339 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
17340 serde_json::json!({"enable_thinking": false, "temperature": 0}),
17341 ];
17342 for extra in &bodies {
17343 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
17345 let mut client = ClientSampling {
17346 seed: Some(3),
17347 ..Default::default()
17348 };
17349 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
17350 client.temperature = Some(t as f32);
17351 }
17352 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
17353 client.top_p = Some(p as f32);
17354 }
17355 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
17356 assert_eq!(
17357 sampler_key(&r.sampler_cfg),
17358 sampler_key(&pre_arm),
17359 "{extra}: single-arm model diverged from the pre-arm resolution law"
17360 );
17361
17362 if r.think != ThinkMode::NoThink {
17365 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
17366 assert_eq!(t.think, r.think, "{extra}");
17367 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
17368 assert_eq!(
17369 sampler_key(&t.sampler_cfg),
17370 sampler_key(&r.sampler_cfg),
17371 "{extra}: a thinking-on request must not feel the non-thinking arm"
17372 );
17373 }
17374 }
17375 }
17376
17377 #[test]
17378 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
17379 let r = build_with_arms(
17385 &qwen38_two_arm_defaults(),
17386 &qwen38_caps(),
17387 None,
17388 serde_json::json!({"response_format": {"type": "json_object"}}),
17389 );
17390 assert_eq!(
17391 r.think,
17392 ThinkMode::NoThink,
17393 "constraint forces the switch off"
17394 );
17395 assert_eq!(
17396 r.sampler_cfg.temperature, 0.7,
17397 "and the arm follows the real mode"
17398 );
17399 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
17400 }
17401
17402 #[test]
17403 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
17404 let caps = ModelCaps {
17409 chat_temperature_default: Some(0.5),
17410 chat_top_p_default: Some(0.9),
17411 chat_ok: true,
17412 ..Default::default()
17413 };
17414 let metadata = OpenRouterModelMetadata {
17415 default_temperature: Some(1.0),
17416 default_top_p: Some(0.95),
17417 default_top_k: Some(64),
17418 ..Default::default()
17419 };
17420
17421 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
17422 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
17423 assert_eq!(caps_only.top_p, Some(0.9));
17424 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
17425
17426 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
17427 assert_eq!(
17428 both.temperature,
17429 Some(1.0),
17430 "metadata outranks the arch cap"
17431 );
17432 assert_eq!(both.top_p, Some(0.95));
17433 assert_eq!(both.top_k, Some(64));
17434
17435 let partial = SamplingDefaults::resolve(
17437 Some(&OpenRouterModelMetadata {
17438 default_temperature: Some(0.7),
17439 ..Default::default()
17440 }),
17441 Some(&caps),
17442 );
17443 assert_eq!(partial.temperature, Some(0.7));
17444 assert_eq!(
17445 partial.top_p,
17446 Some(0.9),
17447 "an undeclared metadata field must fall through to the cap, not to 1.0"
17448 );
17449
17450 assert_eq!(
17452 SamplingDefaults::resolve(None, None),
17453 SamplingDefaults::default()
17454 );
17455 }
17456
17457 #[test]
17458 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
17459 let resolved = |d: &SamplingDefaults| {
17472 resolve_sampler_config(
17473 ClientSampling {
17474 seed: Some(1),
17475 ..Default::default()
17476 },
17477 d,
17478 )
17479 };
17480
17481 assert!(
17483 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
17484 .is_spec_sampling(),
17485 "the API-standard default must stay in the fast pure-temp regime"
17486 );
17487
17488 for (name, d) in [
17489 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
17490 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
17491 ] {
17492 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
17493 assert!(
17494 !sampler.is_greedy(),
17495 "{name}: vendor default must not be greedy"
17496 );
17497 assert!(
17498 !sampler.is_spec_sampling(),
17499 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
17500 starts passing, either the vendor numbers changed or the in-graph draft \
17501 learned filters, and the perf note in docs/SERVING.md needs revisiting"
17502 );
17503 }
17504
17505 let opted_out = resolve_sampler_config(
17507 ClientSampling {
17508 top_p: Some(1.0),
17509 top_k: Some(0),
17510 seed: Some(1),
17511 ..Default::default()
17512 },
17513 &qwen38_vendor_defaults(),
17514 );
17515 assert!(
17516 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
17517 "explicitly disabling the filters must restore the pure-temp regime"
17518 );
17519 }
17520
17521 #[test]
17522 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
17523 let comp_seed = |body: serde_json::Value| {
17530 let req: CompletionReq = serde_json::from_value(body).unwrap();
17531 let (tx, _rx) = worker::event_channel();
17532 build_request(&req, tx, lanes::Lane::Interactive, None)
17533 .sampler_cfg
17534 .seed
17535 };
17536 let chat_seed = |body: serde_json::Value| {
17537 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17538 let (tx, _rx) = worker::event_channel();
17539 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17540 .unwrap()
17541 .request
17542 .sampler_cfg
17543 .seed
17544 };
17545
17546 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
17549 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
17550 let c = chat_seed(serde_json::json!({
17551 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
17552 assert_ne!(
17553 a, 0,
17554 "omitted seed must not be the pinned 0 that caused the loop"
17555 );
17556 assert_ne!(b, 0);
17557 assert_ne!(c, 0);
17558 assert_ne!(
17559 a, b,
17560 "two seed-omitting requests must get DIFFERENT streams"
17561 );
17562 assert_ne!(a, c);
17563
17564 assert_eq!(
17567 comp_seed(serde_json::json!({
17568 "model": "m", "prompt": "t", "seed": 0})),
17569 0,
17570 "explicit seed 0 must stay 0 — the determinism gates depend on it"
17571 );
17572 assert_eq!(
17573 comp_seed(serde_json::json!({
17574 "model": "m", "prompt": "t", "seed": 12345})),
17575 12345
17576 );
17577 assert_eq!(
17578 chat_seed(serde_json::json!({
17579 "model": "m", "messages": [{"role": "user", "content": "t"}],
17580 "seed": 777})),
17581 777
17582 );
17583 assert_eq!(
17585 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
17586 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
17587 );
17588
17589 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
17592 assert_eq!(
17593 seeds.len(),
17594 256,
17595 "fresh_seed must not collide across rapid calls"
17596 );
17597 assert!(!seeds.contains(&0));
17598 }
17599
17600 #[test]
17601 fn response_format_builds_grammar_only_when_present() {
17602 let mk = |rf: Option<serde_json::Value>| {
17606 let mut body = serde_json::json!({
17607 "model": "m", "messages": [{"role": "user", "content": "t"}]});
17608 if let Some(rf) = rf {
17609 body["response_format"] = rf;
17610 }
17611 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
17612 let (tx, _rx) = worker::event_channel();
17613 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
17614 };
17615 assert!(mk(None).unwrap().request.grammar.is_none());
17616 assert!(
17617 mk(Some(serde_json::json!({"type": "text"})))
17618 .unwrap()
17619 .request
17620 .grammar
17621 .is_none()
17622 );
17623 assert!(matches!(
17624 mk(Some(serde_json::json!({"type": "json_object"})))
17625 .unwrap()
17626 .request
17627 .grammar,
17628 Some(constrained::GrammarSpec::JsonObject)
17629 ));
17630 assert!(matches!(
17631 mk(Some(serde_json::json!({"type": "json_schema",
17632 "json_schema": {"schema": {"type": "object"}}})))
17633 .unwrap()
17634 .request
17635 .grammar,
17636 Some(constrained::GrammarSpec::JsonSchema(_))
17637 ));
17638 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
17640 }
17641
17642 #[test]
17652 fn response_format_think_table_switch_postthink_refusal() {
17653 let mk = |caps: &ModelCaps| {
17654 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17655 "model": "m", "messages": [{"role": "user", "content": "t"}],
17656 "response_format": {"type": "json_object"}}))
17657 .unwrap();
17658 let (tx, _rx) = worker::event_channel();
17659 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
17660 };
17661 let switch = ModelCaps {
17663 chat_ok: true,
17664 qwen_think: true,
17665 think_switch: true,
17666 ..Default::default()
17667 };
17668 let plan = mk(&switch).unwrap();
17669 assert_eq!(
17670 plan.request.think,
17671 memra_tokenizer::chat::ThinkMode::NoThink,
17672 "switch-carrying template must keep the grammar-from-token-1 path"
17673 );
17674 assert!(plan.request.grammar.is_some());
17675
17676 let postthink = ModelCaps {
17678 chat_ok: true,
17679 qwen_think: true,
17680 think_switch: false,
17681 think_close: vec![128799],
17682 ..Default::default()
17683 };
17684 let plan = mk(&postthink).unwrap();
17685 assert_ne!(
17686 plan.request.think,
17687 memra_tokenizer::chat::ThinkMode::NoThink,
17688 "post-think constrained request must keep the think channel ON"
17689 );
17690 assert!(plan.request.grammar.is_some());
17691
17692 let no_contract = ModelCaps {
17694 chat_ok: true,
17695 qwen_think: true,
17696 think_switch: false,
17697 think_close: Vec::new(),
17698 ..Default::default()
17699 };
17700 let err = match mk(&no_contract) {
17701 Err(err) => err,
17702 Ok(_) => panic!("think-forced template with no close contract must refuse"),
17703 };
17704 assert!(
17705 err.contains("think-close"),
17706 "refusal must name the missing close contract: {err}"
17707 );
17708 }
17709
17710 #[test]
17711 fn unsupported_semantic_params_are_named_rejections() {
17712 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17714 "model": "m", "messages": [{"role": "user", "content": "t"}],
17715 "response_format": {"type": "json_object"}
17716 }))
17717 .unwrap();
17718 assert!(req.response_format.is_some());
17719 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
17720 "model": "m", "messages": [{"role": "user", "content": "t"}],
17721 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
17722 "user": "u-1", "stream_options": {"include_usage": true}
17723 }))
17724 .unwrap();
17725 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
17727 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
17728 assert_eq!(req.n, Some(1));
17729 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
17731 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
17732 assert_eq!(param, "logit_bias");
17733 assert_eq!(msg, "logit_bias is not supported (why)");
17734 }
17735
17736 #[test]
17737 fn completions_accept_openai_stop_forms() {
17738 for (value, expected) in [
17739 (serde_json::json!("Problem:"), vec!["Problem:"]),
17740 (
17741 serde_json::json!(["Question:", "Problem:"]),
17742 vec!["Question:", "Problem:"],
17743 ),
17744 (serde_json::Value::Null, Vec::<&str>::new()),
17745 ] {
17746 let req: CompletionReq = serde_json::from_value(serde_json::json!({
17747 "model": "plain_quant", "prompt": "task", "stop": value
17748 }))
17749 .unwrap();
17750 assert_eq!(req.stop.into_vec(), expected);
17751 }
17752 }
17753
17754 fn fake_worker_state() -> AppState {
17761 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
17762 }
17763
17764 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
17765 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
17766 }
17767
17768 struct WorkerSaw {
17773 sampler_cfg: SamplerConfig,
17774 think: ThinkMode,
17775 reasoning_effort: Option<String>,
17776 }
17777
17778 fn fake_worker_state_full(
17785 steps: usize,
17786 step_delay: std::time::Duration,
17787 caps: HashMap<String, ModelCaps>,
17788 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
17789 ) -> AppState {
17790 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
17791 let health = health::WorkerHealth::new();
17792 let h = health.clone();
17793 std::thread::spawn(move || {
17794 h.mark_ready();
17795 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
17796 if let Some(tx) = &saw_tx {
17797 let _ = tx.send(WorkerSaw {
17798 sampler_cfg: req.sampler_cfg.clone(),
17799 think: req.think,
17800 reasoning_effort: req.reasoning_effort.clone(),
17801 });
17802 }
17803 worker::release_pending_admit();
17807 worker::release_admission_reservation(req.lane);
17808 h.beat_busy();
17809 if let Some(ready) = req.constraint_ready.take() {
17810 let _ = ready.send(Ok(()));
17811 }
17812 let _ = req.tx.send(Event::PromptUsage {
17813 n_prompt: 1,
17814 n_cached: 0,
17815 });
17816 for step in 0..steps {
17817 h.beat_busy();
17818 let text = if steps == 1 { "ok" } else { "x" };
17819 let _ = req.tx.send(Event::Token {
17820 id: step as u32 + 1,
17821 text: text.into(),
17822 });
17823 if !step_delay.is_zero() {
17824 std::thread::sleep(step_delay);
17825 }
17826 }
17827 let _ = req.tx.send(Event::Done {
17828 stop_reason: "Eos".into(),
17829 n_tokens: steps,
17830 n_prompt: 1,
17831 n_cached: 0,
17832 elapsed_s: 0.01,
17833 spec: None,
17834 });
17835 h.set_phase(health::PHASE_IDLE);
17836 }
17837 });
17838 for _ in 0..2000 {
17841 if health.live().is_ok() {
17842 break;
17843 }
17844 std::thread::sleep(std::time::Duration::from_millis(1));
17845 }
17846 AppState {
17847 cmd_tx,
17848 models: Arc::new(vec!["m".into()]),
17849 caps: Arc::new(caps),
17850 openrouter_metadata: Arc::new(HashMap::new()),
17851 provider_metadata: Arc::new(None),
17852 metering: None,
17853
17854 budget_tokenizers: None,
17855 api_auth: ApiAuth::default(),
17856 metrics_auth: MetricsAuth::default(),
17857 metrics: SharedMetrics::default(),
17858 inflight: Arc::new(Default::default()),
17859 tenant_inflight: Arc::new(Default::default()),
17860 health,
17861 bg: None,
17862 }
17863 }
17864
17865 #[tokio::test]
17866 #[allow(clippy::await_holding_lock)] async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
17868 let _l = drain_lock();
17869 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
17870 let normal_state = st.clone();
17871 let normal = tokio::spawn(async move {
17872 chat_completions(
17873 State(normal_state),
17874 axum::http::HeaderMap::new(),
17875 None,
17876 Json(
17877 serde_json::from_value(serde_json::json!({
17878 "model": "m",
17879 "messages": [{"role": "user", "content": "keep decoding"}],
17880 }))
17881 .unwrap(),
17882 ),
17883 )
17884 .await
17885 });
17886 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
17887
17888 let mut deep = serde_json::json!({"type": "string"});
17889 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
17890 deep = serde_json::json!({"allOf": [deep]});
17891 }
17892 let bad = chat_completions(
17893 State(st.clone()),
17894 axum::http::HeaderMap::new(),
17895 None,
17896 Json(
17897 serde_json::from_value(serde_json::json!({
17898 "model": "m",
17899 "messages": [{"role": "user", "content": "bad schema"}],
17900 "response_format": {
17901 "type": "json_schema",
17902 "json_schema": {"schema": deep},
17903 },
17904 }))
17905 .unwrap(),
17906 ),
17907 )
17908 .await;
17909 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
17910 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
17911 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
17912 .await
17913 .unwrap();
17914 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17915 assert!(
17916 payload["error"]["message"]
17917 .as_str()
17918 .unwrap()
17919 .contains("maximum nesting depth")
17920 );
17921 assert!(
17922 !normal.is_finished(),
17923 "bad schema stalled or replaced the normal decode"
17924 );
17925
17926 let normal_response = normal.await.unwrap();
17927 assert_eq!(normal_response.status(), StatusCode::OK);
17928 let snapshot = st.health.snapshot();
17929 assert!(
17930 st.health.live().is_ok(),
17931 "normal decode left health stalled"
17932 );
17933 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
17934 }
17935
17936 #[tokio::test]
17937 #[allow(clippy::await_holding_lock)] async fn valid_response_format_preflight_preserves_generation() {
17939 let _l = drain_lock();
17940 let response = chat_completions(
17941 State(fake_worker_state()),
17942 axum::http::HeaderMap::new(),
17943 None,
17944 Json(
17945 serde_json::from_value(serde_json::json!({
17946 "model": "m",
17947 "messages": [{"role": "user", "content": "valid schema"}],
17948 "response_format": {"type": "json_object"},
17949 }))
17950 .unwrap(),
17951 ),
17952 )
17953 .await;
17954 assert_eq!(response.status(), StatusCode::OK);
17955 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
17956 .await
17957 .unwrap();
17958 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17959 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
17960 }
17961
17962 #[tokio::test]
17963 #[allow(clippy::await_holding_lock)] async fn unknown_model_refuses_model_not_found_before_admission() {
17965 let _l = drain_lock();
17966 let response = chat_completions(
17971 State(fake_worker_state()),
17972 axum::http::HeaderMap::new(),
17973 None,
17974 Json(
17975 serde_json::from_value(serde_json::json!({
17976 "model": "qwen/qwen3.8-27b-typo",
17977 "messages": [{"role": "user", "content": "hi"}],
17978 }))
17979 .unwrap(),
17980 ),
17981 )
17982 .await;
17983 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
17984 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
17985 .await
17986 .unwrap();
17987 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17988 assert_eq!(payload["error"]["code"], "model_not_found");
17989 assert_eq!(payload["error"]["type"], "invalid_request_error");
17990
17991 let response = completions(
17993 State(fake_worker_state()),
17994 axum::http::HeaderMap::new(),
17995 None,
17996 Json(
17997 serde_json::from_value(serde_json::json!({
17998 "model": "nope",
17999 "prompt": "hi",
18000 }))
18001 .unwrap(),
18002 ),
18003 )
18004 .await;
18005 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
18006 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18007 .await
18008 .unwrap();
18009 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18010 assert_eq!(payload["error"]["code"], "model_not_found");
18011 }
18012
18013 const METRICS_KEY_ACME: &str = "completion-acme-secret";
18014 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
18015
18016 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
18017 let spec = format!(
18018 "acme:{},blue:{}",
18019 auth::sha256_hex(METRICS_KEY_ACME),
18020 auth::sha256_hex(METRICS_KEY_BLUE),
18021 );
18022 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
18023 let mut st = fake_worker_state();
18024 st.api_auth.keyring = Some(keyring);
18025 st.metrics_auth = MetricsAuth::new(
18026 true,
18027 st.api_auth.configured(),
18028 metrics_token.map(str::to_string),
18029 );
18030 {
18031 let mut metrics = st.metrics.lock().unwrap();
18032 metrics.admitted = 17;
18033 metrics.prompt_tokens_in = 400;
18034 metrics.cached_tokens_in = 60;
18035 metrics.prefix_hits = 2;
18036 metrics.prefix_misses = 3;
18037 metrics.prefix_inserts = 5;
18038 metrics.prefix_evictions = 7;
18039 metrics.prefix_skips_budget = 9;
18040 metrics.prefix_skips_pinned = 10;
18041 metrics.prefix_hit_tokens = 11;
18042 metrics.lcp_hist[4] = 13;
18043 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
18044 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
18045 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
18046 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
18047 metrics.prefix_entries = 29;
18048 metrics.prefix_bytes = 31;
18049 metrics.active_sessions = 3;
18050 metrics.queued_requests = 5;
18051 metrics.admission_inflight.insert("m".into(), 4);
18052 metrics
18053 .admission_booked_bytes
18054 .insert("m".into(), 41_000_000);
18055 metrics.continuation_pool_entries = 7;
18056 metrics.spec_pool_entries = 11;
18057 metrics.cuda_driver_free_bytes = 13;
18058 metrics.cuda_pool_reserved_bytes = 17;
18059 metrics.cuda_pool_used_bytes = 19;
18060 metrics.cuda_pool_cached_bytes = 23;
18061 metrics.batch_size_last = 37;
18062 metrics.spec.insert(
18063 "m".into(),
18064 memra_engine::spec::SpecTelemetry {
18065 rounds: 2,
18066 drafted: 6,
18067 accepted: 4,
18068 ..Default::default()
18069 },
18070 );
18071 let mut spec_window = memra_engine::spec::SpecTelemetry {
18072 rounds: 4,
18073 drafted: 12,
18074 accepted: 6,
18075 ..Default::default()
18076 };
18077 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
18078 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
18079 metrics.spec_window.insert("m".into(), spec_window);
18080 metrics.constraint_compiler_fail_closed.insert(
18081 "m".into(),
18082 Arc::new(std::sync::atomic::AtomicBool::new(true)),
18083 );
18084 }
18085 st
18086 }
18087
18088 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
18089 let mut headers = HeaderMap::new();
18090 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
18091 let response = get_metrics(State(st), headers).await;
18092 assert_eq!(response.status(), StatusCode::OK);
18093 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18094 .await
18095 .unwrap();
18096 serde_json::from_slice(&bytes).unwrap()
18097 }
18098
18099 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
18100 let mut headers = HeaderMap::new();
18101 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
18102 let response = yield_metrics(State(st), headers).await;
18103 assert_eq!(response.status(), StatusCode::OK);
18104 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18105 .await
18106 .unwrap();
18107 serde_json::from_slice(&bytes).unwrap()
18108 }
18109
18110 #[test]
18111 fn exposed_open_bind_is_refused_before_server_start() {
18112 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
18113 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
18114
18115 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
18116 assert!(err.contains("refusing unauthenticated non-loopback bind"));
18117 assert!(err.contains("MEMRA_API_KEY"));
18118 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
18119 assert!(validate_bind_security("[::]:8000", false, false).is_err());
18120
18121 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
18122 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
18123 }
18124
18125 #[tokio::test]
18126 async fn keyed_metrics_require_and_accept_api_bearer() {
18127 let mut st = fake_worker_state();
18128 st.api_auth.single_key = Some(Arc::from("completion-secret"));
18129 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
18130
18131 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
18132 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
18133 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
18134 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
18135
18136 let mut headers = HeaderMap::new();
18137 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
18138 assert_eq!(
18139 get_metrics(State(st.clone()), headers.clone())
18140 .await
18141 .status(),
18142 StatusCode::OK,
18143 );
18144 let body = metrics_json(st.clone(), "completion-secret").await;
18145 assert!(
18146 body.get("admitted").is_some(),
18147 "the legacy single-key domain keeps cumulative counters",
18148 );
18149 assert!(
18150 body.get("active_sessions").is_none(),
18151 "a static completion key is not an operator metrics principal",
18152 );
18153 assert_eq!(
18154 yield_metrics(State(st), headers).await.status(),
18155 StatusCode::OK
18156 );
18157 }
18158
18159 #[tokio::test]
18160 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
18161 let st = multi_key_metrics_state(None);
18162 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
18163 assert_eq!(
18164 body.as_object().unwrap().len(),
18165 2,
18166 "completion metrics must contain only tenant-scoped rows",
18167 );
18168 let tenants = body["tenants"].as_object().unwrap();
18169 assert_eq!(tenants.len(), 1);
18170 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
18171 assert!(!tenants.contains_key("t:blue"));
18172 let adsd = body["adsd_suspect_total"].as_object().unwrap();
18173 assert_eq!(adsd.len(), 1);
18174 assert_eq!(adsd["t:acme"], 1);
18175 assert!(!adsd.contains_key("t:blue"));
18176
18177 let mut headers = HeaderMap::new();
18178 headers.insert(
18179 "authorization",
18180 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
18181 );
18182 assert_eq!(
18183 yield_metrics(State(st), headers).await.status(),
18184 StatusCode::FORBIDDEN,
18185 "the process-wide yield view requires an operator metrics token",
18186 );
18187 }
18188
18189 #[tokio::test]
18190 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
18191 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
18192 for operator_only in [
18193 "prefix_cache_entries",
18194 "prefix_cache_bytes",
18195 "prefix_cache_skips_budget",
18196 "prefix_cache_skips_pinned",
18197 "active_sessions",
18198 "queued_requests",
18199 "admission_inflight",
18200 "admission_booked_bytes",
18201 "continuation_pool_entries",
18202 "spec_pool_entries",
18203 "cuda_driver_free_bytes",
18204 "cuda_pool_reserved_bytes",
18205 "cuda_pool_used_bytes",
18206 "cuda_pool_cached_bytes",
18207 "constraint_compiler_fail_closed",
18208 "serve_idle_seconds",
18209 "spec",
18210 "spec_tau",
18211 "spec_accept_by_position",
18212 "dual_pp",
18213 "pp_wave",
18214 "peer_probe_bypassed",
18215 "peer_probe_boundary_copies",
18216 "peer_probe_runtime_reprobes",
18217 "peer_probe_runtime_failures",
18218 "peer_probe_deferred_total",
18219 "peer_probe_integrity_degraded",
18220 "peer_probe_degraded_to_host_bounce",
18221 ] {
18222 assert!(
18223 body.get(operator_only).is_none(),
18224 "tenant metrics must not expose operator field {operator_only}",
18225 );
18226 }
18227 }
18228
18229 #[test]
18230 fn populated_spec_acceptance_metrics_are_operator_only() {
18231 for scope in [
18232 MetricsScope::CompletionDomain,
18233 MetricsScope::Tenant("t:acme".into()),
18234 ] {
18235 let mut body = json!({});
18236 insert_spec_acceptance_metrics(&mut body, &scope, || {
18237 panic!("tenant scope evaluated the process-wide spec snapshot")
18238 });
18239 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
18240 assert!(
18241 body.get("spec_accept_by_position").is_none(),
18242 "{scope:?} leaked the accept histogram"
18243 );
18244 }
18245
18246 let mut telemetry = memra_engine::spec::SpecTelemetry {
18247 rounds: 4,
18248 drafted: 12,
18249 accepted: 6,
18250 ..Default::default()
18251 };
18252 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
18253 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
18254 let mut body = json!({});
18255 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
18256 HashMap::from([("model-a".to_string(), telemetry)])
18257 });
18258 assert_eq!(body["spec_tau"]["model-a"], 1.5);
18259 let histogram = &body["spec_accept_by_position"]["model-a"];
18260 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
18261 assert_eq!(histogram["rounds"], 4);
18262 assert_eq!(histogram["offered"], json!([4, 4, 4]));
18263 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
18264 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
18265 }
18266
18267 #[test]
18268 fn populated_dual_pp_metrics_are_operator_only() {
18269 let populated = DualPpMetricsSnapshot {
18270 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
18271 stage_samples: [1, 1, 1, 1],
18272 dropped_timing_samples: 0,
18273 overlaps: 17,
18274 slot_pairs: 19,
18275 slot_uses: [19, 19],
18276 slot_collisions: 0,
18277 };
18278 for scope in [
18279 MetricsScope::CompletionDomain,
18280 MetricsScope::Tenant("t:acme".into()),
18281 ] {
18282 let mut body = json!({});
18283 insert_dual_pp_metrics(&mut body, &scope, || populated);
18284 assert!(
18285 body.get("dual_pp").is_none(),
18286 "{scope:?} leaked dual PP topology"
18287 );
18288 }
18289
18290 let mut body = json!({});
18291 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
18292 assert_eq!(body["dual_pp"]["overlaps"], 17);
18293 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
18294 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
18295 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
18296 assert_eq!(
18297 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
18298 1.0
18299 );
18300 }
18301
18302 #[test]
18303 fn populated_pp_wave_metrics_are_operator_only() {
18304 let populated = PpWaveMetricsSnapshot {
18305 ticks: 11,
18306 cells: 96,
18307 overlaps: 37,
18308 };
18309 for scope in [
18310 MetricsScope::CompletionDomain,
18311 MetricsScope::Tenant("t:acme".into()),
18312 ] {
18313 let mut body = json!({});
18314 insert_pp_wave_metrics(&mut body, &scope, || populated);
18315 assert!(
18316 body.get("pp_wave").is_none(),
18317 "{scope:?} leaked PP wave topology"
18318 );
18319 }
18320
18321 let mut body = json!({});
18322 insert_pp_wave_metrics(&mut body, &MetricsScope::All, || populated);
18323 assert_eq!(body["pp_wave"]["ticks"], 11);
18324 assert_eq!(body["pp_wave"]["cells"], 96);
18325 assert_eq!(body["pp_wave"]["overlaps"], 37);
18326 }
18327
18328 #[test]
18329 fn peer_probe_metrics_are_operator_only() {
18330 let populated = memra_engine::pp::PeerProbeMetrics {
18331 bypassed: 1,
18332 boundary_copies: 8_192,
18333 runtime_probes: 1,
18334 runtime_failures: 0,
18335 deferred_total: 4,
18336 integrity_degraded: true,
18337 degraded_to_host_bounce: true,
18338 };
18339 for scope in [
18340 MetricsScope::CompletionDomain,
18341 MetricsScope::Tenant("t:acme".into()),
18342 ] {
18343 let mut body = json!({});
18344 insert_peer_probe_metrics(&mut body, &scope, || populated);
18345 assert!(body.get("peer_probe_bypassed").is_none());
18346 }
18347
18348 let mut body = json!({});
18349 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
18350 assert_eq!(body["peer_probe_bypassed"], 1);
18351 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
18352 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
18353 assert_eq!(body["peer_probe_runtime_failures"], 0);
18354 assert_eq!(body["peer_probe_deferred_total"], 4);
18355 assert_eq!(body["peer_probe_integrity_degraded"], true);
18356 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
18357 }
18358
18359 #[tokio::test]
18360 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
18361 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
18362 for operator_only in [
18363 "lcp_histogram",
18364 "cache_hit_token_ratio",
18365 "prefix_cache_hits",
18366 "prefix_cache_misses",
18367 "prefix_cache_inserts",
18368 "prefix_cache_evictions",
18369 "prefix_cache_skips_budget",
18370 "prefix_cache_skips_pinned",
18371 "prefix_cache_hit_tokens",
18372 ] {
18373 assert!(
18374 tenant_body.get(operator_only).is_none(),
18375 "tenant metrics must not expose global prefix field {operator_only}",
18376 );
18377 }
18378 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
18379 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
18380 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
18381 assert_eq!(
18382 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
18383 0.4
18384 );
18385
18386 let operator_body = metrics_json(
18387 multi_key_metrics_state(Some("scrape-secret")),
18388 "scrape-secret",
18389 )
18390 .await;
18391 assert_eq!(operator_body["prefix_cache_hits"], 2);
18392 assert_eq!(operator_body["prefix_cache_misses"], 3);
18393 assert_eq!(operator_body["prefix_cache_inserts"], 5);
18394 assert_eq!(operator_body["prefix_cache_evictions"], 7);
18395 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
18396 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
18397 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
18398 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
18399 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
18400 }
18401
18402 #[tokio::test]
18403 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
18404 let st = multi_key_metrics_state(Some("scrape-secret"));
18405 let mut completion_headers = HeaderMap::new();
18406 completion_headers.insert(
18407 "authorization",
18408 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
18409 );
18410 assert_eq!(
18411 get_metrics(State(st.clone()), completion_headers.clone())
18412 .await
18413 .status(),
18414 StatusCode::FORBIDDEN,
18415 );
18416 assert_eq!(
18417 yield_metrics(State(st.clone()), completion_headers)
18418 .await
18419 .status(),
18420 StatusCode::FORBIDDEN,
18421 );
18422
18423 let body = metrics_json(st.clone(), "scrape-secret").await;
18424 let tenants = body["tenants"].as_object().unwrap();
18425 assert_eq!(tenants.len(), 2);
18426 assert!(tenants.contains_key("t:acme"));
18427 assert!(tenants.contains_key("t:blue"));
18428 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
18429 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
18430 assert_eq!(body["active_sessions"], 3);
18431 assert_eq!(body["queued_requests"], 5);
18432 assert_eq!(body["admission_inflight"]["m"], 4);
18434 assert_eq!(body["admission_booked_bytes"]["m"], 41_000_000);
18435 assert_eq!(body["prefix_cache_bytes"], 31);
18436 assert_eq!(body["cuda_driver_free_bytes"], 13);
18437 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
18438 assert_eq!(body["spec"]["m"]["drafted"], 6);
18439 assert_eq!(body["spec_tau"]["m"], 1.5);
18440 assert_eq!(
18441 body["spec_accept_by_position"]["m"]["accepted"],
18442 json!([3, 2, 1])
18443 );
18444 let yield_body = yield_metrics_json(st, "scrape-secret").await;
18445 assert_eq!(yield_body["batch_size_last"], 37);
18446 }
18447
18448 #[tokio::test]
18449 async fn metrics_token_protects_public_override_without_api_keys() {
18450 let mut st = fake_worker_state();
18451 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
18452
18453 assert_eq!(
18454 get_metrics(State(st.clone()), HeaderMap::new())
18455 .await
18456 .status(),
18457 StatusCode::UNAUTHORIZED,
18458 );
18459 let mut headers = HeaderMap::new();
18460 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
18461 assert_eq!(
18462 get_metrics(State(st.clone()), headers.clone())
18463 .await
18464 .status(),
18465 StatusCode::OK,
18466 );
18467 assert_eq!(
18468 yield_metrics(State(st), headers).await.status(),
18469 StatusCode::OK
18470 );
18471 }
18472
18473 #[tokio::test]
18474 async fn no_key_loopback_metrics_remain_open_for_development() {
18475 let mut st = fake_worker_state();
18476 st.metrics_auth = MetricsAuth::new(true, false, None);
18477 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
18478 assert_eq!(response.status(), StatusCode::OK);
18479 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18480 .await
18481 .unwrap();
18482 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18483 assert!(
18484 body.get("active_sessions").is_some(),
18485 "no-key loopback development keeps full operator visibility",
18486 );
18487 assert_eq!(
18488 yield_metrics(State(st), HeaderMap::new()).await.status(),
18489 StatusCode::OK,
18490 );
18491 }
18492
18493 #[test]
18494 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
18495 let metrics = SharedMetrics::default();
18496 let rl = RateLimit::compute(4, 1, &metrics);
18498 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
18499 let rl = RateLimit::compute(4, 3, &metrics);
18500 assert_eq!(rl.remaining, 1);
18501 let rl = RateLimit::compute(4, 4, &metrics);
18503 assert_eq!(rl.remaining, 0);
18504 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
18505 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
18507 let m = worker::Metrics {
18509 completed: 2,
18510 tokens_out: 200,
18511 step_p50_ms: 20.0,
18512 ..Default::default()
18513 };
18514 assert_eq!(reset_estimate_s(&m), 2); }
18516
18517 #[test]
18518 fn inflight_guard_counts_up_and_frees_on_drop() {
18519 let counts: InflightCounts = Arc::new(Default::default());
18520 let tenants: TenantGauge = Arc::new(Default::default());
18521 let (g1, n1, t1) = InflightGuard::try_acquire(
18522 counts.clone(),
18523 lanes::Lane::Interactive,
18524 tenants.clone(),
18525 "acme",
18526 None,
18527 )
18528 .unwrap();
18529 let (g2, n2, t2) = InflightGuard::try_acquire(
18530 counts.clone(),
18531 lanes::Lane::Interactive,
18532 tenants.clone(),
18533 "acme",
18534 None,
18535 )
18536 .unwrap();
18537 assert_eq!((n1, n2), (1, 2));
18538 assert_eq!((t1, t2), (1, 2));
18540 let (gj, nj, tj) = InflightGuard::try_acquire(
18542 counts.clone(),
18543 lanes::Lane::Judge,
18544 tenants.clone(),
18545 "blue",
18546 None,
18547 )
18548 .unwrap();
18549 assert_eq!((nj, tj), (1, 1));
18550 drop(g1);
18551 drop(gj);
18552 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
18553 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
18554 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
18555 assert!(tenants.lock().unwrap().get("blue").is_none());
18557 drop(g2);
18558 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
18559 assert!(tenants.lock().unwrap().is_empty());
18560 }
18561
18562 #[test]
18563 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
18564 let counts: InflightCounts = Arc::new(Default::default());
18565 let tenants: TenantGauge = Arc::new(Default::default());
18566 let start = Arc::new(std::sync::Barrier::new(3));
18567 let attempted = Arc::new(std::sync::Barrier::new(3));
18568 let mut joins = Vec::new();
18569 for _ in 0..2 {
18570 let counts = counts.clone();
18571 let tenants = tenants.clone();
18572 let start = start.clone();
18573 let attempted = attempted.clone();
18574 joins.push(std::thread::spawn(move || {
18575 start.wait();
18576 let result = InflightGuard::try_acquire(
18577 counts,
18578 lanes::Lane::Interactive,
18579 tenants,
18580 "preview_001",
18581 Some(1),
18582 );
18583 let won = result.is_ok();
18584 attempted.wait(); drop(result);
18586 won
18587 }));
18588 }
18589 start.wait();
18590 attempted.wait();
18591 let wins = joins
18592 .into_iter()
18593 .map(|join| join.join().unwrap())
18594 .filter(|won| *won)
18595 .count();
18596 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
18597 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
18598 assert!(tenants.lock().unwrap().is_empty());
18599 }
18600
18601 #[tokio::test]
18602 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
18603 let st = fake_worker_state();
18604 let tenant = auth::TenantCtx {
18605 tenant: "preview_001".into(),
18606 lane_class: auth::LaneClass::Interactive,
18607 rate_limit: Some(1),
18608 key_prefix: None,
18609 };
18610 let first_env = Envelope::new(true);
18611 let (guard, first_rl) =
18612 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
18613 Ok(slot) => slot,
18614 Err(_) => panic!("the first request must acquire the tenant slot"),
18615 };
18616 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
18617
18618 let second_env = Envelope::new(true);
18619 let response =
18620 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
18621 Err(response) => response,
18622 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
18623 };
18624 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
18625 assert_eq!(response.headers()["retry-after"], "2");
18626 assert_eq!(response.headers()["retry-after-ms"], "2000");
18627 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
18628 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
18629 assert_eq!(response.headers()["x-request-id"], second_env.id);
18630 assert_eq!(
18631 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18632 1,
18633 "rejected request must not consume a lane slot"
18634 );
18635 assert_eq!(
18636 st.tenant_inflight
18637 .lock()
18638 .unwrap()
18639 .get("preview_001")
18640 .copied(),
18641 Some(1),
18642 "rejected request must not increment the tenant gauge"
18643 );
18644 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18645 .await
18646 .unwrap();
18647 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18648 assert_eq!(payload["error"]["type"], "rate_limit_error");
18649 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
18650 assert!(
18651 payload["error"]["message"]
18652 .as_str()
18653 .unwrap()
18654 .contains("concurrent request limit")
18655 );
18656
18657 drop(guard);
18658 let _ = InflightGuard::try_acquire(
18659 st.inflight.clone(),
18660 lanes::Lane::Interactive,
18661 st.tenant_inflight.clone(),
18662 "preview_001",
18663 Some(1),
18664 )
18665 .expect("slot must reopen after the in-flight request completes");
18666 }
18667
18668 #[test]
18669 fn tenant_rate_limit_override_is_min_with_global_cap() {
18670 let metrics = SharedMetrics::default();
18671 let unlimited = auth::TenantCtx::default_tenant();
18672 let capped = auth::TenantCtx {
18673 tenant: "acme".into(),
18674 lane_class: auth::LaneClass::Interactive,
18675 rate_limit: Some(2),
18676 key_prefix: None,
18677 };
18678 let global = lane_cap(lanes::Lane::Interactive);
18679 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
18681 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
18682 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
18684 assert_eq!((rl.limit, rl.remaining), (2, 1));
18685 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
18686 assert_eq!(rl.remaining, 0);
18687 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
18688 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
18692 assert_eq!(rl.remaining, 0);
18693 let wide = auth::TenantCtx {
18694 rate_limit: Some(global + 100),
18695 ..capped.clone()
18696 };
18697 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
18698 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
18699 }
18700
18701 #[test]
18702 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
18703 let batch = auth::TenantCtx {
18704 tenant: "bulk".into(),
18705 lane_class: auth::LaneClass::Batch,
18706 rate_limit: None,
18707 key_prefix: None,
18708 };
18709 let interactive = auth::TenantCtx::default_tenant();
18710 let hdr = |v: Option<&str>| {
18711 let mut h = axum::http::HeaderMap::new();
18712 if let Some(v) = v {
18713 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
18714 }
18715 h
18716 };
18717 assert_eq!(
18719 lane_for_tenant(&hdr(None), &interactive).unwrap(),
18720 lanes::Lane::Interactive
18721 );
18722 assert_eq!(
18723 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
18724 lanes::Lane::Judge
18725 );
18726 assert_eq!(
18728 lane_for_tenant(&hdr(None), &batch).unwrap(),
18729 lanes::Lane::Harvest
18730 );
18731 assert_eq!(
18732 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
18733 lanes::Lane::Judge
18734 );
18735 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
18736 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
18737 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
18739 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
18740 }
18741
18742 #[tokio::test]
18743 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
18744 let hdr = |v: &str| {
18749 let mut h = axum::http::HeaderMap::new();
18750 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
18751 h
18752 };
18753 let body = |resp: Response| async move {
18754 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18755 .await
18756 .unwrap();
18757 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
18758 };
18759
18760 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
18761 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
18762 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
18763 let payload = body(resp).await;
18764 assert!(
18765 payload["error"].is_object(),
18766 "bare-string error body: {payload}"
18767 );
18768 assert_eq!(payload["error"]["type"], "invalid_request_error");
18769 assert_eq!(payload["error"]["param"], "x-lane");
18770 assert_eq!(payload["error"]["code"], "invalid_lane");
18771
18772 let batch = auth::TenantCtx {
18773 tenant: "bulk".into(),
18774 lane_class: auth::LaneClass::Batch,
18775 rate_limit: None,
18776 key_prefix: None,
18777 };
18778 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
18779 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
18780 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
18781 let payload = body(resp).await;
18782 assert_eq!(payload["error"]["type"], "authentication_error");
18783 assert_eq!(payload["error"]["param"], "x-lane");
18784 }
18785
18786 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
18789
18790 fn drain_lock() -> std::sync::MutexGuard<'static, ()> {
18803 let guard = DRAIN_LOCK.lock().unwrap_or_else(|poisoned| {
18804 DRAIN_LOCK.clear_poison();
18807 poisoned.into_inner()
18808 });
18809 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18810 guard
18811 }
18812
18813 struct DrainingRestore;
18819 impl Drop for DrainingRestore {
18820 fn drop(&mut self) {
18821 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18822 }
18823 }
18824
18825 #[tokio::test]
18826 #[allow(clippy::await_holding_lock)] async fn responses_carry_rate_limit_headers_and_slot_frees() {
18828 let _l = drain_lock();
18829 let st = fake_worker_state();
18830 let resp = chat_completions(
18833 State(st.clone()),
18834 axum::http::HeaderMap::new(),
18835 None,
18836 Json(
18837 serde_json::from_value(serde_json::json!({
18838 "model": "m", "messages": [{"role": "user", "content": "t"}]
18839 }))
18840 .unwrap(),
18841 ),
18842 )
18843 .await;
18844 assert_eq!(resp.status(), StatusCode::OK);
18845 let h = resp.headers();
18846 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
18847 let remaining: usize = h["x-ratelimit-remaining"]
18848 .to_str()
18849 .unwrap()
18850 .parse()
18851 .unwrap();
18852 assert_eq!(remaining, limit - 1);
18853 assert_eq!(h["x-ratelimit-reset"], "0");
18854 assert_eq!(
18855 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18856 0,
18857 "slot must free at completion"
18858 );
18859 let resp = completions(
18862 State(st.clone()),
18863 axum::http::HeaderMap::new(),
18864 None,
18865 Json(
18866 serde_json::from_value(serde_json::json!({
18867 "model": "m", "prompt": "t", "stream": true
18868 }))
18869 .unwrap(),
18870 ),
18871 )
18872 .await;
18873 assert_eq!(resp.status(), StatusCode::OK);
18874 assert!(resp.headers().contains_key("x-ratelimit-limit"));
18875 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
18876 assert!(resp.headers().contains_key("x-ratelimit-reset"));
18877 assert_eq!(
18878 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18879 1,
18880 "stream in flight holds the slot"
18881 );
18882 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
18883 .await
18884 .unwrap();
18885 assert_eq!(
18886 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18887 0,
18888 "slot must free when the stream completes"
18889 );
18890 }
18891
18892 #[tokio::test]
18893 #[allow(clippy::await_holding_lock)] async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
18895 let _l = drain_lock();
18896 let mut st = fake_worker_state();
18897 let mock = MockMetering::admit_all();
18898 st.metering = Some(mock.clone());
18899
18900 let nonstream = chat_completions(
18901 State(st.clone()),
18902 HeaderMap::new(),
18903 None,
18904 Json(
18905 serde_json::from_value(json!({
18906 "model": "m",
18907 "messages": [{"role": "user", "content": "t"}],
18908 }))
18909 .unwrap(),
18910 ),
18911 )
18912 .await;
18913 assert_eq!(nonstream.status(), StatusCode::OK);
18914 let nonstream_id = nonstream.headers()["x-request-id"]
18915 .to_str()
18916 .unwrap()
18917 .to_string();
18918
18919 let stream = completions(
18920 State(st),
18921 HeaderMap::new(),
18922 None,
18923 Json(
18924 serde_json::from_value(json!({
18925 "model": "m",
18926 "prompt": "t",
18927 "stream": true,
18928 }))
18929 .unwrap(),
18930 ),
18931 )
18932 .await;
18933 assert_eq!(stream.status(), StatusCode::OK);
18934 let stream_id = stream.headers()["x-request-id"]
18935 .to_str()
18936 .unwrap()
18937 .to_string();
18938 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
18939 .await
18940 .unwrap();
18941
18942 let events = mock.events();
18946 let opened: Vec<&str> = events
18947 .iter()
18948 .filter_map(|e| match e {
18949 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
18950 _ => None,
18951 })
18952 .collect();
18953 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
18954 let completes = events
18955 .iter()
18956 .filter(|e| {
18957 matches!(
18958 e,
18959 MeterEvent::Complete {
18960 prompt: 1,
18961 cached: 0,
18962 completion: 1,
18963 }
18964 )
18965 })
18966 .count();
18967 assert_eq!(
18968 completes, 2,
18969 "both surfaces settle complete with worker-truth usage: {events:?}"
18970 );
18971 }
18972
18973 #[tokio::test]
18974 #[allow(clippy::await_holding_lock)] async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
18976 let _l = drain_lock();
18977 let mock = MockMetering::with_limits(vec![
18983 ReserveScript::Insufficient,
18984 ReserveScript::Admit { with_permit: false },
18985 ReserveScript::Blocked,
18986 ReserveScript::Admit { with_permit: true },
18987 ]);
18988 let mut st = fake_worker_state();
18989 st.metering = Some(mock.clone());
18990
18991 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
18993 assert_eq!(metrics.status(), StatusCode::OK);
18994 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
18995 .await
18996 .unwrap();
18997 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
18998 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
18999 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
19000 assert_eq!(metrics_body["budget_source_available"], true);
19001
19002 let request = || {
19003 Json(
19004 serde_json::from_value::<CompletionReq>(json!({
19005 "model": "m",
19006 "prompt_ids": [1],
19007 "max_tokens": 1,
19008 }))
19009 .unwrap(),
19010 )
19011 };
19012
19013 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19014 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
19015 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
19016 .await
19017 .unwrap();
19018 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
19019 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
19020 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
19021
19022 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19023 assert_eq!(included.status(), StatusCode::OK);
19024
19025 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19028 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
19029
19030 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
19031 assert_eq!(admitted.status(), StatusCode::OK);
19032
19033 let events = mock.events();
19034 let terminal: Vec<&MeterEvent> = events
19035 .iter()
19036 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
19037 .collect();
19038 assert_eq!(
19039 terminal.len(),
19040 4,
19041 "four requests, four terminal settles: {events:?}"
19042 );
19043 assert!(matches!(
19044 terminal[0],
19045 MeterEvent::Reject { status: 402, .. }
19046 ));
19047 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
19048 assert!(matches!(
19049 terminal[2],
19050 MeterEvent::Reject { status: 402, .. }
19051 ));
19052 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
19053 let permits: Vec<bool> = events
19055 .iter()
19056 .filter_map(|e| match e {
19057 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
19058 _ => None,
19059 })
19060 .collect();
19061 assert_eq!(
19062 permits,
19063 vec![false, false, false, true],
19064 "the permit rides the receipt exactly when reserve minted one: {events:?}"
19065 );
19066 }
19067
19068 #[tokio::test]
19072 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
19073 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
19074 let mut st = fake_worker_state();
19075 st.metering = Some(mock.clone());
19076 let tenant = auth::TenantCtx {
19077 tenant: "acme".into(),
19078 lane_class: auth::LaneClass::Interactive,
19079 rate_limit: None,
19080 key_prefix: Some("mk-acme-testprefix00".into()),
19081 };
19082 let mut request = gate_request(1, 1);
19083 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
19084 .expect_err("a capped key must be refused at admission");
19085 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
19086 let (response, outcome) = rejection.into_response();
19087 assert_eq!(outcome, "key_spend_cap_reached");
19088 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
19089 let body = body_value(response).await;
19090 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
19091 assert!(
19092 body["error"]["message"].as_str().unwrap().contains("cap"),
19093 "the 402 must point at the KEY's cap, not tenant credit: {body}"
19094 );
19095 let events = mock.events();
19096 assert!(
19097 events.contains(&MeterEvent::Reserve {
19098 tenant: "acme".into(),
19099 principal: Some("mk-acme-testprefix00".into()),
19100 model: "qwen/qwen3.8-27b".into(),
19101 }),
19102 "the key prefix must reach reserve: {events:?}"
19103 );
19104 }
19105
19106 #[tokio::test]
19107 #[allow(clippy::await_holding_lock)] async fn streaming_client_disconnect_records_partial_usage_and_cost() {
19109 let _l = drain_lock();
19110 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
19111 let mock = MockMetering::admit_all();
19112 st.metering = Some(mock.clone());
19113
19114 let response = completions(
19115 State(st),
19116 HeaderMap::new(),
19117 None,
19118 Json(
19119 serde_json::from_value(json!({
19120 "model": "m",
19121 "prompt": "disconnect after one delta",
19122 "stream": true,
19123 }))
19124 .unwrap(),
19125 ),
19126 )
19127 .await;
19128 assert_eq!(response.status(), StatusCode::OK);
19129 let request_id = response.headers()["x-request-id"]
19130 .to_str()
19131 .unwrap()
19132 .to_string();
19133 let mut body = Box::pin(response.into_body().into_data_stream());
19134 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
19135 .await
19136 .expect("stream ended before first delta")
19137 .expect("stream body failed");
19138 assert!(
19139 is_sse_data_frame(&first),
19140 "first frame was not SSE data: {first:?}"
19141 );
19142 drop(body);
19143
19144 let mut dropped = None;
19147 for _ in 0..500 {
19148 if let Some(event) = mock
19149 .events()
19150 .into_iter()
19151 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
19152 {
19153 dropped = Some(event);
19154 break;
19155 }
19156 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
19157 }
19158 let events = mock.events();
19159 assert!(
19160 events
19161 .iter()
19162 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
19163 "the receipt was opened under the caller-visible request id: {events:?}"
19164 );
19165 assert_eq!(
19166 dropped,
19167 Some(MeterEvent::Dropped {
19168 prompt: 1,
19169 cached: 0,
19170 completion: 1,
19171 }),
19172 "a client disconnect must leave the partial counts on the dropped receipt \
19173 (the implementation prices that drop): {events:?}"
19174 );
19175 }
19176
19177 #[tokio::test]
19178 #[allow(clippy::await_holding_lock)] async fn draining_rejects_new_requests_with_503_and_retry_after() {
19180 let _l = drain_lock();
19181 let st = fake_worker_state();
19182 let _down = DrainingRestore;
19185 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
19186 let resp = chat_completions(
19188 State(st.clone()),
19189 axum::http::HeaderMap::new(),
19190 None,
19191 Json(
19192 serde_json::from_value(serde_json::json!({
19193 "model": "m", "messages": [{"role": "user", "content": "t"}]
19194 }))
19195 .unwrap(),
19196 ),
19197 )
19198 .await;
19199 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19200 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
19205 let ra_s: u64 = ra
19206 .parse()
19207 .expect("Retry-After must be integer delay-seconds");
19208 assert!(
19209 ra_s > 0 && ra_s <= 60,
19210 "Retry-After {ra_s}s is outside the honored window"
19211 );
19212 let ra_ms: u64 = resp.headers()["retry-after-ms"]
19213 .to_str()
19214 .unwrap()
19215 .parse()
19216 .unwrap();
19217 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
19218 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19219 .await
19220 .unwrap();
19221 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19222 assert!(
19223 payload["error"]["message"]
19224 .as_str()
19225 .unwrap()
19226 .contains("draining")
19227 );
19228 assert_eq!(payload["error"]["type"], "server_error");
19229 assert_eq!(payload["error"]["code"], "draining");
19230 let resp = completions(
19231 State(st.clone()),
19232 axum::http::HeaderMap::new(),
19233 None,
19234 Json(
19235 serde_json::from_value(serde_json::json!({
19236 "model": "m", "prompt": "t"
19237 }))
19238 .unwrap(),
19239 ),
19240 )
19241 .await;
19242 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19243 assert!(resp.headers().contains_key("retry-after"));
19244 assert_eq!(
19245 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
19246 0,
19247 "rejected requests must not hold slots"
19248 );
19249 let resp = health_live(State(st.clone())).await.into_response();
19252 assert_eq!(
19253 resp.status(),
19254 StatusCode::OK,
19255 "a drain must not look like a liveness fault"
19256 );
19257 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19258 .await
19259 .unwrap();
19260 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19261 assert_eq!(payload["status"], "draining");
19262 let resp = health_ready(State(st.clone())).await.into_response();
19264 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19265 let retry_s = drain_deadline_s().clamp(1, 60);
19266 let retry_s_text = retry_s.to_string();
19267 let retry_ms_text = (retry_s * 1000).to_string();
19268 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
19269 assert_eq!(
19270 resp.headers().get("retry-after-ms").unwrap(),
19271 retry_ms_text.as_str()
19272 );
19273 assert_ne!(
19274 resp.headers()
19275 .get("x-should-retry")
19276 .and_then(|v| v.to_str().ok()),
19277 Some("false")
19278 );
19279 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19280 .await
19281 .unwrap();
19282 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19283 assert_eq!(payload["status"], "not_ready");
19284 assert!(payload["detail"].as_str().unwrap().contains("draining"));
19285 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
19286 let resp = chat_completions(
19288 State(st.clone()),
19289 axum::http::HeaderMap::new(),
19290 None,
19291 Json(
19292 serde_json::from_value(serde_json::json!({
19293 "model": "m", "messages": [{"role": "user", "content": "t"}]
19294 }))
19295 .unwrap(),
19296 ),
19297 )
19298 .await;
19299 assert_eq!(resp.status(), StatusCode::OK);
19300 }
19301
19302 #[tokio::test]
19305 #[allow(clippy::await_holding_lock)] async fn health_is_green_only_while_the_worker_is_alive() {
19307 let _l = drain_lock();
19310 let st = fake_worker_state();
19311 let resp = health_live(State(st.clone())).await.into_response();
19314 assert_eq!(resp.status(), StatusCode::OK);
19315 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19316 .await
19317 .unwrap();
19318 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19319 assert_eq!(payload["status"], "ok");
19320 assert_eq!(payload["worker"]["phase"], "idle");
19321 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
19322 let ready = health_ready(State(st.clone())).await.into_response();
19323 assert_eq!(ready.status(), StatusCode::OK);
19324
19325 st.health.mark_dead("worker thread panicked: test-injected");
19329 let resp = health_live(State(st.clone())).await.into_response();
19330 assert_eq!(
19331 resp.status(),
19332 StatusCode::SERVICE_UNAVAILABLE,
19333 "a dead worker MUST NOT report a healthy liveness"
19334 );
19335 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19336 .await
19337 .unwrap();
19338 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19339 assert_eq!(payload["status"], "unhealthy");
19340 assert!(
19342 payload["detail"]
19343 .as_str()
19344 .unwrap()
19345 .contains("test-injected"),
19346 "cause not surfaced: {payload}"
19347 );
19348 let ready = health_ready(State(st.clone())).await.into_response();
19349 assert_eq!(
19350 ready.status(),
19351 StatusCode::SERVICE_UNAVAILABLE,
19352 "dead is also not ready"
19353 );
19354
19355 st.health.mark_ready();
19358 assert_eq!(
19359 health_live(State(st.clone()))
19360 .await
19361 .into_response()
19362 .status(),
19363 StatusCode::OK,
19364 "mark_ready must clear the latch (a successful respawn)"
19365 );
19366 }
19367
19368 #[tokio::test]
19369 #[allow(clippy::await_holding_lock)] async fn readyz_peer_probe_integrity_is_present_and_advisory() {
19371 let _l = drain_lock();
19372 let st = fake_worker_state();
19373
19374 let ready = health_ready(State(st.clone())).await.into_response();
19375 assert_eq!(ready.status(), StatusCode::OK);
19376 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
19377 .await
19378 .unwrap();
19379 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19380 assert_eq!(payload["peer_probe_integrity"], "ok");
19381
19382 st.health.note_peer_probe_deferral(2, false);
19383 let deferred = health_ready(State(st.clone())).await.into_response();
19384 assert_eq!(deferred.status(), StatusCode::OK);
19385 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
19386 .await
19387 .unwrap();
19388 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19389 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
19390
19391 st.health.note_peer_probe_deferral(4, true);
19392 let degraded = health_ready(State(st.clone())).await.into_response();
19393 assert_eq!(
19394 degraded.status(),
19395 StatusCode::OK,
19396 "peer degradation is advisory while plain serving remains healthy"
19397 );
19398 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
19399 .await
19400 .unwrap();
19401 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19402 assert_eq!(payload["peer_probe_integrity"], "degraded");
19403
19404 st.health.mark_dead("test-injected worker failure");
19405 let unready = health_ready(State(st)).await.into_response();
19406 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
19407 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
19408 .await
19409 .unwrap();
19410 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19411 assert_eq!(
19412 payload["peer_probe_integrity"], "degraded",
19413 "the advisory field must also survive an unrelated readiness failure"
19414 );
19415 }
19416
19417 #[tokio::test]
19418 #[allow(clippy::await_holding_lock)] async fn liveness_failure_obeys_the_retry_contract() {
19420 let _l = drain_lock();
19425 let st = fake_worker_state();
19426 st.health
19427 .mark_dead("worker thread panicked: retry-contract-test");
19428
19429 let resp = health_live(State(st)).await.into_response();
19430 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19431 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
19432 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
19433 assert_ne!(
19434 resp.headers()
19435 .get("x-should-retry")
19436 .and_then(|v| v.to_str().ok()),
19437 Some("false")
19438 );
19439 }
19440
19441 #[tokio::test]
19442 #[allow(clippy::await_holding_lock)] async fn readiness_failure_obeys_the_retry_contract() {
19444 let _l = drain_lock();
19445 let st = fake_worker_state();
19446 st.health
19447 .mark_dead("worker thread panicked: retry-contract-test");
19448
19449 let resp = health_ready(State(st)).await.into_response();
19450 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19451 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
19452 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
19453 assert_ne!(
19454 resp.headers()
19455 .get("x-should-retry")
19456 .and_then(|v| v.to_str().ok()),
19457 Some("false")
19458 );
19459 }
19460
19461 #[tokio::test]
19462 #[allow(clippy::await_holding_lock)] async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
19464 let _l = drain_lock();
19474 let st = fake_worker_state();
19475 assert_eq!(
19476 health_live(State(st.clone()))
19477 .await
19478 .into_response()
19479 .status(),
19480 StatusCode::OK
19481 );
19482 st.health
19483 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
19484 let resp = health_live(State(st.clone())).await.into_response();
19485 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
19486 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
19487 .await
19488 .unwrap();
19489 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19490 assert!(
19491 payload["detail"]
19492 .as_str()
19493 .unwrap()
19494 .contains("probe exceeded")
19495 );
19496 st.health.mark_ready();
19499 assert_eq!(
19500 health_live(State(st.clone()))
19501 .await
19502 .into_response()
19503 .status(),
19504 StatusCode::SERVICE_UNAVAILABLE,
19505 "a GPU fault must not be cleared by an in-process respawn"
19506 );
19507 }
19508
19509 #[test]
19510 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
19511 let caps = ModelCaps {
19513 tools_branch: true,
19514 hy3: false,
19515 qwen_think: true,
19516 think_switch: true,
19517 chat_ok: true,
19518 context_length: 262144,
19519 tokenizer: "qwen2".into(),
19520 instruct_type: Some("chatml".into()),
19521 effort_levels: false,
19522 qwen_effort: false,
19523 gemma_think: false,
19524 dsv4: false,
19525 glm5: false,
19526 chat_temperature_default: None,
19527 chat_top_p_default: None,
19528 n_vocab: 151_936,
19529 think_close: Vec::new(),
19530 };
19531 let e = model_entry_v1("main", Some(&caps), None);
19532 assert_eq!(e["id"], "main");
19533 assert_eq!(e["name"], "main");
19534 assert_eq!(e["object"], "model");
19535 assert_eq!(e["context_length"], 262144);
19536 assert!(e["pricing"]["input"].is_null());
19538 assert!(e["pricing"]["output"].is_null());
19539
19540 let meta = OpenRouterModelMetadata {
19543 pricing: OpenRouterPricing {
19544 prompt: Some("0.00000038".into()),
19545 cached_prompt: Some("0.0000002".into()),
19546 completion: Some("0.0000026".into()),
19547 ..Default::default()
19548 },
19549 input_modalities: vec!["image".into(), "video".into()],
19550 max_output_length: Some(32768),
19551 ..Default::default()
19552 };
19553 let e = model_entry_v1("main", Some(&caps), Some(&meta));
19554 assert_eq!(e["pricing"]["currency"], "USD");
19557 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
19558 assert_eq!(e["pricing"]["input"], "0.38");
19559 assert_eq!(e["pricing"]["output"], "2.60");
19560 assert_eq!(e["pricing"]["cached_input"], "0.20");
19561 assert!(e["pricing"]["cache_write"].is_null());
19562 assert_eq!(e["pricing"]["minimum_request"], "0");
19563 assert_eq!(e["owned_by"], "main");
19564 assert_eq!(e["type"], "chat");
19565 assert_eq!(e["max_output_tokens"], 32768);
19566 assert_eq!(e["endpoints"], json!(["chat/completions"]));
19567 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
19568 assert_eq!(e["output_modalities"], json!(["text"]));
19569 assert_eq!(e["capabilities"]["streaming"], true);
19570 assert_eq!(e["capabilities"]["tools"], true);
19571 assert_eq!(e["lifecycle"]["status"], "active");
19572 assert!(e["lifecycle"]["deprecation_at"].is_null());
19573 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
19574 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
19575 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
19579 keys.sort_unstable();
19580 assert_eq!(
19581 keys,
19582 [
19583 "capabilities",
19584 "context_length",
19585 "endpoints",
19586 "id",
19587 "input_modalities",
19588 "lifecycle",
19589 "max_output_tokens",
19590 "name",
19591 "object",
19592 "output_modalities",
19593 "owned_by",
19594 "pricing",
19595 "reliability",
19596 "type",
19597 ],
19598 "unexpected /v1/models entry keys"
19599 );
19600 let mut price_keys: Vec<&str> = e["pricing"]
19601 .as_object()
19602 .unwrap()
19603 .keys()
19604 .map(String::as_str)
19605 .collect();
19606 price_keys.sort_unstable();
19607 assert_eq!(
19608 price_keys,
19609 [
19610 "cache_write",
19611 "cached_input",
19612 "currency",
19613 "input",
19614 "minimum_request",
19615 "output",
19616 "unit",
19617 ],
19618 "unexpected /v1/models pricing keys"
19619 );
19620
19621 let e = model_entry_v1("m", None, None);
19623 assert!(e["context_length"].is_null());
19624 assert!(e["max_output_tokens"].is_null());
19625 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
19627 assert!(e["context_length"].is_null());
19628 }
19629
19630 #[test]
19636 fn catalog_row_follows_the_declared_surface() {
19637 let caps = ModelCaps {
19638 tools_branch: true,
19639 ..Default::default()
19640 };
19641
19642 let embed = OpenRouterModelMetadata {
19643 surface: Some("embedding".into()),
19644 max_output_length: Some(1),
19645 ..Default::default()
19646 };
19647 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
19648 assert_eq!(e["type"], "embedding");
19649 assert_eq!(e["endpoints"], json!(["embeddings"]));
19650 assert_eq!(e["output_modalities"], json!(["embeddings"]));
19651 assert_eq!(e["capabilities"]["streaming"], false);
19652 assert_eq!(
19653 e["capabilities"]["tools"], false,
19654 "an embedder has no tools"
19655 );
19656 assert_eq!(e["capabilities"]["reasoning"], false);
19657 assert_eq!(e["capabilities"]["structured_output"], false);
19658 assert_eq!(e["capabilities"]["prompt_caching"], false);
19659 assert!(
19660 e["max_output_tokens"].is_null(),
19661 "a surface that emits no completion tokens must not advertise a ceiling"
19662 );
19663
19664 let rerank = OpenRouterModelMetadata {
19665 surface: Some("rerank".into()),
19666 ..Default::default()
19667 };
19668 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
19669 assert_eq!(r["type"], "rerank");
19670 assert_eq!(r["endpoints"], json!(["rerank"]));
19671 assert_eq!(r["output_modalities"], json!(["rerank"]));
19672 assert_eq!(r["capabilities"]["tools"], false);
19673 assert_eq!(r["capabilities"]["reasoning"], false);
19674
19675 let chat = OpenRouterModelMetadata {
19678 max_output_length: Some(32768),
19679 ..Default::default()
19680 };
19681 let c = model_entry_v1("main", Some(&caps), Some(&chat));
19682 assert_eq!(c["type"], "chat");
19683 assert_eq!(c["endpoints"], json!(["chat/completions"]));
19684 assert_eq!(c["output_modalities"], json!(["text"]));
19685 assert_eq!(c["capabilities"]["tools"], true);
19686 assert_eq!(c["max_output_tokens"], 32768);
19687 }
19688
19689 #[test]
19692 fn unknown_surface_is_rejected_at_config_load() {
19693 let bad = OpenRouterModelMetadata {
19694 surface: Some("embeddings".into()), ..Default::default()
19696 };
19697 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
19698 .expect_err("an unknown surface must not load");
19699 assert!(err.contains("surface"), "{err}");
19700
19701 for good in ["chat", "embedding", "rerank"] {
19702 let ok = OpenRouterModelMetadata {
19703 surface: Some(good.into()),
19704 ..Default::default()
19705 };
19706 assert!(
19707 validate_openrouter_metadata("m", &ok).is_ok(),
19708 "{good} must load"
19709 );
19710 }
19711 }
19712
19713 #[test]
19714 fn per_million_price_is_exact_decimal_shift() {
19715 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
19717 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
19718 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
19719 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
19720 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
19721 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
19722 assert_eq!(per_million_price("not-a-price"), None);
19723 assert_eq!(per_million_price(""), None);
19724 }
19725
19726 #[test]
19727 fn metadata_provider_block_parses_and_validates() {
19728 let (_, provider) = OpenRouterMetadataFile::parse(
19729 r#"
19730 [provider]
19731 id = "tiyuvta"
19732 status_url = "https://status.tiyuvta.ai"
19733 support_contact = "mailto:support@tiyuvta.ai"
19734 incident_contact = "mailto:incidents@tiyuvta.ai"
19735 regions = ["eu-central"]
19736 "#,
19737 )
19738 .unwrap();
19739 let provider = provider.unwrap();
19740 assert_eq!(provider.id, "tiyuvta");
19741 assert_eq!(provider.regions, vec!["eu-central"]);
19742 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
19744 assert!(err.contains("provider.id"), "{err}");
19745 let err = OpenRouterMetadataFile::parse(
19747 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
19748 )
19749 .unwrap_err();
19750 assert!(err.contains("must be a URI"), "{err}");
19751 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
19753 assert!(provider.is_none());
19754 }
19755
19756 #[test]
19757 fn models_openai_default_body_stays_byte_identical() {
19758 let body = models_openai_body(&["main".into(), "judge".into()]);
19759 let bytes = serde_json::to_vec(&body).unwrap();
19760 assert_eq!(
19761 bytes,
19762 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
19763 );
19764 }
19765
19766 #[test]
19767 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
19768 let loaded = vec![
19770 "qwen/qwen3.6-27b".to_string(),
19771 "qwen/qwen3.6-35b-a3b".to_string(),
19772 ];
19773 assert_eq!(
19774 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
19775 Some("qwen/qwen3.6-35b-a3b"),
19776 );
19777 assert_eq!(
19778 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
19779 Some("qwen/qwen3.6-27b"),
19780 );
19781 assert_eq!(
19783 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
19784 Some("qwen/qwen3.6-35b-a3b"),
19785 );
19786 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
19788 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
19789 assert_eq!(canonical_model_id(&loaded, ""), None);
19790 }
19791
19792 #[test]
19793 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
19794 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
19797 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
19798 assert_eq!(
19800 canonical_model_id(&loaded, "a/shared-name").as_deref(),
19801 Some("a/shared-name")
19802 );
19803 assert_eq!(
19804 canonical_model_id(&loaded, "b/shared-name").as_deref(),
19805 Some("b/shared-name")
19806 );
19807 let bare = vec!["solo".to_string()];
19809 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
19810 }
19811
19812 #[test]
19813 fn openrouter_models_entry_serializes_complete_metadata() {
19814 let metadata = OpenRouterMetadataFile::from_toml(
19815 r#"
19816[models.main]
19817hugging_face_id = "Qwen/Qwen3.6-27B"
19818created = 1786032000
19819quantization = "nvfp4"
19820description = "Qwen3.6 27B served by memra."
19821max_prompt_length = 245760
19822max_output_length = 16384
19823default_output_length = 4096
19824is_ready = true
19825is_free = false
19826discount_to_user = 0.1
19827openrouter_slug = "qwen/qwen3.6-27b"
19828datacenters = [{ country_code = "US", region = "us-east" }]
19829zdr = true
19830hipaa = false
19831
19832[models.main.pricing]
19833prompt = "0.000000234"
19834cached_prompt = "0.0000000585"
19835cache_write = "0.000000234"
19836completion = "0.000001872"
19837internal_reasoning = "0.000001872"
19838request = "0.01"
19839
19840[models.main.capacity]
19841prompt_tpm = 1000000
19842cached_prompt_tpm = 2000000
19843completion_tpm = 500000
19844request_rpm = 1000
19845concurrency = 64
19846"#,
19847 )
19848 .unwrap();
19849 let caps = ModelCaps {
19850 tools_branch: true,
19851 qwen_think: true,
19852 think_switch: true,
19853 chat_ok: true,
19854 context_length: 262144,
19855 tokenizer: "qwen2".into(),
19856 instruct_type: Some("chatml".into()),
19857 ..Default::default()
19858 };
19859 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
19860
19861 assert_eq!(entry["schema_version"], "2.4");
19862 assert_eq!(entry["id"], "main");
19863 assert_eq!(entry["name"], "main");
19864 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
19865 assert_eq!(entry["created"], 1786032000u64);
19866 assert_eq!(entry["quantization"], "nvfp4");
19867 assert_eq!(entry["tokenizer"], "qwen2");
19868 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
19869 assert!(
19870 entry.get("object").is_none(),
19871 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
19872 );
19873
19874 let input = &entry["input_modalities"][0];
19875 assert_eq!(input["type"], "text");
19876 assert_eq!(
19877 input["supported_inputs"]["max_context_length"]["value"],
19878 262144
19879 );
19880 assert_eq!(
19881 input["supported_inputs"]["max_prompt_length"]["value"],
19882 245760
19883 );
19884 let input_prices = input["pricing"].as_array().unwrap();
19885 let input_price = |kind: &str| {
19886 input_prices
19887 .iter()
19888 .find(|price| price["type"] == kind)
19889 .unwrap()
19890 };
19891 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
19892 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
19893 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
19894 assert_eq!(input["capacity"][0]["value"], 1000000);
19895 assert_eq!(input["capacity"][1]["value"], 2000000);
19896
19897 let output = &entry["output_modalities"][0];
19898 assert_eq!(output["type"], "text");
19899 assert_eq!(output["max_length"]["value"], 16384);
19900 assert_eq!(output["streaming"], true);
19901 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
19902 assert_eq!(
19903 output["supported_parameters"]["structured_outputs"]["type"],
19904 "boolean"
19905 );
19906 assert_eq!(
19907 output["supported_parameters"]["reasoning"]["type"],
19908 "boolean"
19909 );
19910 assert_eq!(output["pricing"][0]["type"], "completion");
19911 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
19912 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
19913 assert_eq!(output["capacity"][0]["value"], 500000);
19914 assert_eq!(output["capacity"][1]["type"], "concurrency");
19915 assert_eq!(output["capacity"][1]["value"], 64);
19916
19917 assert_eq!(entry["pricing"][0]["type"], "request");
19918 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
19919 assert_eq!(entry["capacity"][0]["value"], 1000);
19920 assert_eq!(entry["is_ready"], true);
19921 assert_eq!(entry["is_free"], false);
19922 assert_eq!(entry["discount_to_user"], 0.1);
19923 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
19924 assert_eq!(entry["datacenters"][0]["country_code"], "US");
19925 assert_eq!(entry["compliance"]["zdr"], true);
19926 assert_eq!(entry["compliance"]["hipaa"], false);
19927 }
19928
19929 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
19934[models."qwen/qwen3.6-35b-a3b"]
19935hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
19936created = 1777260255
19937quantization = "int4"
19938description = "Qwen3.6 35B-A3B fixture entry."
19939max_prompt_length = 262144
19940max_output_length = 262144
19941default_output_length = 8192
19942is_ready = true
19943is_free = false
19944discount_to_user = 0.0
19945openrouter_slug = "qwen/qwen3.6-35b-a3b"
19946zdr = false
19947hipaa = false
19948
19949[[models."qwen/qwen3.6-35b-a3b".datacenters]]
19950country_code = "CA"
19951region = "Ontario"
19952
19953[models."qwen/qwen3.6-35b-a3b".pricing]
19954prompt = "0.0000000931"
19955cached_prompt = "0.0000000652"
19956completion = "0.0000009025"
19957
19958[models."qwen/qwen3.6-35b-a3b".capacity]
19959prompt_tpm = 780000
19960cached_prompt_tpm = 310000
19961completion_tpm = 9600
19962request_rpm = 160
19963concurrency = 16
19964
19965[planned_models."qwen/qwen3.8-27b"]
19966description = "Planned fixture entry; must never be emitted."
19967max_prompt_length = 262144
19968max_output_length = 262144
19969default_output_length = 8192
19970is_ready = false
19971is_free = false
19972discount_to_user = 0.0
19973openrouter_slug = "qwen/qwen3.8-27b"
19974zdr = false
19975hipaa = false
19976
19977[planned_models."qwen/qwen3.8-27b".pricing]
19978prompt = "0.0000002745"
19979cached_prompt = "0.0000001922"
19980completion = "0.0000022800"
19981
19982[planned_models."google/gemma-4-26b-a4b-it"]
19983hugging_face_id = "google/gemma-4-26B-A4B-it"
19984created = 1775227989
19985quantization = "int4"
19986description = "Planned fixture entry; must never be emitted."
19987max_prompt_length = 262144
19988max_output_length = 262144
19989default_output_length = 8192
19990is_ready = false
19991is_free = false
19992discount_to_user = 0.0
19993openrouter_slug = "google/gemma-4-26b-a4b-it"
19994zdr = false
19995hipaa = false
19996
19997[planned_models."google/gemma-4-26b-a4b-it".pricing]
19998prompt = "0.0000000665"
19999cached_prompt = "0.0000000466"
20000completion = "0.0000003230"
20001"#;
20002
20003 #[test]
20004 fn gateway_registry_generates_the_staged_active_shape() {
20005 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
20006 let caps = ModelCaps {
20007 tools_branch: true,
20008 qwen_think: true,
20009 think_switch: true,
20010 chat_ok: true,
20011 context_length: 262144,
20012 tokenizer: "qwen2".into(),
20013 instruct_type: Some("chatml".into()),
20014 ..Default::default()
20015 };
20016 let q35_entry = model_entry_openrouter(
20017 "qwen/qwen3.6-35b-a3b",
20018 Some(&caps),
20019 metadata.get("qwen/qwen3.6-35b-a3b"),
20020 );
20021 assert_eq!(q35_entry["created"], 1777260255u64);
20022 assert_eq!(q35_entry["quantization"], "int4");
20023 assert_eq!(q35_entry["is_ready"], true);
20024 assert_eq!(
20025 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
20026 262144
20027 );
20028 assert_eq!(
20029 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
20030 262144
20031 );
20032 assert_eq!(
20033 q35_entry["output_modalities"][0]["max_length"]["value"],
20034 262144
20035 );
20036 let prices = q35_entry["input_modalities"][0]["pricing"]
20037 .as_array()
20038 .unwrap();
20039 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
20040 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
20041 assert_eq!(
20045 q35_entry["input_modalities"][0]["capacity"][0]["value"],
20046 780000
20047 );
20048 assert_eq!(
20049 q35_entry["input_modalities"][0]["capacity"][1]["value"],
20050 310000
20051 );
20052 assert_eq!(
20053 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
20054 262144
20055 );
20056 assert_eq!(
20057 q35_entry["output_modalities"][0]["capacity"][0]["value"],
20058 9600
20059 );
20060 assert_eq!(
20061 q35_entry["output_modalities"][0]["capacity"][1]["value"],
20062 16
20063 );
20064 assert_eq!(
20065 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
20066 "0.0000009025"
20067 );
20068 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
20070
20071 assert_eq!(
20072 metadata.len(),
20073 1,
20074 "planned models must never enter the active map"
20075 );
20076 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
20077 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
20078 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
20079
20080 let openmodels = model_entry_openmodels(
20081 "qwen/qwen3.6-35b-a3b",
20082 Some(&caps),
20083 metadata.get("qwen/qwen3.6-35b-a3b"),
20084 )
20085 .unwrap();
20086 assert_eq!(openmodels["currency"], "USD");
20087 assert_eq!(openmodels["max_output_length"], 262144);
20088 assert_eq!(openmodels["is_ready"], true);
20089 assert_eq!(openmodels["is_free"], false);
20090 assert_eq!(openmodels["discount_to_user"], 0.0);
20091 }
20092
20093 #[test]
20094 fn gateway_registry_limits_are_live_request_limits() {
20095 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
20096 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
20097 let caps = ModelCaps {
20098 context_length: 262_144,
20099 ..Default::default()
20100 };
20101 let build = |value: serde_json::Value| {
20102 let req: CompletionReq = serde_json::from_value(value).unwrap();
20103 let (tx, _rx) = worker::event_channel();
20104 build_request(&req, tx, lanes::Lane::Interactive, None)
20105 };
20106
20107 let mut omitted = build(json!({
20108 "model": "qwen/qwen3.6-35b-a3b",
20109 "prompt_ids": [1, 2, 3]
20110 }));
20111 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
20112 assert_eq!(omitted.params.max_new, 8_192);
20113 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
20114
20115 let mut field_top = build(json!({
20116 "model": "qwen/qwen3.6-35b-a3b",
20117 "prompt_ids": [1],
20118 "max_tokens": 262144
20119 }));
20120 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
20121 assert_eq!(field_top.params.max_new, 262_144);
20122 assert_eq!(
20123 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
20124 262_044,
20125 "the field-top output request is accepted but bounded by remaining trained context",
20126 );
20127
20128 let mut too_much_output = build(json!({
20129 "model": "qwen/qwen3.6-35b-a3b",
20130 "prompt_ids": [1],
20131 "max_tokens": 262145
20132 }));
20133 let (message, param) =
20134 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
20135 .unwrap_err();
20136 assert_eq!(param, "max_tokens");
20137 assert!(message.contains("262145"));
20138
20139 let mut oversized_allocation = build(json!({
20140 "model": "qwen/qwen3.6-35b-a3b",
20141 "prompt_ids": [1],
20142 "max_tokens": 1,
20143 "max_ctx": 262145
20144 }));
20145 let (_, param) =
20146 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
20147 .unwrap_err();
20148 assert_eq!(param, "max_ctx");
20149 }
20150
20151 #[test]
20152 fn planned_registry_entries_are_validated_but_never_activated() {
20153 let parsed = OpenRouterMetadataFile::from_toml(
20154 r#"
20155[planned_models.future]
20156max_output_length = 262144
20157default_output_length = 8192
20158
20159[planned_models.future.pricing]
20160prompt = "0.0000001"
20161"#,
20162 )
20163 .unwrap();
20164 assert!(parsed.is_empty());
20165
20166 let error = OpenRouterMetadataFile::from_toml(
20167 r#"
20168[planned_models.future]
20169default_output_length = 8192
20170"#,
20171 )
20172 .unwrap_err();
20173 assert!(error.contains("requires max_output_length"));
20174 }
20175
20176 #[test]
20181 fn every_catalog_feed_honours_the_declared_surface() {
20182 let metadata = OpenRouterMetadataFile::from_toml(
20183 r#"
20184[models."qwen/qwen3-embedding-8b"]
20185surface = "embedding"
20186created = 1787961600
20187max_output_length = 1
20188is_ready = true
20189is_free = false
20190discount_to_user = 0.0
20191
20192[models."qwen/qwen3-embedding-8b".pricing]
20193prompt = "0.00000001"
20194cached_prompt = "0.0"
20195completion = "0.0"
20196
20197[models."main"]
20198created = 1787443200
20199max_output_length = 32768
20200is_ready = true
20201is_free = false
20202discount_to_user = 0.0
20203
20204[models."main".pricing]
20205prompt = "0.00000025"
20206cached_prompt = "0.00000009"
20207completion = "0.0000012"
20208"#,
20209 )
20210 .unwrap();
20211 let caps = ModelCaps {
20212 tools_branch: true,
20213 qwen_think: true,
20214 think_switch: true,
20219 chat_ok: true,
20220 context_length: 32768,
20221 ..Default::default()
20222 };
20223 let embed = metadata.get("qwen/qwen3-embedding-8b");
20224 let chat = metadata.get("main");
20225
20226 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
20228 let out = &or["output_modalities"][0];
20229 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
20230 assert!(
20231 out.get("streaming").is_none(),
20232 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
20233 );
20234 let params = &out["supported_parameters"];
20239 assert_eq!(
20240 params.as_object().map(|o| o.len()),
20241 Some(0),
20242 "no completion parameter belongs on an embedder row: {params}"
20243 );
20244 for field in [
20245 "tools",
20246 "tool_choice",
20247 "reasoning",
20248 "max_tokens",
20249 "json_mode",
20250 "structured_outputs",
20251 "stop",
20252 "temperature",
20253 "seed",
20254 ] {
20255 assert!(params[field].is_null(), "{field} leaked onto an embedder");
20256 }
20257 assert!(
20258 out["max_length"].is_null(),
20259 "a surface emitting no completion tokens advertises no ceiling: {out}"
20260 );
20261
20262 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
20264 .expect("openmodels entry builds");
20265 assert_eq!(om["output_modalities"], json!(["embeddings"]));
20266 let features = om["supported_features"].as_array().unwrap();
20267 assert!(
20268 !features
20269 .iter()
20270 .any(|f| f == "tool_calling" || f == "reasoning"),
20271 "chat-only features leaked onto an embedder: {features:?}"
20272 );
20273
20274 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
20276 assert_eq!(v1["type"], "embedding");
20277 assert_eq!(v1["capabilities"]["tools"], false);
20278
20279 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
20281 let out_chat = &or_chat["output_modalities"][0];
20282 assert_eq!(out_chat["type"], "text");
20283 assert_eq!(out_chat["streaming"], true);
20284 assert!(!out_chat["supported_parameters"]["tools"].is_null());
20285 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
20286 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
20287 assert_eq!(out_chat["max_length"]["value"], 32768u64);
20288 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
20289 assert_eq!(om_chat["output_modalities"], json!(["text"]));
20290 assert!(
20291 om_chat["supported_features"]
20292 .as_array()
20293 .unwrap()
20294 .iter()
20295 .any(|f| f == "tool_calling")
20296 );
20297 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
20298 }
20299
20300 #[test]
20307 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
20308 let raw = std::fs::read_to_string(concat!(
20309 env!("CARGO_MANIFEST_DIR"),
20310 "/../../research/gateway-20260812/raw/sources/",
20311 "openrouter-provider-schema-v2.4-20260812.json"
20312 ))
20313 .expect("vendored Provider Monitor 2.4 schema is in-tree");
20314 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
20315 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
20316 .as_array()
20317 .expect("OutputModality is a oneOf");
20318
20319 let metadata = OpenRouterMetadataFile::from_toml(
20320 r#"
20321[models."embed"]
20322surface = "embedding"
20323created = 1787961600
20324max_output_length = 1
20325is_ready = true
20326is_free = false
20327discount_to_user = 0.0
20328
20329[models."embed".pricing]
20330prompt = "0.00000001"
20331cached_prompt = "0.0"
20332completion = "0.0"
20333
20334[models."rr"]
20335surface = "rerank"
20336created = 1787961600
20337max_output_length = 1
20338is_ready = true
20339is_free = false
20340discount_to_user = 0.0
20341
20342[models."rr".pricing]
20343prompt = "0.00000003"
20344cached_prompt = "0.0"
20345completion = "0.0"
20346
20347[models."chatty"]
20348created = 1787443200
20349max_output_length = 32768
20350is_ready = true
20351is_free = false
20352discount_to_user = 0.0
20353
20354[models."chatty".pricing]
20355prompt = "0.00000025"
20356cached_prompt = "0.00000009"
20357completion = "0.0000012"
20358"#,
20359 )
20360 .unwrap();
20361 let caps = ModelCaps {
20362 tools_branch: true,
20363 qwen_think: true,
20364 chat_ok: true,
20365 context_length: 32768,
20366 ..Default::default()
20367 };
20368
20369 for (alias, want_type) in [
20370 ("embed", "embeddings"),
20371 ("rr", "rerank"),
20372 ("chatty", "text"),
20373 ] {
20374 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
20375 let modality = &row["output_modalities"][0];
20376 assert_eq!(modality["type"], want_type, "{alias}: {row}");
20377
20378 let branch = branches
20380 .iter()
20381 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
20382 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
20383 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
20384 .as_object()
20385 .expect("branch properties")
20386 .keys()
20387 .map(String::as_str)
20388 .collect();
20389 for key in modality.as_object().expect("modality object").keys() {
20390 assert!(
20391 allowed.contains(key.as_str()),
20392 "{alias}: {key:?} is not a property of the {want_type:?} branch \
20393 (additionalProperties:false); allowed = {allowed:?}"
20394 );
20395 }
20396 for req in branch["required"].as_array().into_iter().flatten() {
20397 let req = req.as_str().expect("required entry is a string");
20398 assert!(
20399 modality.get(req).is_some(),
20400 "{alias}: required property {req:?} missing from the {want_type:?} branch"
20401 );
20402 }
20403 }
20404 }
20405
20406 #[test]
20407 fn openrouter_models_entry_omits_undeclared_optional_fields() {
20408 let entry = model_entry_openrouter("minimal", None, None);
20409 let object = entry.as_object().unwrap();
20410 for field in [
20411 "hugging_face_id",
20412 "created",
20413 "quantization",
20414 "tokenizer",
20415 "description",
20416 "pricing",
20417 "capacity",
20418 "is_ready",
20419 "is_free",
20420 "discount_to_user",
20421 "openrouter",
20422 "datacenters",
20423 "compliance",
20424 ] {
20425 assert!(
20426 !object.contains_key(field),
20427 "optional field {field} must be absent, not null"
20428 );
20429 }
20430 assert_eq!(entry["schema_version"], "2.4");
20431 assert_eq!(entry["input_modalities"][0]["type"], "text");
20432 assert!(
20433 entry["input_modalities"][0]
20434 .get("supported_inputs")
20435 .is_none()
20436 );
20437 assert!(entry["input_modalities"][0].get("pricing").is_none());
20438 assert!(entry["input_modalities"][0].get("capacity").is_none());
20439 assert_eq!(entry["output_modalities"][0]["type"], "text");
20440 assert_eq!(entry["output_modalities"][0]["streaming"], true);
20441 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
20442 assert!(entry["output_modalities"][0].get("max_length").is_none());
20443 assert!(entry["output_modalities"][0].get("pricing").is_none());
20444 assert!(entry["output_modalities"][0].get("capacity").is_none());
20445 }
20446
20447 #[test]
20448 fn openmodels_entry_serializes_standard_provider_shape() {
20449 let metadata = OpenRouterMetadataFile::from_toml(
20450 r#"
20451[models."qwen/qwen3.6-27b"]
20452created = 1786032000
20453max_output_length = 16384
20454is_ready = true
20455is_free = false
20456discount_to_user = 0.05
20457
20458[models."qwen/qwen3.6-27b".pricing]
20459prompt = "0.000000291"
20460cached_prompt = "0.000000291"
20461completion = "0.000002763"
20462request = "0"
20463"#,
20464 )
20465 .unwrap();
20466 let caps = ModelCaps {
20467 tools_branch: true,
20468 qwen_think: true,
20469 chat_ok: true,
20470 context_length: 262144,
20471 ..Default::default()
20472 };
20473 let entry = model_entry_openmodels(
20474 "qwen/qwen3.6-27b",
20475 Some(&caps),
20476 metadata.get("qwen/qwen3.6-27b"),
20477 )
20478 .unwrap();
20479
20480 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
20481 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
20482 assert_eq!(entry["created"], 1786032000u64);
20483 assert_eq!(entry["input_modalities"], json!(["text"]));
20484 assert_eq!(entry["output_modalities"], json!(["text"]));
20485 assert_eq!(entry["context_length"], 262144u64);
20486 assert_eq!(entry["max_output_length"], 16384u64);
20487 assert_eq!(entry["currency"], "USD");
20488 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
20489 assert_eq!(entry["pricing"]["completion"], "0.000002763");
20490 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
20491 assert_eq!(entry["pricing"]["request"], "0");
20492 assert_eq!(
20493 entry["supported_features"],
20494 json!(["tool_calling", "reasoning"])
20495 );
20496 assert_eq!(entry["is_ready"], true);
20497 assert_eq!(entry["is_free"], false);
20498 assert_eq!(entry["discount_to_user"], 0.05);
20499 assert!(entry.get("schema_version").is_none());
20500 assert!(entry.get("quantization").is_none());
20501 }
20502
20503 #[test]
20504 fn openmodels_entry_rejects_missing_operator_metadata() {
20505 let caps = ModelCaps {
20506 context_length: 262144,
20507 ..Default::default()
20508 };
20509 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
20510 assert_eq!(
20511 error,
20512 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
20513 );
20514 }
20515
20516 #[tokio::test]
20517 async fn blocking_response_excludes_stop_text_across_token_events() {
20518 let (tx, rx) = worker::event_channel();
20519 tx.send(Event::Token {
20520 id: 1,
20521 text: "answer\nPro".into(),
20522 })
20523 .unwrap();
20524 tx.send(Event::Token {
20525 id: 2,
20526 text: "blem: leaked prompt".into(),
20527 })
20528 .unwrap();
20529 tx.send(Event::Done {
20530 stop_reason: "Callback".into(),
20531 n_tokens: 2,
20532 n_prompt: 8,
20533 n_cached: 0,
20534 elapsed_s: 0.5,
20535 spec: None,
20536 })
20537 .unwrap();
20538 drop(tx);
20539 let response = blocking_response(
20540 rx,
20541 "plain_quant".into(),
20542 false,
20543 vec!["Problem:".into()],
20544 None,
20545 Envelope::new(false),
20546 )
20547 .await;
20548 assert_eq!(response.status(), StatusCode::OK);
20549 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
20550 .await
20551 .unwrap();
20552 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
20553 assert_eq!(payload["text"], "answer\n");
20554 assert_eq!(payload["stop_reason"], "Callback");
20555 }
20556
20557 #[test]
20561 fn step_walker_expansion_and_separator_law() {
20562 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
20564 let uri = format!("data:image/png;base64,{PNG64}");
20565 let content = serde_json::json!([
20566 {"type": "text", "text": "look at"},
20567 {"type": "text", "text": "this:"},
20568 {"type": "image_url", "image_url": {"url": uri}},
20569 {"type": "text", "text": "what is it?"},
20570 ]);
20571 let mut pending: Vec<PendingStepImage> = Vec::new();
20572 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
20573 let mut expansion = String::from("<im_start>");
20574 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
20575 expansion.push_str("<im_patch>");
20576 }
20577 expansion.push_str("<im_end>");
20578 assert_eq!(out, format!("look at this:{expansion}what is it?"));
20581 assert_eq!(pending.len(), 1);
20582 assert_eq!(pending[0].plan.n_tiles, 0);
20583 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
20584
20585 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
20587 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
20588 let http = serde_json::json!([
20589 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
20590 ]);
20591 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
20592 }
20593}
20594
20595#[cfg(test)]
20601mod build_identity_tests {
20602 use super::{BUILD_GIT_SHA, BUILD_ID_NOTE, BUILD_ID_SRC, SYSTEM_FINGERPRINT, build_id};
20603
20604 #[test]
20606 fn baked_fingerprint_is_real_and_well_formed() {
20607 assert!(!SYSTEM_FINGERPRINT.is_empty());
20608 assert_ne!(SYSTEM_FINGERPRINT, "memra-unknown");
20609 assert!(
20610 !SYSTEM_FINGERPRINT.contains("unknown"),
20611 "fingerprint {SYSTEM_FINGERPRINT:?} still carries the degraded literal"
20612 );
20613 assert!(
20614 build_id::fingerprint_is_well_formed(SYSTEM_FINGERPRINT),
20615 "fingerprint {SYSTEM_FINGERPRINT:?} is not memra-<version>-<12 hex>"
20616 );
20617 assert!(
20620 SYSTEM_FINGERPRINT.starts_with(concat!("memra-", env!("CARGO_PKG_VERSION"), "-")),
20621 "fingerprint {SYSTEM_FINGERPRINT:?} does not name this crate version"
20622 );
20623 }
20624
20625 #[test]
20628 fn the_shape_check_rejects_what_shipped_to_prod() {
20629 assert!(!build_id::fingerprint_is_well_formed("memra-unknown"));
20630 assert!(!build_id::fingerprint_is_well_formed(
20631 "memra-0.123.0-unknown"
20632 ));
20633 let old_form = format!("memra-{}", "0".repeat(12));
20638 assert!(!build_id::fingerprint_is_well_formed(&old_form));
20639 assert!(!build_id::fingerprint_is_well_formed(""));
20640 assert!(!build_id::fingerprint_is_well_formed("memra-"));
20641 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-"));
20642 assert!(!build_id::fingerprint_is_well_formed("memra-0.123.0-abc"));
20644 assert!(!build_id::fingerprint_is_well_formed(
20645 "memra-0.123.0-ABCDEF012345"
20646 ));
20647 assert!(!build_id::fingerprint_is_well_formed(
20648 "memra-0.123.0-zzzzzzzzzzzz"
20649 ));
20650 assert!(build_id::fingerprint_is_well_formed(
20652 "memra-0.123.0-4b1f9c02d7a3"
20653 ));
20654 }
20655
20656 #[test]
20664 fn build_id_is_rederivable_from_the_source_tree() {
20665 let root = build_id::workspace_root(env!("CARGO_MANIFEST_DIR"));
20666 let scan = root.as_deref().and_then(build_id::content_id);
20667 match scan {
20668 Some(scan) => {
20669 assert_eq!(
20670 BUILD_ID_SRC,
20671 build_id::BUILD_ID_SRC_TREE,
20672 "the source tree is readable, so the baked id must come from it"
20673 );
20674 assert!(BUILD_ID_NOTE.is_empty(), "note set on a non-degraded build");
20675 let expected =
20676 format!(concat!("memra-", env!("CARGO_PKG_VERSION"), "-{}"), scan.id);
20677 assert_eq!(
20678 SYSTEM_FINGERPRINT,
20679 expected,
20680 "baked fingerprint disagrees with a re-derivation over {} files: the id \
20681 is not a pure function of the source tree, or the build script did not \
20682 re-run after an edit",
20683 scan.files.len()
20684 );
20685 assert!(scan.files.len() > 100, "suspiciously small hashed file set");
20686 }
20687 None => {
20688 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
20692 assert!(
20693 !BUILD_ID_NOTE.is_empty(),
20694 "a degraded build must state its reason so the boot WARN can print it"
20695 );
20696 }
20697 }
20698 }
20699
20700 #[test]
20703 fn identity_is_independent_of_git_history() {
20704 let id = SYSTEM_FINGERPRINT.rsplit_once('-').unwrap().1;
20705 assert_ne!(
20706 id, BUILD_GIT_SHA,
20707 "the content id equals the git sha; the identity must not be history, it has to \
20708 survive a rewrite that changes every commit"
20709 );
20710 assert!(
20711 !SYSTEM_FINGERPRINT.contains(BUILD_GIT_SHA),
20712 "the git sha leaked into the customer-visible fingerprint {SYSTEM_FINGERPRINT:?}"
20713 );
20714 assert!(!BUILD_GIT_SHA.is_empty());
20717 }
20718
20719 #[test]
20722 fn content_digest_is_deterministic_and_change_sensitive() {
20723 let a = build_id::degraded_build_id("memra-server", "0.123.0");
20724 let b = build_id::degraded_build_id("memra-server", "0.123.0");
20725 assert_eq!(a, b, "the digest is not deterministic");
20726 assert_eq!(a.len(), build_id::BUILD_ID_HEX);
20727 assert!(
20728 a.chars()
20729 .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
20730 );
20731 assert_ne!(a, build_id::degraded_build_id("memra-server", "0.123.1"));
20732 assert_ne!(a, build_id::degraded_build_id("memra-serve", "r0.123.0"));
20733 assert_eq!(build_id::render_build_id(0).len(), build_id::BUILD_ID_HEX);
20735 assert_eq!(
20736 build_id::render_build_id(0),
20737 "0".repeat(build_id::BUILD_ID_HEX)
20738 );
20739 }
20740
20741 #[test]
20744 fn two_scans_of_one_tree_agree() {
20745 let Some(root) = build_id::workspace_root(env!("CARGO_MANIFEST_DIR")) else {
20746 assert_eq!(BUILD_ID_SRC, build_id::BUILD_ID_SRC_DEGRADED);
20747 return;
20748 };
20749 let first = build_id::content_id(&root).expect("first scan");
20750 let second = build_id::content_id(&root).expect("second scan");
20751 assert_eq!(first.id, second.id);
20752 assert_eq!(first.files.len(), second.files.len());
20753 }
20754}
20755
20756#[cfg(test)]
20762mod vision_placement_gate_tests {
20763 use super::vision_media_admissible;
20764
20765 #[test]
20766 fn a_media_part_is_admitted_only_when_the_placement_admits() {
20767 assert_eq!(vision_media_admissible(true, "image"), Ok(()));
20768 assert_eq!(vision_media_admissible(true, "video"), Ok(()));
20769 let err = vision_media_admissible(false, "image").unwrap_err();
20770 assert!(
20771 err.starts_with("image input is not enabled on this deployment"),
20772 "same named refusal the armed-off path gives, so clients see one contract: {err}"
20773 );
20774 assert!(
20775 err.contains("placement"),
20776 "the refusal names its cause: {err}"
20777 );
20778 let err = vision_media_admissible(false, "video").unwrap_err();
20779 assert!(
20780 err.starts_with("video input is not enabled on this deployment"),
20781 "{err}"
20782 );
20783 }
20784
20785 fn live_src() -> String {
20786 let src: String = include_str!("lib.rs")
20787 .lines()
20788 .map(|l| match l.find("//") {
20789 Some(i) => &l[..i],
20790 None => l,
20791 })
20792 .collect::<Vec<_>>()
20793 .join("\n");
20794 let end = src
20795 .find("\nmod vision_placement_gate_tests")
20796 .expect("this test module exists");
20797 src[..end].to_string()
20798 }
20799
20800 fn item_body<'a>(live: &'a str, head: &str) -> &'a str {
20802 let start = live
20803 .find(head)
20804 .unwrap_or_else(|| panic!("{head} not found — did it get renamed?"));
20805 let body = &live[start..];
20806 let end = body.find("\n}\n").expect("item body closes");
20807 &body[..end]
20808 }
20809
20810 fn head_of(s: &str, n: usize) -> &str {
20812 match s.char_indices().nth(n) {
20813 Some((i, _)) => &s[..i],
20814 None => s,
20815 }
20816 }
20817
20818 #[test]
20823 fn no_family_switch_reads_the_placement_decision() {
20824 let live = live_src();
20825 for switch in [
20826 "fn vision_enabled()",
20827 "fn gemma_vision_enabled()",
20828 "fn step_vision_enabled()",
20829 ] {
20830 let body = item_body(&live, switch);
20831 assert!(
20832 !body.contains("vision_placement_serving")
20833 && !body.contains("vision_placement_admits"),
20834 "{switch} routes text rendering; it must stay keyed on the operator knobs alone"
20835 );
20836 }
20837 let walker = item_body(&live, "fn content_to_text_vision(");
20838 assert!(
20839 walker.contains(
20840 "if step_vision_enabled() {\n return content_to_text_vision_step(v, step_images);"
20841 ),
20842 "the step walker dispatch is keyed on the armed switch alone"
20843 );
20844 }
20845
20846 #[test]
20849 fn every_media_accepting_arm_passes_the_placement_gate() {
20850 let live = live_src();
20851 let step = item_body(&live, "fn content_to_text_vision_step(");
20852 let arm = step
20853 .split("Some(\"image_url\") => {")
20854 .nth(1)
20855 .expect("the step walker has an image arm");
20856 assert!(
20857 head_of(arm, 120).contains("vision_placement_admits(\"image\")?;"),
20858 "the step image arm must pass the placement gate first: {}",
20859 head_of(arm, 120)
20860 );
20861 let walker = item_body(&live, "fn content_to_text_vision(");
20862 for (head, kind) in [
20863 (
20864 "Some(\"image_url\") if gemma_vision_enabled() => {",
20865 "image",
20866 ),
20867 ("Some(\"image_url\") => {", "image"),
20868 ("Some(\"video_url\") => {", "video"),
20869 ] {
20870 let arm = walker
20871 .split(head)
20872 .nth(1)
20873 .unwrap_or_else(|| panic!("{head} is not an arm of the walker"));
20874 let window = head_of(arm, 400);
20875 assert!(
20876 window.contains(&format!("vision_placement_admits(\"{kind}\")?;")),
20877 "{head} must pass the placement gate before planning anything: {window}"
20878 );
20879 }
20880 assert!(live.contains("GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)"));
20884 let gate = item_body(&live, "fn vision_placement_admits(");
20886 assert!(gate.contains("vision_media_admissible(vision_placement_serving(), kind)"));
20887 }
20888}