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;
80mod anthropic;
86mod dsv4_serve;
87mod embed_api;
88pub mod metering;
95mod responses_api;
96mod surfaces;
97mod toolcall;
98mod ttft;
99mod worker;
100
101use std::collections::HashMap;
102use std::net::{SocketAddr, ToSocketAddrs};
103use std::sync::Arc;
104use std::sync::mpsc::Sender;
105
106use axum::{
107 Extension, Json, Router,
108 body::Body,
109 extract::{DefaultBodyLimit, Query, Request as AxumRequest, State},
110 http::{
111 HeaderMap, StatusCode,
112 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
113 },
114 middleware::{self, Next},
115 response::{
116 IntoResponse, Response,
117 sse::{Event as SseEvent, Sse},
118 },
119 routing::{get, post},
120};
121use futures_core::Stream as _;
122use serde::{Deserialize, Serialize};
123use serde_json::json;
124
125use memra_engine::decode::GenParams;
126use memra_engine::sampler::SamplerConfig;
127use memra_tokenizer::{
128 Tokenizer,
129 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
130};
131use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
132use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
133
134const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
158const MAX_BODY_ADMISSIONS: usize = 4;
159const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
160#[allow(clippy::identity_op)] const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
165const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
166const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
167const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
168const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
169const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
170
171fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
172 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
173 SEMAPHORE
174 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
175 .clone()
176}
177
178fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
179 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
180 SEMAPHORE
181 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
182 .clone()
183}
184
185fn declared_body_length(req: &AxumRequest) -> Option<usize> {
186 req.headers()
187 .get(CONTENT_LENGTH)
188 .and_then(|value| value.to_str().ok())
189 .and_then(|value| value.parse().ok())
190}
191
192fn body_requires_admission(req: &AxumRequest) -> bool {
193 if req.headers().contains_key(TRANSFER_ENCODING) {
197 return true;
198 }
199 declared_body_length(req).is_none_or(|length| length > BODY_ADMISSION_BYPASS_BYTES)
200}
201
202fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
206 let Some(length) = declared_body_length(req) else {
207 return BODY_READ_TIMEOUT;
208 };
209 let bytes = length as u64;
210 let extra_seconds =
211 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
212 let seconds = BODY_READ_TIMEOUT
213 .as_secs()
214 .saturating_add(extra_seconds)
215 .min(BODY_READ_TIMEOUT_MAX.as_secs());
216 std::time::Duration::from_secs(seconds)
217}
218
219async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
224 let resp = next.run(req).await;
225 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
226 return resp;
227 }
228 error_response_coded(
229 StatusCode::PAYLOAD_TOO_LARGE,
230 &format!(
231 "request body exceeds the {} MiB limit",
232 MAX_BODY_BYTES / (1024 * 1024)
233 ),
234 "invalid_request_error",
235 None,
236 Some("request_too_large"),
237 )
238}
239
240fn apply_body_limit(app: Router) -> Router {
243 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
244 .layer(middleware::from_fn(shape_payload_too_large))
245}
246
247fn protected_inference_path(path: &str) -> bool {
248 matches!(
249 path,
250 "/v1/auth/check"
251 | "/v1/completions"
252 | "/v1/chat/completions"
253 | "/v1/messages"
254 | "/v1/responses"
255 | "/v1/embeddings"
256 | "/v1/rerank"
257 )
258}
259
260async fn shape_inference_early_response(path: &str, response: Response) -> Response {
264 let request_id = Envelope::new(path != "/v1/completions");
265 if path == "/v1/messages" {
266 anthropic::with_anthropic_request_id(
267 &request_id.id,
268 anthropic::reshape_error(response, &request_id.id).await,
269 )
270 } else {
271 with_request_id(&request_id.id, response)
272 }
273}
274
275async fn authenticate_inference_before_body(
280 State(st): State<AppState>,
281 mut req: AxumRequest,
282 next: Next,
283) -> Response {
284 if !protected_inference_path(req.uri().path()) {
285 return next.run(req).await;
286 }
287 let path = req.uri().path().to_string();
288 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
292 return shape_inference_early_response(
293 &path,
294 error_response_coded(
295 StatusCode::PAYLOAD_TOO_LARGE,
296 &format!(
297 "request body exceeds the {} MiB limit",
298 MAX_BODY_BYTES / (1024 * 1024)
299 ),
300 "invalid_request_error",
301 None,
302 Some("request_too_large"),
303 ),
304 )
305 .await;
306 }
307 let headers = req.headers();
308 let bearer = bearer_token(headers);
309 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
310 let api_key = headers
311 .get("x-api-key")
312 .and_then(|value| value.to_str().ok());
313 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
314 } else {
315 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
316 };
317 if let Err(why) = auth {
318 return shape_inference_early_response(&path, authentication_error(why)).await;
319 }
320 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
328 let body_admission = if body_requires_admission(&req) {
329 body_admission_semaphore()
330 } else {
331 small_body_admission_semaphore()
332 };
333 let body_permit = match body_admission.try_acquire_owned() {
334 Ok(permit) => Some(permit),
335 Err(tokio::sync::TryAcquireError::Closed) => {
336 let response = retry_contract_response(
337 error_response_coded(
338 StatusCode::SERVICE_UNAVAILABLE,
339 "request body admission is unavailable",
340 "server_error",
341 None,
342 Some("body_admission_unavailable"),
343 ),
344 Some(BODY_ADMISSION_RETRY_AFTER_S),
345 );
346 return shape_inference_early_response(&path, response).await;
347 }
348 Err(tokio::sync::TryAcquireError::NoPermits) => {
349 let response = retry_contract_response(
350 error_response_coded(
351 StatusCode::TOO_MANY_REQUESTS,
352 "request body admission is busy",
353 "rate_limit_error",
354 None,
355 Some("body_admission_busy"),
356 ),
357 Some(BODY_ADMISSION_RETRY_AFTER_S),
358 );
359 return shape_inference_early_response(&path, response).await;
360 }
361 };
362 let body = std::mem::replace(req.body_mut(), Body::empty());
366 let mut body = Box::pin(body.into_data_stream());
367 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
368 let body_timed_out_flag = body_timed_out.clone();
369 let guarded_body = async_stream::stream! {
370 loop {
371 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
372 if remaining.is_zero() {
373 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
374 yield Err(std::io::Error::new(
375 std::io::ErrorKind::TimedOut,
376 "request body read deadline exceeded",
377 ));
378 break;
379 }
380 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
381 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
382 Ok(frame) => frame,
383 Err(_) => {
384 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
385 yield Err(std::io::Error::new(
386 std::io::ErrorKind::TimedOut,
387 "request body idle timeout exceeded",
388 ));
389 break;
390 }
391 };
392 match frame {
393 Some(Ok(bytes)) => yield Ok(bytes),
394 Some(Err(error)) => {
395 yield Err(std::io::Error::other(error.to_string()));
396 break;
397 }
398 None => break,
399 }
400 }
401 drop(body_permit);
402 };
403 *req.body_mut() = Body::from_stream(guarded_body);
404 let response = next.run(req).await;
405 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
406 let request_id = Envelope::new(path != "/v1/completions");
407 let timeout = error_response_coded(
408 StatusCode::REQUEST_TIMEOUT,
409 "request body read timed out",
410 "invalid_request_error",
411 None,
412 Some("request_body_timeout"),
413 );
414 return if path == "/v1/messages" {
415 anthropic::with_anthropic_request_id(
416 &request_id.id,
417 anthropic::reshape_error(timeout, &request_id.id).await,
418 )
419 } else {
420 with_request_id(&request_id.id, timeout)
421 };
422 }
423 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
424 let request_id = Envelope::new(true);
425 return anthropic::with_anthropic_request_id(
426 &request_id.id,
427 anthropic::reshape_error(response, &request_id.id).await,
428 );
429 }
430 response
431}
432
433#[cfg(test)]
434mod body_limit_tests {
435 use super::*;
436 use tower::ServiceExt as _;
437
438 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
439
440 fn test_app() -> Router {
444 let app = Router::new()
445 .route(
446 "/bytes",
447 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
448 )
449 .route(
450 "/json",
451 post(|Json(v): Json<serde_json::Value>| async move {
452 v["pad"].as_str().unwrap_or("").len().to_string()
453 }),
454 );
455 apply_body_limit(app)
456 }
457
458 fn streamed_body(chunks: usize) -> Body {
459 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
462 Body::from_stream(async_stream::stream! {
463 for _ in 0..chunks {
464 yield Ok::<_, std::io::Error>(chunk.clone());
465 }
466 })
467 }
468
469 #[tokio::test]
470 async fn bodies_past_the_old_2mib_default_are_accepted() {
471 for (path, body) in [
474 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
475 (
476 "/json",
477 Body::from(
478 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
479 ),
480 ),
481 ] {
482 let resp = test_app()
483 .oneshot(
484 axum::http::Request::post(path)
485 .header(CONTENT_TYPE, "application/json")
486 .body(body)
487 .unwrap(),
488 )
489 .await
490 .unwrap();
491 assert_eq!(resp.status(), StatusCode::OK, "{path}");
492 }
493 }
494
495 #[tokio::test]
496 async fn body_at_exactly_the_limit_is_accepted() {
497 let resp = test_app()
498 .oneshot(
499 axum::http::Request::post("/bytes")
500 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
501 .unwrap(),
502 )
503 .await
504 .unwrap();
505 assert_eq!(resp.status(), StatusCode::OK);
506 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
507 .await
508 .unwrap();
509 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
510 }
511
512 #[tokio::test]
513 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
514 for path in ["/bytes", "/json"] {
518 let resp = test_app()
519 .oneshot(
520 axum::http::Request::post(path)
521 .header(CONTENT_TYPE, "application/json")
522 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
523 .unwrap(),
524 )
525 .await
526 .unwrap();
527 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
528 assert_eq!(
529 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
530 Some(b"false".as_ref()),
531 "{path}: retrying identical bytes cannot fix a 413"
532 );
533 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
534 .await
535 .unwrap();
536 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
537 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
538 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
539 assert!(
540 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
541 "{path}: message names the limit"
542 );
543 }
544 }
545
546 #[tokio::test]
547 async fn authenticated_body_admission_is_finite() {
548 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
549 let semaphore = body_admission_semaphore();
550 let mut permits = Vec::new();
551 for _ in 0..MAX_BODY_ADMISSIONS {
552 permits.push(semaphore.clone().acquire_owned().await.unwrap());
553 }
554 assert!(
555 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
556 .await
557 .is_err(),
558 "body parser admission must not be unbounded"
559 );
560 drop(permits);
561 assert!(semaphore.acquire().await.is_ok());
562 }
563
564 #[tokio::test]
565 async fn small_body_admission_is_finite_and_separate() {
566 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
567 let large = body_admission_semaphore();
568 let small = small_body_admission_semaphore();
569 let mut small_permits = Vec::new();
570 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
571 small_permits.push(small.clone().acquire_owned().await.unwrap());
572 }
573 assert!(
574 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
575 .await
576 .is_err(),
577 "small body parser admission must be bounded"
578 );
579 assert!(
580 large.clone().try_acquire().is_ok(),
581 "small uploads must not consume large-upload permits"
582 );
583 drop(small_permits);
584 assert!(small.acquire().await.is_ok());
585 }
586
587 #[test]
588 fn small_declared_bodies_bypass_large_upload_admission() {
589 let request = axum::http::Request::post("/v1/chat/completions")
590 .header(CONTENT_LENGTH, "2048")
591 .body(Body::empty())
592 .unwrap();
593 assert!(!body_requires_admission(&request));
594
595 let request = axum::http::Request::post("/v1/chat/completions")
596 .header(
597 CONTENT_LENGTH,
598 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
599 )
600 .body(Body::empty())
601 .unwrap();
602 assert!(body_requires_admission(&request));
603
604 let request = axum::http::Request::post("/v1/chat/completions")
605 .header(CONTENT_LENGTH, "2048")
606 .header(TRANSFER_ENCODING, "chunked")
607 .body(Body::empty())
608 .unwrap();
609 assert!(body_requires_admission(&request));
610 }
611
612 #[test]
613 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
614 let unknown = axum::http::Request::post("/v1/chat/completions")
615 .body(Body::empty())
616 .unwrap();
617 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
618
619 let large = axum::http::Request::post("/v1/chat/completions")
620 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
621 .body(Body::empty())
622 .unwrap();
623 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
624 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
625
626 let absurd = axum::http::Request::post("/v1/chat/completions")
627 .header(CONTENT_LENGTH, u64::MAX.to_string())
628 .body(Body::empty())
629 .unwrap();
630 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
631 }
632
633 #[tokio::test]
634 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
635 let too_large = shape_inference_early_response(
636 "/v1/messages",
637 error_response_coded(
638 StatusCode::PAYLOAD_TOO_LARGE,
639 "request body exceeds the 192 MiB limit",
640 "invalid_request_error",
641 None,
642 Some("request_too_large"),
643 ),
644 )
645 .await;
646 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
647 let house_id = too_large.headers()["x-request-id"].clone();
648 assert_eq!(too_large.headers()["request-id"], house_id);
649 assert_eq!(too_large.headers()["x-should-retry"], "false");
650 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
651 .await
652 .unwrap();
653 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
654 assert_eq!(payload["type"], "error");
655 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
656
657 let busy = shape_inference_early_response(
658 "/v1/chat/completions",
659 retry_contract_response(
660 error_response_coded(
661 StatusCode::TOO_MANY_REQUESTS,
662 "request body admission is busy",
663 "rate_limit_error",
664 None,
665 Some("body_admission_busy"),
666 ),
667 Some(BODY_ADMISSION_RETRY_AFTER_S),
668 ),
669 )
670 .await;
671 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
672 assert!(!busy.headers()["x-request-id"].is_empty());
673 assert_eq!(busy.headers()["retry-after"], "1");
674 assert_eq!(busy.headers()["retry-after-ms"], "1000");
675 assert!(busy.headers().get("x-should-retry").is_none());
676 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
677 .await
678 .unwrap();
679 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
680 assert_eq!(payload["error"]["code"], "body_admission_busy");
681 }
682}
683
684#[derive(Clone, Default)]
685struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
686
687fn is_sse_data_frame(bytes: &[u8]) -> bool {
688 bytes
689 .windows(b"data:".len())
690 .any(|window| window == b"data:")
691}
692
693async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
694 let trace = ttft::start(req.uri().path());
695 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
696 let response = next.run(req).await;
697 let Some(trace) = trace else {
698 return response;
699 };
700 let is_sse = response
701 .headers()
702 .get(CONTENT_TYPE)
703 .and_then(|value| value.to_str().ok())
704 .is_some_and(|value| value.starts_with("text/event-stream"));
705 if !is_sse {
706 return response;
707 }
708
709 let (parts, body) = response.into_parts();
712 let mut body = Box::pin(body.into_data_stream());
713 let stream = async_stream::stream! {
714 while let Some(frame) =
715 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
716 {
717 if frame
718 .as_ref()
719 .is_ok_and(|bytes| is_sse_data_frame(bytes))
720 {
721 trace.mark_first_sse_byte();
722 }
723 yield frame;
724 }
725 };
726 Response::from_parts(parts, Body::from_stream(stream))
727}
728
729const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
730const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
731
732#[derive(Debug, Clone, Default, Deserialize)]
733#[serde(deny_unknown_fields)]
734struct OpenRouterMetadataFile {
735 #[serde(default)]
736 models: HashMap<String, OpenRouterModelMetadata>,
737 #[serde(default)]
740 planned_models: HashMap<String, OpenRouterModelMetadata>,
741 #[serde(default)]
744 provider: Option<ProviderMetadata>,
745}
746
747#[derive(Debug, Clone, Deserialize)]
751#[serde(deny_unknown_fields)]
752struct ProviderMetadata {
753 id: String,
754 #[serde(default)]
755 status_url: Option<String>,
756 #[serde(default)]
757 support_contact: Option<String>,
758 #[serde(default)]
759 incident_contact: Option<String>,
760 #[serde(default)]
761 regions: Vec<String>,
762}
763
764#[derive(Debug, Clone, Default, Deserialize)]
766#[serde(deny_unknown_fields)]
767struct LifecycleMetadata {
768 #[serde(default)]
769 status: Option<String>,
770 #[serde(default)]
771 deprecation_at: Option<String>,
772 #[serde(default)]
773 retirement_at: Option<String>,
774 #[serde(default)]
775 replacement_model_id: Option<String>,
776}
777
778#[derive(Debug, Clone, Default, Deserialize)]
780#[serde(deny_unknown_fields)]
781struct ReliabilityMetadata {
782 #[serde(default)]
783 first_token_timeout_seconds: Option<u64>,
784 #[serde(default)]
785 completion_timeout_seconds: Option<u64>,
786 #[serde(default)]
787 stream_idle_timeout_seconds: Option<u64>,
788 #[serde(default)]
789 capacity_scope: Option<String>,
790}
791
792#[derive(Debug, Clone, Default, Deserialize)]
793#[serde(deny_unknown_fields)]
794struct OpenRouterModelMetadata {
795 #[serde(default)]
797 owned_by: Option<String>,
798 #[serde(default)]
799 lifecycle: Option<LifecycleMetadata>,
800 #[serde(default)]
801 reliability: Option<ReliabilityMetadata>,
802 #[serde(default)]
803 hugging_face_id: Option<String>,
804 #[serde(default)]
805 created: Option<u64>,
806 #[serde(default)]
807 quantization: Option<String>,
808 #[serde(default)]
809 description: Option<String>,
810 #[serde(default)]
811 max_prompt_length: Option<u64>,
812 #[serde(default)]
813 max_output_length: Option<u64>,
814 #[serde(default)]
817 default_output_length: Option<u64>,
818 #[serde(default)]
819 pricing: OpenRouterPricing,
820 #[serde(default)]
821 capacity: OpenRouterCapacity,
822 #[serde(default)]
823 is_ready: Option<bool>,
824 #[serde(default)]
825 is_free: Option<bool>,
826 #[serde(default)]
827 discount_to_user: Option<f64>,
828 #[serde(default)]
829 openrouter_slug: Option<String>,
830 #[serde(default)]
831 datacenters: Vec<OpenRouterDatacenter>,
832 #[serde(default)]
836 input_modalities: Vec<String>,
837 #[serde(default)]
852 surface: Option<String>,
853 #[serde(default)]
854 zdr: Option<bool>,
855 #[serde(default)]
856 hipaa: Option<bool>,
857 #[serde(default)]
867 default_reasoning_effort: Option<String>,
868 #[serde(default)]
886 default_temperature: Option<f32>,
887 #[serde(default)]
888 default_top_p: Option<f32>,
889 #[serde(default)]
891 default_top_k: Option<usize>,
892 #[serde(default)]
893 default_min_p: Option<f32>,
894 #[serde(default)]
895 default_presence_penalty: Option<f32>,
896 #[serde(default)]
897 default_frequency_penalty: Option<f32>,
898 #[serde(default)]
900 default_repetition_penalty: Option<f32>,
901 #[serde(default)]
922 non_thinking_sampling: Option<SamplingArmMetadata>,
923}
924
925#[derive(Debug, Clone, Default, Deserialize)]
931#[serde(deny_unknown_fields)]
932struct SamplingArmMetadata {
933 #[serde(default)]
934 temperature: Option<f32>,
935 #[serde(default)]
936 top_p: Option<f32>,
937 #[serde(default)]
938 top_k: Option<usize>,
939 #[serde(default)]
940 min_p: Option<f32>,
941 #[serde(default)]
942 presence_penalty: Option<f32>,
943 #[serde(default)]
944 frequency_penalty: Option<f32>,
945 #[serde(default)]
946 repetition_penalty: Option<f32>,
947}
948
949impl SamplingArmMetadata {
950 fn is_empty(&self) -> bool {
951 self.temperature.is_none()
952 && self.top_p.is_none()
953 && self.top_k.is_none()
954 && self.min_p.is_none()
955 && self.presence_penalty.is_none()
956 && self.frequency_penalty.is_none()
957 && self.repetition_penalty.is_none()
958 }
959}
960
961#[derive(Debug, Clone, Default, Deserialize)]
962#[serde(deny_unknown_fields)]
963struct OpenRouterPricing {
964 #[serde(default)]
965 prompt: Option<String>,
966 #[serde(default)]
967 cached_prompt: Option<String>,
968 #[serde(default)]
969 cache_write: Option<String>,
970 #[serde(default)]
971 completion: Option<String>,
972 #[serde(default)]
973 internal_reasoning: Option<String>,
974 #[serde(default)]
975 request: Option<String>,
976}
977
978#[derive(Debug, Clone, Default, Deserialize)]
979#[serde(deny_unknown_fields)]
980struct OpenRouterCapacity {
981 #[serde(default)]
982 prompt_tpm: Option<u64>,
983 #[serde(default)]
984 cached_prompt_tpm: Option<u64>,
985 #[serde(default)]
986 completion_tpm: Option<u64>,
987 #[serde(default)]
988 request_rpm: Option<u64>,
989 #[serde(default)]
990 concurrency: Option<u64>,
991}
992
993#[derive(Debug, Clone, Deserialize, Serialize)]
994#[serde(deny_unknown_fields)]
995struct OpenRouterDatacenter {
996 country_code: String,
997 #[serde(default, skip_serializing_if = "Option::is_none")]
998 region: Option<String>,
999}
1000
1001impl OpenRouterMetadataFile {
1002 fn parse(
1003 text: &str,
1004 ) -> Result<
1005 (
1006 HashMap<String, OpenRouterModelMetadata>,
1007 Option<ProviderMetadata>,
1008 ),
1009 String,
1010 > {
1011 let file: Self =
1012 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
1013 for (alias, metadata) in &file.models {
1014 validate_openrouter_metadata(alias, metadata)?;
1015 }
1016 for (alias, metadata) in &file.planned_models {
1017 validate_openrouter_metadata(alias, metadata)?;
1018 if file.models.contains_key(alias) {
1019 return Err(format!(
1020 "model alias {alias:?} appears in both models and planned_models"
1021 ));
1022 }
1023 }
1024 if let Some(provider) = &file.provider {
1025 if provider.id.is_empty() {
1026 return Err("provider.id must be a non-empty slug".into());
1027 }
1028 for (field, value) in [
1030 ("provider.support_contact", &provider.support_contact),
1031 ("provider.incident_contact", &provider.incident_contact),
1032 ] {
1033 if let Some(value) = value
1034 && !value.contains(':')
1035 {
1036 return Err(format!(
1037 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1038 ));
1039 }
1040 }
1041 }
1042 Ok((file.models, file.provider))
1043 }
1044
1045 #[cfg(test)]
1046 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1047 Self::parse(text).map(|(models, _)| models)
1048 }
1049}
1050
1051fn per_million_price(per_token: &str) -> Option<String> {
1055 if !valid_price_string(per_token) {
1056 return None;
1057 }
1058 let (whole, frac) = match per_token.split_once('.') {
1059 Some((whole, frac)) => (whole, frac),
1060 None => (per_token, ""),
1061 };
1062 let mut digits = format!("{whole}{frac}");
1063 let point = whole.len() + 6;
1064 while digits.len() < point {
1065 digits.push('0');
1066 }
1067 let (int_part, frac_part) = digits.split_at(point);
1068 let int_part = int_part.trim_start_matches('0');
1069 let int_part = if int_part.is_empty() { "0" } else { int_part };
1070 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1071 while frac_out.len() < 2 {
1072 frac_out.push('0');
1073 }
1074 Some(format!("{int_part}.{frac_out}"))
1075}
1076
1077fn valid_price_string(value: &str) -> bool {
1078 let mut parts = value.split('.');
1079 let whole = parts.next().unwrap_or_default();
1080 let fraction = parts.next();
1081 !whole.is_empty()
1082 && whole.bytes().all(|b| b.is_ascii_digit())
1083 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1084 && parts.next().is_none()
1085}
1086
1087fn validate_openrouter_metadata(
1088 alias: &str,
1089 metadata: &OpenRouterModelMetadata,
1090) -> Result<(), String> {
1091 if alias.is_empty() {
1092 return Err("models metadata contains an empty model alias".into());
1093 }
1094 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1097 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1098 {
1099 return Err(format!(
1100 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1101 reasoning_effort level (none|minimal|low|medium|high)"
1102 ));
1103 }
1104 validate_sampling_defaults(alias, metadata)?;
1105 for m in &metadata.input_modalities {
1106 if m != "image" && m != "video" {
1107 return Err(format!(
1108 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1109 ));
1110 }
1111 }
1112 if let Some(sfc) = metadata.surface.as_deref()
1113 && !matches!(sfc, "chat" | "embedding" | "rerank")
1114 {
1115 return Err(format!(
1116 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1117 ));
1118 }
1119 if let Some(q) = metadata.quantization.as_deref()
1120 && !matches!(
1121 q,
1122 "int4"
1123 | "int8"
1124 | "fp4"
1125 | "mxfp4"
1126 | "nvfp4"
1127 | "fp6"
1128 | "fp8"
1129 | "mxfp8"
1130 | "fp16"
1131 | "bf16"
1132 | "fp32"
1133 )
1134 {
1135 return Err(format!(
1136 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1137 ));
1138 }
1139 for (field, value) in [
1140 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1141 (
1142 "pricing.cached_prompt",
1143 metadata.pricing.cached_prompt.as_deref(),
1144 ),
1145 (
1146 "pricing.cache_write",
1147 metadata.pricing.cache_write.as_deref(),
1148 ),
1149 ("pricing.completion", metadata.pricing.completion.as_deref()),
1150 (
1151 "pricing.internal_reasoning",
1152 metadata.pricing.internal_reasoning.as_deref(),
1153 ),
1154 ("pricing.request", metadata.pricing.request.as_deref()),
1155 ] {
1156 if let Some(value) = value
1157 && !valid_price_string(value)
1158 {
1159 return Err(format!(
1160 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1161 ));
1162 }
1163 }
1164 for (field, value) in [
1165 ("created", metadata.created),
1166 ("max_prompt_length", metadata.max_prompt_length),
1167 ("max_output_length", metadata.max_output_length),
1168 ("default_output_length", metadata.default_output_length),
1169 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1170 (
1171 "capacity.cached_prompt_tpm",
1172 metadata.capacity.cached_prompt_tpm,
1173 ),
1174 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1175 ("capacity.request_rpm", metadata.capacity.request_rpm),
1176 ("capacity.concurrency", metadata.capacity.concurrency),
1177 ] {
1178 if let Some(value) = value
1179 && value > JSON_SAFE_INTEGER_MAX
1180 {
1181 return Err(format!(
1182 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1183 ));
1184 }
1185 }
1186 for (field, value) in [
1187 ("max_prompt_length", metadata.max_prompt_length),
1188 ("max_output_length", metadata.max_output_length),
1189 ("default_output_length", metadata.default_output_length),
1190 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1191 (
1192 "capacity.cached_prompt_tpm",
1193 metadata.capacity.cached_prompt_tpm,
1194 ),
1195 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1196 ("capacity.request_rpm", metadata.capacity.request_rpm),
1197 ("capacity.concurrency", metadata.capacity.concurrency),
1198 ] {
1199 if value == Some(0) {
1200 return Err(format!(
1201 "model {alias:?}: {field} must be greater than zero when declared"
1202 ));
1203 }
1204 }
1205 if let (Some(default), Some(maximum)) =
1206 (metadata.default_output_length, metadata.max_output_length)
1207 && default > maximum
1208 {
1209 return Err(format!(
1210 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1211 ));
1212 }
1213 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1214 return Err(format!(
1215 "model {alias:?}: default_output_length requires max_output_length"
1216 ));
1217 }
1218 if let Some(discount) = metadata.discount_to_user
1219 && (!discount.is_finite() || discount >= 1.0)
1220 {
1221 return Err(format!(
1222 "model {alias:?}: discount_to_user must be finite and less than 1"
1223 ));
1224 }
1225 if metadata
1226 .openrouter_slug
1227 .as_deref()
1228 .is_some_and(str::is_empty)
1229 {
1230 return Err(format!(
1231 "model {alias:?}: openrouter_slug must not be empty when declared"
1232 ));
1233 }
1234 for dc in &metadata.datacenters {
1235 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1236 return Err(format!(
1237 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1238 dc.country_code
1239 ));
1240 }
1241 }
1242 Ok(())
1243}
1244
1245fn validate_sampling_defaults(
1260 alias: &str,
1261 metadata: &OpenRouterModelMetadata,
1262) -> Result<(), String> {
1263 validate_sampling_arm(
1264 alias,
1265 &[
1266 "default_temperature",
1267 "default_top_p",
1268 "default_min_p",
1269 "default_presence_penalty",
1270 "default_frequency_penalty",
1271 "default_repetition_penalty",
1272 ],
1273 metadata.default_temperature,
1274 metadata.default_top_p,
1275 metadata.default_min_p,
1276 metadata.default_presence_penalty,
1277 metadata.default_frequency_penalty,
1278 metadata.default_repetition_penalty,
1279 )?;
1280 if let Some(arm) = &metadata.non_thinking_sampling {
1281 if arm.is_empty() {
1285 return Err(format!(
1286 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1287 least one vendor recommendation or delete the table"
1288 ));
1289 }
1290 validate_sampling_arm(
1291 alias,
1292 &[
1293 "non_thinking_sampling.temperature",
1294 "non_thinking_sampling.top_p",
1295 "non_thinking_sampling.min_p",
1296 "non_thinking_sampling.presence_penalty",
1297 "non_thinking_sampling.frequency_penalty",
1298 "non_thinking_sampling.repetition_penalty",
1299 ],
1300 arm.temperature,
1301 arm.top_p,
1302 arm.min_p,
1303 arm.presence_penalty,
1304 arm.frequency_penalty,
1305 arm.repetition_penalty,
1306 )?;
1307 }
1308 Ok(())
1309}
1310
1311#[allow(clippy::too_many_arguments)]
1317fn validate_sampling_arm(
1318 alias: &str,
1319 keys: &[&str; 6],
1320 temperature: Option<f32>,
1321 top_p: Option<f32>,
1322 min_p: Option<f32>,
1323 presence_penalty: Option<f32>,
1324 frequency_penalty: Option<f32>,
1325 repetition_penalty: Option<f32>,
1326) -> Result<(), String> {
1327 if let Some(t) = temperature
1328 && (!t.is_finite() || t <= 0.0 || t > 2.0)
1329 {
1330 return Err(format!(
1331 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1332 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1333 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1334 vendor recommendation, not greedy); clients reach greedy by sending an \
1335 explicit temperature 0.",
1336 keys[0]
1337 ));
1338 }
1339 if let Some(p) = top_p
1340 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1341 {
1342 return Err(format!(
1343 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1344 keys[1]
1345 ));
1346 }
1347 if let Some(m) = min_p
1348 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1349 {
1350 return Err(format!(
1351 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1352 keys[2]
1353 ));
1354 }
1355 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1356 if let Some(v) = value
1357 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1358 {
1359 return Err(format!(
1360 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1361 ));
1362 }
1363 }
1364 if let Some(r) = repetition_penalty
1365 && (!r.is_finite() || r <= 0.0)
1366 {
1367 return Err(format!(
1368 "model {alias:?}: {} {r} must be finite and \
1369 greater than zero (1.0 = off)",
1370 keys[5]
1371 ));
1372 }
1373 Ok(())
1374}
1375
1376fn load_openrouter_metadata(
1377 models: &[(String, String, Option<String>)],
1378) -> Result<
1379 (
1380 HashMap<String, OpenRouterModelMetadata>,
1381 Option<ProviderMetadata>,
1382 ),
1383 String,
1384> {
1385 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1386 Ok(path) => path,
1387 Err(_) => return Ok((HashMap::new(), None)),
1388 };
1389 let p = std::path::Path::new(&path);
1390 if !p.is_file() {
1391 return Err(format!(
1392 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1393 ));
1394 }
1395 let text =
1396 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1397 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1398 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1399 for alias in metadata.keys() {
1400 if !models.iter().any(|(name, _, _)| name == alias) {
1401 return Err(format!(
1402 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1403 ));
1404 }
1405 }
1406 eprintln!(
1407 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1408 metadata.len()
1409 );
1410 Ok((metadata, provider))
1411}
1412
1413#[derive(Clone)]
1414struct AppState {
1415 cmd_tx: Sender<Cmd>,
1416 models: Arc<Vec<String>>,
1417 caps: Arc<HashMap<String, ModelCaps>>,
1418 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1419 provider_metadata: Arc<Option<ProviderMetadata>>,
1421 metering: Option<Arc<dyn metering::Metering>>,
1427 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1430 api_auth: ApiAuth,
1433 metrics_auth: MetricsAuth,
1435 metrics: SharedMetrics,
1436 inflight: InflightCounts,
1440 tenant_inflight: TenantGauge,
1443 health: health::SharedHealth,
1447 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1451}
1452
1453impl AppState {
1454 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1468 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1469 }
1470}
1471
1472#[derive(Clone, Default)]
1473struct ApiAuth {
1474 keyring: Option<&'static auth::KeyStore>,
1475 single_key: Option<Arc<str>>,
1476}
1477
1478impl ApiAuth {
1479 fn from_env() -> Result<ApiAuth, String> {
1480 let single_key = match std::env::var("MEMRA_API_KEY") {
1481 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1482 Ok(key) => Some(Arc::from(key)),
1483 Err(std::env::VarError::NotPresent) => None,
1484 Err(std::env::VarError::NotUnicode(_)) => {
1485 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1486 }
1487 };
1488 Ok(ApiAuth {
1489 keyring: auth::global(),
1490 single_key,
1491 })
1492 }
1493
1494 fn configured(&self) -> bool {
1495 self.keyring.is_some() || self.single_key.is_some()
1496 }
1497}
1498
1499#[derive(Clone, Default)]
1500struct MetricsAuth {
1501 required: bool,
1502 token: Option<Arc<str>>,
1503}
1504
1505impl MetricsAuth {
1506 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1507 let token = token.map(Arc::from);
1508 MetricsAuth {
1509 required: !bind_loopback || api_auth_configured || token.is_some(),
1510 token,
1511 }
1512 }
1513}
1514
1515fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1516 let mut resolved = addr
1517 .to_socket_addrs()
1518 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1519 let first = resolved
1520 .next()
1521 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1522 let mut loopback = first.ip().to_canonical().is_loopback();
1523 for socket in resolved {
1524 loopback &= socket.ip().to_canonical().is_loopback();
1525 }
1526 Ok((first, loopback))
1527}
1528
1529fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1530 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1531}
1532
1533fn validate_bind_security(
1534 addr: &str,
1535 api_auth_configured: bool,
1536 allow_open_bind: bool,
1537) -> Result<bool, String> {
1538 let loopback = bind_is_loopback(addr)?;
1539 if !loopback && !api_auth_configured && !allow_open_bind {
1540 return Err(format!(
1541 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1542 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1543 ));
1544 }
1545 Ok(loopback)
1546}
1547
1548type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1566
1567type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1570
1571struct InflightGuard {
1575 counts: InflightCounts,
1576 idx: usize,
1577 tenants: TenantGauge,
1578 tenant: String,
1579}
1580
1581impl InflightGuard {
1582 fn try_acquire(
1586 counts: InflightCounts,
1587 lane: lanes::Lane,
1588 tenants: TenantGauge,
1589 tenant: &str,
1590 tenant_cap: Option<usize>,
1591 ) -> Result<(Self, usize, usize), usize> {
1592 let idx = lane.idx();
1593 let nt = {
1594 let mut m = tenants.lock().unwrap();
1595 let e = m.entry(tenant.to_string()).or_insert(0);
1596 if tenant_cap.is_some_and(|cap| *e >= cap) {
1597 return Err(*e);
1598 }
1599 *e += 1;
1600 *e
1601 };
1602 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1603 Ok((
1604 InflightGuard {
1605 counts,
1606 idx,
1607 tenants,
1608 tenant: tenant.to_string(),
1609 },
1610 n,
1611 nt,
1612 ))
1613 }
1614}
1615
1616impl Drop for InflightGuard {
1617 fn drop(&mut self) {
1618 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1619 let mut m = self.tenants.lock().unwrap();
1620 if let Some(e) = m.get_mut(&self.tenant) {
1621 *e -= 1;
1622 if *e == 0 {
1623 m.remove(&self.tenant);
1624 }
1625 }
1626 }
1627}
1628
1629fn lane_cap(lane: lanes::Lane) -> usize {
1633 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1634 CAPS.get_or_init(|| {
1635 let batching = std::env::var("MEMRA_SERVE_BATCH")
1636 .map(|v| v != "0")
1637 .unwrap_or(true);
1638 let interactive = if batching {
1639 std::env::var("MEMRA_MAX_SESSIONS")
1640 .ok()
1641 .and_then(|v| v.parse().ok())
1642 .unwrap_or(64)
1643 } else {
1644 worker::MAX_ACTIVE
1645 };
1646 let p = lanes::LanePolicy::from_env();
1647 [interactive, p.max_sessions[1], p.max_sessions[2]]
1648 })[lane.idx()]
1649}
1650
1651fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1654 if m.completed > 0 && m.step_p50_ms > 0.0 {
1655 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1656 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1657 }
1658 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1659 *D.get_or_init(|| {
1660 std::env::var("MEMRA_RL_RESET_S")
1661 .ok()
1662 .and_then(|v| v.parse().ok())
1663 .unwrap_or(2)
1664 })
1665}
1666
1667pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1684pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1685pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1686
1687pub(crate) fn timeout_ms_max() -> u64 {
1697 static V: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1698 *V.get_or_init(|| {
1699 std::env::var("MEMRA_TIMEOUT_MS_MAX")
1700 .ok()
1701 .and_then(|s| s.parse::<u64>().ok())
1702 .filter(|&ms| ms >= TIMEOUT_MS_MIN)
1703 .unwrap_or(TIMEOUT_MS_MAX)
1704 })
1705}
1706
1707pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1711 let max = timeout_ms_max();
1712 let Some(v) = v.filter(|v| !v.is_null()) else {
1713 return Ok(max);
1715 };
1716 let Some(ms) = v.as_u64() else {
1717 return Err(format!(
1718 "timeout_ms must be an integer number of milliseconds in \
1719 {TIMEOUT_MS_MIN}..={max}, got {v}; for work longer than \
1720 {max} ms use \"stream\": true — the deadline then bounds only the \
1721 time to first token and the stream may run as long as it needs"
1722 ));
1723 };
1724 if !(TIMEOUT_MS_MIN..=max).contains(&ms) {
1725 return Err(format!(
1726 "timeout_ms {ms} is outside the accepted range \
1727 {TIMEOUT_MS_MIN}..={max} (milliseconds). {max} is a \
1728 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1729 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1730 lie. For work longer than {max} ms use \"stream\": true — the \
1731 deadline then bounds only the time to first token and the stream may run as \
1732 long as it needs"
1733 ));
1734 }
1735 Ok(ms)
1736}
1737
1738#[derive(Clone, Copy)]
1741pub(crate) struct RequestDeadline {
1742 pub(crate) at: tokio::time::Instant,
1743 pub(crate) ms: u64,
1744}
1745
1746impl RequestDeadline {
1747 pub(crate) fn starting_now(ms: u64) -> Self {
1748 Self {
1749 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
1750 ms,
1751 }
1752 }
1753
1754 pub(crate) fn remaining(&self) -> std::time::Duration {
1755 self.at
1756 .saturating_duration_since(tokio::time::Instant::now())
1757 }
1758}
1759
1760pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
1766 let what = if stream {
1767 "the first token was produced"
1768 } else {
1769 "the response completed"
1770 };
1771 let msg = format!(
1772 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
1773 {what}; generation was cancelled and this request is not billed"
1774 );
1775 error_response_coded(
1776 StatusCode::REQUEST_TIMEOUT,
1777 &msg,
1778 "timeout",
1779 Some("timeout_ms"),
1780 Some("deadline_exceeded"),
1781 )
1782}
1783
1784pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
1824
1825pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
1830
1831pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
1835
1836fn env_flag_on(name: &'static str, default_on: bool) -> bool {
1842 match std::env::var(name) {
1843 Ok(v) => !matches!(
1844 v.trim().to_ascii_lowercase().as_str(),
1845 "0" | "off" | "false"
1846 ),
1847 Err(_) => default_on,
1848 }
1849}
1850
1851fn env_u64(name: &'static str, default: u64) -> u64 {
1854 std::env::var(name)
1855 .ok()
1856 .and_then(|v| v.parse::<u64>().ok())
1857 .filter(|v| *v > 0)
1858 .unwrap_or(default)
1859}
1860
1861const CHARS_PER_TOKEN_FLOOR: usize = 6;
1872
1873pub(crate) fn prompt_tokens_estimate(
1874 request: &worker::Request,
1875 tokenizer: Option<&Tokenizer>,
1876) -> u64 {
1877 if !request.prompt_ids.is_empty() {
1878 return request.prompt_ids.len() as u64;
1879 }
1880 let mut text = String::new();
1881 text.push_str(&request.prompt_text);
1882 for turn in &request.chat_turns {
1883 text.push_str(&turn.content);
1884 }
1885 for tool in &request.tools_json {
1886 text.push_str(tool);
1887 }
1888 if let Some(tokenizer) = tokenizer {
1889 return tokenizer.encode(text.as_str(), false).len() as u64;
1890 }
1891 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
1892}
1893
1894pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
1898 let prefill_ms = prompt_tokens
1899 .saturating_mul(1_000)
1900 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
1901 .unwrap_or(u64::MAX);
1902 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
1903 if decode_ms == 0 {
1904 return None;
1905 }
1906 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
1907}
1908
1909pub(crate) fn nonstream_deadline_gate(
1919 request: &worker::Request,
1920 stream: bool,
1921 deadline: RequestDeadline,
1922 caller_declared_max_tokens: bool,
1923 tokenizer: Option<&Tokenizer>,
1924) -> Result<(), String> {
1925 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
1926 return Ok(());
1927 }
1928 let max_new = request.params.max_new as u64;
1929 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
1936 {
1937 return Ok(());
1938 }
1939 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
1940 let remaining_ms = deadline.remaining().as_millis() as u64;
1941 let prefill_ms = prompt_tokens.saturating_mul(1_000)
1942 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
1943 let decode_ms = max_new.saturating_mul(1_000)
1944 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
1945 let est_ms = prefill_ms.saturating_add(decode_ms);
1946 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
1947 if est_ms <= bound_ms {
1948 return Ok(());
1949 }
1950 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
1951 let advice = match fits {
1952 Some(fits) if fits > 0 => format!(
1953 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
1954 stream's deadline bounds only the time to first token, so it may run as long \
1955 as it needs"
1956 ),
1957 _ => format!(
1958 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
1959 token, so no max_tokens fits: set \"stream\": true"
1960 ),
1961 };
1962 Err(format!(
1963 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
1964 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
1965 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
1966 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
1967 rather than after the deadline: {advice}",
1968 est_ms / 1_000,
1969 ))
1970}
1971
1972fn max_queue_depth(cap: usize) -> usize {
1976 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1977 D.get_or_init(|| {
1978 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
1979 .ok()
1980 .and_then(|v| v.parse().ok())
1981 })
1982 .unwrap_or(cap.saturating_mul(4))
1983}
1984
1985pub(crate) struct PendingAdmissionGuard {
2002 reserved: bool,
2003 lane: lanes::Lane,
2004}
2005
2006impl PendingAdmissionGuard {
2007 pub(crate) fn commit(mut self) {
2011 self.reserved = false;
2012 std::mem::forget(self);
2013 }
2014}
2015
2016impl Drop for PendingAdmissionGuard {
2017 fn drop(&mut self) {
2018 if self.reserved {
2019 worker::release_pending_admit();
2020 worker::release_admission_reservation(self.lane);
2021 }
2022 }
2023}
2024
2025#[allow(clippy::result_large_err)] pub(crate) fn reserve_pending_admit(
2027 st: &AppState,
2028 lane: lanes::Lane,
2029 rl: &RateLimit,
2030 deadline: RequestDeadline,
2031) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2032 let cap = lane_cap(lane).max(1);
2037 let bound = max_queue_depth(cap);
2038 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2039 loop {
2040 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2041 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2042 let backlog = reservations;
2046 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2047 if backlog >= bound {
2048 let msg = format!(
2049 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2050 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2051 coarse estimate, not a promise)",
2052 lane.as_str()
2053 );
2054 let resp = retry_contract_response(
2055 (
2056 StatusCode::TOO_MANY_REQUESTS,
2057 Json(error_body(
2058 &msg,
2059 "rate_limit_error",
2060 None,
2061 Some("shed_queue"),
2062 )),
2063 )
2064 .into_response(),
2065 Some(est_wait_s),
2066 );
2067 return Err((resp, "shed_queue"));
2068 }
2069 let remaining_ms = deadline.remaining().as_millis() as u64;
2070 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2074 if lane == lanes::Lane::Interactive
2075 && waits_for_capacity
2076 && est_wait_s.saturating_mul(1_000) > remaining_ms
2077 {
2078 let msg = format!(
2079 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2080 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2081 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2082 estimate, not a promise)"
2083 );
2084 let resp = retry_contract_response(
2085 (
2086 StatusCode::TOO_MANY_REQUESTS,
2087 Json(error_body(
2088 &msg,
2089 "rate_limit_error",
2090 None,
2091 Some("shed_deadline"),
2092 )),
2093 )
2094 .into_response(),
2095 Some(est_wait_s),
2096 );
2097 return Err((resp, "shed_deadline"));
2098 }
2099 if reservations_for_lane
2100 .compare_exchange(
2101 reservations,
2102 reservations.saturating_add(1),
2103 std::sync::atomic::Ordering::AcqRel,
2104 std::sync::atomic::Ordering::Acquire,
2105 )
2106 .is_ok()
2107 {
2108 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2112 return Ok(PendingAdmissionGuard {
2113 reserved: true,
2114 lane,
2115 });
2116 }
2117 }
2118}
2119
2120static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2131
2132fn draining() -> bool {
2133 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2134}
2135
2136fn drain_deadline_s() -> u64 {
2138 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2139 *D.get_or_init(|| {
2140 std::env::var("MEMRA_DRAIN_S")
2141 .ok()
2142 .and_then(|v| v.parse().ok())
2143 .unwrap_or(30)
2144 })
2145}
2146
2147fn drain_response() -> Response {
2157 let resp = (
2158 StatusCode::SERVICE_UNAVAILABLE,
2159 Json(error_body(
2160 "server is draining (shutdown in progress); retry",
2161 "server_error",
2162 None,
2163 Some("draining"),
2164 )),
2165 )
2166 .into_response();
2167 retry_contract_response(resp, Some(drain_deadline_s()))
2168}
2169
2170struct RateLimit {
2172 limit: usize,
2173 remaining: usize,
2174 reset_s: u64,
2175}
2176
2177impl RateLimit {
2178 fn at_admit(
2183 lane: lanes::Lane,
2184 n_inflight: usize,
2185 metrics: &SharedMetrics,
2186 tenant: &auth::TenantCtx,
2187 n_tenant: usize,
2188 ) -> Self {
2189 let global = lane_cap(lane);
2190 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2191 return Self::compute(global, n_inflight, metrics);
2192 };
2193 let headroom = t
2194 .saturating_sub(n_tenant)
2195 .min(global.saturating_sub(n_inflight));
2196 Self::compute(t, t - headroom, metrics)
2198 }
2199
2200 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2201 let remaining = limit.saturating_sub(n_inflight);
2202 let reset_s = if remaining > 0 {
2203 0
2204 } else {
2205 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2206 reset_estimate_s(&m)
2207 };
2208 RateLimit {
2209 limit,
2210 remaining,
2211 reset_s,
2212 }
2213 }
2214
2215 fn attach(&self, mut resp: Response) -> Response {
2217 let h = resp.headers_mut();
2218 for (k, v) in [
2219 ("x-ratelimit-limit", self.limit as u64),
2220 ("x-ratelimit-remaining", self.remaining as u64),
2221 ("x-ratelimit-reset", self.reset_s),
2222 ] {
2223 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2224 h.insert(axum::http::HeaderName::from_static(k), v);
2225 }
2226 }
2227 resp
2228 }
2229}
2230
2231#[allow(clippy::result_large_err)] fn acquire_request_slot(
2236 st: &AppState,
2237 lane: lanes::Lane,
2238 tenant: &auth::TenantCtx,
2239 env: &Envelope,
2240) -> Result<(InflightGuard, RateLimit), Response> {
2241 let global = lane_cap(lane);
2242 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2243 match InflightGuard::try_acquire(
2244 st.inflight.clone(),
2245 lane,
2246 st.tenant_inflight.clone(),
2247 &tenant.tenant,
2248 tenant_cap,
2249 ) {
2250 Ok((guard, n_inflight, n_tenant)) => {
2251 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2252 Ok((guard, rl))
2253 }
2254 Err(n_tenant) => {
2255 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2256 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2257 let error =
2258 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2259 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2260 }
2261 }
2262}
2263
2264#[derive(Deserialize)]
2266struct CompletionReq {
2267 model: String,
2268 #[serde(default)]
2269 prompt: String,
2270 #[serde(default)]
2272 prompt_ids: Vec<u32>,
2273 #[serde(default)]
2276 max_tokens: Option<usize>,
2277 #[serde(default)]
2290 temperature: Option<f32>,
2291 #[serde(default)]
2292 top_p: Option<f32>,
2293 #[serde(default)]
2295 top_k: Option<usize>,
2296 #[serde(default)]
2298 min_p: Option<f32>,
2299 #[serde(default)]
2301 frequency_penalty: Option<f32>,
2302 #[serde(default)]
2303 presence_penalty: Option<f32>,
2304 #[serde(default)]
2306 repetition_penalty: Option<f32>,
2307 #[serde(default)]
2313 seed: Option<u64>,
2314 #[serde(default)]
2315 stop: StopSequences,
2316 #[serde(default)]
2319 logit_bias: Option<serde_json::Value>,
2320 #[serde(default)]
2321 logprobs: Option<serde_json::Value>,
2322 #[serde(default)]
2323 n: Option<usize>,
2324 #[serde(default)]
2325 best_of: Option<usize>,
2326 #[serde(default)]
2328 chat: bool,
2329 #[serde(default)]
2331 stream: bool,
2332 #[serde(default)]
2334 max_ctx: Option<usize>,
2335 #[serde(default)]
2337 trace_id: Option<String>,
2338 #[serde(default)]
2342 cache_salt: Option<String>,
2343 #[serde(default)]
2347 session_id: Option<String>,
2348 #[serde(default)]
2349 user: Option<String>,
2350 #[serde(default)]
2354 timeout_ms: Option<serde_json::Value>,
2355}
2356
2357#[derive(Deserialize)]
2358struct ChatMessage {
2359 role: String,
2360 #[serde(default)]
2362 content: serde_json::Value,
2363 #[serde(default)]
2365 tool_calls: Vec<ReqToolCall>,
2366 #[serde(default)]
2369 tool_call_id: Option<String>,
2370 #[serde(default)]
2373 name: Option<String>,
2374 #[serde(default, alias = "reasoning_content")]
2384 reasoning: Option<String>,
2385}
2386
2387#[derive(Deserialize)]
2388struct ReqToolCall {
2389 #[serde(default)]
2390 #[allow(dead_code)]
2391 id: Option<String>,
2392 function: ReqToolFunction,
2393}
2394
2395#[derive(Deserialize)]
2396struct ReqToolFunction {
2397 name: String,
2398 #[serde(default)]
2400 arguments: serde_json::Value,
2401}
2402
2403#[derive(Clone, Default, Deserialize)]
2404#[serde(untagged)]
2405enum StopSequences {
2406 One(String),
2407 Many(Vec<String>),
2408 #[default]
2409 None,
2410}
2411
2412impl StopSequences {
2413 fn into_vec(self) -> Vec<String> {
2419 let stops = match self {
2420 Self::One(stop) => vec![stop],
2421 Self::Many(stops) => stops,
2422 Self::None => Vec::new(),
2423 };
2424 stops.into_iter().filter(|s| !s.is_empty()).collect()
2425 }
2426}
2427
2428#[derive(Deserialize)]
2434struct ChatCompletionReq {
2435 model: String,
2436 messages: Vec<ChatMessage>,
2437 #[serde(default, alias = "max_completion_tokens")]
2440 max_tokens: Option<usize>,
2441 #[serde(default)]
2444 temperature: Option<f32>,
2445 #[serde(default)]
2446 top_p: Option<f32>,
2447 #[serde(default)]
2451 top_k: Option<usize>,
2452 #[serde(default)]
2454 min_p: Option<f32>,
2455 #[serde(default)]
2457 frequency_penalty: Option<f32>,
2458 #[serde(default)]
2459 presence_penalty: Option<f32>,
2460 #[serde(default)]
2462 repetition_penalty: Option<f32>,
2463 #[serde(default)]
2465 seed: Option<u64>,
2466 #[serde(default)]
2467 stop: StopSequences,
2468 #[serde(default)]
2469 stream: bool,
2470 #[serde(default)]
2471 max_ctx: Option<usize>,
2472 #[serde(default)]
2477 response_format: Option<serde_json::Value>,
2478 #[serde(default)]
2479 logit_bias: Option<serde_json::Value>,
2480 #[serde(default)]
2481 logprobs: Option<serde_json::Value>,
2482 #[serde(default)]
2483 top_logprobs: Option<usize>,
2484 #[serde(default)]
2485 n: Option<usize>,
2486 #[serde(default)]
2488 tools: Vec<serde_json::Value>,
2489 #[serde(default)]
2491 tool_choice: Option<serde_json::Value>,
2492 #[serde(default)]
2498 reasoning_effort: Option<String>,
2499 #[serde(default)]
2508 reasoning: Option<serde_json::Value>,
2509 #[serde(default)]
2524 include_reasoning: Option<bool>,
2525 #[serde(default)]
2534 enable_thinking: Option<bool>,
2535 #[serde(default)]
2541 chat_template_kwargs: Option<serde_json::Value>,
2542 #[serde(default)]
2546 cache_salt: Option<String>,
2547 #[serde(default)]
2549 session_id: Option<String>,
2550 #[serde(default)]
2551 user: Option<String>,
2552 #[serde(default)]
2557 timeout_ms: Option<serde_json::Value>,
2558}
2559fn one() -> f32 {
2560 1.0
2561}
2562fn default_temperature() -> f32 {
2568 1.0
2569}
2570
2571#[derive(Debug, Clone, Copy, Default, PartialEq)]
2591struct SamplingDefaults {
2592 temperature: Option<f32>,
2593 top_p: Option<f32>,
2594 top_k: Option<usize>,
2595 min_p: Option<f32>,
2596 frequency_penalty: Option<f32>,
2597 presence_penalty: Option<f32>,
2598 repetition_penalty: Option<f32>,
2599}
2600
2601impl SamplingDefaults {
2602 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2605 SamplingDefaults {
2606 temperature: metadata
2607 .and_then(|m| m.default_temperature)
2608 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2609 top_p: metadata
2610 .and_then(|m| m.default_top_p)
2611 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2612 top_k: metadata.and_then(|m| m.default_top_k),
2613 min_p: metadata.and_then(|m| m.default_min_p),
2614 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2615 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2616 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2617 }
2618 }
2619}
2620
2621#[derive(Debug, Clone, Copy, Default, PartialEq)]
2634struct ModelSamplingDefaults {
2635 thinking: SamplingDefaults,
2636 non_thinking: Option<SamplingDefaults>,
2637}
2638
2639impl ModelSamplingDefaults {
2640 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2641 ModelSamplingDefaults {
2642 thinking: SamplingDefaults::resolve(metadata, caps),
2643 non_thinking: metadata
2648 .and_then(|m| m.non_thinking_sampling.as_ref())
2649 .map(|arm| SamplingDefaults {
2650 temperature: arm.temperature,
2651 top_p: arm.top_p,
2652 top_k: arm.top_k,
2653 min_p: arm.min_p,
2654 frequency_penalty: arm.frequency_penalty,
2655 presence_penalty: arm.presence_penalty,
2656 repetition_penalty: arm.repetition_penalty,
2657 }),
2658 }
2659 }
2660
2661 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
2674 match (think, &self.non_thinking) {
2675 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
2676 _ => &self.thinking,
2677 }
2678 }
2679
2680 #[cfg(test)] fn single(thinking: SamplingDefaults) -> Self {
2684 ModelSamplingDefaults {
2685 thinking,
2686 non_thinking: None,
2687 }
2688 }
2689}
2690
2691#[derive(Debug, Clone, Copy, Default)]
2697struct ClientSampling {
2698 temperature: Option<f32>,
2699 top_p: Option<f32>,
2700 top_k: Option<usize>,
2701 min_p: Option<f32>,
2702 frequency_penalty: Option<f32>,
2703 presence_penalty: Option<f32>,
2704 repetition_penalty: Option<f32>,
2705 seed: Option<u64>,
2706}
2707
2708impl From<&CompletionReq> for ClientSampling {
2709 fn from(r: &CompletionReq) -> Self {
2710 ClientSampling {
2711 temperature: r.temperature,
2712 top_p: r.top_p,
2713 top_k: r.top_k,
2714 min_p: r.min_p,
2715 frequency_penalty: r.frequency_penalty,
2716 presence_penalty: r.presence_penalty,
2717 repetition_penalty: r.repetition_penalty,
2718 seed: r.seed,
2719 }
2720 }
2721}
2722
2723impl From<&ChatCompletionReq> for ClientSampling {
2724 fn from(r: &ChatCompletionReq) -> Self {
2725 ClientSampling {
2726 temperature: r.temperature,
2727 top_p: r.top_p,
2728 top_k: r.top_k,
2729 min_p: r.min_p,
2730 frequency_penalty: r.frequency_penalty,
2731 presence_penalty: r.presence_penalty,
2732 repetition_penalty: r.repetition_penalty,
2733 seed: r.seed,
2734 }
2735 }
2736}
2737
2738fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
2744 sampler_config(
2745 client
2746 .temperature
2747 .or(defaults.temperature)
2748 .unwrap_or_else(default_temperature),
2749 client.top_k.or(defaults.top_k).unwrap_or(0),
2750 client.top_p.or(defaults.top_p).unwrap_or_else(one),
2751 client.min_p.or(defaults.min_p).unwrap_or(0.0),
2752 client
2753 .frequency_penalty
2754 .or(defaults.frequency_penalty)
2755 .unwrap_or(0.0),
2756 client
2757 .presence_penalty
2758 .or(defaults.presence_penalty)
2759 .unwrap_or(0.0),
2760 client
2761 .repetition_penalty
2762 .or(defaults.repetition_penalty)
2763 .unwrap_or_else(one),
2764 client.seed,
2765 )
2766}
2767
2768#[derive(Serialize)]
2769struct CompletionResp {
2770 model: String,
2771 text: String,
2772 tokens: Vec<u32>,
2773 stop_reason: String,
2777 #[serde(default, skip_serializing_if = "Option::is_none")]
2782 error: Option<serde_json::Value>,
2783 n_tokens: usize,
2784 prompt_tokens: usize,
2787 cached_tokens: usize,
2788 elapsed_s: f64,
2789}
2790
2791fn usage_json(
2799 n_prompt: usize,
2800 n_tokens: usize,
2801 n_cached: usize,
2802 elapsed_s: f64,
2803 spec: Option<worker::SpecUsage>,
2804) -> serde_json::Value {
2805 let mut u = json!({
2806 "prompt_tokens": n_prompt,
2807 "completion_tokens": n_tokens,
2808 "total_tokens": n_prompt + n_tokens,
2809 "prompt_tokens_details": { "cached_tokens": n_cached },
2810 "elapsed_s": elapsed_s,
2811 });
2812 if let Some(sp) = spec {
2813 u["spec"] = json!({
2814 "rounds": sp.rounds,
2815 "drafted": sp.drafted,
2816 "accepted": sp.accepted,
2817 "acceptance_rate": if sp.drafted > 0 {
2818 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
2819 });
2820 }
2821 u
2822}
2823
2824const SYSTEM_FINGERPRINT: &str = concat!("memra-", env!("MEMRA_BUILD_SHA"));
2838
2839fn gen_hex128() -> String {
2842 use std::hash::{BuildHasher, Hasher};
2843 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2844 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2845 let t = std::time::SystemTime::now()
2846 .duration_since(std::time::UNIX_EPOCH)
2847 .map(|d| d.as_nanos() as u64)
2848 .unwrap_or(0);
2849 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
2850 h1.write_u64(n);
2851 h1.write_u64(t);
2852 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
2853 h2.write_u64(t.rotate_left(17));
2854 h2.write_u64(n);
2855 format!("{:016x}{:016x}", h1.finish(), h2.finish())
2856}
2857
2858#[derive(Clone)]
2861struct Envelope {
2862 id: String,
2863 created: u64,
2864}
2865
2866impl Envelope {
2867 fn new(chat: bool) -> Self {
2868 Envelope {
2869 id: format!(
2870 "{}-{}",
2871 if chat { "chatcmpl" } else { "cmpl" },
2872 gen_hex128()
2873 ),
2874 created: std::time::SystemTime::now()
2875 .duration_since(std::time::UNIX_EPOCH)
2876 .map(|d| d.as_secs())
2877 .unwrap_or(0),
2878 }
2879 }
2880
2881 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
2883 v["id"] = json!(self.id);
2884 v["created"] = json!(self.created);
2885 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
2886 v
2887 }
2888}
2889
2890fn with_request_id(id: &str, mut resp: Response) -> Response {
2892 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
2893 resp.headers_mut()
2894 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
2895 }
2896 resp
2897}
2898
2899fn openai_compat() -> bool {
2907 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2908 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
2909 Ok("openai") => true,
2910 Ok(_) => false,
2911 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
2912 })
2913}
2914
2915fn cache_namespace(cache_salt: &Option<String>) -> String {
2927 cache_salt.clone().unwrap_or_default()
2928}
2929
2930const CACHE_SALT_MAX_BYTES: usize = 64;
2931
2932fn validate_cache_namespace(
2933 cache_salt: &Option<String>,
2934 keyring_configured: bool,
2935) -> Result<String, &'static str> {
2936 let raw = cache_namespace(cache_salt);
2937 if raw.len() > CACHE_SALT_MAX_BYTES {
2938 return Err("cache_salt must be at most 64 bytes");
2939 }
2940 if !keyring_configured && raw.starts_with("t:") {
2941 return Err("cache_salt must not use the reserved t: prefix without a keyring");
2942 }
2943 if !raw
2944 .bytes()
2945 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
2946 {
2947 return Err("cache_salt contains unsupported characters");
2948 }
2949 Ok(raw)
2950}
2951
2952fn affinity_key(
2968 session_id: &Option<String>,
2969 user: &Option<String>,
2970 headers: &axum::http::HeaderMap,
2971) -> Option<String> {
2972 let clean = |s: &str| -> Option<String> {
2973 let t = s.trim();
2974 if t.is_empty() {
2975 None
2976 } else {
2977 Some(t.to_string())
2978 }
2979 };
2980 session_id
2981 .as_deref()
2982 .and_then(clean)
2983 .or_else(|| user.as_deref().and_then(clean))
2984 .or_else(|| {
2985 headers
2986 .get("x-session-id")
2987 .and_then(|v| v.to_str().ok())
2988 .and_then(clean)
2989 })
2990}
2991
2992fn error_body(
2997 message: &str,
2998 etype: &str,
2999 param: Option<&str>,
3000 code: Option<&str>,
3001) -> serde_json::Value {
3002 json!({ "error": {
3003 "message": message,
3004 "type": etype,
3005 "param": param,
3006 "code": code,
3007 } })
3008}
3009
3010fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3011 error_response_coded(status, message, etype, param, None)
3012}
3013
3014fn error_response_coded(
3019 status: StatusCode,
3020 message: &str,
3021 etype: &str,
3022 param: Option<&str>,
3023 code: Option<&str>,
3024) -> Response {
3025 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3026 if status.is_client_error()
3027 && status != StatusCode::TOO_MANY_REQUESTS
3028 && status != StatusCode::REQUEST_TIMEOUT
3029 && status != StatusCode::CONFLICT
3030 {
3031 resp.headers_mut().insert(
3032 "x-should-retry",
3033 axum::http::HeaderValue::from_static("false"),
3034 );
3035 }
3036 resp
3037}
3038
3039fn bad_request(message: &str, param: Option<&str>) -> Response {
3040 error_response(
3041 StatusCode::BAD_REQUEST,
3042 message,
3043 "invalid_request_error",
3044 param,
3045 )
3046}
3047
3048const 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>) {
3077 use worker::ErrClass as C;
3078 match class {
3079 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3080 C::ContextLength => (
3081 StatusCode::BAD_REQUEST,
3082 "invalid_request_error",
3083 Some("context_length_exceeded"),
3084 ),
3085 C::ModelNotFound => (
3086 StatusCode::BAD_REQUEST,
3087 "invalid_request_error",
3088 Some("model_not_found"),
3089 ),
3090 C::RateLimit => (
3091 StatusCode::TOO_MANY_REQUESTS,
3092 "rate_limit_error",
3093 Some("rate_limit_exceeded"),
3094 ),
3095 C::Overloaded => (
3096 StatusCode::SERVICE_UNAVAILABLE,
3097 "server_error",
3098 Some("overloaded"),
3099 ),
3100 C::Engine => (
3101 StatusCode::INTERNAL_SERVER_ERROR,
3102 "server_error",
3103 Some("engine_error"),
3104 ),
3105 }
3106}
3107
3108fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3110 use worker::ErrClass as C;
3111 match class {
3112 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3113 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3114 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3118 }
3119}
3120
3121fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3124 let (_, etype, code) = class_http(e.class);
3125 error_body(&e.message, etype, e.param, code)
3126}
3127
3128fn engine_error_response(e: &worker::EngineError) -> Response {
3134 engine_error_response_with_retry_after(
3135 e,
3136 e.retry_after_s.or_else(|| class_retry_after_s(e.class)),
3137 )
3138}
3139
3140fn engine_error_response_with_retry_after(
3141 e: &worker::EngineError,
3142 retry_after_s: Option<u64>,
3143) -> Response {
3144 let (status, _, _) = class_http(e.class);
3145 let resp = (status, Json(engine_error_body(e))).into_response();
3146 retry_contract_response(resp, retry_after_s)
3147}
3148
3149fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3151 let status = resp.status();
3152 let h = resp.headers_mut();
3153 match retry_after_s {
3154 Some(secs) => {
3155 let secs = secs.clamp(1, 60);
3157 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3158 h.insert(axum::http::header::RETRY_AFTER, v);
3159 }
3160 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3161 h.insert("retry-after-ms", v);
3162 }
3163 }
3164 None if status.is_client_error() => {
3165 h.insert(
3168 "x-should-retry",
3169 axum::http::HeaderValue::from_static("false"),
3170 );
3171 }
3172 None => {}
3173 }
3174 resp
3175}
3176
3177fn worker_unavailable_response() -> Response {
3178 engine_error_response_with_retry_after(
3179 &worker::EngineError::overloaded("worker unavailable"),
3180 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3181 )
3182}
3183
3184fn stop_reason_to_finish(r: &str) -> &'static str {
3185 match r {
3186 "Eos" | "Callback" => "stop",
3187 "MaxNew" | "ContextFull" => "length",
3188 _ => "stop",
3189 }
3190}
3191
3192fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3196 match v {
3197 serde_json::Value::Null => Ok(String::new()),
3198 serde_json::Value::String(s) => Ok(s.clone()),
3199 serde_json::Value::Array(parts) => {
3200 let mut out = String::new();
3201 for p in parts {
3202 match p.get("type").and_then(|t| t.as_str()) {
3203 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3204 Some(t) => out.push_str(t),
3205 None => return Err("content part has no text field".into()),
3206 },
3207 Some(other) => {
3208 return Err(format!(
3209 "unsupported content part type {other:?} (text only)"
3210 ));
3211 }
3212 }
3213 }
3214 Ok(out)
3215 }
3216 _ => Err("content must be a string, null, or an array of text parts".into()),
3217 }
3218}
3219
3220fn vision_enabled() -> bool {
3223 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3224 *ON.get_or_init(|| {
3225 std::env::var("MEMRA_VISION_DIR").is_ok()
3226 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3227 })
3228}
3229
3230fn gemma_vision_enabled() -> bool {
3235 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3236 *ON.get_or_init(|| {
3237 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3238 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3239 })
3240}
3241
3242pub(crate) static GLM5_VISION_SERVING: std::sync::atomic::AtomicBool =
3250 std::sync::atomic::AtomicBool::new(false);
3251
3252fn glm5_vision_enabled() -> bool {
3255 GLM5_VISION_SERVING.load(std::sync::atomic::Ordering::Acquire)
3256}
3257
3258fn step_vision_enabled() -> bool {
3264 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3265 *ON.get_or_init(|| {
3266 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3267 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3268 })
3269}
3270
3271const VISION_MAX_IMAGES: usize = 8;
3273
3274pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3280 std::sync::atomic::AtomicUsize::new(0);
3281pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3285 tokio::sync::Semaphore::const_new(1);
3286
3287pub(crate) struct VisionMemoryPermit {
3288 bytes: usize,
3289}
3290
3291#[derive(Debug)]
3292pub(crate) enum VisionMemoryError {
3293 Request(String),
3294 Capacity(String),
3295}
3296
3297impl Drop for VisionMemoryPermit {
3298 fn drop(&mut self) {
3299 if self.bytes != 0 {
3300 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3301 }
3302 }
3303}
3304
3305fn try_reserve_vision_memory(
3306 bytes: usize,
3307) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3308 if bytes == 0 {
3309 return Ok(None);
3310 }
3311 if bytes > MAX_VISION_PATCH_BYTES {
3312 return Err(VisionMemoryError::Request(format!(
3313 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3314 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3315 )));
3316 }
3317 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3318 loop {
3319 let Some(next) = in_use.checked_add(bytes) else {
3320 return Err(VisionMemoryError::Capacity(
3321 "vision patch memory reservation overflowed".into(),
3322 ));
3323 };
3324 if next > MAX_VISION_PATCH_BYTES {
3325 return Err(VisionMemoryError::Capacity(format!(
3326 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3327 in_use / (1024 * 1024),
3328 bytes / (1024 * 1024)
3329 )));
3330 }
3331 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3332 in_use,
3333 next,
3334 std::sync::atomic::Ordering::AcqRel,
3335 std::sync::atomic::Ordering::Acquire,
3336 ) {
3337 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3338 Err(actual) => in_use = actual,
3339 }
3340 }
3341}
3342
3343pub(crate) fn vision_memory_error_response(
3344 error: VisionMemoryError,
3345 param: Option<&str>,
3346) -> Response {
3347 match error {
3348 VisionMemoryError::Request(message) => bad_request(&message, param),
3349 VisionMemoryError::Capacity(message) => retry_contract_response(
3350 error_response_coded(
3351 StatusCode::SERVICE_UNAVAILABLE,
3352 &message,
3353 "server_error",
3354 None,
3355 Some("vision_memory_busy"),
3356 ),
3357 Some(RETRY_AFTER_S_OVERLOADED),
3358 ),
3359 }
3360}
3361
3362enum PendingVisionUnit {
3368 Still {
3369 bytes: Vec<u8>,
3370 gh: usize,
3371 gw: usize,
3372 },
3373 Video {
3374 bytes: Vec<u8>,
3375 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3376 video: usize,
3377 },
3378}
3379
3380struct PendingGemmaImage {
3382 bytes: Vec<u8>,
3383 gw: usize,
3384 gh: usize,
3385}
3386
3387struct PendingGlm5Image {
3390 bytes: Vec<u8>,
3391 gh: usize,
3392 gw: usize,
3393}
3394
3395struct PendingStepImage {
3398 bytes: Vec<u8>,
3399 plan: memra_engine::vision_step::StepImagePlan,
3400}
3401
3402fn content_to_text_vision_step(
3412 v: &serde_json::Value,
3413 step_images: &mut Vec<PendingStepImage>,
3414) -> Result<String, String> {
3415 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
3416 let parts = match v {
3417 serde_json::Value::Array(parts) => parts,
3418 _ => return content_to_text(v),
3419 };
3420 let mut out = String::new();
3421 let mut needs_sep = false;
3422 for p in parts {
3423 match p.get("type").and_then(|t| t.as_str()) {
3424 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3425 Some(t) => {
3426 if needs_sep {
3427 out.push(' ');
3428 }
3429 out.push_str(t);
3430 needs_sep = true;
3431 }
3432 None => return Err("content part has no text field".into()),
3433 },
3434 Some("image_url") => {
3435 let url = p
3436 .get("image_url")
3437 .and_then(|u| {
3438 if u.is_string() {
3439 u.as_str()
3440 } else {
3441 u.get("url").and_then(|x| x.as_str())
3442 }
3443 })
3444 .ok_or("image_url part has no url")?;
3445 if !url.starts_with("data:") {
3446 return Err(
3447 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3448 );
3449 }
3450 if step_images.len() >= VISION_MAX_IMAGES {
3451 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3452 }
3453 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3457 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3458 let plan = memra_engine::vision_step::step_plan_image(&bytes)
3459 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3460 for i in 0..plan.n_tiles {
3461 out.push_str("<patch_start>");
3462 for _ in 0..SV_TILE_ROWS {
3463 out.push_str("<im_patch>");
3464 }
3465 out.push_str("<patch_end>");
3466 if plan.newline_mask[i] {
3467 out.push_str("<patch_newline>");
3468 }
3469 }
3470 out.push_str("<im_start>");
3471 for _ in 0..SV_MAIN_ROWS {
3472 out.push_str("<im_patch>");
3473 }
3474 out.push_str("<im_end>");
3475 step_images.push(PendingStepImage { bytes, plan });
3476 needs_sep = false;
3477 }
3478 Some("video_url") => {
3479 return Err("step37 has no video input (image-only processor)".into());
3480 }
3481 Some(other) => {
3482 return Err(format!("unsupported content part type {other:?}"));
3483 }
3484 }
3485 }
3486 Ok(out)
3487}
3488
3489fn content_to_text_vision(
3498 v: &serde_json::Value,
3499 images: &mut Vec<PendingVisionUnit>,
3500 gemma_images: &mut Vec<PendingGemmaImage>,
3501 glm5_images: &mut Vec<PendingGlm5Image>,
3502 step_images: &mut Vec<PendingStepImage>,
3503 next_video: &mut usize,
3504) -> Result<String, String> {
3505 if step_vision_enabled() {
3509 return content_to_text_vision_step(v, step_images);
3510 }
3511 let parts = match v {
3512 serde_json::Value::Array(parts) => parts,
3513 _ => return content_to_text(v),
3514 };
3515 let mut out = String::new();
3516 for p in parts {
3517 match p.get("type").and_then(|t| t.as_str()) {
3518 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3519 Some(t) => out.push_str(t),
3520 None => return Err("content part has no text field".into()),
3521 },
3522 Some("image_url") if glm5_vision_enabled() => {
3523 let url = p
3524 .get("image_url")
3525 .and_then(|u| {
3526 if u.is_string() {
3527 u.as_str()
3528 } else {
3529 u.get("url").and_then(|x| x.as_str())
3530 }
3531 })
3532 .ok_or("image_url part has no url")?;
3533 if !url.starts_with("data:") {
3534 return Err(
3535 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3536 );
3537 }
3538 if glm5_images.len() >= VISION_MAX_IMAGES {
3539 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3540 }
3541 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3545 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
3546 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes)
3547 .map_err(|e| format!("image {}: {e}", glm5_images.len() + 1))?;
3548 out.push_str("<|begin_of_image|>");
3552 for _ in 0..memra_engine::vision_glm5::n_merged_for_grid(gh, gw) {
3553 out.push_str("<|image|>");
3554 }
3555 out.push_str("<|end_of_image|>");
3556 glm5_images.push(PendingGlm5Image { bytes, gh, gw });
3557 }
3558 Some("video_url") if glm5_vision_enabled() => {
3559 return Err(
3560 "glm5 video input is not served (tensor census only; image input is the \
3561 supported surface)"
3562 .into(),
3563 );
3564 }
3565 Some("image_url") if gemma_vision_enabled() => {
3566 let url = p
3567 .get("image_url")
3568 .and_then(|u| {
3569 if u.is_string() {
3570 u.as_str()
3571 } else {
3572 u.get("url").and_then(|x| x.as_str())
3573 }
3574 })
3575 .ok_or("image_url part has no url")?;
3576 if !url.starts_with("data:") {
3577 return Err(
3578 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3579 );
3580 }
3581 if gemma_images.len() >= VISION_MAX_IMAGES {
3582 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3583 }
3584 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
3588 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3589 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
3590 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3591 out.push_str("<|image>");
3593 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
3594 out.push_str("<|image|>");
3595 }
3596 out.push_str("<image|>");
3597 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
3598 }
3599 Some("image_url") => {
3600 if !vision_enabled() {
3601 return Err("image input is not enabled on this deployment".into());
3602 }
3603 let url = p
3604 .get("image_url")
3605 .and_then(|u| {
3606 if u.is_string() {
3607 u.as_str()
3608 } else {
3609 u.get("url").and_then(|x| x.as_str())
3610 }
3611 })
3612 .ok_or("image_url part has no url")?;
3613 if !url.starts_with("data:") {
3614 return Err(
3615 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3616 );
3617 }
3618 if images
3619 .iter()
3620 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
3621 .count()
3622 >= VISION_MAX_IMAGES
3623 {
3624 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3625 }
3626 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3631 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3632 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
3633 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3634 out.push_str("<|vision_start|>");
3635 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
3636 out.push_str("<|image_pad|>");
3637 }
3638 out.push_str("<|vision_end|>");
3639 images.push(PendingVisionUnit::Still { bytes, gh, gw });
3640 }
3641 Some("video_url") if gemma_vision_enabled() => {
3642 return Err("gemma-4 has no video input (image-only projector)".into());
3643 }
3644 Some("video_url") => {
3645 if !vision_enabled() {
3646 return Err("video input is not enabled on this deployment".into());
3647 }
3648 let url = p
3649 .get("video_url")
3650 .and_then(|u| {
3651 if u.is_string() {
3652 u.as_str()
3653 } else {
3654 u.get("url").and_then(|x| x.as_str())
3655 }
3656 })
3657 .ok_or("video_url part has no url")?;
3658 if !url.starts_with("data:") {
3659 return Err(
3660 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3661 );
3662 }
3663 if *next_video >= 2 {
3664 return Err("too many videos (max 2)".into());
3665 }
3666 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
3669 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
3670 .map_err(|e| format!("video: {e}"))?;
3671 let vidx = *next_video;
3672 *next_video += 1;
3673 for group in &vid.groups {
3675 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
3676 out.push_str("<|vision_start|>");
3677 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
3678 out.push_str("<|video_pad|>");
3679 }
3680 out.push_str("<|vision_end|>");
3681 }
3682 images.push(PendingVisionUnit::Video {
3685 bytes,
3686 groups: vid.groups,
3687 video: vidx,
3688 });
3689 }
3690 Some(other) => {
3691 return Err(format!("unsupported content part type {other:?}"));
3692 }
3693 }
3694 }
3695 Ok(out)
3696}
3697
3698fn pyjson(v: &serde_json::Value, out: &mut String) {
3702 match v {
3703 serde_json::Value::Object(m) => {
3704 out.push('{');
3705 for (i, (k, val)) in m.iter().enumerate() {
3706 if i > 0 {
3707 out.push_str(", ");
3708 }
3709 out.push_str(&serde_json::Value::String(k.clone()).to_string());
3710 out.push_str(": ");
3711 pyjson(val, out);
3712 }
3713 out.push('}');
3714 }
3715 serde_json::Value::Array(a) => {
3716 out.push('[');
3717 for (i, val) in a.iter().enumerate() {
3718 if i > 0 {
3719 out.push_str(", ");
3720 }
3721 pyjson(val, out);
3722 }
3723 out.push(']');
3724 }
3725 scalar => out.push_str(&scalar.to_string()),
3726 }
3727}
3728
3729fn pyjson_str(v: &serde_json::Value) -> String {
3730 let mut s = String::new();
3731 pyjson(v, &mut s);
3732 s
3733}
3734
3735#[allow(clippy::too_many_arguments)] fn sampler_config(
3743 temperature: f32,
3744 top_k: usize,
3745 top_p: f32,
3746 min_p: f32,
3747 frequency_penalty: f32,
3748 presence_penalty: f32,
3749 repetition_penalty: f32,
3750 seed: Option<u64>,
3751) -> SamplerConfig {
3752 let penalties_on =
3753 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
3754 SamplerConfig {
3755 temperature,
3756 top_k,
3757 top_p,
3758 min_p,
3759 penalty_last_n: if penalties_on {
3760 memra_engine::spec::PEN_WINDOW_MAX
3761 } else {
3762 0
3763 },
3764 penalty_repeat: repetition_penalty,
3765 penalty_freq: frequency_penalty,
3766 penalty_present: presence_penalty,
3767 seed: seed.unwrap_or_else(fresh_seed),
3770 }
3771}
3772
3773fn fresh_seed() -> u64 {
3778 use std::sync::atomic::{AtomicU64, Ordering};
3779 static COUNTER: AtomicU64 = AtomicU64::new(0);
3780 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
3781 let nanos = std::time::SystemTime::now()
3782 .duration_since(std::time::UNIX_EPOCH)
3783 .map(|d| d.as_nanos() as u64)
3784 .unwrap_or(0);
3785 let mut z = nanos
3786 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
3787 .wrapping_add(0x9E3779B97F4A7C15);
3788 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
3789 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
3790 z ^= z >> 31;
3791 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
3794}
3795
3796fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
3801 for (param, present, why) in fields {
3802 if *present {
3803 return Err((format!("{param} is not supported{why}"), param.to_string()));
3804 }
3805 }
3806 Ok(())
3807}
3808
3809#[derive(PartialEq)]
3810enum ToolChoice {
3811 Auto,
3812 None,
3813}
3814
3815fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
3816 match v {
3817 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
3818 Some(serde_json::Value::String(s)) => match s.as_str() {
3819 "auto" => Ok(ToolChoice::Auto),
3820 "none" => Ok(ToolChoice::None),
3821 "required" => Err("tool_choice \"required\" is not supported (no constrained \
3822 decoding); use \"auto\""
3823 .into()),
3824 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
3825 },
3826 Some(serde_json::Value::Object(_)) => {
3827 Err("named-function tool_choice is not supported; use \"auto\"".into())
3828 }
3829 Some(other) => Err(format!("bad tool_choice: {other}")),
3830 }
3831}
3832
3833fn parse_think(
3883 reasoning_effort: &Option<String>,
3884 reasoning: &Option<serde_json::Value>,
3885 vllm_switch: Option<bool>,
3886 suppress_switch: Option<bool>,
3887 default_effort: Option<&str>,
3888 max_tier: bool,
3889) -> Result<(ThinkMode, Option<String>, bool), String> {
3890 let mut effort = reasoning_effort.clone();
3891 let ReasoningObject {
3892 mut enabled,
3893 effort: object_effort,
3894 exclude,
3895 } = parse_reasoning_object(reasoning)?;
3896 if let Some(e) = object_effort {
3897 effort = Some(e);
3898 }
3899 match (enabled, vllm_switch) {
3904 (Some(a), Some(b)) if a != b => {
3905 return Err(format!(
3906 "contradictory reasoning switches: reasoning.enabled={a} and \
3907 enable_thinking={b} — send one"
3908 ));
3909 }
3910 (None, Some(b)) => enabled = Some(b),
3911 _ => {}
3912 }
3913 let suppress = match (exclude, suppress_switch) {
3925 (Some(true), _) | (_, Some(false)) => Some(false),
3926 _ => None,
3927 };
3928 match (enabled, suppress) {
3929 (Some(true), Some(false)) => {
3930 return Err(
3931 "contradictory reasoning switches: reasoning is enabled but \
3932 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
3933 on this server not delivering reasoning means not generating it, so send one"
3934 .into(),
3935 );
3936 }
3937 (None, Some(b)) => enabled = Some(b),
3938 _ => {}
3939 }
3940 let client_explicit = effort.is_some() || enabled.is_some();
3944 if effort.is_none() && enabled.is_none() {
3949 effort = default_effort.map(str::to_string);
3950 }
3951 let effort_arm = match effort.as_deref() {
3956 None => None,
3957 Some(raw) => {
3958 let level = canonical_effort_for(raw, max_tier).ok_or_else(|| {
3959 format!(
3960 "bad reasoning_effort {raw:?} \
3961 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
3962 highest level this model's template distinguishes)"
3963 )
3964 })?;
3965 Some(match level {
3966 "none" | "minimal" => (ThinkMode::NoThink, "low"),
3967 "low" => (ThinkMode::Think, "low"),
3968 "medium" => (ThinkMode::Think, "medium"),
3969 "max" => (ThinkMode::Think, "max"),
3970 _ => (ThinkMode::Think, "high"),
3971 })
3972 }
3973 };
3974 let (think, level) = match (enabled, effort_arm) {
3975 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
3978 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
3979 (None, Some((think, level))) => (think, Some(level.to_string())),
3980 (None, None) => (ThinkMode::Default, None),
3981 };
3982 Ok((think, level, client_explicit))
3983}
3984
3985struct ReasoningObject {
3987 enabled: Option<bool>,
3988 effort: Option<String>,
3989 exclude: Option<bool>,
3990}
3991
3992fn parse_reasoning_object(
4011 reasoning: &Option<serde_json::Value>,
4012) -> Result<ReasoningObject, String> {
4013 let mut out = ReasoningObject {
4014 enabled: None,
4015 effort: None,
4016 exclude: None,
4017 };
4018 let Some(v) = reasoning else { return Ok(out) };
4019 let obj = match v {
4020 serde_json::Value::Null => return Ok(out),
4021 serde_json::Value::Object(obj) => obj,
4022 _ => return Err("reasoning must be an object".into()),
4023 };
4024 for (key, value) in obj {
4025 match key.as_str() {
4033 "enabled" => {
4034 if !value.is_null() {
4035 out.enabled = Some(
4036 value
4037 .as_bool()
4038 .ok_or("reasoning.enabled must be true or false")?,
4039 );
4040 }
4041 }
4042 "exclude" => {
4043 if !value.is_null() {
4044 out.exclude = Some(
4045 value
4046 .as_bool()
4047 .ok_or("reasoning.exclude must be true or false")?,
4048 );
4049 }
4050 }
4051 "effort" => {
4052 if !value.is_null() {
4053 out.effort = Some(
4054 value
4055 .as_str()
4056 .ok_or("reasoning.effort must be a string")?
4057 .to_string(),
4058 );
4059 }
4060 }
4061 "max_tokens" => {
4062 return Err(
4063 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4064 are output tokens here, and max_tokens is the ONE output budget covering \
4065 reasoning and content together — there is no separate reasoning budget to \
4066 spend against, so honouring this field is impossible rather than merely \
4067 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4068 reasoning.enabled:false) to spend less of it on reasoning"
4069 .into(),
4070 );
4071 }
4072 other => {
4073 return Err(format!(
4074 "reasoning.{other} is not a field this server implements (it would change \
4075 nothing about the request); the supported keys are enabled, effort and \
4076 exclude"
4077 ));
4078 }
4079 }
4080 }
4081 Ok(out)
4082}
4083
4084fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4105 let Some(v) = kwargs else { return Ok(None) };
4106 let obj = match v {
4107 serde_json::Value::Null => return Ok(None),
4108 serde_json::Value::Object(obj) => obj,
4109 _ => return Err("chat_template_kwargs must be an object".into()),
4110 };
4111 let mut switch = None;
4112 for (key, value) in obj {
4113 match key.as_str() {
4114 "enable_thinking" => {
4115 switch = Some(
4116 value
4117 .as_bool()
4118 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4119 );
4120 }
4121 "preserve_thinking" => {
4122 let preserve = value
4123 .as_bool()
4124 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4125 if !preserve {
4126 return Err(
4127 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4128 server: the renderer implements the vendor DEFAULT (replay every prior \
4129 assistant turn's <think> block, empty when no reasoning was sent) but \
4130 not the strip arm — serving replay bytes under a strip request would \
4131 misdescribe the prompt. Omit the flag or send true"
4132 .into(),
4133 );
4134 }
4135 }
4137 other => {
4138 return Err(format!(
4139 "chat_template_kwargs.{other} is not supported by this server's \
4140 template renderer (it would change nothing about the prompt); the only \
4141 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4142 refuses in both directions — see its own message)"
4143 ));
4144 }
4145 }
4146 }
4147 Ok(switch)
4148}
4149
4150fn resolve_vllm_think_switch(
4154 enable_thinking: Option<bool>,
4155 kwargs: &Option<serde_json::Value>,
4156) -> Result<Option<bool>, String> {
4157 let from_kwargs = parse_template_kwargs(kwargs)?;
4158 match (enable_thinking, from_kwargs) {
4159 (Some(a), Some(b)) if a != b => Err(format!(
4160 "contradictory reasoning switches: enable_thinking={a} and \
4161 chat_template_kwargs.enable_thinking={b} — send one"
4162 )),
4163 (Some(a), _) => Ok(Some(a)),
4164 (None, b) => Ok(b),
4165 }
4166}
4167
4168pub(crate) fn canonical_effort_for(value: &str, max_tier: bool) -> Option<&'static str> {
4188 match value {
4189 "none" => Some("none"),
4190 "minimal" => Some("minimal"),
4191 "low" => Some("low"),
4192 "medium" => Some("medium"),
4193 "high" => Some("high"),
4194 "xhigh" | "max" | "ultra" => Some(if max_tier { "max" } else { "high" }),
4200 _ => None,
4201 }
4202}
4203
4204pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4207 canonical_effort_for(value, false)
4208}
4209
4210fn json_to_val(v: &serde_json::Value) -> chat::Val {
4213 match v {
4214 serde_json::Value::Null => chat::Val::Null,
4215 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4216 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4217 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4218 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4219 serde_json::Value::Object(o) => chat::Val::Obj(
4222 o.iter()
4223 .map(|(k, val)| (k.clone(), json_to_val(val)))
4224 .collect(),
4225 ),
4226 }
4227}
4228
4229#[allow(clippy::type_complexity)]
4233fn prepare_tools(
4234 tools: &[serde_json::Value],
4235) -> Result<
4236 (
4237 Vec<String>,
4238 Vec<chat::Val>,
4239 HashMap<String, HashMap<String, String>>,
4240 ),
4241 String,
4242> {
4243 let mut tools_json = Vec::with_capacity(tools.len());
4244 let mut tools_struct = Vec::with_capacity(tools.len());
4245 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4246 for t in tools {
4247 let f = t
4248 .get("function")
4249 .ok_or("each tool needs a function object")?;
4250 let name = f
4251 .get("name")
4252 .and_then(|n| n.as_str())
4253 .ok_or("each tool needs function.name")?;
4254 let mut params: HashMap<String, String> = HashMap::new();
4255 if let Some(props) = f
4256 .get("parameters")
4257 .and_then(|p| p.get("properties"))
4258 .and_then(|p| p.as_object())
4259 {
4260 for (p, def) in props {
4261 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4262 params.insert(p.clone(), ty.to_string());
4263 }
4264 }
4265 }
4266 schemas.insert(name.to_string(), params);
4267 tools_json.push(pyjson_str(t));
4268 tools_struct.push(json_to_val(f));
4270 }
4271 Ok((tools_json, tools_struct, schemas))
4272}
4273
4274fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4279 let parsed: serde_json::Value = match &tc.function.arguments {
4280 serde_json::Value::Null => json!({}),
4281 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4282 serde_json::Value::String(s) => serde_json::from_str(s)
4283 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4284 v @ serde_json::Value::Object(_) => v.clone(),
4285 _ => return Err("tool_calls arguments must be a JSON object".into()),
4286 };
4287 let obj = parsed
4288 .as_object()
4289 .ok_or("tool_calls arguments must decode to a JSON object")?;
4290 let params = obj
4291 .iter()
4292 .map(|(k, v)| {
4293 let rendered = match v {
4294 serde_json::Value::String(s) => s.clone(),
4295 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4296 scalar => scalar.to_string(),
4297 };
4298 (k.clone(), rendered)
4299 })
4300 .collect();
4301 let args = obj
4304 .iter()
4305 .map(|(k, v)| (k.clone(), json_to_val(v)))
4306 .collect();
4307 Ok(TmplToolCall {
4308 name: tc.function.name.clone(),
4309 params,
4310 args,
4311 id: tc.id.clone(),
4312 })
4313}
4314
4315fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4317 json!({ "id": c.id, "type": "function",
4318 "function": { "name": c.name, "arguments": c.arguments } })
4319}
4320
4321#[tokio::main]
4325pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4326 serve_with(ServerWiring::stock()).await
4327}
4328
4329enum MeteringWiring {
4331 Stock,
4335 Custom(metering::MeteringFactory),
4340}
4341
4342pub struct ServerWiring {
4346 metering: MeteringWiring,
4347 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
4350 claimed_env: Vec<&'static str>,
4354}
4355
4356impl ServerWiring {
4357 pub fn stock() -> Self {
4359 ServerWiring {
4360 metering: MeteringWiring::Stock,
4361 on_ready: None,
4362 claimed_env: Vec::new(),
4363 }
4364 }
4365
4366 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
4369 ServerWiring {
4370 metering: MeteringWiring::Custom(factory),
4371 on_ready: None,
4372 claimed_env: Vec::new(),
4373 }
4374 }
4375
4376 pub fn claiming(mut self, var: &'static str) -> Self {
4380 self.claimed_env.push(var);
4381 self
4382 }
4383
4384 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
4385 self.on_ready = Some(Box::new(hook));
4386 self
4387 }
4388}
4389
4390pub struct RuntimeHandles {
4393 pub trim: TrimHandle,
4394 pub purge: PurgeHandle,
4397 pub shutdown: tokio::sync::watch::Receiver<bool>,
4402}
4403
4404#[derive(Clone)]
4407pub struct TrimHandle {
4408 cmd_tx: Sender<Cmd>,
4409}
4410
4411impl TrimHandle {
4412 pub async fn trim(&self) -> Result<serde_json::Value, String> {
4414 let (tx, rx) = tokio::sync::oneshot::channel();
4415 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
4416 return Err("worker is down".into());
4417 }
4418 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4419 Ok(Ok(report)) => Ok(json!(report)),
4420 _ => Err("worker did not answer the trim within 30s".into()),
4421 }
4422 }
4423}
4424
4425#[derive(Clone)]
4434pub struct PurgeHandle {
4435 cmd_tx: Sender<Cmd>,
4436}
4437
4438impl PurgeHandle {
4439 pub async fn purge_tenant(&self, tenant: &str) -> Result<serde_json::Value, String> {
4441 let (tx, rx) = tokio::sync::oneshot::channel();
4442 let cmd = Cmd::PurgeTenantHost {
4443 tenant: tenant.to_string(),
4444 tx,
4445 };
4446 if self.cmd_tx.send(cmd).is_err() {
4447 return Err("worker is down".into());
4448 }
4449 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4450 Ok(Ok(report)) => Ok(json!(report)),
4451 _ => Err("worker did not answer the purge within 30s".into()),
4452 }
4453 }
4454}
4455
4456pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
4457 let args: Vec<String> = std::env::args().skip(1).collect();
4460 if let Some(code) = auth::run_cli(&args) {
4461 std::process::exit(code);
4462 }
4463 auth::init_from_env();
4466 let api_auth = match ApiAuth::from_env() {
4467 Ok(auth) => auth,
4468 Err(err) => {
4469 eprintln!("[server] FATAL: {err}");
4470 std::process::exit(1);
4471 }
4472 };
4473 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
4474 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
4475 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
4476 Ok(resolved) => resolved,
4477 Err(err) => {
4478 eprintln!("[server] FATAL: {err}");
4479 std::process::exit(1);
4480 }
4481 };
4482 if let Err(message) = validate_bind_security(&addr, api_auth.configured(), allow_open_bind) {
4487 eprintln!("[server] FATAL: {message}");
4488 std::process::exit(1);
4489 }
4490 if !bind_loopback && !api_auth.configured() {
4491 eprintln!(
4492 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
4493 metrics remain bearer-protected"
4494 );
4495 }
4496 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
4497 Ok(token) if token.is_empty() => {
4498 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
4499 std::process::exit(1);
4500 }
4501 Ok(token) => Some(token),
4502 Err(std::env::VarError::NotPresent) => None,
4503 Err(std::env::VarError::NotUnicode(_)) => {
4504 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
4505 std::process::exit(1);
4506 }
4507 };
4508 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
4509
4510 let models = parse_models_config();
4511 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
4512 Ok(loaded) => loaded,
4513 Err(err) => {
4514 eprintln!("[server] FATAL: {err}");
4515 std::process::exit(1);
4516 }
4517 };
4518 let metering_obj: Option<Arc<dyn metering::Metering>> = {
4524 let factory = match wiring.metering {
4525 MeteringWiring::Stock => None,
4526 MeteringWiring::Custom(factory) => Some(factory),
4527 };
4528 for deployment_only in [
4529 "MEMRA_REQUEST_LEDGER",
4530 "MEMRA_TENANT_BUDGETS",
4531 "MEMRA_ADMIN_ADDR",
4532 "MEMRA_ADMIN_TOKEN_FILE",
4533 "MEMRA_CAPTURE_DIR",
4534 ] {
4535 if std::env::var_os(deployment_only).is_some()
4536 && !wiring.claimed_env.contains(&deployment_only)
4537 {
4538 eprintln!(
4539 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
4540 build ships no accounting/admin/capture. Wire a Metering implementation \
4541 through ServerWiring and claim the vars it consumes."
4542 );
4543 std::process::exit(1);
4544 }
4545 }
4546 match factory {
4547 None => None,
4548 Some(factory) => {
4549 let model_ids: Vec<String> =
4550 models.iter().map(|(name, _, _)| name.clone()).collect();
4551 match factory(&metering::MeteringInit { models: &model_ids }) {
4552 Ok(metering_obj) => metering_obj,
4553 Err(err) => {
4554 eprintln!("[server] FATAL: metering wiring: {err}");
4555 std::process::exit(1);
4556 }
4557 }
4558 }
4559 }
4560 };
4561 let budget_tokenizers = if metering_obj
4562 .as_ref()
4563 .is_some_and(|manager| manager.enforces_limits())
4564 {
4565 match load_budget_tokenizers(&models) {
4566 Ok(tokenizers) => Some(tokenizers),
4567 Err(err) => {
4568 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
4569 std::process::exit(1);
4570 }
4571 }
4572 } else {
4573 None
4574 };
4575 eprintln!("[server] starting; models config = {models:?}");
4576
4577 let health_state = health::WorkerHealth::new();
4582 health::spawn_gpu_watch(health_state.clone());
4586 health::spawn_sd_watchdog(health_state.clone());
4587
4588 let (cmd_tx, model_names, caps, metrics, worker_thread) =
4590 match worker::spawn(models, health_state.clone()) {
4591 Ok(v) => v,
4592 Err(err) => {
4593 eprintln!("[server] FATAL: worker init failed: {err}");
4594 health_state.mark_dead(format!("worker init failed: {err}"));
4595 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
4596 std::process::exit(1);
4597 }
4598 };
4599 eprintln!("[server] worker ready; serving models: {model_names:?}");
4600
4601 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
4608 if let Some(on_ready) = wiring.on_ready {
4609 on_ready(RuntimeHandles {
4610 trim: TrimHandle {
4611 cmd_tx: cmd_tx.clone(),
4612 },
4613 purge: PurgeHandle {
4614 cmd_tx: cmd_tx.clone(),
4615 },
4616 shutdown: drain_shutdown_rx.clone(),
4617 });
4618 }
4619
4620 let bg_handle = darklane::spawn_from_env(health_state.clone());
4623 let bg_state = bg_handle.as_ref().map(|h| {
4624 let mode = darklane::BgConfig::from_env()
4625 .map(|c| c.yield_mode.as_str())
4626 .unwrap_or("stop");
4627 (h.state.clone(), mode)
4628 });
4629
4630 let state = AppState {
4631 cmd_tx,
4632 models: model_names,
4633 caps,
4634 openrouter_metadata: Arc::new(openrouter_metadata),
4635 provider_metadata: Arc::new(provider_metadata),
4636 metering: metering_obj,
4637 budget_tokenizers,
4638 api_auth,
4639 metrics_auth,
4640 metrics,
4641 inflight: Arc::new(Default::default()),
4642 tenant_inflight: Arc::new(Default::default()),
4643 health: health_state.clone(),
4644 bg: bg_state,
4645 };
4646 let inflight_handle = state.inflight.clone();
4647 let drain_metering = state.metering.clone();
4650 worker::register_http_inflight(state.inflight.clone());
4655 let app = Router::new()
4656 .route("/health", get(health_live))
4661 .route("/livez", get(health_live))
4662 .route("/readyz", get(health_ready))
4663 .route("/models", get(list_models))
4664 .route("/v1/models", get(list_models_v1))
4665 .route("/v1/auth/check", get(auth_check))
4666 .route("/v1/completions", post(completions))
4667 .route("/v1/embeddings", post(embed_api::embeddings))
4668 .route("/v1/rerank", post(embed_api::rerank))
4669 .route("/v1/chat/completions", post(chat_completions))
4670 .route("/v1/messages", post(anthropic::messages))
4674 .route("/v1/responses", post(responses_api::responses))
4675 .route("/metrics", get(get_metrics))
4676 .route("/yield/metrics", get(yield_metrics))
4677 .with_state(state.clone());
4678 let app = apply_body_limit(app);
4681 let app = app.layer(middleware::from_fn_with_state(
4684 state,
4685 authenticate_inference_before_body,
4686 ));
4687 let app = if ttft::enabled() {
4688 app.layer(middleware::from_fn(ttft_request_start))
4689 } else {
4690 app
4691 };
4692
4693 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
4694 eprintln!("[server] listening on http://{bind_addr}");
4695 drop(drain_shutdown_rx);
4696 health::sd_notify("READY=1\nSTATUS=serving");
4700 let inflight = inflight_handle;
4707 let signal_admin_shutdown = drain_shutdown_tx.clone();
4708 let serve_result = axum::serve(listener, app)
4709 .with_graceful_shutdown(async move {
4710 let mut sigterm =
4711 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
4712 Ok(s) => s,
4713 Err(err) => {
4714 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
4715 std::future::pending::<()>().await;
4716 unreachable!()
4717 }
4718 };
4719 sigterm.recv().await;
4720 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
4721 let _ = signal_admin_shutdown.send(true);
4722 health::sd_notify(&format!(
4726 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
4727 (drain_deadline_s() + 5) * 1_000_000
4728 ));
4729 let n: usize = inflight
4730 .iter()
4731 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4732 .sum();
4733 eprintln!(
4734 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
4735 drain_deadline_s()
4736 );
4737 let deadline = std::time::Duration::from_secs(drain_deadline_s());
4738 let t0 = std::time::Instant::now();
4739 loop {
4740 let n: usize = inflight
4741 .iter()
4742 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4743 .sum();
4744 if n == 0 {
4745 eprintln!(
4746 "[server] drain complete in {:.1}s; exiting",
4747 t0.elapsed().as_secs_f64()
4748 );
4749 break;
4750 }
4751 if t0.elapsed() >= deadline {
4752 eprintln!(
4753 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
4754 drain_deadline_s()
4755 );
4756 if let Some(metering) = drain_metering.as_ref() {
4763 metering.drain_kill();
4764 }
4765 break;
4766 }
4767 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4768 }
4769 })
4770 .await;
4771 let _ = drain_shutdown_tx.send(true);
4774 serve_result?;
4775 if let Some(h) = bg_handle {
4779 h.shutdown();
4780 }
4781 worker_thread.join().map_err(|_| {
4790 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
4791 })?;
4792 eprintln!("[server] GPU worker shutdown complete");
4793 Ok(())
4794}
4795
4796fn validate_model_path(path: &str) -> Result<(), String> {
4802 let p = std::path::Path::new(path);
4803 if !p.exists() {
4804 return Err(format!("model path {path:?} does not exist"));
4805 }
4806 if p.is_file() {
4807 return Ok(()); }
4809 if p.join("manifest.json").exists() {
4810 return Ok(()); }
4812 let has_st =
4813 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
4814 if !has_st {
4815 return Err(format!(
4816 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
4817 model.safetensors.index.json + config.json (HF safetensors dir), or \
4818 manifest.json (memra repack dir)"
4819 ));
4820 }
4821 if !p.join("config.json").exists() {
4822 return Err(format!(
4823 "model dir {path:?} has safetensors weights but no config.json"
4824 ));
4825 }
4826 Ok(())
4827}
4828
4829fn parse_models_config() -> Vec<(String, String, Option<String>)> {
4837 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
4838 let mut out = Vec::new();
4839 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
4840 if let Some((name, path)) = entry.split_once('=') {
4841 let (mpath, dpath) = match path.trim().split_once('+') {
4844 Some((m, d)) => (m.trim(), Some(d.trim())),
4845 None => (path.trim(), None),
4846 };
4847 let resolve = |p: &str| {
4848 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
4849 eprintln!("[server] FATAL: model {name:?}: {err}");
4850 std::process::exit(1);
4851 })
4852 };
4853 let mpath = resolve(mpath);
4854 if let Err(err) = validate_model_path(&mpath) {
4855 eprintln!("[server] FATAL: model {name:?}: {err}");
4856 std::process::exit(1);
4857 }
4858 let dpath = dpath.map(|d| {
4868 let d = resolve(d);
4869 let p = std::path::Path::new(&d);
4870 if !p.exists() {
4871 eprintln!(
4872 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
4873 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
4874 rather than serving plain decode under a config that asked \
4875 for speculative decoding."
4876 );
4877 std::process::exit(1);
4878 }
4879 if !p.is_file() {
4880 eprintln!(
4881 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
4882 file — a '+draft' attach must be a NextN/MTP GGUF file."
4883 );
4884 std::process::exit(1);
4885 }
4886 d
4887 });
4888 out.push((name.trim().to_string(), mpath, dpath));
4889 } else {
4890 eprintln!(
4891 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
4892 );
4893 }
4894 }
4895 if !out.is_empty() {
4896 return out;
4897 }
4898 }
4899 vec![
4901 (
4902 "main".into(),
4903 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
4904 None,
4905 ),
4906 (
4907 "judge".into(),
4908 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
4909 None,
4910 ),
4911 ]
4912}
4913
4914fn load_budget_tokenizers(
4915 models: &[(String, String, Option<String>)],
4916) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
4917 let mut tokenizers = HashMap::new();
4918 for (alias, path, _) in models {
4919 let path = std::path::Path::new(path);
4920 let tokenizer = if path.is_dir() {
4921 let tokenizer_dir = if path.join("manifest.json").exists() {
4922 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
4923 format!("model {alias:?}: open repack tokenizer source: {err}")
4924 })?;
4925 repack
4926 .source_dir()
4927 .filter(|source| source.join("tokenizer.json").exists())
4928 .unwrap_or(path)
4929 .to_path_buf()
4930 } else {
4931 path.to_path_buf()
4932 };
4933 Tokenizer::from_hf_dir(&tokenizer_dir)
4934 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4935 } else {
4936 let gguf = memra_gguf::GgufFile::open(path)
4937 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
4938 Tokenizer::from_gguf(&gguf)
4939 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4940 };
4941 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
4942 }
4943 Ok(Arc::new(tokenizers))
4944}
4945
4946fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4948 let s = st.health.snapshot();
4949 let mut v = json!({
4950 "status": status,
4951 "models": *st.models,
4952 "worker": {
4953 "phase": health::phase_name(s.phase),
4954 "beat_age_ms": s.beat_age_ms,
4955 "tick_max_ms": s.tick_max_ms,
4956 "stall_threshold_ms": s.stall_threshold_ms,
4957 "generation": s.generation,
4958 "xid_warnings": s.xid_warns,
4959 },
4960 });
4961 if let Some(d) = detail {
4962 v["detail"] = json!(d);
4963 }
4964 v
4965}
4966
4967fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4971 let mut v = health_payload(st, status, detail);
4972 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
4973 v
4974}
4975
4976async fn auth_check() -> impl IntoResponse {
4980 StatusCode::NO_CONTENT
4981}
4982
4983async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
4999 if draining() {
5000 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
5003 }
5004 match st.health.live() {
5005 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
5006 Err(why) => retry_contract_response(
5007 (
5008 StatusCode::SERVICE_UNAVAILABLE,
5009 Json(health_payload(&st, "unhealthy", Some(&why))),
5010 )
5011 .into_response(),
5012 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
5013 ),
5014 }
5015}
5016
5017async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
5028 let is_draining = draining();
5029 match st.health.ready(is_draining) {
5030 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
5031 Err(why) => retry_contract_response(
5032 (
5033 StatusCode::SERVICE_UNAVAILABLE,
5034 Json(readiness_payload(&st, "not_ready", Some(&why))),
5035 )
5036 .into_response(),
5037 Some(if is_draining {
5038 drain_deadline_s()
5039 } else {
5040 worker::WORKER_RESPAWN_BACKOFF_BASE_S
5041 }),
5042 ),
5043 }
5044}
5045
5046#[derive(Clone, Copy)]
5047struct DualPpMetricsSnapshot {
5048 stage_ns: [u64; 4],
5049 stage_samples: [usize; 4],
5050 dropped_timing_samples: usize,
5051 overlaps: usize,
5052 slot_pairs: usize,
5053 slot_uses: [usize; 2],
5054 slot_collisions: usize,
5055}
5056
5057impl DualPpMetricsSnapshot {
5058 fn current() -> Self {
5059 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
5060 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
5061 Self {
5062 stage_ns,
5063 stage_samples,
5064 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
5065 overlaps: memra_engine::pp::dual_pp_overlaps(),
5066 slot_pairs,
5067 slot_uses,
5068 slot_collisions,
5069 }
5070 }
5071
5072 fn populated(self) -> bool {
5073 self.stage_samples.iter().any(|&n| n > 0)
5074 || self.dropped_timing_samples > 0
5075 || self.slot_pairs > 0
5076 || self.slot_collisions > 0
5077 }
5078}
5079
5080fn insert_dual_pp_metrics(
5081 body: &mut serde_json::Value,
5082 metrics_scope: &MetricsScope,
5083 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
5084) {
5085 if !metrics_scope.operator() {
5088 return;
5089 }
5090 let snapshot = snapshot();
5091 if !snapshot.populated() {
5092 return;
5093 }
5094 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
5095 .iter()
5096 .enumerate()
5097 .map(|(i, name)| {
5098 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
5099 (
5100 name.to_string(),
5101 json!({
5102 "samples": snapshot.stage_samples[i],
5103 "total_ms": total_ms,
5104 "mean_ms": if snapshot.stage_samples[i] > 0 {
5105 total_ms / snapshot.stage_samples[i] as f64
5106 } else { 0.0 },
5107 }),
5108 )
5109 })
5110 .collect();
5111 body["dual_pp"] = json!({
5112 "overlaps": snapshot.overlaps,
5113 "slot_pairs": snapshot.slot_pairs,
5114 "slot_uses": snapshot.slot_uses,
5115 "slot_collisions": snapshot.slot_collisions,
5116 "cuda_event_spans": timings,
5117 "dropped_timing_samples": snapshot.dropped_timing_samples,
5118 });
5119}
5120
5121#[derive(Clone, Copy)]
5122struct PpWaveMetricsSnapshot {
5123 ticks: usize,
5124 cells: usize,
5125 overlaps: usize,
5126}
5127
5128impl PpWaveMetricsSnapshot {
5129 fn current() -> Self {
5130 let (ticks, cells, overlaps) = memra_engine::pp::pp_wave_snapshot();
5131 Self {
5132 ticks,
5133 cells,
5134 overlaps,
5135 }
5136 }
5137}
5138
5139fn insert_pp_wave_metrics(
5140 body: &mut serde_json::Value,
5141 metrics_scope: &MetricsScope,
5142 snapshot: impl FnOnce() -> PpWaveMetricsSnapshot,
5143) {
5144 if !metrics_scope.operator() {
5145 return;
5146 }
5147 let snapshot = snapshot();
5148 if snapshot.ticks == 0 && snapshot.cells == 0 {
5149 return;
5150 }
5151 body["pp_wave"] = json!({
5152 "ticks": snapshot.ticks,
5153 "cells": snapshot.cells,
5154 "overlaps": snapshot.overlaps,
5155 });
5156}
5157
5158fn insert_spec_acceptance_metrics(
5159 body: &mut serde_json::Value,
5160 metrics_scope: &MetricsScope,
5161 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
5162) {
5163 if !metrics_scope.operator() {
5166 return;
5167 }
5168 let snapshot = snapshot();
5169 if snapshot.is_empty() {
5170 return;
5171 }
5172
5173 let mut tau = serde_json::Map::new();
5174 let mut by_position = serde_json::Map::new();
5175 for (model, telemetry) in snapshot {
5176 if telemetry.rounds == 0 {
5177 continue;
5178 }
5179 let n_pos = telemetry
5180 .pos_drafted
5181 .iter()
5182 .rposition(|&n| n > 0)
5183 .map_or(0, |position| position + 1);
5184 tau.insert(model.clone(), json!(telemetry.tau()));
5185 by_position.insert(
5186 model,
5187 json!({
5188 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
5189 "rounds": telemetry.rounds,
5190 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
5191 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
5192 "accept_rate": (0..n_pos).map(|position| {
5193 let offered = telemetry.pos_drafted[position];
5194 if offered > 0 {
5195 telemetry.pos_accepted[position] as f64 / offered as f64
5196 } else {
5197 0.0
5198 }
5199 }).collect::<Vec<f64>>(),
5200 }),
5201 );
5202 }
5203 if !tau.is_empty() {
5204 body["spec_tau"] = serde_json::Value::Object(tau);
5205 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
5206 }
5207}
5208
5209fn insert_peer_probe_metrics(
5210 body: &mut serde_json::Value,
5211 metrics_scope: &MetricsScope,
5212 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
5213) {
5214 if !metrics_scope.operator() {
5217 return;
5218 }
5219 let snapshot = snapshot();
5220 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
5221 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
5222 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
5223 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
5224 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
5225 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
5226 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
5227}
5228
5229async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5231 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5232 Ok(scope) => scope,
5233 Err(response) => return response,
5234 };
5235 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5236 let mut body = if metrics_scope.process_wide() {
5240 json!({
5241 "admitted": m.admitted,
5242 "completed": m.completed,
5243 "tokens_out": m.tokens_out,
5244 "step_p50_ms": m.step_p50_ms,
5245 "step_p99_ms": m.step_p99_ms,
5246 "prompt_tokens_in": m.prompt_tokens_in,
5248 "cached_tokens_in": m.cached_tokens_in,
5249 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5252 "admission_session_defers": m.admission_session_defers,
5255 "admission_vram_defers": m.admission_vram_defers,
5256 "step_oom_parks": m.step_oom_parks,
5257 "continuation_pool_hits": m.continuation_pool_hits,
5258 "continuation_pool_evictions": m.continuation_pool_evictions,
5259 "plain_affinity_rewinds": m.plain_affinity_rewinds,
5260 "served_dspark": m.served_dspark,
5261 "served_spec": m.served_spec,
5262 "served_plain": m.served_plain,
5263 "spec_pool_hits": m.spec_pool_hits,
5264 "spec_pool_misses": m.spec_pool_misses,
5265 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
5266 "spec_pool_evictions": m.spec_pool_evictions,
5267 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
5270 })
5271 } else {
5272 json!({})
5273 };
5274 if metrics_scope.operator() {
5278 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
5279 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
5280 body["budget_source_reload_consecutive"] =
5281 json!(budget_health.source_reload_consecutive);
5282 body["budget_source_available"] = json!(budget_health.source_available);
5283 }
5284 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
5286 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
5287 } else {
5288 0.0
5289 });
5290 body["prefix_cache_hits"] = json!(m.prefix_hits);
5291 body["prefix_cache_misses"] = json!(m.prefix_misses);
5292 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
5293 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
5294 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
5295 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
5296 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
5297 body["prefix_host_entries"] = json!(m.prefix_host_entries);
5301 body["prefix_host_bytes"] = json!(m.prefix_host_bytes);
5302 body["prefix_host_demotions"] = json!(m.prefix_host_demotions);
5303 body["prefix_host_promotions"] = json!(m.prefix_host_promotions);
5304 body["prefix_host_demote_ms"] = json!(m.prefix_host_demote_ms);
5305 body["prefix_host_promote_ms"] = json!(m.prefix_host_promote_ms);
5306 body["prefix_host_rejected_allocs"] = json!(m.prefix_host_rejected_allocs);
5307 body["prefix_host_purges"] = json!(m.prefix_host_purges);
5308 body["prefix_host_purged_entries"] = json!(m.prefix_host_purged_entries);
5309 body["prefix_host_purged_bytes"] = json!(m.prefix_host_purged_bytes);
5310 body["prefix_host_tenant_rejects"] = json!(m.prefix_host_tenant_rejects);
5311 body["prefix_host_pause_demotes"] = json!(m.prefix_host_pause_demotes);
5315 body["prefix_host_pause_cancels"] = json!(m.prefix_host_pause_cancels);
5316 body["kv_flex_borrowed_bytes"] = json!(m.kv_flex_borrowed_bytes);
5321 body["kv_flex_sheds"] = json!(m.kv_flex_sheds);
5322 body["kv_flex_shed_ms"] = json!(m.kv_flex_shed_ms);
5323 body["lcp_histogram"] = json!({
5326 "edges": worker::LCP_HIST_EDGES.to_vec(),
5327 "counts": m.lcp_hist.to_vec(),
5328 });
5329 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
5333 body["prefix_cache_entries"] = json!(m.prefix_entries);
5334 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
5335 body["active_sessions"] = json!(m.active_sessions);
5336 body["queued_requests"] = json!(m.queued_requests);
5337 body["admission_inflight"] = json!(m.admission_inflight);
5341 body["admission_booked_bytes"] = json!(m.admission_booked_bytes);
5342 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
5343 body["spec_pool_entries"] = json!(m.spec_pool_entries);
5344 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
5345 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
5346 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
5347 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
5348 if !m.constraint_compiler_fail_closed.is_empty() {
5349 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
5350 m.constraint_compiler_fail_closed
5351 .iter()
5352 .map(|(model, gauge)| {
5353 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
5354 (model.clone(), json!(value))
5355 })
5356 .collect(),
5357 );
5358 }
5359 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
5360 }
5361 if !m.ns_tokens.is_empty() {
5366 let tenants: serde_json::Map<String, serde_json::Value> = m
5367 .ns_tokens
5368 .iter()
5369 .filter(|(ns, _)| metrics_scope.includes(ns))
5370 .map(|(ns, [p, c])| {
5371 (
5372 ns.clone(),
5373 json!({
5374 "prompt_tokens_in": p,
5375 "cached_tokens_in": c,
5376 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
5377 }),
5378 )
5379 })
5380 .collect();
5381 if !tenants.is_empty() {
5382 body["tenants"] = serde_json::Value::Object(tenants);
5383 }
5384 }
5385 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
5386 .adsd_suspect_total
5387 .iter()
5388 .filter(|(tenant, _)| metrics_scope.includes(tenant))
5389 .map(|(tenant, total)| (tenant.clone(), json!(total)))
5390 .collect();
5391 if !adsd_suspect_total.is_empty() {
5392 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
5393 }
5394 if metrics_scope.operator()
5396 && let Some((bg, mode)) = &st.bg
5397 {
5398 body["bg"] = bg.to_json(mode);
5399 }
5400 if metrics_scope.operator() {
5407 let spec: serde_json::Map<String, serde_json::Value> = m
5408 .spec
5409 .iter()
5410 .map(|(model, t)| {
5411 let n_pos = t
5412 .pos_drafted
5413 .iter()
5414 .rposition(|&d| d > 0)
5415 .map_or(0, |p| p + 1);
5416 (
5417 model.clone(),
5418 json!({
5419 "rounds": t.rounds,
5420 "drafted": t.drafted,
5421 "accepted": t.accepted,
5422 "acceptance_rate": if t.drafted > 0 {
5423 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
5424 "tokens_per_round": if t.rounds > 0 {
5425 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
5426 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
5427 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
5428 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
5429 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
5430 .collect::<Vec<f64>>(),
5431 }),
5432 )
5433 })
5434 .collect();
5435 if !spec.is_empty() {
5436 body["spec"] = serde_json::Value::Object(spec);
5437 }
5438 }
5439 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
5440 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
5441 insert_pp_wave_metrics(&mut body, &metrics_scope, PpWaveMetricsSnapshot::current);
5442 insert_peer_probe_metrics(
5443 &mut body,
5444 &metrics_scope,
5445 memra_engine::pp::peer_probe_metrics,
5446 );
5447 Json(body).into_response()
5448}
5449
5450#[derive(Debug, Default, Deserialize)]
5451struct ModelsQuery {
5452 #[serde(default)]
5453 schema: Option<String>,
5454}
5455
5456fn models_openai_body(models: &[String]) -> serde_json::Value {
5457 let data: Vec<_> = models
5458 .iter()
5459 .map(|m| json!({ "id": m, "object": "model" }))
5460 .collect();
5461 json!({ "object": "list", "data": data })
5462}
5463
5464fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
5469 match metadata.and_then(|m| m.surface.as_deref()) {
5470 Some("embedding") => "embedding",
5471 Some("rerank") => "rerank",
5472 _ => "chat",
5473 }
5474}
5475
5476fn openrouter_supported_parameters(
5477 caps: Option<&ModelCaps>,
5478 max_output_length: Option<u64>,
5479 is_chat: bool,
5480) -> serde_json::Value {
5481 let mut parameters = serde_json::Map::new();
5482 if !is_chat {
5490 return serde_json::Value::Object(parameters);
5491 }
5492 for name in [
5493 "temperature",
5494 "top_p",
5495 "min_p",
5496 "frequency_penalty",
5497 "presence_penalty",
5498 "repetition_penalty",
5499 "stop",
5500 ] {
5501 parameters.insert(name.into(), json!({ "type": "unknown" }));
5502 }
5503 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
5504 parameters.insert(
5505 "seed".into(),
5506 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
5507 );
5508 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
5509 if let Some(max) = max_output_length {
5510 max_tokens["max"] = json!(max);
5511 }
5512 parameters.insert("max_tokens".into(), max_tokens);
5513 if is_chat && caps.is_some_and(|c| !c.dsv4 && !(c.qwen_think && !c.think_switch)) {
5521 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
5522 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
5523 }
5524 if is_chat && caps.is_some_and(|c| c.tools_branch) {
5525 parameters.insert("tools".into(), json!({ "type": "boolean" }));
5526 parameters.insert(
5527 "tool_choice".into(),
5528 json!({ "type": "enum", "values": ["auto", "none"] }),
5529 );
5530 }
5531 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5532 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
5533 }
5534 serde_json::Value::Object(parameters)
5535}
5536
5537fn published_context_length(
5553 caps: Option<&ModelCaps>,
5554 metadata: Option<&OpenRouterModelMetadata>,
5555) -> Option<u64> {
5556 let trained = caps
5557 .map(|c| c.context_length as u64)
5558 .filter(|&value| value > 0)?;
5559 let envelope = metadata.and_then(|m| {
5560 let prompt = m.max_prompt_length?;
5561 let output = m.max_output_length?;
5562 prompt.checked_add(output)
5563 });
5564 Some(envelope.map_or(trained, |envelope| trained.min(envelope)))
5565}
5566
5567fn model_entry_openrouter(
5568 name: &str,
5569 caps: Option<&ModelCaps>,
5570 metadata: Option<&OpenRouterModelMetadata>,
5571) -> serde_json::Value {
5572 let empty = OpenRouterModelMetadata::default();
5573 let metadata = metadata.unwrap_or(&empty);
5574 let context_length =
5575 published_context_length(caps, Some(metadata)).filter(|&v| v <= JSON_SAFE_INTEGER_MAX);
5576 let tokenizer = caps
5577 .map(|c| c.tokenizer.as_str())
5578 .filter(|tokenizer| !tokenizer.is_empty());
5579
5580 let mut input = serde_json::Map::new();
5581 input.insert("type".into(), json!("text"));
5582 let mut supported_inputs = serde_json::Map::new();
5583 if let Some(value) = context_length {
5584 supported_inputs.insert(
5585 "max_context_length".into(),
5586 json!({ "value": value, "unit": "token" }),
5587 );
5588 }
5589 if let Some(value) = metadata.max_prompt_length {
5590 supported_inputs.insert(
5591 "max_prompt_length".into(),
5592 json!({ "value": value, "unit": "token" }),
5593 );
5594 }
5595 if !supported_inputs.is_empty() {
5596 input.insert(
5597 "supported_inputs".into(),
5598 serde_json::Value::Object(supported_inputs),
5599 );
5600 }
5601 let mut input_pricing = Vec::new();
5602 for (kind, cost) in [
5603 ("prompt", metadata.pricing.prompt.as_deref()),
5604 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
5605 ("cache_write", metadata.pricing.cache_write.as_deref()),
5606 ] {
5607 if let Some(cost) = cost {
5608 input_pricing.push(json!({
5609 "type": kind,
5610 "unit": "token",
5611 "cost_usd": cost,
5612 }));
5613 }
5614 }
5615 if !input_pricing.is_empty() {
5616 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
5617 }
5618 let mut input_capacity = Vec::new();
5619 for (kind, value) in [
5620 ("prompt", metadata.capacity.prompt_tpm),
5621 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
5622 ] {
5623 if let Some(value) = value {
5624 input_capacity.push(json!({
5625 "type": kind,
5626 "unit": "token",
5627 "per": "minute",
5628 "value": value,
5629 }));
5630 }
5631 }
5632 if !input_capacity.is_empty() {
5633 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
5634 }
5635
5636 let or_surface = declared_surface(Some(metadata));
5637 let or_is_chat = or_surface == "chat";
5638 let mut output = serde_json::Map::new();
5639 output.insert(
5647 "type".into(),
5648 json!(match or_surface {
5649 "embedding" => "embeddings",
5650 "rerank" => "rerank",
5651 _ => "text",
5652 }),
5653 );
5654 output.insert(
5655 "supported_parameters".into(),
5656 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
5657 );
5658 if or_is_chat {
5662 output.insert("streaming".into(), json!(true));
5663 }
5664 if let Some(value) = metadata.max_output_length
5667 && or_is_chat
5668 {
5669 output.insert(
5670 "max_length".into(),
5671 json!({ "value": value, "unit": "token" }),
5672 );
5673 }
5674 let mut output_pricing = Vec::new();
5675 for (kind, cost) in [
5676 ("completion", metadata.pricing.completion.as_deref()),
5677 (
5678 "internal_reasoning",
5679 metadata.pricing.internal_reasoning.as_deref(),
5680 ),
5681 ] {
5682 if let Some(cost) = cost {
5683 output_pricing.push(json!({
5684 "type": kind,
5685 "unit": "token",
5686 "cost_usd": cost,
5687 }));
5688 }
5689 }
5690 if !output_pricing.is_empty() {
5691 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
5692 }
5693 let mut output_capacity = Vec::new();
5694 if let Some(value) = metadata.capacity.completion_tpm {
5695 output_capacity.push(json!({
5696 "type": "completion",
5697 "unit": "token",
5698 "per": "minute",
5699 "value": value,
5700 }));
5701 }
5702 if let Some(value) = metadata.capacity.concurrency {
5703 output_capacity.push(json!({
5704 "type": "concurrency",
5705 "unit": "request",
5706 "value": value,
5707 }));
5708 }
5709 if !output_capacity.is_empty() {
5710 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
5711 }
5712
5713 let mut entry = serde_json::Map::new();
5714 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
5715 entry.insert("id".into(), json!(name));
5716 entry.insert("name".into(), json!(name));
5717 if let Some(value) = metadata.hugging_face_id.as_deref() {
5718 entry.insert("hugging_face_id".into(), json!(value));
5719 }
5720 if let Some(value) = metadata.created {
5721 entry.insert("created".into(), json!(value));
5722 }
5723 if let Some(value) = metadata.quantization.as_deref() {
5724 entry.insert("quantization".into(), json!(value));
5725 }
5726 if let Some(value) = tokenizer {
5727 entry.insert("tokenizer".into(), json!(value));
5728 }
5729 if let Some(value) = metadata.description.as_deref() {
5730 entry.insert("description".into(), json!(value));
5731 }
5732 let mut input_modalities = vec![serde_json::Value::Object(input)];
5733 for m in &metadata.input_modalities {
5734 let mut extra = serde_json::Map::new();
5735 extra.insert("type".into(), json!(m));
5736 if let Some(cost) = metadata.pricing.prompt.as_deref() {
5737 extra.insert(
5739 "pricing".into(),
5740 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
5741 );
5742 }
5743 input_modalities.push(serde_json::Value::Object(extra));
5744 }
5745 entry.insert(
5746 "input_modalities".into(),
5747 serde_json::Value::Array(input_modalities),
5748 );
5749 entry.insert(
5750 "output_modalities".into(),
5751 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
5752 );
5753 if let Some(cost) = metadata.pricing.request.as_deref() {
5754 entry.insert(
5755 "pricing".into(),
5756 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
5757 );
5758 }
5759 if let Some(value) = metadata.capacity.request_rpm {
5760 entry.insert(
5761 "capacity".into(),
5762 json!([{
5763 "type": "request",
5764 "unit": "request",
5765 "per": "minute",
5766 "value": value,
5767 }]),
5768 );
5769 }
5770 if let Some(value) = metadata.is_ready {
5771 entry.insert("is_ready".into(), json!(value));
5772 }
5773 if let Some(value) = metadata.is_free {
5774 entry.insert("is_free".into(), json!(value));
5775 }
5776 if let Some(value) = metadata.discount_to_user {
5777 entry.insert("discount_to_user".into(), json!(value));
5778 }
5779 if let Some(value) = metadata.openrouter_slug.as_deref() {
5780 entry.insert("openrouter".into(), json!({ "slug": value }));
5781 }
5782 if !metadata.datacenters.is_empty() {
5783 entry.insert("datacenters".into(), json!(metadata.datacenters));
5784 }
5785 let mut compliance = serde_json::Map::new();
5786 if let Some(value) = metadata.zdr {
5787 compliance.insert("zdr".into(), json!(value));
5788 }
5789 if let Some(value) = metadata.hipaa {
5790 compliance.insert("hipaa".into(), json!(value));
5791 }
5792 if !compliance.is_empty() {
5793 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
5794 }
5795 serde_json::Value::Object(entry)
5796}
5797
5798fn models_openrouter_body(st: &AppState) -> serde_json::Value {
5799 let data: Vec<_> = st
5800 .models
5801 .iter()
5802 .map(|model| {
5803 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
5804 })
5805 .collect();
5806 json!({ "data": data })
5807}
5808
5809fn model_entry_openmodels(
5810 name: &str,
5811 caps: Option<&ModelCaps>,
5812 metadata: Option<&OpenRouterModelMetadata>,
5813) -> Result<serde_json::Value, String> {
5814 let metadata = metadata.ok_or_else(|| {
5815 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
5816 })?;
5817 let context_length = published_context_length(caps, Some(metadata))
5818 .filter(|&value| value <= JSON_SAFE_INTEGER_MAX)
5819 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
5820 let created = metadata
5821 .created
5822 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
5823 let max_output_length = metadata
5824 .max_output_length
5825 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
5826 let prompt = metadata
5827 .pricing
5828 .prompt
5829 .as_deref()
5830 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
5831 let completion =
5832 metadata.pricing.completion.as_deref().ok_or_else(|| {
5833 format!("OpenModels feed requires pricing.completion for model {name:?}")
5834 })?;
5835 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
5836 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
5837 })?;
5838 let is_ready = metadata
5839 .is_ready
5840 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
5841 let is_free = metadata
5842 .is_free
5843 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
5844 let discount_to_user = metadata
5845 .discount_to_user
5846 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
5847
5848 let mut pricing = serde_json::Map::new();
5849 pricing.insert("prompt".into(), json!(prompt));
5850 pricing.insert("completion".into(), json!(completion));
5851 pricing.insert("input_cache_read".into(), json!(input_cache_read));
5852 if let Some(value) = metadata.pricing.request.as_deref() {
5853 pricing.insert("request".into(), json!(value));
5854 }
5855
5856 let om_surface = declared_surface(Some(metadata));
5857 let om_is_chat = om_surface == "chat";
5858 let mut supported_features = Vec::new();
5859 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
5860 supported_features.push("tool_calling");
5861 }
5862 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5863 supported_features.push("reasoning");
5864 }
5865
5866 let mut entry = serde_json::Map::new();
5867 entry.insert("id".into(), json!(name));
5868 entry.insert("name".into(), json!(name));
5869 entry.insert("created".into(), json!(created));
5870 entry.insert("input_modalities".into(), json!(["text"]));
5871 entry.insert(
5872 "output_modalities".into(),
5873 json!(match om_surface {
5874 "embedding" => ["embeddings"],
5875 "rerank" => ["rerank"],
5876 _ => ["text"],
5877 }),
5878 );
5879 entry.insert("context_length".into(), json!(context_length));
5880 entry.insert("max_output_length".into(), json!(max_output_length));
5881 entry.insert("currency".into(), json!("USD"));
5884 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
5885 entry.insert("supported_features".into(), json!(supported_features));
5886 entry.insert("is_ready".into(), json!(is_ready));
5887 entry.insert("is_free".into(), json!(is_free));
5888 entry.insert("discount_to_user".into(), json!(discount_to_user));
5889 Ok(serde_json::Value::Object(entry))
5890}
5891
5892fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
5893 let data: Result<Vec<_>, _> = st
5894 .models
5895 .iter()
5896 .map(|model| {
5897 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
5898 })
5899 .collect();
5900 Ok(json!({ "data": data? }))
5901}
5902
5903async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
5904 match query.schema.as_deref() {
5905 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
5906 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
5907 Some("openmodels") => match models_openmodels_body(&st) {
5908 Ok(body) => Json(body).into_response(),
5909 Err(error) => bad_request(&error, Some("schema")),
5910 },
5911 Some(schema) => bad_request(
5912 &format!(
5913 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
5914 ),
5915 Some("schema"),
5916 ),
5917 }
5918}
5919
5920fn model_entry_v1(
5928 name: &str,
5929 caps: Option<&ModelCaps>,
5930 metadata: Option<&OpenRouterModelMetadata>,
5931) -> serde_json::Value {
5932 let ctx = published_context_length(caps, metadata);
5933 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
5937 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
5940 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
5941 Some(p) => json!(p),
5942 None => serde_json::Value::Null,
5943 };
5944 let owned_by = metadata
5945 .and_then(|m| m.owned_by.as_deref())
5946 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
5947 let mut input_modalities = vec!["text"];
5948 if let Some(meta) = metadata {
5949 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
5950 }
5951 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
5952 let reliability = metadata.and_then(|m| m.reliability.as_ref());
5953 let surface = declared_surface(metadata);
5959 let (model_type, endpoints, output_modalities) = match surface {
5960 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
5964 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
5965 _ => ("chat", vec!["chat/completions"], vec!["text"]),
5966 };
5967 let is_chat = surface == "chat";
5968 json!({
5969 "id": name,
5970 "name": name,
5971 "object": "model",
5972 "owned_by": owned_by,
5973 "type": model_type,
5974 "context_length": ctx,
5975 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
5978 "endpoints": endpoints,
5979 "input_modalities": input_modalities,
5980 "output_modalities": output_modalities,
5981 "capabilities": {
5982 "streaming": is_chat,
5985 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
5986 "structured_output": is_chat
5992 && !is_dsv4
5993 && !caps.is_some_and(|c| c.qwen_think && !c.think_switch),
5994 "reasoning": is_chat && thinking,
5995 "prompt_caching": is_chat && !is_dsv4,
5996 },
5997 "pricing": {
5998 "currency": "USD",
5999 "unit": "per_1m_tokens",
6000 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
6001 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
6002 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
6003 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
6004 "minimum_request": metadata
6006 .and_then(|m| m.pricing.request.as_deref())
6007 .unwrap_or("0"),
6008 },
6009 "lifecycle": {
6010 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
6011 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
6012 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
6013 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
6014 },
6015 "reliability": {
6016 "first_token_timeout_seconds":
6017 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
6018 "completion_timeout_seconds":
6019 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
6020 "stream_idle_timeout_seconds":
6021 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
6022 "capacity_scope":
6023 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
6024 },
6025 })
6026}
6027
6028async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
6031 let data: Vec<_> = st
6032 .models
6033 .iter()
6034 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
6035 .collect();
6036 let mut body = json!({
6037 "object": "list",
6038 "contract_version": "2.0",
6039 "data": data,
6040 });
6041 if let Some(provider) = st.provider_metadata.as_ref() {
6046 body["provider"] = json!({
6047 "id": provider.id,
6048 "status_url": provider.status_url,
6049 "support_contact": provider.support_contact,
6050 "incident_contact": provider.incident_contact,
6051 "regions": provider.regions,
6052 "request_id_header": "x-request-id",
6053 "error_contract": {
6054 "rate_limit_status": 429,
6055 "overload_status": 503,
6056 "retry_after_header": "Retry-After",
6057 "account_quota_error_codes": ["insufficient_balance"],
6058 },
6059 });
6060 }
6061 Json(body)
6062}
6063
6064async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
6067 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
6068 Ok(scope) => scope,
6069 Err(response) => return response,
6070 };
6071 if !metrics_scope.process_wide() {
6072 return error_response(
6073 StatusCode::FORBIDDEN,
6074 "completion api keys do not authorize process-wide yield metrics; configure \
6075 MEMRA_METRICS_TOKEN",
6076 "authentication_error",
6077 None,
6078 );
6079 }
6080 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
6081 let lane = |i: usize| {
6082 json!({
6083 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
6084 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
6085 })
6086 };
6087 let mut body = json!({
6088 "lanes": {
6089 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
6090 },
6091 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
6092 });
6093 if metrics_scope.operator() {
6094 body["batch_size_last"] = json!(m.batch_size_last);
6095 }
6096 Json(body).into_response()
6097}
6098
6099async fn peek_admission(
6112 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
6113) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, (Response, &'static str)> {
6114 match rx.recv().await {
6115 Some(Event::Error(e)) => {
6120 let error_code = engine_error_code(e.class);
6121 Err((engine_error_response(&e), error_code))
6122 }
6123 first => {
6124 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
6125 if let Some(ev) = first {
6126 let _ = tx2.send(ev);
6127 }
6128 tokio::spawn(forward_events(rx, tx2));
6129 Ok(rx2)
6130 }
6131 }
6132}
6133
6134async fn forward_events(
6145 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
6146 tx2: tokio::sync::mpsc::UnboundedSender<Event>,
6147) {
6148 loop {
6149 tokio::select! {
6150 biased;
6151 () = tx2.closed() => break,
6152 ev = rx.recv() => match ev {
6153 Some(ev) => {
6154 if tx2.send(ev).is_err() {
6155 break;
6156 }
6157 }
6158 None => break,
6159 },
6160 }
6161 }
6162}
6163
6164async fn peek_first_token(
6179 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
6180 deadline: RequestDeadline,
6181) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, ()> {
6182 let mut buffered: Vec<Event> = Vec::new();
6183 loop {
6184 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
6185 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
6188 let first_delivery = matches!(
6189 ev,
6190 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
6191 );
6192 buffered.push(ev);
6193 if first_delivery {
6194 break;
6195 }
6196 }
6197 }
6198 }
6199 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
6200 for ev in buffered {
6201 let _ = tx2.send(ev);
6202 }
6203 tokio::spawn(forward_events(rx, tx2));
6204 Ok(rx2)
6205}
6206
6207#[cfg(test)]
6209fn build_request(
6213 req: &CompletionReq,
6214 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6215 lane: lanes::Lane,
6216 affinity: Option<String>,
6217) -> Request {
6218 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
6219}
6220
6221fn build_request_with_trace(
6222 req: &CompletionReq,
6223 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6224 lane: lanes::Lane,
6225 affinity: Option<String>,
6226 ttft: Option<Arc<ttft::Trace>>,
6227 sampling_defaults: &SamplingDefaults,
6228) -> Request {
6229 let params = GenParams {
6230 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6231 max_ctx: req.max_ctx,
6232 eos: Vec::new(), };
6234 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
6239 Request {
6240 model: req.model.clone(),
6241 prompt_ids: req.prompt_ids.clone(),
6242 prompt_text: req.prompt.clone(),
6243 chat: req.chat,
6244 chat_turns: Vec::new(),
6245 tools_json: Vec::new(),
6246 tools_struct: Vec::new(),
6247 think: ThinkMode::Default,
6248 reasoning_effort: None, params,
6250 sampler_cfg,
6251 stop_strings: req.stop.clone().into_vec(),
6252 trace_id: req.trace_id.clone(),
6253 request_id: String::new(),
6256 admit_predict_logged: false,
6257 max_prompt_tokens: None,
6258 cache_ns: cache_namespace(&req.cache_salt),
6259 affinity,
6260 lane,
6261 grammar: None, prepared_constraint: None,
6263 constraint_ready: None,
6264 oom_retries: 0, spec_k_replay: None,
6266 prepared_prompt: None,
6267 capture: None, images: Vec::new(), gemma_images: Vec::new(),
6270 glm5_images: Vec::new(),
6271 step_images: Vec::new(),
6272 vision_memory: None,
6273 ttft,
6274 tx,
6275 }
6276}
6277
6278struct ChatPlan {
6281 request: Request,
6282 parser: Option<ToolStreamParser>,
6285 pending_images: Vec<PendingVisionUnit>,
6288 pending_gemma: Vec<PendingGemmaImage>,
6289 pending_glm5: Vec<PendingGlm5Image>,
6290 pending_step: Vec<PendingStepImage>,
6291 vision_memory: Option<VisionMemoryPermit>,
6295}
6296
6297pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
6298 req.messages.iter().any(|message| {
6299 message.content.as_array().is_some_and(|parts| {
6300 parts.iter().any(|part| {
6301 matches!(
6302 part.get("type").and_then(serde_json::Value::as_str),
6303 Some("image_url" | "video_url")
6304 )
6305 })
6306 })
6307 })
6308}
6309
6310fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
6311 let mut total = 0usize;
6312 let mut add = |bytes: usize| {
6313 total = total.checked_add(bytes).ok_or_else(|| {
6314 "vision patch memory reservation overflowed while planning".to_string()
6315 })?;
6316 Ok::<(), String>(())
6317 };
6318 for unit in &plan.pending_images {
6319 let bytes = match unit {
6320 PendingVisionUnit::Still { gh, gw, .. } => gh
6321 .checked_mul(*gw)
6322 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6323 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6324 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
6325 PendingVisionUnit::Video { groups, .. } => {
6326 groups.iter().try_fold(0usize, |total, group| {
6327 let bytes = group
6328 .gh
6329 .checked_mul(group.gw)
6330 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6331 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6332 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6333 total.checked_add(bytes).ok_or_else(|| {
6334 "vision patch memory reservation overflowed while planning".to_string()
6335 })
6336 })?
6337 }
6338 };
6339 add(bytes)?;
6340 }
6341 for unit in &plan.pending_gemma {
6342 let bytes = unit
6343 .gw
6344 .checked_mul(unit.gh)
6345 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
6346 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6347 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6348 add(bytes)?;
6349 }
6350 for unit in &plan.pending_glm5 {
6351 let bytes = unit
6352 .gh
6353 .checked_mul(unit.gw)
6354 .and_then(|n| n.checked_mul(memra_engine::vision_glm5::G5V_PATCH_IN))
6355 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6356 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6357 add(bytes)?;
6358 }
6359 for unit in &plan.pending_step {
6360 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
6361 let patches = unit
6363 .plan
6364 .n_tiles
6365 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
6366 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
6367 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6368 let bytes = patches
6369 .checked_mul(SV_PATCH_IN)
6370 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6371 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6372 add(bytes)?;
6373 }
6374 Ok(total)
6375}
6376
6377pub(crate) fn reserve_vision_memory(
6378 plan: &ChatPlan,
6379) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
6380 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
6381 try_reserve_vision_memory(bytes)
6382}
6383
6384#[cfg(test)]
6385fn build_chat_request(
6386 req: ChatCompletionReq,
6387 caps: Option<&ModelCaps>,
6388 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6389 lane: lanes::Lane,
6390 affinity: Option<String>,
6391) -> Result<ChatPlan, String> {
6392 let defaults = ModelSamplingDefaults::resolve(None, caps);
6395 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
6396}
6397
6398#[allow(clippy::too_many_arguments)]
6412fn build_chat_request_with_trace(
6413 req: ChatCompletionReq,
6414 caps: Option<&ModelCaps>,
6415 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6416 lane: lanes::Lane,
6417 affinity: Option<String>,
6418 ttft: Option<Arc<ttft::Trace>>,
6419 default_effort: Option<&str>,
6420 sampling_defaults: &ModelSamplingDefaults,
6421) -> Result<ChatPlan, String> {
6422 let client_sampling: ClientSampling = (&req).into();
6425 let tool_choice = parse_tool_choice(&req.tool_choice)?;
6426 if let Some(c) = caps
6431 && !c.chat_ok
6432 {
6433 return Err(format!(
6434 "model {:?} has no chat template (checkpoint carries neither \
6435 tokenizer_config.json chat_template nor chat_template.jinja) — \
6436 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
6437 req.model
6438 ));
6439 }
6440 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
6441 let (mut think, effort_level, think_client_explicit) = parse_think(
6442 &req.reasoning_effort,
6443 &req.reasoning,
6444 vllm_switch,
6445 req.include_reasoning,
6446 default_effort,
6447 caps.is_some_and(|c| c.dsv4 || c.glm5),
6451 )?;
6452 let level_template = caps
6456 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort || c.glm5)
6457 .unwrap_or(false);
6458 if think_client_explicit
6479 && think == ThinkMode::NoThink
6480 && let Some(c) = caps
6481 && c.qwen_think
6482 && !c.think_switch
6483 && !c.dsv4
6484 {
6485 return Err(format!(
6486 "model {:?} cannot disable reasoning: its chat template opens a think \
6487 tail unconditionally and carries no enable_thinking switch, so \
6488 reasoning_effort/enable_thinking cannot turn it off on this model",
6489 req.model
6490 ));
6491 }
6492 let reasoning_effort = if level_template { effort_level } else { None };
6525 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
6528 if grammar.is_some()
6542 && let Some(c) = caps
6543 && c.qwen_think
6544 && think != ThinkMode::NoThink
6545 {
6546 if c.think_switch {
6547 think = ThinkMode::NoThink;
6548 } else if c.think_close.is_empty() {
6549 return Err(
6550 "response_format requires the model's think channel to close \
6551 before the grammar can engage, but this chat template has \
6552 neither an enable_thinking switch nor a recognizable \
6553 think-close token sequence"
6554 .into(),
6555 );
6556 }
6557 }
6561
6562 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
6569
6570 let (tools_json, tools_struct, schemas) =
6573 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
6574 prepare_tools(&req.tools)?
6575 } else {
6576 (Vec::new(), Vec::new(), HashMap::new())
6577 };
6578
6579 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
6580 let mut images: Vec<PendingVisionUnit> = Vec::new();
6581 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
6582 let mut glm5_images: Vec<PendingGlm5Image> = Vec::new();
6583 let mut step_images: Vec<PendingStepImage> = Vec::new();
6584 let mut next_video = 0usize;
6585 for msg in &req.messages {
6586 let content = content_to_text_vision(
6587 &msg.content,
6588 &mut images,
6589 &mut gemma_images,
6590 &mut glm5_images,
6591 &mut step_images,
6592 &mut next_video,
6593 )
6594 .map_err(|e| format!("{} message: {e}", msg.role))?;
6595 let tool_calls = msg
6596 .tool_calls
6597 .iter()
6598 .map(render_req_tool_call)
6599 .collect::<Result<Vec<_>, _>>()?;
6600 if !tool_calls.is_empty() && msg.role != "assistant" {
6601 return Err("tool_calls are only valid on assistant messages".into());
6602 }
6603 let role = if msg.role == "developer" {
6606 "system".to_string()
6607 } else {
6608 msg.role.clone()
6609 };
6610 turns.push(TmplTurn {
6611 role,
6612 content,
6613 tool_calls,
6614 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
6616 tool_call_id: msg.tool_call_id.clone(),
6617 tool_name: msg.name.clone(),
6618 tool_responses: Vec::new(),
6619 task: None,
6623 tools: Vec::new(),
6624 });
6625 }
6626
6627 let has_tool_features = !tools_json.is_empty()
6630 || turns
6631 .iter()
6632 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
6633 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
6634 return Err(format!(
6635 "model {:?} chat template has no tools branch",
6636 req.model
6637 ));
6638 }
6639
6640 let think_open = caps
6643 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
6644 .unwrap_or(false);
6645 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
6663 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
6669 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
6670 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
6671 let glm5 = caps.map(|c| c.glm5).unwrap_or(false);
6677 let is_hy3 = caps.map(|c| c.hy3).unwrap_or(false);
6680 let hy3_think_open = is_hy3 && think == ThinkMode::Think;
6681 let hy3_tools = is_hy3 && !tools_json.is_empty();
6682 let parser = if glm5 {
6683 Some(ToolStreamParser::glm5(think_open, schemas))
6684 } else if is_hy3 && (hy3_tools || hy3_think_open) {
6685 Some(ToolStreamParser::hy3(schemas, hy3_think_open))
6686 } else if is_dsv4 && (dsv4_tools || dsv4_think_open) {
6687 Some(ToolStreamParser::dsv4(dsv4_think_open))
6688 } else if gemma_tools {
6689 Some(ToolStreamParser::gemma_tools())
6690 } else if !tools_json.is_empty() {
6691 Some(ToolStreamParser::new(schemas, think_open))
6692 } else if think_open {
6693 Some(ToolStreamParser::reasoning_only())
6694 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
6695 Some(ToolStreamParser::gemma_thought())
6703 } else {
6704 None
6705 };
6706
6707 Ok(ChatPlan {
6708 request: Request {
6709 model: req.model,
6710 prompt_ids: Vec::new(),
6711 prompt_text: String::new(),
6712 chat: false,
6713 chat_turns: turns,
6714 tools_json,
6715 tools_struct,
6716 think,
6717 reasoning_effort,
6718 params: GenParams {
6719 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6720 max_ctx: req.max_ctx,
6721 eos: Vec::new(),
6722 },
6723 sampler_cfg,
6724 stop_strings: {
6725 let mut stops = req.stop.into_vec();
6730 if gemma_tools {
6731 stops.push("<tool_call|>".to_string());
6732 }
6733 if dsv4_tools {
6738 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
6739 }
6740 if hy3_tools {
6743 stops.push("</tool_calls:opensource>".to_string());
6744 }
6745 stops
6746 },
6747 trace_id: None,
6748 request_id: String::new(),
6751 admit_predict_logged: false,
6752 max_prompt_tokens: None,
6753 cache_ns: cache_namespace(&req.cache_salt),
6754 affinity,
6755 lane,
6756 grammar,
6757 prepared_constraint: None,
6758 constraint_ready: None,
6759 oom_retries: 0, spec_k_replay: None,
6761 prepared_prompt: None,
6762 images: Vec::new(),
6767 gemma_images: Vec::new(),
6768 glm5_images: Vec::new(),
6769 step_images: Vec::new(),
6770 capture: None, vision_memory: None,
6772 ttft,
6773 tx,
6774 },
6775 parser,
6776 pending_images: images,
6777 pending_gemma: gemma_images,
6778 pending_glm5: glm5_images,
6779 pending_step: step_images,
6780 vision_memory: None,
6781 })
6782}
6783
6784fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
6790 for (i, unit) in plan.pending_images.drain(..).enumerate() {
6791 match unit {
6792 PendingVisionUnit::Still { bytes, gh, gw } => {
6793 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
6794 .map_err(|e| format!("image {}: {e}", i + 1))?;
6795 if (prep.gh, prep.gw) != (gh, gw) {
6796 return Err(format!(
6797 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
6798 i + 1,
6799 prep.gh,
6800 prep.gw
6801 ));
6802 }
6803 plan.request
6804 .images
6805 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
6806 }
6807 PendingVisionUnit::Video {
6808 bytes,
6809 groups,
6810 video,
6811 } => {
6812 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
6813 .map_err(|e| format!("video {}: {e}", i + 1))?;
6814 if prepared.groups.len() != groups.len() {
6815 return Err(format!(
6816 "video {}: decoded {} groups differ from its header-planned {} groups",
6817 i + 1,
6818 prepared.groups.len(),
6819 groups.len()
6820 ));
6821 }
6822 for ((group, prep), timestamp) in
6823 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
6824 {
6825 if (prep.gh, prep.gw) != (group.gh, group.gw) {
6826 return Err(format!(
6827 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
6828 i + 1,
6829 prep.gh,
6830 prep.gw,
6831 group.gh,
6832 group.gw
6833 ));
6834 }
6835 if (timestamp - group.timestamp).abs() > 0.001 {
6836 return Err(format!(
6837 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
6838 i + 1,
6839 group.timestamp
6840 ));
6841 }
6842 plan.request
6843 .images
6844 .push(memra_engine::vision_pre::VisionUnit {
6845 prep,
6846 video: Some(video),
6847 });
6848 }
6849 }
6850 }
6851 }
6852 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
6853 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
6854 .map_err(|e| format!("image {}: {e}", i + 1))?;
6855 if (gw, gh) != (unit.gw, unit.gh) {
6856 return Err(format!(
6857 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
6858 i + 1,
6859 unit.gw,
6860 unit.gh
6861 ));
6862 }
6863 plan.request
6864 .gemma_images
6865 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
6866 }
6867 for (i, unit) in plan.pending_glm5.drain(..).enumerate() {
6868 let (patches, gh, gw) = memra_engine::vision_glm5::glm5_prep_image(&unit.bytes)
6869 .map_err(|e| format!("image {}: {e}", i + 1))?;
6870 if (gh, gw) != (unit.gh, unit.gw) {
6871 return Err(format!(
6872 "image {}: decoded grid {gh}x{gw} differs from its header-planned grid {}x{} — refusing (placeholder runs already rendered)",
6873 i + 1,
6874 unit.gh,
6875 unit.gw
6876 ));
6877 }
6878 plan.request
6879 .glm5_images
6880 .push(memra_engine::vision_glm5::Glm5VisionUnit { patches, gh, gw });
6881 }
6882 for (i, unit) in plan.pending_step.drain(..).enumerate() {
6883 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
6884 .map_err(|e| format!("image {}: {e}", i + 1))?;
6885 if prepped.tiles.len() != unit.plan.n_tiles
6886 || prepped.newline_mask != unit.plan.newline_mask
6887 {
6888 return Err(format!(
6889 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
6890 ({} tiles) — refusing (pad runs already rendered)",
6891 i + 1,
6892 prepped.tiles.len(),
6893 unit.plan.n_tiles
6894 ));
6895 }
6896 plan.request.step_images.push(prepped);
6897 }
6898 Ok(())
6899}
6900
6901fn bearer_token(headers: &HeaderMap) -> Option<&str> {
6909 headers
6910 .get("authorization")
6911 .and_then(|value| value.to_str().ok())
6912 .and_then(|value| value.strip_prefix("Bearer "))
6913}
6914
6915fn authentication_error(why: auth::AuthDenied) -> Response {
6916 match why {
6917 auth::AuthDenied::Unknown => error_response(
6918 StatusCode::UNAUTHORIZED,
6919 "invalid api key",
6920 "authentication_error",
6921 None,
6922 ),
6923 auth::AuthDenied::Disabled => error_response(
6924 StatusCode::FORBIDDEN,
6925 "api key is disabled",
6926 "authentication_error",
6927 None,
6928 ),
6929 }
6930}
6931
6932#[allow(clippy::result_large_err)] fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
6934 auth::authenticate_with(
6935 api_auth.keyring,
6936 api_auth.single_key.as_deref(),
6937 bearer_token(headers),
6938 )
6939 .map_err(authentication_error)
6940}
6941
6942#[derive(Debug, Clone, PartialEq, Eq)]
6943enum MetricsScope {
6944 All,
6945 CompletionDomain,
6946 Tenant(String),
6947}
6948
6949impl MetricsScope {
6950 fn operator(&self) -> bool {
6951 matches!(self, MetricsScope::All)
6952 }
6953
6954 fn process_wide(&self) -> bool {
6955 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
6956 }
6957
6958 fn includes(&self, tenant_row: &str) -> bool {
6959 match self {
6960 MetricsScope::All | MetricsScope::CompletionDomain => true,
6961 MetricsScope::Tenant(tenant) => tenant == tenant_row,
6962 }
6963 }
6964}
6965
6966#[allow(clippy::result_large_err)] fn authorize_metrics(
6968 api_auth: &ApiAuth,
6969 metrics_auth: &MetricsAuth,
6970 headers: &HeaderMap,
6971) -> Result<MetricsScope, Response> {
6972 if !metrics_auth.required {
6973 return Ok(MetricsScope::All);
6974 }
6975 let Some(candidate) = bearer_token(headers) else {
6976 return Err(authentication_error(auth::AuthDenied::Unknown));
6977 };
6978 if let Some(token) = metrics_auth.token.as_deref() {
6979 if auth::constant_time_secret_eq(token, candidate) {
6980 return Ok(MetricsScope::All);
6981 }
6982 if api_auth.configured() {
6983 return match auth::authenticate_with(
6984 api_auth.keyring,
6985 api_auth.single_key.as_deref(),
6986 Some(candidate),
6987 ) {
6988 Ok(_) => Err(error_response(
6989 StatusCode::FORBIDDEN,
6990 "completion api keys do not authorize metrics while \
6991 MEMRA_METRICS_TOKEN is configured",
6992 "authentication_error",
6993 None,
6994 )),
6995 Err(why) => Err(authentication_error(why)),
6996 };
6997 }
6998 return Err(authentication_error(auth::AuthDenied::Unknown));
6999 }
7000 if api_auth.configured() {
7001 let tenant = authenticate(api_auth, headers)?;
7002 return Ok(if api_auth.keyring.is_some() {
7003 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
7004 } else {
7005 MetricsScope::CompletionDomain
7009 });
7010 }
7011 Err(authentication_error(auth::AuthDenied::Unknown))
7012}
7013
7014#[allow(clippy::result_large_err)] fn lane_for_tenant(
7021 headers: &axum::http::HeaderMap,
7022 tenant: &auth::TenantCtx,
7023) -> Result<lanes::Lane, Response> {
7024 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
7025 None => None,
7026 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
7031 error_response_coded(
7032 StatusCode::BAD_REQUEST,
7033 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
7034 "invalid_request_error",
7035 Some("x-lane"),
7036 Some("invalid_lane"),
7037 )
7038 })?),
7039 };
7040 match tenant.lane_class {
7041 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
7042 auth::LaneClass::Batch => match requested {
7043 None => Ok(lanes::Lane::Harvest),
7044 Some(lanes::Lane::Interactive) => Err(error_response(
7045 StatusCode::FORBIDDEN,
7046 "this api key is batch-class: x-lane interactive is not permitted \
7047 (use judge or harvest)",
7048 "authentication_error",
7049 Some("x-lane"),
7050 )),
7051 Some(l) => Ok(l),
7052 },
7053 }
7054}
7055
7056fn tenant_namespace(
7060 tenant: &auth::TenantCtx,
7061 cache_salt: &Option<String>,
7062) -> Result<String, &'static str> {
7063 let keyring_configured = auth::global().is_some();
7064 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
7065 if keyring_configured {
7066 Ok(auth::scope_namespace(&tenant.tenant, &raw))
7067 } else {
7068 Ok(raw)
7069 }
7070}
7071
7072fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
7077 eprintln!(
7078 "[meter] admit id={} tenant={} lane={} model={:?}",
7079 env.id,
7080 tenant.tenant,
7081 lane.as_str(),
7082 model
7083 );
7084}
7085
7086fn apply_model_request_limits(
7087 request: &mut Request,
7088 metadata: Option<&OpenRouterModelMetadata>,
7089 caps: Option<&ModelCaps>,
7090) -> Result<(), (String, &'static str)> {
7091 let Some(metadata) = metadata else {
7092 return Ok(());
7093 };
7094 let max_prompt = metadata
7095 .max_prompt_length
7096 .map(usize::try_from)
7097 .transpose()
7098 .map_err(|_| {
7099 (
7100 "configured model prompt limit does not fit this platform".into(),
7101 "model",
7102 )
7103 })?;
7104 let max_output = metadata
7105 .max_output_length
7106 .map(usize::try_from)
7107 .transpose()
7108 .map_err(|_| {
7109 (
7110 "configured model output limit does not fit this platform".into(),
7111 "model",
7112 )
7113 })?;
7114
7115 request.max_prompt_tokens = max_prompt;
7116 if let Some(max_output) = max_output {
7117 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
7118 request.params.max_new = metadata
7119 .default_output_length
7120 .map(usize::try_from)
7121 .transpose()
7122 .map_err(|_| {
7123 (
7124 "configured default output length does not fit this platform".into(),
7125 "model",
7126 )
7127 })?
7128 .unwrap_or(max_output);
7129 } else if request.params.max_new > max_output {
7130 return Err((
7131 format!(
7132 "max_tokens {} exceeds configured model maximum {max_output}",
7133 request.params.max_new
7134 ),
7135 "max_tokens",
7136 ));
7137 }
7138 }
7139
7140 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
7144 (max_prompt, max_output, request.params.max_ctx)
7145 {
7146 let operational_ctx = max_prompt
7147 .checked_add(max_output)
7148 .and_then(|value| value.checked_add(8))
7149 .ok_or_else(|| {
7150 (
7151 "configured model context envelope overflowed".into(),
7152 "model",
7153 )
7154 })?;
7155 let operational_ctx = caps
7156 .map(|caps| caps.context_length)
7157 .filter(|&context| context > 0)
7158 .map_or(operational_ctx, |context| operational_ctx.min(context));
7159 if requested_ctx > operational_ctx {
7160 return Err((
7161 format!(
7162 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
7163 ),
7164 "max_ctx",
7165 ));
7166 }
7167 }
7168 Ok(())
7169}
7170
7171fn effective_max_tokens(request: &worker::Request) -> Option<u64> {
7175 (request.params.max_new != worker::MAX_NEW_CTX_BOUNDED).then_some(request.params.max_new as u64)
7176}
7177
7178#[allow(clippy::too_many_arguments)]
7179fn start_request_receipt(
7180 st: &AppState,
7181 env: &Envelope,
7182 tenant: &auth::TenantCtx,
7183 model: &str,
7184 route: &'static str,
7185 lane: lanes::Lane,
7186 stream: bool,
7187 max_tokens: Option<u64>,
7188 reserved_ctx: Option<u64>,
7189 budget_permit: Option<metering::Permit>,
7190) -> Option<Box<dyn metering::Receipt>> {
7191 st.metering.as_ref().map(|accounting| {
7192 accounting.open(
7193 &metering::RequestMeta {
7194 request_id: &env.id,
7195 tenant: &tenant.tenant,
7196 principal: tenant.key_prefix.as_deref(),
7197 model,
7198 route,
7199 lane: lane.as_str(),
7200 stream,
7201 max_tokens,
7202 reserved_ctx,
7203 },
7204 budget_permit,
7205 )
7206 })
7207}
7208
7209fn arm_capture(
7215 mut receipt: Option<Box<dyn metering::Receipt>>,
7216 prompt: impl FnOnce() -> serde_json::Value,
7217) -> Option<Box<dyn metering::Receipt>> {
7218 if let Some(receipt) = receipt.as_mut()
7219 && receipt.wants_capture()
7220 {
7221 receipt.arm_capture(prompt());
7222 }
7223 receipt
7224}
7225
7226fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
7230 serde_json::Value::Array(
7231 messages
7232 .iter()
7233 .map(|message| {
7234 let mut row = json!({ "role": message.role, "content": message.content });
7235 if !message.tool_calls.is_empty() {
7236 row["tool_calls"] = serde_json::Value::Array(
7237 message
7238 .tool_calls
7239 .iter()
7240 .map(|call| {
7241 json!({
7242 "id": call.id,
7243 "function": {
7244 "name": call.function.name,
7245 "arguments": call.function.arguments,
7246 },
7247 })
7248 })
7249 .collect(),
7250 );
7251 }
7252 row
7253 })
7254 .collect(),
7255 )
7256}
7257
7258enum BudgetRejection {
7259 Invalid(String),
7260 Insufficient,
7261 Unenrolled,
7262 PrincipalCapped,
7265 Unavailable(String),
7266}
7267
7268impl BudgetRejection {
7269 fn into_response(self) -> (Response, &'static str) {
7270 match self {
7271 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
7272 Self::Insufficient => (
7273 error_response_coded(
7274 StatusCode::PAYMENT_REQUIRED,
7275 "tenant prepaid balance is insufficient for this request",
7276 "insufficient_balance",
7277 None,
7278 Some("insufficient_balance"),
7279 ),
7280 "insufficient_balance",
7281 ),
7282 Self::Unenrolled => (
7283 error_response_coded(
7284 StatusCode::PAYMENT_REQUIRED,
7285 "tenant is not enrolled for prepaid billing",
7286 "tenant_not_enrolled",
7287 None,
7288 Some("tenant_not_enrolled"),
7289 ),
7290 "tenant_not_enrolled",
7291 ),
7292 Self::PrincipalCapped => (
7293 error_response_coded(
7294 StatusCode::PAYMENT_REQUIRED,
7295 "this API key's spend cap is reached; raise or clear the key's cap to continue",
7296 "key_spend_cap_reached",
7297 None,
7298 Some("key_spend_cap_reached"),
7299 ),
7300 "key_spend_cap_reached",
7301 ),
7302 Self::Unavailable(err) => {
7303 eprintln!("[budget] ERROR: admission unavailable: {err}");
7304 (
7305 error_response_coded(
7306 StatusCode::SERVICE_UNAVAILABLE,
7307 "tenant budget accounting is unavailable",
7308 "server_error",
7309 None,
7310 Some("tenant_budget_unavailable"),
7311 ),
7312 "tenant_budget_unavailable",
7313 )
7314 }
7315 }
7316 }
7317}
7318
7319fn prepare_budget_prompt(
7320 request: &mut Request,
7321 tokenizer: Option<&Tokenizer>,
7322) -> Result<usize, String> {
7323 if let Some(error) = worker::prompt_source_limit_error(request) {
7324 return Err(error);
7325 }
7326 if request.prepared_prompt.is_none() {
7327 if let Some(trace) = request.ttft.as_ref() {
7328 trace.mark_tokenize_start();
7329 }
7330 let prompt = if !request.prompt_ids.is_empty() {
7331 request.prompt_ids.clone()
7332 } else if !request.chat_turns.is_empty() {
7333 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7334 let plain = worker::plain_chat_render_path(
7341 &request.tools_json,
7342 &request.think,
7343 request.reasoning_effort.as_deref(),
7344 &request.chat_turns,
7345 tokenizer.has_qwen_effort_ladder(),
7346 );
7347 let rendered = if plain {
7348 let messages: Vec<_> = request
7349 .chat_turns
7350 .iter()
7351 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
7352 .collect();
7353 tokenizer.apply_chat_template(&messages, true)
7354 } else {
7355 tokenizer
7356 .apply_chat_template_tools_ex(
7357 &request.chat_turns,
7358 true,
7359 &request.tools_json,
7360 &request.tools_struct,
7361 request.think,
7362 request.reasoning_effort.as_deref(),
7363 )
7364 .map_err(|err| format!("chat template: {err}"))?
7365 };
7366 tokenizer.encode(&rendered, true)
7367 } else if request.chat {
7368 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7369 let rendered =
7370 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
7371 tokenizer.encode(&rendered, true)
7372 } else {
7373 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7374 tokenizer.encode(&request.prompt_text, true)
7375 };
7376 if prompt.is_empty() {
7377 return Err("empty prompt after tokenization".into());
7378 }
7379 if let Some(trace) = request.ttft.as_ref() {
7380 trace.mark_tokenize_end(prompt.len());
7381 }
7382 request.prepared_prompt = Some(prompt);
7383 }
7384 let prompt_tokens = request
7385 .prepared_prompt
7386 .as_ref()
7387 .expect("budget prompt was prepared")
7388 .len();
7389 if let Some(limit) = request.max_prompt_tokens
7390 && prompt_tokens > limit
7391 {
7392 return Err(format!(
7393 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
7394 ));
7395 }
7396 Ok(prompt_tokens)
7397}
7398
7399fn budget_completion_bound(
7400 request: &Request,
7401 prompt_tokens: usize,
7402 caps: Option<&ModelCaps>,
7403) -> Result<usize, String> {
7404 let max_new = request.params.max_new;
7405 let requested_ctx = match (request.params.max_ctx, max_new) {
7406 (Some(cap), _) => cap,
7407 (None, worker::MAX_NEW_CTX_BOUNDED) => {
7408 let server_ctx = std::env::var("MEMRA_CTX")
7409 .ok()
7410 .and_then(|value| value.parse().ok())
7411 .unwrap_or(8192usize);
7412 let mut cap = server_ctx;
7413 if prompt_tokens.saturating_add(16) > cap {
7414 cap = prompt_tokens.saturating_add(server_ctx);
7415 }
7416 cap
7417 }
7418 (None, max_new) => prompt_tokens
7419 .checked_add(max_new)
7420 .and_then(|value| value.checked_add(8))
7421 .ok_or_else(|| "request context bound overflowed".to_string())?,
7422 };
7423 let ctx_cap = caps
7424 .map(|caps| caps.context_length)
7425 .filter(|&context| context > 0)
7426 .map_or(requested_ctx, |context| requested_ctx.min(context));
7427 if prompt_tokens >= ctx_cap {
7428 return Err(format!(
7429 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
7430 ));
7431 }
7432 Ok(max_new.min(ctx_cap - prompt_tokens))
7433}
7434
7435struct BudgetAdmission {
7440 permit: Option<metering::Permit>,
7441 reserved_ctx: Option<u64>,
7442}
7443
7444impl std::fmt::Debug for BudgetAdmission {
7446 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
7447 f.debug_struct("BudgetAdmission")
7448 .field("permit", &self.permit.is_some())
7449 .field("reserved_ctx", &self.reserved_ctx)
7450 .finish()
7451 }
7452}
7453
7454fn admit_tenant_budget(
7455 st: &AppState,
7456 tenant: &auth::TenantCtx,
7457 request: &mut Request,
7458) -> Result<BudgetAdmission, BudgetRejection> {
7459 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
7460 return Ok(BudgetAdmission {
7461 permit: None,
7462 reserved_ctx: None,
7463 });
7464 };
7465 match accounting.is_limited(&tenant.tenant) {
7466 Ok(false) => return Err(BudgetRejection::Unenrolled),
7467 Ok(true) => {}
7468 Err(metering::AdmitError::Unavailable(err)) => {
7469 return Err(BudgetRejection::Unavailable(err));
7470 }
7471 Err(other) => {
7472 return Err(BudgetRejection::Unavailable(format!(
7473 "unexpected budget enrollment result: {other:?}"
7474 )));
7475 }
7476 }
7477 let tokenizer = st
7478 .budget_tokenizers
7479 .as_ref()
7480 .and_then(|tokenizers| tokenizers.get(&request.model))
7481 .map(Arc::as_ref);
7482 if request.prompt_ids.is_empty() && tokenizer.is_none() {
7483 return Err(BudgetRejection::Unavailable(format!(
7484 "no reservation tokenizer for model {:?}",
7485 request.model
7486 )));
7487 }
7488 let prompt_tokens =
7489 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
7490 let completion_tokens =
7491 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
7492 .map_err(BudgetRejection::Invalid)?;
7493 let prompt_tokens = u64::try_from(prompt_tokens)
7494 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
7495 let completion_tokens = u64::try_from(completion_tokens)
7496 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
7497 match accounting.reserve(
7498 &tenant.tenant,
7499 tenant.key_prefix.as_deref(),
7500 &request.model,
7501 prompt_tokens,
7502 completion_tokens,
7503 ) {
7504 Ok(permit) => Ok(BudgetAdmission {
7505 permit,
7506 reserved_ctx: Some(prompt_tokens.saturating_add(completion_tokens)),
7507 }),
7508 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
7509 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
7510 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
7514 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
7515 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
7516 }
7517}
7518
7519fn request_ledger_error_response() -> Response {
7520 error_response_coded(
7521 StatusCode::INTERNAL_SERVER_ERROR,
7522 "request completion could not be committed to the billing ledger",
7523 "server_error",
7524 None,
7525 Some("request_ledger_unavailable"),
7526 )
7527}
7528
7529fn request_ledger_error_body() -> serde_json::Value {
7530 error_body(
7531 "request completion could not be committed to the billing ledger",
7532 "server_error",
7533 None,
7534 Some("request_ledger_unavailable"),
7535 )
7536}
7537
7538fn ledger_rejected(
7539 mut receipt: Option<Box<dyn metering::Receipt>>,
7540 response: Response,
7541 error_code: &str,
7542 request_id: &str,
7543) -> Response {
7544 let status = response.status().as_u16();
7545 if let Some(receipt) = receipt.as_mut()
7546 && let Err(err) = receipt.reject(status, error_code)
7547 {
7548 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
7549 return with_request_id(request_id, request_ledger_error_response());
7550 }
7551 with_request_id(request_id, response)
7552}
7553
7554fn ledger_unbilled(
7559 mut receipt: Option<Box<dyn metering::Receipt>>,
7560 response: Response,
7561 outcome: &'static str,
7562 error_code: &str,
7563 request_id: &str,
7564) -> Response {
7565 let status = response.status().as_u16();
7566 if let Some(receipt) = receipt.as_mut()
7567 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
7568 {
7569 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
7570 return with_request_id(request_id, request_ledger_error_response());
7571 }
7572 with_request_id(request_id, response)
7573}
7574
7575fn engine_error_code(class: worker::ErrClass) -> &'static str {
7576 use worker::ErrClass as C;
7577 match class {
7578 C::InvalidRequest => "invalid_request",
7579 C::ContextLength => "context_length_exceeded",
7580 C::ModelNotFound => "model_not_found",
7581 C::RateLimit => "rate_limit_exceeded",
7582 C::Overloaded => "overloaded",
7583 C::Engine => "engine_error",
7584 }
7585}
7586
7587fn model_not_found_response(models: &[String], requested: &str) -> Response {
7607 error_response_coded(
7608 StatusCode::BAD_REQUEST,
7609 &format!("unknown model {requested:?}; loaded: {models:?}"),
7610 "invalid_request_error",
7611 Some("model"),
7612 Some("model_not_found"),
7613 )
7614}
7615
7616fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
7625 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
7626 return Ok(());
7627 };
7628 if let Some((pos, &id)) = ids
7629 .iter()
7630 .enumerate()
7631 .find(|&(_, &id)| id as usize >= n_vocab)
7632 {
7633 return Err(format!(
7634 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
7635 ));
7636 }
7637 Ok(())
7638}
7639
7640#[cfg(test)]
7641mod prompt_ids_tests {
7642 use super::*;
7643
7644 #[test]
7645 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
7646 let caps = ModelCaps {
7647 n_vocab: 8,
7648 ..Default::default()
7649 };
7650 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
7652 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
7653 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
7655 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
7656 assert!(err.contains("vocab size 8"), "{err}");
7657 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
7658 assert!(err.contains("4294967295"), "{err}");
7659 let unknown = ModelCaps::default();
7661 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
7662 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
7663 }
7664}
7665
7666fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
7667 if models.iter().any(|m| m == requested) {
7668 return Some(requested.to_string());
7669 }
7670 if requested.is_empty() || requested.contains('/') {
7671 return None;
7672 }
7673 let mut matches = models.iter().filter(|m| {
7674 m.rsplit('/')
7675 .next()
7676 .is_some_and(|suffix| suffix == requested)
7677 });
7678 match (matches.next(), matches.next()) {
7679 (Some(only), None) => Some(only.clone()),
7680 _ => None,
7681 }
7682}
7683
7684async fn completions(
7685 State(st): State<AppState>,
7686 headers: axum::http::HeaderMap,
7687 trace: Option<Extension<TtftRequestTrace>>,
7688 Json(mut req): Json<CompletionReq>,
7689) -> Response {
7690 let env = Envelope::new(false);
7691 match canonical_model_id(&st.models, &req.model) {
7692 Some(canonical) => req.model = canonical,
7693 None => {
7694 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7695 }
7696 }
7697 let ttft = trace.and_then(|Extension(trace)| trace.0);
7700 if let Some(trace) = ttft.as_ref() {
7701 trace.mark_parsed();
7702 trace.bind_request(&env.id, &req.model);
7703 }
7704 let tenant = match authenticate(&st.api_auth, &headers) {
7705 Ok(t) => t,
7706 Err(resp) => return with_request_id(&env.id, resp),
7707 };
7708 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7709 Ok(ns) => ns,
7710 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7711 };
7712 if let Err((msg, param)) = reject_unsupported(&[
7714 (
7715 "logit_bias",
7716 req.logit_bias.is_some(),
7717 " (device-side sampling has no bias hook yet)",
7718 ),
7719 ("logprobs", req.logprobs.is_some(), ""),
7720 (
7721 "n",
7722 req.n.is_some_and(|n| n != 1),
7723 " for n != 1 (single choice only)",
7724 ),
7725 (
7726 "best_of",
7727 req.best_of.is_some_and(|n| n != 1),
7728 " (single choice only)",
7729 ),
7730 ]) {
7731 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7732 }
7733 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
7736 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
7737 }
7738 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7741 Ok(ms) => RequestDeadline::starting_now(ms),
7742 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7743 };
7744 let lane = match lane_for_tenant(&headers, &tenant) {
7745 Ok(l) => l,
7746 Err(resp) => return resp,
7747 };
7748 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7749 let model = req.model.clone();
7750 let stream = req.stream;
7751 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7752 let mut request = build_request_with_trace(
7753 &req,
7754 tx,
7755 lane,
7756 affinity,
7757 ttft.clone(),
7758 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
7762 );
7763 request.cache_ns = cache_ns;
7764 request.request_id = env.id.clone();
7765 if let Err((message, param)) = apply_model_request_limits(
7766 &mut request,
7767 st.openrouter_metadata.get(&model),
7768 st.caps.get(&model),
7769 ) {
7770 return with_request_id(&env.id, bad_request(&message, Some(param)));
7771 }
7772 if let Err(msg) = nonstream_deadline_gate(
7777 &request,
7778 req.stream,
7779 deadline,
7780 req.max_tokens.is_some(),
7781 st.budget_tokenizers
7782 .as_ref()
7783 .and_then(|t| t.get(&req.model))
7784 .map(Arc::as_ref),
7785 ) {
7786 return with_request_id(
7787 &env.id,
7788 error_response_coded(
7789 StatusCode::BAD_REQUEST,
7790 &msg,
7791 "invalid_request_error",
7792 Some("max_tokens"),
7793 Some("nonstream_deadline_infeasible"),
7794 ),
7795 );
7796 }
7797 if draining() {
7800 let receipt = start_request_receipt(
7801 &st,
7802 &env,
7803 &tenant,
7804 &req.model,
7805 "/v1/completions",
7806 lane,
7807 req.stream,
7808 effective_max_tokens(&request),
7809 None,
7810 None,
7811 );
7812 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7813 }
7814 let budget = match admit_tenant_budget(&st, &tenant, &mut request) {
7815 Ok(budget) => budget,
7816 Err(rejection) => {
7817 let (response, error_code) = rejection.into_response();
7818 let receipt = start_request_receipt(
7819 &st,
7820 &env,
7821 &tenant,
7822 &req.model,
7823 "/v1/completions",
7824 lane,
7825 req.stream,
7826 effective_max_tokens(&request),
7827 None,
7828 None,
7829 );
7830 return ledger_rejected(receipt, response, error_code, &env.id);
7831 }
7832 };
7833 let receipt = start_request_receipt(
7834 &st,
7835 &env,
7836 &tenant,
7837 &req.model,
7838 "/v1/completions",
7839 lane,
7840 req.stream,
7841 effective_max_tokens(&request),
7842 budget.reserved_ctx,
7843 budget.permit,
7844 );
7845 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
7846 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7849 Ok(slot) => slot,
7850 Err(resp) => {
7851 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7852 }
7853 };
7854 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7857 Ok(guard) => guard,
7858 Err((resp, outcome)) => {
7859 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7860 }
7861 };
7862 meter_admit(&env, &tenant, &model, lane);
7863 let stop_strings = request.stop_strings.clone();
7864
7865 if let Some(trace) = ttft.as_ref() {
7870 trace.mark_submitted();
7871 }
7872 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
7873 drop(pending_admit);
7874 return ledger_rejected(
7875 receipt,
7876 rl.attach(worker_unavailable_response()),
7877 "worker_unavailable",
7878 &env.id,
7879 );
7880 }
7881 pending_admit.commit();
7882 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7886 Ok(Ok(rx)) => rx,
7887 Ok(Err((resp, error_code))) => {
7888 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7889 }
7890 Err(_) => {
7891 return ledger_unbilled(
7892 receipt,
7893 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7894 "deadline_exceeded",
7895 "deadline_exceeded",
7896 &env.id,
7897 );
7898 }
7899 };
7900
7901 let resp = if stream {
7902 let rx = match peek_first_token(rx, deadline).await {
7906 Ok(rx) => rx,
7907 Err(()) => {
7908 return ledger_unbilled(
7909 receipt,
7910 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7911 "deadline_exceeded",
7912 "deadline_exceeded",
7913 &env.id,
7914 );
7915 }
7916 };
7917 sse_response_with_receipt(
7918 rx,
7919 model,
7920 false,
7921 None,
7922 env.clone(),
7923 stop_strings,
7924 Some(guard),
7925 receipt,
7926 )
7927 .into_response()
7928 } else {
7929 let mut receipt = receipt;
7935 let resp = blocking_response_with_receipt(
7936 rx,
7937 model,
7938 false,
7939 stop_strings,
7940 None,
7941 env.clone(),
7942 &mut receipt,
7943 Some(deadline),
7944 )
7945 .await;
7946 drop(guard); resp.into_response()
7948 };
7949 rl.attach(with_request_id(&env.id, resp))
7950}
7951
7952async fn chat_completions(
7953 State(st): State<AppState>,
7954 headers: axum::http::HeaderMap,
7955 trace: Option<Extension<TtftRequestTrace>>,
7956 Json(mut req): Json<ChatCompletionReq>,
7957) -> Response {
7958 let env = Envelope::new(true);
7959 match canonical_model_id(&st.models, &req.model) {
7964 Some(canonical) => req.model = canonical,
7965 None => {
7966 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7967 }
7968 }
7969 let ttft = trace.and_then(|Extension(trace)| trace.0);
7970 if let Some(trace) = ttft.as_ref() {
7971 trace.mark_parsed();
7972 trace.bind_request(&env.id, &req.model);
7973 }
7974 let tenant = match authenticate(&st.api_auth, &headers) {
7975 Ok(t) => t,
7976 Err(resp) => return with_request_id(&env.id, resp),
7977 };
7978 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7979 Ok(ns) => ns,
7980 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7981 };
7982 if req.messages.is_empty()
7983 || req.messages.iter().any(|message| {
7984 !matches!(
7985 message.role.as_str(),
7986 "system" | "developer" | "user" | "assistant" | "tool"
7987 )
7988 })
7989 {
7990 return with_request_id(
7991 &env.id,
7992 bad_request(
7993 "messages must use system/developer/user/assistant/tool roles",
7994 Some("messages"),
7995 ),
7996 );
7997 }
7998 if let Err((msg, param)) = reject_unsupported(&[
8003 (
8004 "logit_bias",
8005 req.logit_bias.is_some(),
8006 " (device-side sampling has no bias hook yet)",
8007 ),
8008 (
8009 "logprobs",
8010 req.logprobs
8011 .as_ref()
8012 .is_some_and(|v| v.as_bool() != Some(false)),
8013 "",
8014 ),
8015 ("top_logprobs", req.top_logprobs.is_some(), ""),
8016 (
8017 "n",
8018 req.n.is_some_and(|n| n != 1),
8019 " for n != 1 (single choice only)",
8020 ),
8021 ]) {
8022 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
8023 }
8024 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
8027 Ok(ms) => RequestDeadline::starting_now(ms),
8028 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
8029 };
8030 let lane = match lane_for_tenant(&headers, &tenant) {
8031 Ok(l) => l,
8032 Err(resp) => return resp,
8033 };
8034 let model = req.model.clone();
8035 let stream = req.stream;
8036 let capture_prompt = st
8039 .metering
8040 .as_ref()
8041 .filter(|m| m.captures(&tenant.tenant))
8042 .map(|_| capture_chat_messages(&req.messages));
8043 let declared_max_tokens = req.max_tokens.is_some();
8047 let vision_preprocess_permit = if request_has_vision(&req) {
8051 match VISION_PREPROCESS_SEMAPHORE.acquire().await {
8052 Ok(permit) => Some(permit),
8053 Err(_) => {
8054 return with_request_id(
8055 &env.id,
8056 error_response(
8057 StatusCode::SERVICE_UNAVAILABLE,
8058 "vision preprocessing is unavailable",
8059 "server_error",
8060 None,
8061 ),
8062 );
8063 }
8064 }
8065 } else {
8066 None
8067 };
8068 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
8069 let affinity = affinity_key(&req.session_id, &req.user, &headers);
8070 let mut plan = match build_chat_request_with_trace(
8071 req,
8072 st.caps.get(&model),
8073 tx,
8074 lane,
8075 affinity,
8076 ttft.clone(),
8077 st.openrouter_metadata
8078 .get(&model)
8079 .and_then(|m| m.default_reasoning_effort.as_deref()),
8080 &st.sampling_defaults(&model),
8081 ) {
8082 Ok(plan) => plan,
8083 Err(err) => {
8084 return with_request_id(&env.id, bad_request(&err, None));
8085 }
8086 };
8087 plan.request.cache_ns = cache_ns;
8088 plan.request.request_id = env.id.clone();
8089 if let Err((message, param)) = apply_model_request_limits(
8090 &mut plan.request,
8091 st.openrouter_metadata.get(&model),
8092 st.caps.get(&model),
8093 ) {
8094 return with_request_id(&env.id, bad_request(&message, Some(param)));
8095 }
8096 if let Err(msg) = nonstream_deadline_gate(
8099 &plan.request,
8100 stream,
8101 deadline,
8102 declared_max_tokens,
8103 st.budget_tokenizers
8104 .as_ref()
8105 .and_then(|t| t.get(&model))
8106 .map(Arc::as_ref),
8107 ) {
8108 return with_request_id(
8109 &env.id,
8110 error_response_coded(
8111 StatusCode::BAD_REQUEST,
8112 &msg,
8113 "invalid_request_error",
8114 Some("max_tokens"),
8115 Some("nonstream_deadline_infeasible"),
8116 ),
8117 );
8118 }
8119 plan.vision_memory = match reserve_vision_memory(&plan) {
8120 Ok(permit) => permit,
8121 Err(err) => {
8122 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
8123 }
8124 };
8125 if draining() {
8128 let receipt = start_request_receipt(
8129 &st,
8130 &env,
8131 &tenant,
8132 &model,
8133 "/v1/chat/completions",
8134 lane,
8135 stream,
8136 effective_max_tokens(&plan.request),
8137 None,
8138 None,
8139 );
8140 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
8141 }
8142 let budget = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
8143 Ok(budget) => budget,
8144 Err(rejection) => {
8145 let (response, error_code) = rejection.into_response();
8146 let receipt = start_request_receipt(
8147 &st,
8148 &env,
8149 &tenant,
8150 &model,
8151 "/v1/chat/completions",
8152 lane,
8153 stream,
8154 effective_max_tokens(&plan.request),
8155 None,
8156 None,
8157 );
8158 return ledger_rejected(receipt, response, error_code, &env.id);
8159 }
8160 };
8161 let receipt = start_request_receipt(
8162 &st,
8163 &env,
8164 &tenant,
8165 &model,
8166 "/v1/chat/completions",
8167 lane,
8168 stream,
8169 effective_max_tokens(&plan.request),
8170 budget.reserved_ctx,
8171 budget.permit,
8172 );
8173 let receipt = if let Some(prompt) = capture_prompt {
8174 arm_capture(receipt, move || prompt)
8175 } else {
8176 receipt
8177 };
8178 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
8182 Ok(slot) => slot,
8183 Err(resp) => {
8184 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
8185 }
8186 };
8187 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
8190 Ok(guard) => guard,
8191 Err((resp, outcome)) => {
8192 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
8193 }
8194 };
8195 if let Err(err) = decode_pending_vision(&mut plan) {
8200 return ledger_rejected(
8201 receipt,
8202 rl.attach(bad_request(&err, Some("messages"))),
8203 "invalid_request_error",
8204 &env.id,
8205 );
8206 }
8207 plan.request.vision_memory = plan.vision_memory.take();
8208 drop(vision_preprocess_permit);
8209 let constraint_ready = if plan.request.grammar.is_some() {
8210 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
8211 plan.request.constraint_ready = Some(ready_tx);
8212 Some(ready_rx)
8213 } else {
8214 None
8215 };
8216 meter_admit(&env, &tenant, &model, lane);
8217 let stop_strings = plan.request.stop_strings.clone();
8218 if let Some(trace) = ttft.as_ref() {
8220 trace.mark_submitted();
8221 }
8222 if st
8223 .cmd_tx
8224 .send(Cmd::Generate(Box::new(plan.request)))
8225 .is_err()
8226 {
8227 drop(pending_admit);
8228 return ledger_rejected(
8229 receipt,
8230 rl.attach(worker_unavailable_response()),
8231 "worker_unavailable",
8232 &env.id,
8233 );
8234 }
8235 pending_admit.commit();
8236 if let Some(ready) = constraint_ready {
8242 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
8243 match tokio::time::timeout(bound, ready).await {
8244 Ok(Ok(Ok(()))) => {}
8245 Ok(Ok(Err(err))) => {
8246 return ledger_rejected(
8247 receipt,
8248 rl.attach(engine_error_response(&err)),
8249 engine_error_code(err.class),
8250 &env.id,
8251 );
8252 }
8253 Ok(Err(_)) => {
8254 return ledger_rejected(
8255 receipt,
8256 rl.attach(worker_unavailable_response()),
8257 "worker_unavailable",
8258 &env.id,
8259 );
8260 }
8261 Err(_) if deadline.remaining().is_zero() => {
8262 return ledger_unbilled(
8263 receipt,
8264 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8265 "deadline_exceeded",
8266 "deadline_exceeded",
8267 &env.id,
8268 );
8269 }
8270 Err(_) => {
8271 return ledger_rejected(
8272 receipt,
8273 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
8274 "constraint_compile_timeout",
8275 &env.id,
8276 );
8277 }
8278 }
8279 }
8280 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
8282 Ok(Ok(rx)) => rx,
8283 Ok(Err((resp, error_code))) => {
8284 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
8285 }
8286 Err(_) => {
8287 return ledger_unbilled(
8288 receipt,
8289 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8290 "deadline_exceeded",
8291 "deadline_exceeded",
8292 &env.id,
8293 );
8294 }
8295 };
8296 let resp = if stream {
8297 let rx = match peek_first_token(rx, deadline).await {
8299 Ok(rx) => rx,
8300 Err(()) => {
8301 return ledger_unbilled(
8302 receipt,
8303 rl.attach(deadline_exceeded_response(deadline.ms, true)),
8304 "deadline_exceeded",
8305 "deadline_exceeded",
8306 &env.id,
8307 );
8308 }
8309 };
8310 sse_response_with_receipt(
8311 rx,
8312 model,
8313 true,
8314 plan.parser,
8315 env.clone(),
8316 stop_strings,
8317 Some(guard),
8318 receipt,
8319 )
8320 .into_response()
8321 } else {
8322 let mut receipt = receipt;
8325 let resp = blocking_response_with_receipt(
8326 rx,
8327 model,
8328 true,
8329 stop_strings,
8330 plan.parser,
8331 env.clone(),
8332 &mut receipt,
8333 Some(deadline),
8334 )
8335 .await;
8336 drop(guard); resp.into_response()
8338 };
8339 rl.attach(with_request_id(&env.id, resp))
8340}
8341
8342#[cfg(test)]
8351fn sse_response(
8352 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8353 model: String,
8354 chat: bool,
8355 parser: Option<ToolStreamParser>,
8356 env: Envelope,
8357 stop_strings: Vec<String>,
8358 guard: Option<InflightGuard>,
8359) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8360 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
8361}
8362
8363#[allow(clippy::too_many_arguments)] fn sse_response_with_receipt(
8365 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8366 model: String,
8367 chat: bool,
8368 mut parser: Option<ToolStreamParser>,
8369 env: Envelope,
8370 stop_strings: Vec<String>,
8371 guard: Option<InflightGuard>,
8372 mut receipt: Option<Box<dyn metering::Receipt>>,
8373) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8374 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
8378 .then(|| StopScrubber::new(stop_strings));
8379 let stream = async_stream::stream! {
8380 let _guard = guard;
8383 let mut call_index: usize = 0;
8384 let mut role_sent = false;
8387 macro_rules! chat_chunk {
8388 ($delta:expr, $finish:expr) => {{
8389 let mut delta = $delta;
8390 if chat && !role_sent {
8391 role_sent = true;
8392 delta["role"] = json!("assistant");
8393 }
8394 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
8395 "choices": [{ "index": 0, "delta": delta,
8396 "finish_reason": $finish }] }))
8397 .to_string()
8398 }};
8399 }
8400 macro_rules! piece_chunks {
8402 ($piece:expr) => {{
8403 let mut payloads: Vec<String> = Vec::new();
8404 match $piece {
8405 Piece::Content(text) => {
8406 let text = match scrub.as_mut() {
8407 Some(sc) => sc.push(&text),
8408 None => text,
8409 };
8410 if !text.is_empty() {
8411 payloads.push(chat_chunk!(json!({ "content": text }),
8412 serde_json::Value::Null));
8413 }
8414 }
8415 Piece::Reasoning(text) => payloads.push(
8419 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
8420 Piece::Call(call) => {
8421 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8422 "index": call_index, "id": call.id, "type": "function",
8423 "function": { "name": call.name, "arguments": "" } }] }),
8424 serde_json::Value::Null));
8425 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8426 "index": call_index,
8427 "function": { "arguments": call.arguments } }] }),
8428 serde_json::Value::Null));
8429 call_index += 1;
8430 }
8431 }
8432 payloads
8433 }};
8434 }
8435 let mut terminal = false;
8439 while let Some(ev) = rx.recv().await {
8440 match ev {
8441 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8443 if let Some(receipt) = receipt.as_mut()
8444 && let Err(err) = receipt.record_prompt_usage(
8445 n_prompt as u64,
8446 n_cached as u64,
8447 )
8448 {
8449 eprintln!(
8450 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8451 env.id
8452 );
8453 let _ = receipt.reject(500, "request_ledger_unavailable");
8456 let payload = request_ledger_error_body().to_string();
8457 if chat || openai_compat() {
8458 yield Ok(SseEvent::default().data(payload));
8459 yield Ok(SseEvent::default().data("[DONE]"));
8460 } else {
8461 yield Ok(SseEvent::default().event("error").data(payload));
8462 }
8463 terminal = true;
8464 break;
8465 }
8466 }
8467 Event::Token { id, text } => {
8468 if let Some(receipt) = receipt.as_mut()
8469 && let Err(err) = receipt.record_completion_token()
8470 {
8471 eprintln!(
8472 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8473 env.id
8474 );
8475 let _ = receipt.reject(500, "request_ledger_unavailable");
8476 let payload = request_ledger_error_body().to_string();
8477 if chat || openai_compat() {
8478 yield Ok(SseEvent::default().data(payload));
8479 yield Ok(SseEvent::default().data("[DONE]"));
8480 } else {
8481 yield Ok(SseEvent::default().event("error").data(payload));
8482 }
8483 terminal = true;
8484 break;
8485 }
8486 if let Some(receipt) = receipt.as_mut() {
8489 receipt.capture_completion_delta(&text);
8490 }
8491 if let Some(p) = parser.as_mut() {
8492 for piece in p.push(&text) {
8493 for payload in piece_chunks!(piece) {
8494 yield Ok(SseEvent::default().data(payload));
8495 }
8496 }
8497 continue;
8498 }
8499 let text = match scrub.as_mut() {
8500 Some(sc) => sc.push(&text),
8501 None => text,
8502 };
8503 if text.is_empty() && scrub.is_some() {
8504 continue; }
8506 let payload = if chat {
8507 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
8508 } else if openai_compat() {
8509 env.stamp(json!({ "object": "text_completion", "model": model,
8510 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
8511 .to_string()
8512 } else {
8513 json!({ "model": model, "id": id, "text": text }).to_string()
8514 };
8515 yield Ok(SseEvent::default().data(payload));
8516 }
8517 Event::TokenSnapshot(_) => {}
8521 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
8522 let mut finish = stop_reason_to_finish(&stop_reason);
8523 if let Some(p) = parser.as_mut() {
8524 for piece in p.finish() {
8525 for payload in piece_chunks!(piece) {
8526 yield Ok(SseEvent::default().data(payload));
8527 }
8528 }
8529 if p.n_calls() > 0 { finish = "tool_calls"; }
8530 }
8531 if let Some(sc) = scrub.as_mut() {
8533 let tail = sc.finish();
8534 if !tail.is_empty() {
8535 let payload = if chat {
8536 chat_chunk!(json!({ "content": tail }),
8537 serde_json::Value::Null)
8538 } else {
8539 env.stamp(json!({ "object": "text_completion",
8540 "model": model,
8541 "choices": [{ "index": 0, "text": tail,
8542 "finish_reason": null }] })).to_string()
8543 };
8544 yield Ok(SseEvent::default().data(payload));
8545 }
8546 }
8547 if let Some(receipt) = receipt.as_mut()
8548 && let Err(err) = receipt.complete(
8549 metering::UsageCounts {
8550 prompt_tokens: n_prompt as u64,
8551 cached_prompt_tokens: n_cached as u64,
8552 completion_tokens: n_tokens as u64,
8553 },
8554 elapsed_s,
8555 )
8556 {
8557 eprintln!(
8558 "[ledger] ERROR: request {} completion receipt failed: {err}",
8559 env.id
8560 );
8561 let _ = receipt.reject(500, "request_ledger_unavailable");
8565 let payload = request_ledger_error_body().to_string();
8566 if chat || openai_compat() {
8567 yield Ok(SseEvent::default().data(payload));
8568 yield Ok(SseEvent::default().data("[DONE]"));
8569 } else {
8570 yield Ok(SseEvent::default().event("error").data(payload));
8571 }
8572 terminal = true;
8573 break;
8574 }
8575 if chat || openai_compat() {
8576 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
8577 let fin = if chat {
8578 let mut v = env.stamp(json!({
8579 "object": "chat.completion.chunk", "model": model,
8580 "choices": [{ "index": 0, "delta": {},
8581 "finish_reason": finish }],
8582 "usage": usage }));
8583 if !role_sent {
8585 v["choices"][0]["delta"]["role"] = json!("assistant");
8586 }
8587 v
8588 } else {
8589 env.stamp(json!({ "object": "text_completion", "model": model,
8590 "choices": [{ "index": 0, "text": "",
8591 "finish_reason": finish }],
8592 "usage": usage }))
8593 }.to_string();
8594 yield Ok(SseEvent::default().data(fin));
8595 yield Ok(SseEvent::default().data("[DONE]"));
8596 } else {
8597 let payload = json!({
8598 "stop_reason": stop_reason, "n_tokens": n_tokens,
8599 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
8600 "elapsed_s": elapsed_s
8601 }).to_string();
8602 yield Ok(SseEvent::default().event("done").data(payload));
8603 }
8604 terminal = true;
8605 break;
8606 }
8607 Event::Error(err) => {
8608 let ledger_error = if let Some(receipt) = receipt.as_mut() {
8618 receipt
8619 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
8620 .err()
8621 } else {
8622 None
8623 };
8624 if let Some(ref ledger_error) = ledger_error {
8625 eprintln!(
8626 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
8627 env.id
8628 );
8629 }
8630 let payload = if ledger_error.is_some() {
8631 request_ledger_error_body().to_string()
8632 } else {
8633 engine_error_body(&err).to_string()
8634 };
8635 if chat || openai_compat() {
8636 yield Ok(SseEvent::default().data(payload));
8639 yield Ok(SseEvent::default().data("[DONE]"));
8640 } else {
8641 yield Ok(SseEvent::default().event("error").data(payload));
8644 }
8645 terminal = true;
8646 break;
8647 }
8648 }
8649 }
8650 if !terminal {
8651 let e = worker::EngineError::overloaded(
8657 "worker closed the stream without completing (worker restart in progress)",
8658 );
8659 if let Some(receipt) = receipt.as_mut()
8660 && let Err(ledger_err) = receipt.reject(
8661 class_http(e.class).0.as_u16(),
8662 engine_error_code(e.class),
8663 )
8664 {
8665 eprintln!(
8666 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8667 env.id
8668 );
8669 }
8670 let payload = engine_error_body(&e).to_string();
8671 if chat || openai_compat() {
8672 yield Ok(SseEvent::default().data(payload));
8673 yield Ok(SseEvent::default().data("[DONE]"));
8674 } else {
8675 yield Ok(SseEvent::default().event("error").data(payload));
8676 }
8677 }
8678 };
8679 Sse::new(stream).keep_alive(
8680 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
8683 )
8684}
8685
8686fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
8688 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
8689 text.truncate(offset);
8690 }
8691}
8692
8693fn partial_stop_suffix(s: &str, tag: &str) -> usize {
8696 let mut best = 0;
8697 for (k, _) in tag.char_indices().skip(1) {
8698 if k <= s.len() && s.ends_with(&tag[..k]) {
8699 best = k;
8700 }
8701 }
8702 best
8703}
8704
8705struct StopScrubber {
8711 stops: Vec<String>,
8712 buf: String,
8713 done: bool,
8714}
8715
8716impl StopScrubber {
8717 fn new(stops: Vec<String>) -> Self {
8718 Self {
8719 stops,
8720 buf: String::new(),
8721 done: false,
8722 }
8723 }
8724
8725 fn push(&mut self, text: &str) -> String {
8727 if self.done {
8728 return String::new();
8729 }
8730 self.buf.push_str(text);
8731 if let Some(i) = self
8732 .stops
8733 .iter()
8734 .filter_map(|s| self.buf.find(s.as_str()))
8735 .min()
8736 {
8737 self.done = true;
8738 let out = self.buf[..i].to_string();
8739 self.buf.clear();
8740 return out;
8741 }
8742 let keep = self
8743 .stops
8744 .iter()
8745 .map(|s| partial_stop_suffix(&self.buf, s))
8746 .max()
8747 .unwrap_or(0);
8748 let emit_to = self.buf.len() - keep;
8749 let out = self.buf[..emit_to].to_string();
8750 self.buf.drain(..emit_to);
8751 out
8752 }
8753
8754 fn finish(&mut self) -> String {
8756 if self.done {
8757 self.buf.clear();
8758 return String::new();
8759 }
8760 std::mem::take(&mut self.buf)
8761 }
8762}
8763
8764#[cfg(test)]
8765async fn blocking_response(
8766 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8767 model: String,
8768 chat: bool,
8769 stop_strings: Vec<String>,
8770 parser: Option<ToolStreamParser>,
8771 env: Envelope,
8772) -> Response {
8773 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
8774 .await
8775}
8776
8777struct BlockingPayload<'a> {
8781 env: &'a Envelope,
8782 model: String,
8783 chat: bool,
8784 finish: &'static str,
8785 text: String,
8786 reasoning: String,
8787 calls: Vec<ParsedToolCall>,
8788 tokens: Vec<u32>,
8789 stop_reason: String,
8790 n_prompt: usize,
8791 n_tokens: usize,
8792 n_cached: usize,
8793 elapsed_s: f64,
8794 spec: Option<worker::SpecUsage>,
8795 deadline_error: Option<serde_json::Value>,
8801}
8802
8803fn blocking_payload(p: BlockingPayload<'_>) -> Response {
8804 let BlockingPayload {
8805 env,
8806 model,
8807 chat,
8808 finish,
8809 text,
8810 reasoning,
8811 calls,
8812 tokens,
8813 stop_reason,
8814 n_prompt,
8815 n_tokens,
8816 n_cached,
8817 elapsed_s,
8818 spec,
8819 deadline_error,
8820 } = p;
8821 if chat {
8822 let content = if !calls.is_empty() && text.is_empty() {
8824 serde_json::Value::Null
8825 } else {
8826 serde_json::Value::String(text)
8827 };
8828 let mut message = json!({ "role": "assistant", "content": content });
8829 if !reasoning.is_empty() {
8832 message["reasoning"] = json!(reasoning);
8833 message["reasoning_details"] = json!([{
8834 "type": "reasoning.text", "text": reasoning }]);
8835 }
8836 if !calls.is_empty() {
8837 message["tool_calls"] =
8838 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
8839 }
8840 let mut body = json!({
8841 "object": "chat.completion", "model": model,
8842 "choices": [{ "index": 0,
8843 "message": message,
8844 "finish_reason": finish }],
8845 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8846 });
8847 if let Some(err) = deadline_error {
8848 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8849 body["error"] = err;
8850 }
8851 return Json(env.stamp(body)).into_response();
8852 }
8853 if openai_compat() {
8854 let mut body = json!({
8855 "object": "text_completion", "model": model,
8856 "choices": [{ "index": 0, "text": text,
8857 "finish_reason": finish }],
8858 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8859 });
8860 if let Some(err) = deadline_error {
8861 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8862 body["error"] = err;
8863 }
8864 return Json(env.stamp(body)).into_response();
8865 }
8866 Json(CompletionResp {
8867 model,
8868 text,
8869 tokens,
8870 stop_reason,
8871 error: deadline_error,
8872 n_tokens,
8873 prompt_tokens: n_prompt,
8874 cached_tokens: n_cached,
8875 elapsed_s,
8876 })
8877 .into_response()
8878}
8879
8880#[allow(clippy::too_many_arguments)] async fn blocking_response_with_receipt(
8906 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8907 model: String,
8908 chat: bool,
8909 stop_strings: Vec<String>,
8910 mut parser: Option<ToolStreamParser>,
8911 env: Envelope,
8912 receipt: &mut Option<Box<dyn metering::Receipt>>,
8913 deadline: Option<RequestDeadline>,
8914) -> Response {
8915 let mut text = String::new();
8916 let mut reasoning = String::new();
8917 let mut tokens: Vec<u32> = Vec::new();
8918 let mut calls: Vec<ParsedToolCall> = Vec::new();
8919 let consume = |pieces: Vec<Piece>,
8920 text: &mut String,
8921 reasoning: &mut String,
8922 calls: &mut Vec<ParsedToolCall>| {
8923 for piece in pieces {
8924 match piece {
8925 Piece::Content(t) => text.push_str(&t),
8926 Piece::Reasoning(t) => reasoning.push_str(&t),
8927 Piece::Call(c) => calls.push(c),
8928 }
8929 }
8930 };
8931 let started = std::time::Instant::now();
8933 let mut seen_prompt: usize = 0;
8934 let mut seen_cached: usize = 0;
8935 let mut seen_tokens: usize = 0;
8936 loop {
8937 let ev = match deadline {
8938 Some(d) => tokio::select! {
8939 biased;
8940 ev = rx.recv() => ev,
8941 () = tokio::time::sleep_until(d.at) => {
8942 drop(rx);
8945 if seen_tokens == 0 {
8946 if let Some(receipt) = receipt.as_mut()
8950 && let Err(err) = receipt.settle_unbilled(
8951 "deadline_exceeded",
8952 StatusCode::REQUEST_TIMEOUT.as_u16(),
8953 "deadline_exceeded",
8954 )
8955 {
8956 eprintln!(
8957 "[ledger] ERROR: request {} deadline receipt failed: {err}",
8958 env.id
8959 );
8960 return request_ledger_error_response();
8961 }
8962 return deadline_exceeded_response(d.ms, false);
8963 }
8964 if let Some(p) = parser.as_mut() {
8965 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8966 }
8967 truncate_at_stop(&mut text, &stop_strings);
8968 let elapsed_s = started.elapsed().as_secs_f64();
8969 if let Some(receipt) = receipt.as_mut()
8972 && let Err(err) = receipt.complete_deadline_partial(
8973 metering::UsageCounts {
8974 prompt_tokens: seen_prompt as u64,
8975 cached_prompt_tokens: seen_cached as u64,
8976 completion_tokens: seen_tokens as u64,
8977 },
8978 elapsed_s,
8979 )
8980 {
8981 eprintln!(
8982 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
8983 env.id
8984 );
8985 let _ = receipt.reject(500, "request_ledger_unavailable");
8986 return request_ledger_error_response();
8987 }
8988 eprintln!(
8989 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
8990 {} ms deadline (prompt {}); non-streaming caller advised to stream",
8991 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
8992 );
8993 let err_obj = json!({
8994 "message": format!(
8995 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
8996 the {} tokens produced before the cut are delivered above and are \
8997 billed. Set \"stream\": true for work this long — a stream's \
8998 deadline bounds only the time to first token — or lower max_tokens.",
8999 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
9000 ),
9001 "code": "deadline_exceeded",
9002 "metadata": { "error_type": "timeout", "provider_name": "memra" }
9003 });
9004 return blocking_payload(BlockingPayload {
9005 env: &env,
9006 model,
9007 chat,
9008 finish: "error",
9009 text,
9010 reasoning,
9011 calls,
9012 tokens,
9013 stop_reason: "Deadline".to_string(),
9014 n_prompt: seen_prompt,
9015 n_tokens: seen_tokens,
9016 n_cached: seen_cached,
9017 elapsed_s,
9018 spec: None,
9019 deadline_error: Some(err_obj),
9020 });
9021 }
9022 },
9023 None => rx.recv().await,
9024 };
9025 let Some(ev) = ev else { break };
9026 match ev {
9027 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
9029 if let Some(receipt) = receipt.as_mut()
9030 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
9031 {
9032 eprintln!(
9033 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
9034 env.id
9035 );
9036 let _ = receipt.reject(500, "request_ledger_unavailable");
9039 return request_ledger_error_response();
9040 }
9041 seen_prompt = n_prompt;
9042 seen_cached = n_cached;
9043 }
9044 Event::Token { id, text: delta } => {
9045 if let Some(receipt) = receipt.as_mut()
9046 && let Err(err) = receipt.record_completion_token()
9047 {
9048 eprintln!(
9049 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
9050 env.id
9051 );
9052 let _ = receipt.reject(500, "request_ledger_unavailable");
9053 return request_ledger_error_response();
9054 }
9055 if let Some(receipt) = receipt.as_mut() {
9057 receipt.capture_completion_delta(&delta);
9058 }
9059 tokens.push(id);
9060 seen_tokens += 1;
9061 match parser.as_mut() {
9062 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
9063 None => text.push_str(&delta),
9064 }
9065 }
9066 Event::TokenSnapshot(ids) => tokens = ids,
9067 Event::Done {
9068 stop_reason,
9069 n_tokens,
9070 n_prompt,
9071 n_cached,
9072 elapsed_s,
9073 spec,
9074 } => {
9075 if let Some(p) = parser.as_mut() {
9076 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
9077 }
9078 truncate_at_stop(&mut text, &stop_strings);
9079 let finish = if calls.is_empty() {
9080 stop_reason_to_finish(&stop_reason)
9081 } else {
9082 "tool_calls"
9083 };
9084 if let Some(receipt) = receipt.as_mut()
9085 && let Err(err) = receipt.complete(
9086 metering::UsageCounts {
9087 prompt_tokens: n_prompt as u64,
9088 cached_prompt_tokens: n_cached as u64,
9089 completion_tokens: n_tokens as u64,
9090 },
9091 elapsed_s,
9092 )
9093 {
9094 eprintln!(
9095 "[ledger] ERROR: request {} completion receipt failed: {err}",
9096 env.id
9097 );
9098 let _ = receipt.reject(500, "request_ledger_unavailable");
9101 return request_ledger_error_response();
9102 }
9103 return blocking_payload(BlockingPayload {
9104 env: &env,
9105 model,
9106 chat,
9107 finish,
9108 text,
9109 reasoning,
9110 calls,
9111 tokens,
9112 stop_reason,
9113 n_prompt,
9114 n_tokens,
9115 n_cached,
9116 elapsed_s,
9117 spec,
9118 deadline_error: None,
9119 });
9120 }
9121 Event::Error(err) => {
9122 if let Some(receipt) = receipt.as_mut()
9126 && let Err(ledger_err) = receipt.reject(
9127 class_http(err.class).0.as_u16(),
9128 engine_error_code(err.class),
9129 )
9130 {
9131 eprintln!(
9132 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
9133 env.id
9134 );
9135 return request_ledger_error_response();
9136 }
9137 return engine_error_response(&err);
9138 }
9139 }
9140 }
9141 let e = worker::EngineError::overloaded(
9146 "worker closed the stream without completing (worker restart in progress)",
9147 );
9148 if let Some(receipt) = receipt.as_mut()
9149 && let Err(ledger_err) =
9150 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
9151 {
9152 eprintln!(
9153 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
9154 env.id
9155 );
9156 return request_ledger_error_response();
9157 }
9158 engine_error_response(&e)
9159}
9160
9161#[cfg(test)]
9162mod tests {
9163 use super::*;
9164
9165 #[derive(Debug, Clone, PartialEq)]
9173 enum MeterEvent {
9174 Reserve {
9175 tenant: String,
9176 principal: Option<String>,
9177 model: String,
9178 },
9179 Open {
9180 request_id: String,
9181 tenant: String,
9182 model: String,
9183 route: &'static str,
9184 stream: bool,
9185 with_permit: bool,
9186 },
9187 PromptUsage {
9188 prompt: u64,
9189 cached: u64,
9190 },
9191 Token,
9192 CapturePrompt(serde_json::Value),
9193 CaptureDelta(String),
9194 Complete {
9195 prompt: u64,
9196 cached: u64,
9197 completion: u64,
9198 },
9199 DeadlinePartial {
9200 prompt: u64,
9201 cached: u64,
9202 completion: u64,
9203 },
9204 Reject {
9205 status: u16,
9206 code: String,
9207 },
9208 Unbilled {
9209 outcome: &'static str,
9210 status: u16,
9211 code: String,
9212 },
9213 Dropped {
9216 prompt: u64,
9217 cached: u64,
9218 completion: u64,
9219 },
9220 }
9221
9222 enum ReserveScript {
9225 Admit { with_permit: bool },
9226 Insufficient,
9227 Blocked,
9228 PrincipalCapped,
9229 }
9230
9231 struct MockMetering {
9232 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
9233 limits: bool,
9234 limited: bool,
9235 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
9236 captures: bool,
9237 }
9238
9239 impl MockMetering {
9240 fn admit_all() -> Arc<Self> {
9241 Arc::new(MockMetering {
9242 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9243 limits: false,
9244 limited: true,
9245 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
9246 captures: false,
9247 })
9248 }
9249
9250 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
9251 Arc::new(MockMetering {
9252 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9253 limits: true,
9254 limited: true,
9255 reserve_script: std::sync::Mutex::new(script.into()),
9256 captures: false,
9257 })
9258 }
9259
9260 fn capturing() -> Arc<Self> {
9261 Arc::new(MockMetering {
9262 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9263 limits: false,
9264 limited: true,
9265 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
9266 captures: true,
9267 })
9268 }
9269
9270 fn events(&self) -> Vec<MeterEvent> {
9271 self.events.lock().unwrap().clone()
9272 }
9273 }
9274
9275 impl metering::Metering for MockMetering {
9276 fn enforces_limits(&self) -> bool {
9277 self.limits
9278 }
9279
9280 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
9281 Ok(self.limited)
9282 }
9283
9284 fn reserve(
9285 &self,
9286 tenant: &str,
9287 principal: Option<&str>,
9288 model: &str,
9289 _prompt_tokens: u64,
9290 _completion_bound: u64,
9291 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
9292 self.events.lock().unwrap().push(MeterEvent::Reserve {
9293 tenant: tenant.into(),
9294 principal: principal.map(str::to_owned),
9295 model: model.into(),
9296 });
9297 match self.reserve_script.lock().unwrap().pop_front() {
9298 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
9299 Some(ReserveScript::Admit { with_permit: true }) => {
9300 Ok(Some(Box::new(()) as metering::Permit))
9301 }
9302 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
9303 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
9304 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
9305 }
9306 }
9307
9308 fn open(
9309 &self,
9310 meta: &metering::RequestMeta<'_>,
9311 permit: Option<metering::Permit>,
9312 ) -> Box<dyn metering::Receipt> {
9313 self.events.lock().unwrap().push(MeterEvent::Open {
9314 request_id: meta.request_id.into(),
9315 tenant: meta.tenant.into(),
9316 model: meta.model.into(),
9317 route: meta.route,
9318 stream: meta.stream,
9319 with_permit: permit.is_some(),
9320 });
9321 Box::new(MockReceipt {
9322 events: self.events.clone(),
9323 wants_capture: self.captures,
9324 prompt: 0,
9325 cached: 0,
9326 completion: 0,
9327 finalized: false,
9328 })
9329 }
9330
9331 fn captures(&self, _tenant: &str) -> bool {
9332 self.captures
9333 }
9334
9335 fn limits_health(&self) -> Option<metering::LimitsHealth> {
9336 self.limits.then_some(metering::LimitsHealth {
9337 source_reload_failed: 0,
9338 source_reload_consecutive: 0,
9339 source_available: true,
9340 })
9341 }
9342 }
9343
9344 struct MockReceipt {
9345 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
9346 wants_capture: bool,
9347 prompt: u64,
9348 cached: u64,
9349 completion: u64,
9350 finalized: bool,
9351 }
9352
9353 impl metering::Receipt for MockReceipt {
9354 fn wants_capture(&self) -> bool {
9355 self.wants_capture
9356 }
9357
9358 fn arm_capture(&mut self, prompt: serde_json::Value) {
9359 self.events
9360 .lock()
9361 .unwrap()
9362 .push(MeterEvent::CapturePrompt(prompt));
9363 }
9364
9365 fn capture_completion_delta(&mut self, text: &str) {
9366 if self.wants_capture {
9367 self.events
9368 .lock()
9369 .unwrap()
9370 .push(MeterEvent::CaptureDelta(text.into()));
9371 }
9372 }
9373
9374 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
9375 self.prompt = prompt;
9376 self.cached = cached;
9377 self.events
9378 .lock()
9379 .unwrap()
9380 .push(MeterEvent::PromptUsage { prompt, cached });
9381 Ok(())
9382 }
9383
9384 fn record_completion_token(&mut self) -> Result<(), String> {
9385 self.completion += 1;
9386 self.events.lock().unwrap().push(MeterEvent::Token);
9387 Ok(())
9388 }
9389
9390 fn complete(
9391 &mut self,
9392 usage: metering::UsageCounts,
9393 _worker_elapsed_s: f64,
9394 ) -> Result<(), String> {
9395 self.finalized = true;
9396 self.events.lock().unwrap().push(MeterEvent::Complete {
9397 prompt: usage.prompt_tokens,
9398 cached: usage.cached_prompt_tokens,
9399 completion: usage.completion_tokens,
9400 });
9401 Ok(())
9402 }
9403
9404 fn complete_deadline_partial(
9405 &mut self,
9406 usage: metering::UsageCounts,
9407 _worker_elapsed_s: f64,
9408 ) -> Result<(), String> {
9409 self.finalized = true;
9410 self.events
9411 .lock()
9412 .unwrap()
9413 .push(MeterEvent::DeadlinePartial {
9414 prompt: usage.prompt_tokens,
9415 cached: usage.cached_prompt_tokens,
9416 completion: usage.completion_tokens,
9417 });
9418 Ok(())
9419 }
9420
9421 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
9422 self.finalized = true;
9423 self.events.lock().unwrap().push(MeterEvent::Reject {
9424 status,
9425 code: error_code.into(),
9426 });
9427 Ok(())
9428 }
9429
9430 fn settle_unbilled(
9431 &mut self,
9432 outcome: &'static str,
9433 status: u16,
9434 error_code: &str,
9435 ) -> Result<(), String> {
9436 self.finalized = true;
9437 self.events.lock().unwrap().push(MeterEvent::Unbilled {
9438 outcome,
9439 status,
9440 code: error_code.into(),
9441 });
9442 Ok(())
9443 }
9444 }
9445
9446 impl Drop for MockReceipt {
9447 fn drop(&mut self) {
9448 if !self.finalized {
9449 self.events.lock().unwrap().push(MeterEvent::Dropped {
9450 prompt: self.prompt,
9451 cached: self.cached,
9452 completion: self.completion,
9453 });
9454 }
9455 }
9456 }
9457
9458 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
9463
9464 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
9467 let req: CompletionReq = serde_json::from_value(json!({
9468 "model": "qwen/qwen3.8-27b",
9469 "prompt_ids": vec![7u32; prompt_ids],
9470 }))
9471 .unwrap();
9472 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9473 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
9474 request.params.max_new = max_new;
9475 request
9476 }
9477
9478 #[test]
9485 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
9486 let prompt = 30_278u64;
9487 let deadline_ms = TIMEOUT_MS_DEFAULT;
9488 let margin = |max_new: u64| {
9489 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
9490 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
9491 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
9492 };
9493 for allowed in [64u64, 2048, 4096, 5120, 6144] {
9494 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
9495 }
9496 for refused in [8192u64, 16384, 262_144] {
9497 assert!(
9498 !margin(refused),
9499 "{refused} measured as a 408 and must be refused"
9500 );
9501 }
9502 }
9503
9504 #[test]
9505 fn the_gate_names_a_max_tokens_that_actually_fits() {
9506 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
9509 assert!(
9510 fits > 0 && fits < 7_800,
9511 "advice {fits} must fit the measured ceiling"
9512 );
9513 assert_eq!(
9515 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
9516 None
9517 );
9518 }
9519
9520 #[test]
9521 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
9522 let req = gate_request(262_144, 30_000);
9523 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
9524 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
9526 assert!(
9527 err.contains("stream"),
9528 "message must name the streaming alternative: {err}"
9529 );
9530 assert!(
9531 err.contains("max_tokens"),
9532 "message must name the knob: {err}"
9533 );
9534 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
9536 let _l = GATE_ENV_LOCK.lock().unwrap(); for off in ["0", "off", "false"] {
9543 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
9544 assert!(
9545 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
9546 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
9547 );
9548 }
9549 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
9550 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
9551 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
9552 assert!(
9553 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
9554 "unset means ON (the documented default)"
9555 );
9556 }
9557
9558 #[test]
9565 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
9566 let strip = |src: &str| -> String {
9573 src.lines()
9574 .map(|line| match line.find("//") {
9575 Some(i) => line[..i].to_string(),
9576 None => line.to_string(),
9577 })
9578 .collect::<Vec<_>>()
9579 .join("\n")
9580 };
9581 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
9583 let start = src
9584 .find(signature)
9585 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
9586 let rest = &src[start + signature.len()..];
9587 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
9588 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
9589 &rest[..end]
9590 }
9591 let main_src = strip(include_str!("lib.rs"));
9592 let surfaces_src = strip(include_str!("surfaces.rs"));
9593 for (surface, src, signature) in [
9594 ("/v1/completions", &main_src, "async fn completions("),
9595 (
9596 "/v1/chat/completions",
9597 &main_src,
9598 "async fn chat_completions(",
9599 ),
9600 (
9601 "/v1/messages + /v1/responses (shared admission)",
9602 &surfaces_src,
9603 "pub(crate) async fn admit_translated(",
9604 ),
9605 ] {
9606 let handler = body(src, signature);
9607 assert!(
9608 handler.contains("nonstream_deadline_gate("),
9609 "{surface} must CALL the feasibility gate inside {signature}"
9610 );
9611 let limits = handler
9614 .find("apply_model_request_limits(")
9615 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
9616 let gate = handler.find("nonstream_deadline_gate(").unwrap();
9617 assert!(
9618 limits < gate,
9619 "{surface}: the gate must run after apply_model_request_limits"
9620 );
9621 }
9622 }
9623
9624 #[test]
9628 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
9629 let err = json!({"code": "deadline_exceeded",
9630 "metadata": {"error_type": "timeout"}});
9631 let cut = CompletionResp {
9632 model: "m".into(),
9633 text: "partial".into(),
9634 tokens: vec![1, 2],
9635 stop_reason: "Deadline".into(),
9636 error: Some(err.clone()),
9637 n_tokens: 2,
9638 prompt_tokens: 9,
9639 cached_tokens: 0,
9640 elapsed_s: 1.0,
9641 };
9642 let v = serde_json::to_value(&cut).unwrap();
9643 assert_eq!(v["stop_reason"], "Deadline");
9644 assert_eq!(v["error"]["code"], "deadline_exceeded");
9645 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
9646 let whole = CompletionResp {
9648 error: None,
9649 stop_reason: "Eos".into(),
9650 ..cut
9651 };
9652 let v = serde_json::to_value(&whole).unwrap();
9653 assert!(
9654 v.get("error").is_none(),
9655 "a complete response must not grow an error key: {v}"
9656 );
9657 }
9658
9659 #[test]
9660 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
9661 let _l = GATE_ENV_LOCK.lock().unwrap();
9662 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
9666 assert!(
9667 nonstream_deadline_gate(
9668 &req,
9669 false,
9670 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9671 false,
9672 None,
9673 )
9674 .is_ok(),
9675 "an omitted max_tokens is never gated — context is its only limit"
9676 );
9677 let resolved = gate_request(32_768, 30_000);
9682 assert!(
9683 nonstream_deadline_gate(
9684 &resolved,
9685 false,
9686 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9687 false,
9688 None,
9689 )
9690 .is_ok(),
9691 "a resolved-but-undeclared cap is not the caller's number to be refused over"
9692 );
9693 assert!(
9695 nonstream_deadline_gate(
9696 &resolved,
9697 false,
9698 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9699 true,
9700 None,
9701 )
9702 .is_err()
9703 );
9704 }
9705
9706 #[test]
9707 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
9708 let req = gate_request(64, 1234);
9709 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
9710 let mut text = gate_request(64, 0);
9711 text.prompt_ids.clear();
9712 text.prompt_text = "x".repeat(6_000);
9713 assert_eq!(
9714 prompt_tokens_estimate(&text, None),
9715 1_000,
9716 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
9717 that would have succeeded"
9718 );
9719 }
9720
9721 #[test]
9722 fn vision_memory_reservation_is_bounded_and_released() {
9723 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
9724 let Err(capacity) = try_reserve_vision_memory(1) else {
9725 panic!("a full process vision budget admitted another request");
9726 };
9727 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
9728 let response = vision_memory_error_response(capacity, Some("messages"));
9729 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
9730 assert_eq!(response.headers()["retry-after"], "5");
9731 assert_eq!(response.headers()["retry-after-ms"], "5000");
9732 drop(permit);
9733 assert!(try_reserve_vision_memory(1).is_ok());
9734 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
9735 panic!("an over-limit vision request was admitted");
9736 };
9737 assert!(matches!(request, VisionMemoryError::Request(_)));
9738 let response = vision_memory_error_response(request, Some("messages"));
9739 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
9740 assert_eq!(response.headers()["x-should-retry"], "false");
9741 let _ = try_reserve_vision_memory(1);
9742 }
9743
9744 #[test]
9745 fn header_auth_gate_covers_only_inference_dialects() {
9746 for path in [
9747 "/v1/auth/check",
9748 "/v1/completions",
9749 "/v1/chat/completions",
9750 "/v1/messages",
9751 "/v1/responses",
9752 "/v1/embeddings",
9753 "/v1/rerank",
9754 ] {
9755 assert!(protected_inference_path(path), "{path}");
9756 }
9757 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
9758 assert!(!protected_inference_path(path), "{path}");
9759 }
9760 }
9761 #[tokio::test]
9768 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
9769 use crate::metering::Metering as _;
9770 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
9771
9772 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
9773 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
9774 tx.send(Event::PromptUsage {
9775 n_prompt: 7,
9776 n_cached: 0,
9777 })
9778 .unwrap();
9779 tx.send(Event::Token {
9780 id: 1,
9781 text: "Hel".into(),
9782 })
9783 .unwrap();
9784 tx.send(Event::Token {
9785 id: 2,
9786 text: "lo".into(),
9787 })
9788 .unwrap();
9789 tx.send(Event::Done {
9790 stop_reason: "eos".into(),
9791 n_tokens: 2,
9792 n_prompt: 7,
9793 n_cached: 0,
9794 elapsed_s: 0.05,
9795 spec: None,
9796 })
9797 .unwrap();
9798 drop(tx);
9799 let mut receipt = receipt;
9800 blocking_response_with_receipt(
9801 rx,
9802 "m".into(),
9803 true,
9804 Vec::new(),
9805 None,
9806 Envelope::new(true),
9807 &mut receipt,
9808 None,
9809 )
9810 .await
9811 };
9812
9813 let plain = MockMetering::admit_all();
9815 let receipt = plain.open(
9816 &metering::RequestMeta {
9817 request_id: "cap-unmarked",
9818 tenant: "unmarked",
9819 principal: None,
9820 model: "m",
9821 route: "/v1/chat/completions",
9822 lane: "interactive",
9823 stream: false,
9824 max_tokens: None,
9825 reserved_ctx: None,
9826 },
9827 None,
9828 );
9829 let response = drive(Some(receipt)).await;
9830 assert_eq!(response.status(), StatusCode::OK);
9831 assert!(
9832 !plain.events().iter().any(|e| matches!(
9833 e,
9834 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
9835 )),
9836 "an unarmed receipt must see no capture traffic: {:?}",
9837 plain.events()
9838 );
9839
9840 let capturing = MockMetering::capturing();
9843 let mut receipt = capturing.open(
9844 &metering::RequestMeta {
9845 request_id: "cap-marked",
9846 tenant: "marked",
9847 principal: None,
9848 model: "m",
9849 route: "/v1/chat/completions",
9850 lane: "interactive",
9851 stream: false,
9852 max_tokens: None,
9853 reserved_ctx: None,
9854 },
9855 None,
9856 );
9857 assert!(receipt.wants_capture());
9858 receipt.arm_capture(prompt.clone());
9859 let response = drive(Some(receipt)).await;
9860 assert_eq!(response.status(), StatusCode::OK);
9861 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
9862 .await
9863 .unwrap();
9864 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
9865 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
9866
9867 let events = capturing.events();
9868 assert!(
9869 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
9870 "prompt must arm byte-exact: {events:?}"
9871 );
9872 let completion: String = events
9873 .iter()
9874 .filter_map(|e| match e {
9875 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
9876 _ => None,
9877 })
9878 .collect();
9879 assert_eq!(
9880 completion, "Hello",
9881 "the deltas must reassemble the served completion byte-exact: {events:?}"
9882 );
9883 assert!(
9884 events.contains(&MeterEvent::Complete {
9885 prompt: 7,
9886 cached: 0,
9887 completion: 2,
9888 }),
9889 "worker-truth usage settles alongside the capture: {events:?}"
9890 );
9891 }
9892
9893 fn tool_caps() -> ModelCaps {
9894 ModelCaps {
9895 tools_branch: true,
9896 qwen_think: true,
9897 think_switch: true,
9898 chat_ok: true,
9899 ..Default::default()
9900 }
9901 }
9902
9903 fn ladder_caps() -> ModelCaps {
9908 ModelCaps {
9909 qwen_effort: true,
9910 ..tool_caps()
9911 }
9912 }
9913
9914 fn gemma_tool_caps() -> ModelCaps {
9915 ModelCaps {
9916 tools_branch: true,
9917 gemma_think: true,
9918 chat_ok: true,
9919 instruct_type: Some("gemma".into()),
9920 ..Default::default()
9921 }
9922 }
9923
9924 fn hy3_tool_caps() -> ModelCaps {
9925 ModelCaps {
9926 tools_branch: true,
9927 hy3: true,
9928 chat_ok: true,
9929 effort_levels: true,
9930 instruct_type: Some("hy3".into()),
9931 ..Default::default()
9932 }
9933 }
9934
9935 fn gemma_template(kind: &str) -> String {
9936 let file = match kind {
9937 "qat" => "qat-trunk-template.jinja",
9938 _ => "official-tooluse-template.jinja",
9939 };
9940 let path = format!(
9941 "{}/../../research/gemma4-tools-20260817/{file}",
9942 env!("CARGO_MANIFEST_DIR")
9943 );
9944 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9945 }
9946
9947 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
9952 let tools_arr = request
9953 .get("tools")
9954 .and_then(|t| t.as_array())
9955 .cloned()
9956 .unwrap_or_default();
9957 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
9958 (Vec::new(), Vec::new(), HashMap::new())
9959 } else {
9960 prepare_tools(&tools_arr).unwrap()
9961 };
9962 let effort = request
9963 .get("reasoning_effort")
9964 .and_then(|v| v.as_str())
9965 .map(String::from);
9966 let (think, _lvl, _explicit) =
9967 parse_think(&effort, &None, None, None, None, false).unwrap();
9968
9969 let mut turns: Vec<TmplTurn> = Vec::new();
9970 for msg in request["messages"].as_array().unwrap() {
9971 let role = msg["role"].as_str().unwrap();
9972 let role = if role == "developer" { "system" } else { role };
9973 let content =
9974 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9975 let tool_calls = msg
9976 .get("tool_calls")
9977 .and_then(|a| a.as_array())
9978 .map(|a| {
9979 a.iter()
9980 .map(|tc| {
9981 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9982 render_req_tool_call(&rtc).unwrap()
9983 })
9984 .collect()
9985 })
9986 .unwrap_or_default();
9987 let tool_responses = msg
9988 .get("tool_responses")
9989 .and_then(|a| a.as_array())
9990 .map(|a| {
9991 a.iter()
9992 .map(|tr| {
9993 (
9994 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
9995 json_to_val(&tr["response"]),
9996 )
9997 })
9998 .collect()
9999 })
10000 .unwrap_or_default();
10001 turns.push(TmplTurn {
10002 role: role.to_string(),
10003 content,
10004 tool_calls,
10005 reasoning: msg
10006 .get("reasoning")
10007 .and_then(|r| r.as_str())
10008 .map(String::from)
10009 .filter(|s| !s.is_empty()),
10010 tool_call_id: msg
10011 .get("tool_call_id")
10012 .and_then(|s| s.as_str())
10013 .map(String::from),
10014 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
10015 tool_responses,
10016 task: None,
10017 tools: Vec::new(),
10018 });
10019 }
10020 chat::apply_chat_template_tools_ex(
10021 Some(template),
10022 &turns,
10023 true,
10024 &tools_json,
10025 &tools_struct,
10026 think,
10027 None,
10028 None,
10029 )
10030 .unwrap()
10031 }
10032
10033 #[test]
10037 fn gemma4_tools_fixtures_match_the_official_jinja() {
10038 let dir = format!(
10039 "{}/../../research/gemma4-tools-20260817/fixtures",
10040 env!("CARGO_MANIFEST_DIR")
10041 );
10042 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10043 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10044 .map(|e| e.unwrap().path())
10045 .filter(|p| p.is_dir())
10046 .collect();
10047 entries.sort();
10048 assert!(
10049 entries.len() >= 14,
10050 "expected >=14 fixtures, found {}",
10051 entries.len()
10052 );
10053 let (mut official, mut qat) = (0u32, 0u32);
10054 for d in entries {
10055 let input: serde_json::Value =
10056 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10057 .unwrap();
10058 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10059 let kind = input
10060 .get("template")
10061 .and_then(|t| t.as_str())
10062 .unwrap_or("official");
10063 match kind {
10064 "qat" => qat += 1,
10065 _ => official += 1,
10066 }
10067 let tmpl = gemma_template(kind);
10068 let got = render_fixture(&input["request"], &tmpl);
10069 assert_eq!(
10070 got, expected,
10071 "fixture {:?} diverged from the jinja oracle",
10072 d
10073 );
10074 }
10075 assert!(
10076 official >= 12 && qat >= 2,
10077 "coverage: {official} official, {qat} qat"
10078 );
10079 }
10080
10081 #[test]
10087 fn gemma4_tools_flow_through_build_chat_request() {
10088 let tmpl = gemma_template("official");
10089 for name in [
10090 "01-system-tools-basic",
10091 "04-single-call-cycle",
10092 "07-multi-cycle-agentic",
10093 ] {
10094 let path = format!(
10095 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
10096 env!("CARGO_MANIFEST_DIR")
10097 );
10098 let input: serde_json::Value =
10099 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
10100 let expected_path = format!(
10101 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
10102 env!("CARGO_MANIFEST_DIR")
10103 );
10104 let expected = std::fs::read_to_string(&expected_path).unwrap();
10105 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
10106 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10107 let plan = build_chat_request(
10108 req,
10109 Some(&gemma_tool_caps()),
10110 tx,
10111 lanes::Lane::Interactive,
10112 None,
10113 )
10114 .unwrap();
10115 let got = chat::apply_chat_template_tools_ex(
10116 Some(&tmpl),
10117 &plan.request.chat_turns,
10118 true,
10119 &plan.request.tools_json,
10120 &plan.request.tools_struct,
10121 plan.request.think,
10122 plan.request.reasoning_effort.as_deref(),
10123 None,
10124 )
10125 .unwrap();
10126 assert_eq!(got, expected, "pipeline render diverged for {name}");
10127 }
10128 }
10129
10130 fn glm5_template() -> String {
10139 let path = format!(
10140 "{}/../../research/glm53-flash-bringup-20260827/chat_template.jinja",
10141 env!("CARGO_MANIFEST_DIR")
10142 );
10143 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
10144 }
10145
10146 fn glm5_caps() -> ModelCaps {
10151 ModelCaps {
10152 tools_branch: true,
10153 qwen_think: true,
10154 think_switch: false,
10155 chat_ok: true,
10156 context_length: 1_048_576,
10157 tokenizer: "glm4".into(),
10158 instruct_type: Some("glm".into()),
10159 effort_levels: true,
10160 glm5: true,
10161 ..Default::default()
10162 }
10163 }
10164
10165 fn glm5_render(body: serde_json::Value) -> Result<String, String> {
10167 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
10168 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10169 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)?;
10170 chat::apply_chat_template_tools_ex(
10171 Some(&glm5_template()),
10172 &plan.request.chat_turns,
10173 true,
10174 &plan.request.tools_json,
10175 &plan.request.tools_struct,
10176 plan.request.think,
10177 plan.request.reasoning_effort.as_deref(),
10178 None,
10179 )
10180 }
10181
10182 #[test]
10187 fn glm5_fixtures_match_the_vendor_jinja() {
10188 let dir = format!(
10189 "{}/../../research/glm53-flash-bringup-20260827/surface-fixtures",
10190 env!("CARGO_MANIFEST_DIR")
10191 );
10192 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10193 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10194 .map(|e| e.unwrap().path())
10195 .filter(|p| p.is_dir())
10196 .collect();
10197 entries.sort();
10198 assert!(
10199 entries.len() >= 22,
10200 "expected >=22 fixtures, found {}",
10201 entries.len()
10202 );
10203 for d in entries {
10204 let input: serde_json::Value =
10205 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10206 .unwrap();
10207 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10208 let got = glm5_render(input["request"].clone())
10209 .unwrap_or_else(|e| panic!("fixture {d:?} refused: {e}"));
10210 assert_eq!(
10211 got, expected,
10212 "fixture {d:?} diverged from the jinja oracle"
10213 );
10214 }
10215 }
10216
10217 #[test]
10223 fn glm5_never_renders_chatml() {
10224 let tmpl = glm5_template();
10225 assert!(tmpl.contains("<think>") && tmpl.contains("add_generation_prompt"));
10227 assert!(tmpl.contains("<tools>"));
10228 assert!(chat::template_is_glm5(&tmpl));
10229 for body in [
10230 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
10231 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
10232 "tools": [{"type": "function", "function": {"name": "f",
10233 "parameters": {"type": "object", "properties": {}}}}]}),
10234 ] {
10235 let got = glm5_render(body).unwrap();
10236 assert!(
10237 !got.contains("<|im_start|>") && !got.contains("<|im_end|>"),
10238 "glm5 rendered ChatML frames: {got:?}"
10239 );
10240 assert!(
10241 got.starts_with("[gMASK]<sop><|system|>Reasoning Effort: "),
10242 "{got:?}"
10243 );
10244 assert!(got.ends_with("<|assistant|><think>"), "{got:?}");
10245 }
10246 }
10247
10248 #[test]
10252 fn glm5_reasoning_effort_renders_and_keeps_its_max_tier() {
10253 for (sent, line) in [
10254 (None, "Max"),
10255 (Some("low"), "Low"),
10256 (Some("medium"), "Low"),
10259 (Some("high"), "High"),
10260 (Some("xhigh"), "Max"),
10261 (Some("max"), "Max"),
10262 (Some("ultra"), "Max"),
10263 ] {
10264 let mut body = json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]});
10265 if let Some(v) = sent {
10266 body["reasoning_effort"] = json!(v);
10267 }
10268 let got = glm5_render(body).unwrap();
10269 assert!(
10270 got.starts_with(&format!("[gMASK]<sop><|system|>Reasoning Effort: {line}<|")),
10271 "reasoning_effort {sent:?} should render {line:?}: {got:?}"
10272 );
10273 }
10274 let sampler_of = |v: &str| {
10277 let req: ChatCompletionReq = serde_json::from_value(
10278 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
10281 "reasoning_effort": v, "seed": 7}),
10282 )
10283 .unwrap();
10284 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10285 let plan =
10286 build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
10287 .unwrap();
10288 format!("{:?}", plan.request.sampler_cfg)
10289 };
10290 assert_eq!(sampler_of("low"), sampler_of("max"));
10291 assert_eq!(canonical_effort_for("max", true), Some("max"));
10293 assert_eq!(canonical_effort_for("xhigh", true), Some("max"));
10294 assert_eq!(canonical_effort_for("max", false), Some("high"));
10295 }
10296
10297 #[test]
10301 fn glm5_refuses_what_its_template_cannot_honour() {
10302 for (value, needle) in [
10303 ("none", "cannot disable reasoning"),
10304 ("minimal", "cannot disable reasoning"),
10305 ("bogus", "bad reasoning_effort"),
10306 ] {
10307 let err = glm5_render(json!({"model": "m",
10308 "messages": [{"role": "user", "content": "hi"}],
10309 "reasoning_effort": value}))
10310 .err()
10311 .unwrap_or_else(|| panic!("reasoning_effort {value:?} must be refused"));
10312 assert!(err.contains(needle), "{value}: {err}");
10313 }
10314 }
10315
10316 #[test]
10320 fn one_glm5_request_renders_identical_bytes_on_all_three_surfaces() {
10321 let chat = json!({
10328 "model": "m",
10329 "reasoning_effort": "high",
10330 "messages": [
10331 {"role": "user", "content": "Weather in Paris and Rome?"},
10332 {"role": "assistant", "content": null,
10333 "tool_calls": [
10334 {"id": "c1", "type": "function",
10335 "function": {"name": "get_weather",
10336 "arguments": "{\"city\": \"Paris\"}"}},
10337 {"id": "c2", "type": "function",
10338 "function": {"name": "get_weather",
10339 "arguments": "{\"city\": \"Rome\"}"}}]},
10340 {"role": "tool", "tool_call_id": "c2", "content": "rome:27"},
10341 {"role": "tool", "tool_call_id": "c1", "content": "paris:21"}
10342 ],
10343 "tools": [{"type": "function", "function": {
10344 "name": "get_weather", "description": "Get the current weather for a city",
10345 "parameters": {"type": "object",
10346 "properties": {"city": {"type": "string"}},
10347 "required": ["city"]}}}]
10348 });
10349 let responses = responses_api::translate(&json!({
10350 "model": "m",
10351 "reasoning": {"effort": "high"},
10352 "input": [
10353 {"type": "message", "role": "user",
10354 "content": [{"type": "input_text", "text": "Weather in Paris and Rome?"}]},
10355 {"type": "function_call", "call_id": "c1", "name": "get_weather",
10356 "arguments": "{\"city\": \"Paris\"}"},
10357 {"type": "function_call", "call_id": "c2", "name": "get_weather",
10358 "arguments": "{\"city\": \"Rome\"}"},
10359 {"type": "function_call_output", "call_id": "c2", "output": "rome:27"},
10360 {"type": "function_call_output", "call_id": "c1", "output": "paris:21"}
10361 ],
10362 "tools": [{"type": "function", "name": "get_weather",
10363 "description": "Get the current weather for a city",
10364 "parameters": {"type": "object",
10365 "properties": {"city": {"type": "string"}},
10366 "required": ["city"]}}]
10367 }))
10368 .expect("/v1/responses translate");
10369 let messages = anthropic::translate(&json!({
10370 "model": "m",
10371 "max_tokens": 256,
10372 "output_config": {"effort": "high"},
10373 "messages": [
10374 {"role": "user", "content": "Weather in Paris and Rome?"},
10375 {"role": "assistant", "content": [
10376 {"type": "tool_use", "id": "c1", "name": "get_weather",
10377 "input": {"city": "Paris"}},
10378 {"type": "tool_use", "id": "c2", "name": "get_weather",
10379 "input": {"city": "Rome"}}]},
10380 {"role": "user", "content": [
10381 {"type": "tool_result", "tool_use_id": "c2", "content": "rome:27"},
10382 {"type": "tool_result", "tool_use_id": "c1", "content": "paris:21"}]}
10383 ],
10384 "tools": [{"name": "get_weather",
10385 "description": "Get the current weather for a city",
10386 "input_schema": {"type": "object",
10387 "properties": {"city": {"type": "string"}},
10388 "required": ["city"]}}]
10389 }))
10390 .expect("/v1/messages translate");
10391 let want = glm5_render(chat).expect("chat");
10392 assert!(
10394 want.contains(
10395 "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value>\
10396 </tool_call><tool_call>get_weather<arg_key>city</arg_key>\
10397 <arg_value>Rome</arg_value></tool_call>"
10398 ),
10399 "{want:?}"
10400 );
10401 assert!(
10405 want.contains(
10406 "<|observation|><tool_response>paris:21</tool_response>\
10407 <tool_response>rome:27</tool_response>"
10408 ),
10409 "{want:?}"
10410 );
10411 assert!(
10412 want.contains("<|system|>Reasoning Effort: High"),
10413 "{want:?}"
10414 );
10415 for (surface, body) in [
10416 ("/v1/responses", responses),
10417 ("/v1/messages", messages.clone()),
10418 ] {
10419 let got = glm5_render(body).unwrap_or_else(|e| panic!("{surface}: {e}"));
10420 assert_eq!(
10421 got, want,
10422 "{surface} rendered DIFFERENT glm5 prompt bytes than /v1/chat/completions"
10423 );
10424 }
10425 let mut idless = messages;
10431 for m in idless["messages"].as_array_mut().unwrap() {
10432 if m["role"] == "tool" {
10433 m.as_object_mut().unwrap().remove("tool_call_id");
10434 }
10435 }
10436 let got = glm5_render(idless).expect("id-less render");
10437 assert_ne!(
10438 got, want,
10439 "dropping tool_call_id must change the rendered order — this test cannot detect \
10440 a surface that loses ids otherwise"
10441 );
10442 assert!(
10443 got.contains(
10444 "<|observation|><tool_response>rome:27</tool_response>\
10445 <tool_response>paris:21</tool_response>"
10446 ),
10447 "{got:?}"
10448 );
10449 }
10450
10451 #[test]
10454 fn glm5_chat_arms_the_native_tool_parser() {
10455 let req: ChatCompletionReq = serde_json::from_value(json!({
10456 "model": "m", "messages": [{"role": "user", "content": "weather?"}],
10457 "tools": [{"type": "function", "function": {"name": "get_weather",
10458 "parameters": {"type": "object",
10459 "properties": {"city": {"type": "string"}}}}}]}))
10460 .unwrap();
10461 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10462 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
10463 .unwrap();
10464 let mut parser = plan.parser.expect("glm5 tools request must carry a parser");
10465 let pieces = parser.push(
10466 "reasoning here</think><tool_call>get_weather<arg_key>city</arg_key>\
10467 <arg_value>Paris</arg_value></tool_call>",
10468 );
10469 let calls: Vec<_> = pieces
10470 .iter()
10471 .filter_map(|p| match p {
10472 toolcall::Piece::Call(c) => Some((c.name.as_str(), c.arguments.as_str())),
10473 _ => None,
10474 })
10475 .collect();
10476 assert_eq!(
10477 calls,
10478 vec![("get_weather", r#"{"city":"Paris"}"#)],
10479 "{pieces:?}"
10480 );
10481 assert!(
10482 pieces
10483 .iter()
10484 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "reasoning here")),
10485 "{pieces:?}"
10486 );
10487 assert!(
10489 !pieces
10490 .iter()
10491 .any(|p| matches!(p, toolcall::Piece::Content(_))),
10492 "{pieces:?}"
10493 );
10494 let req: ChatCompletionReq = serde_json::from_value(
10498 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
10499 )
10500 .unwrap();
10501 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10502 let plan = build_chat_request(req, Some(&glm5_caps()), tx, lanes::Lane::Interactive, None)
10503 .unwrap();
10504 let mut parser = plan
10505 .parser
10506 .expect("glm5 non-tools request must still split reasoning");
10507 let pieces = parser.push("weighing it</think>The answer.");
10508 assert!(
10509 pieces
10510 .iter()
10511 .any(|p| matches!(p, toolcall::Piece::Reasoning(r) if r == "weighing it")),
10512 "{pieces:?}"
10513 );
10514 assert!(
10515 pieces
10516 .iter()
10517 .any(|p| matches!(p, toolcall::Piece::Content(c) if c == "The answer.")),
10518 "{pieces:?}"
10519 );
10520 }
10521
10522 #[test]
10529 fn glm5_plain_fast_path_never_drops_replayed_reasoning() {
10530 let with_reasoning = vec![
10531 chat::Turn {
10532 role: "user".into(),
10533 content: "a".into(),
10534 ..Default::default()
10535 },
10536 chat::Turn {
10537 role: "assistant".into(),
10538 content: "A".into(),
10539 reasoning: Some("I considered a.".into()),
10540 ..Default::default()
10541 },
10542 chat::Turn {
10543 role: "user".into(),
10544 content: "b".into(),
10545 ..Default::default()
10546 },
10547 ];
10548 assert!(!worker::plain_chat_render_path(
10550 &[],
10551 &chat::ThinkMode::Default,
10552 None,
10553 &with_reasoning,
10554 false,
10555 ));
10556 let plain_turns: Vec<chat::Turn> = with_reasoning
10559 .iter()
10560 .cloned()
10561 .map(|mut t| {
10562 t.reasoning = None;
10563 t
10564 })
10565 .collect();
10566 assert!(worker::plain_chat_render_path(
10567 &[],
10568 &chat::ThinkMode::Default,
10569 None,
10570 &plain_turns,
10571 false,
10572 ));
10573 let tmpl = glm5_template();
10576 let via_tools = chat::apply_chat_template_tools_ex(
10577 Some(&tmpl),
10578 &with_reasoning,
10579 true,
10580 &[],
10581 &[],
10582 chat::ThinkMode::Default,
10583 None,
10584 None,
10585 )
10586 .unwrap();
10587 let msgs: Vec<(&str, &str)> = with_reasoning
10588 .iter()
10589 .map(|t| (t.role.as_str(), t.content.as_str()))
10590 .collect();
10591 let via_plain = chat::apply_chat_template_str(Some(&tmpl), &msgs, true);
10592 assert!(
10593 via_tools.contains("<think>I considered a.</think>"),
10594 "{via_tools:?}"
10595 );
10596 assert_ne!(via_tools, via_plain);
10597 let plain_msgs: Vec<(&str, &str)> = plain_turns
10600 .iter()
10601 .map(|t| (t.role.as_str(), t.content.as_str()))
10602 .collect();
10603 assert_eq!(
10604 chat::apply_chat_template_tools_ex(
10605 Some(&tmpl),
10606 &plain_turns,
10607 true,
10608 &[],
10609 &[],
10610 chat::ThinkMode::Default,
10611 None,
10612 None,
10613 )
10614 .unwrap(),
10615 chat::apply_chat_template_str(Some(&tmpl), &plain_msgs, true)
10616 );
10617 }
10618
10619 #[test]
10623 fn glm5_model_row_does_not_claim_structured_output() {
10624 let caps = glm5_caps();
10625 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), None);
10626 assert_eq!(row["capabilities"]["structured_output"], json!(false));
10627 assert_eq!(row["capabilities"]["tools"], json!(true));
10628 assert_eq!(row["capabilities"]["reasoning"], json!(true));
10629 let err = glm5_render(json!({"model": "m",
10631 "messages": [{"role": "user", "content": "hi"}],
10632 "response_format": {"type": "json_object"}}))
10633 .expect_err("response_format must be refused on a switchless think template");
10634 assert!(
10638 err.contains("neither an enable_thinking switch nor a recognizable"),
10639 "{err}"
10640 );
10641 let switchable = model_entry_v1("q", Some(&tool_caps()), None);
10643 assert_eq!(switchable["capabilities"]["structured_output"], json!(true));
10644 let glm_params = openrouter_supported_parameters(Some(&caps), None, true);
10647 assert!(
10648 glm_params.get("structured_outputs").is_none(),
10649 "{glm_params}"
10650 );
10651 assert!(glm_params.get("json_mode").is_none(), "{glm_params}");
10652 assert!(glm_params.get("tools").is_some(), "{glm_params}");
10653 let qwen_params = openrouter_supported_parameters(Some(&tool_caps()), None, true);
10654 assert!(
10655 qwen_params.get("structured_outputs").is_some(),
10656 "{qwen_params}"
10657 );
10658 assert!(qwen_params.get("json_mode").is_some(), "{qwen_params}");
10659 }
10660
10661 #[test]
10669 fn catalog_context_claim_is_capped_by_the_deployment_envelope() {
10670 let caps = glm5_caps();
10671 assert_eq!(caps.context_length, 1_048_576);
10672 let metadata = OpenRouterModelMetadata {
10673 max_prompt_length: Some(126_976),
10674 max_output_length: Some(4_096),
10675 ..Default::default()
10676 };
10677 let row = model_entry_v1("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
10679 assert_eq!(row["context_length"], json!(131_072));
10680 let or_row = model_entry_openrouter("zai/glm-5.3-flash", Some(&caps), Some(&metadata));
10681 assert_eq!(
10682 or_row["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
10683 json!(131_072)
10684 );
10685 assert_eq!(
10686 published_context_length(Some(&caps), Some(&metadata)),
10687 Some(131_072)
10688 );
10689 assert_eq!(published_context_length(Some(&caps), None), Some(1_048_576));
10691 let half = OpenRouterModelMetadata {
10692 max_output_length: Some(4_096),
10693 ..Default::default()
10694 };
10695 assert_eq!(
10696 published_context_length(Some(&caps), Some(&half)),
10697 Some(1_048_576)
10698 );
10699 let wide = OpenRouterModelMetadata {
10701 max_prompt_length: Some(2_000_000),
10702 max_output_length: Some(2_000_000),
10703 ..Default::default()
10704 };
10705 assert_eq!(
10706 published_context_length(Some(&caps), Some(&wide)),
10707 Some(1_048_576)
10708 );
10709 }
10710
10711 fn dsv4_sentinel() -> String {
10719 let path = format!(
10720 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
10721 env!("CARGO_MANIFEST_DIR")
10722 );
10723 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
10724 }
10725
10726 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
10731 let role = msg["role"].as_str().unwrap().to_string();
10732 let content =
10733 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
10734 let reasoning = msg
10735 .get("reasoning")
10736 .or_else(|| msg.get("reasoning_content"))
10737 .and_then(|r| r.as_str())
10738 .map(String::from)
10739 .filter(|s| !s.is_empty());
10740 let tool_calls = msg
10741 .get("tool_calls")
10742 .and_then(|a| a.as_array())
10743 .map(|a| {
10744 a.iter()
10745 .map(|tc| {
10746 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
10747 render_req_tool_call(&rtc).unwrap()
10748 })
10749 .collect()
10750 })
10751 .unwrap_or_default();
10752 let tools = msg
10753 .get("tools")
10754 .and_then(|a| a.as_array())
10755 .map(|a| {
10756 a.iter()
10757 .filter_map(|t| t.get("function").map(json_to_val))
10758 .collect()
10759 })
10760 .unwrap_or_default();
10761 TmplTurn {
10762 role,
10763 content,
10764 tool_calls,
10765 reasoning,
10766 tool_call_id: msg
10767 .get("tool_call_id")
10768 .and_then(|s| s.as_str())
10769 .map(String::from),
10770 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
10771 tool_responses: Vec::new(),
10772 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
10773 tools,
10774 }
10775 }
10776
10777 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
10778 v.and_then(|t| t.as_array())
10779 .map(|a| {
10780 a.iter()
10781 .filter_map(|t| t.get("function").map(json_to_val))
10782 .collect()
10783 })
10784 .unwrap_or_default()
10785 }
10786
10787 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
10791 let dir = format!(
10792 "{}/../../research/dsv4-template-20260818/{subdir}",
10793 env!("CARGO_MANIFEST_DIR")
10794 );
10795 let tmpl = dsv4_sentinel();
10796 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10797 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10798 .map(|e| e.unwrap().path())
10799 .filter(|p| p.is_dir())
10800 .collect();
10801 entries.sort();
10802 assert!(
10803 entries.len() >= min_fixtures,
10804 "expected >={min_fixtures} fixtures, found {}",
10805 entries.len()
10806 );
10807 for d in &entries {
10808 let input: serde_json::Value =
10809 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10810 .unwrap();
10811 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10812 let turns: Vec<TmplTurn> = input["turns"]
10813 .as_array()
10814 .unwrap()
10815 .iter()
10816 .map(dsv4_turn)
10817 .collect();
10818 let think = match input["think"].as_str().unwrap() {
10819 "chat" => ThinkMode::NoThink,
10820 _ => ThinkMode::Think,
10821 };
10822 let effort = input
10823 .get("reasoning_effort")
10824 .and_then(|v| v.as_str())
10825 .map(String::from);
10826 let req_tools = dsv4_req_tools(input.get("req_tools"));
10827 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
10828 let got = chat::apply_chat_template_tools_ex(
10829 Some(&tmpl),
10830 &turns,
10831 agp,
10832 &[],
10833 &req_tools,
10834 think,
10835 effort.as_deref(),
10836 Some(encoding),
10837 )
10838 .unwrap();
10839 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
10840 }
10841 }
10842
10843 #[test]
10844 fn dsv4_template_fixtures_match_the_oracle() {
10845 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
10846 }
10847
10848 #[test]
10854 fn dsv4_0731_fixtures_match_the_oracle() {
10855 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
10856 }
10857
10858 #[test]
10859 fn dsv4_artifact_fixtures_are_byte_identical() {
10860 let base = format!(
10864 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
10865 env!("CARGO_MANIFEST_DIR")
10866 );
10867 let tmpl = dsv4_sentinel();
10868 for (n, think) in [
10869 (1u32, ThinkMode::Think),
10870 (2, ThinkMode::Think),
10871 (3, ThinkMode::Think),
10872 (4, ThinkMode::NoThink),
10873 ] {
10874 let td: serde_json::Value = serde_json::from_str(
10875 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
10876 )
10877 .unwrap();
10878 let (messages, tools) = if td.is_object() {
10879 (td["messages"].clone(), td.get("tools").cloned())
10880 } else {
10881 (td.clone(), None)
10882 };
10883 let mut turns: Vec<TmplTurn> = Vec::new();
10884 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
10885 let mut t = dsv4_turn(msg);
10886 if i == 0
10887 && let Some(tl) = &tools
10888 {
10889 t.tools = tl
10890 .as_array()
10891 .unwrap()
10892 .iter()
10893 .filter_map(|x| x.get("function").map(json_to_val))
10894 .collect();
10895 }
10896 turns.push(t);
10897 }
10898 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
10899 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
10903 let got = chat::apply_chat_template_tools_ex(
10904 Some(&tmpl),
10905 &turns,
10906 true,
10907 &[],
10908 &[],
10909 think,
10910 None,
10911 Some(encoding),
10912 )
10913 .unwrap();
10914 assert_eq!(
10915 got, expected,
10916 "artifact fixture {n} diverged from the oracle under {encoding:?}"
10917 );
10918 }
10919 }
10920 }
10921
10922 #[test]
10923 fn dsv4_default_thinkmode_renders_thinking() {
10924 let tmpl = dsv4_sentinel();
10927 let turns = vec![TmplTurn {
10928 role: "user".into(),
10929 content: "Hi".into(),
10930 ..Default::default()
10931 }];
10932 let dflt = chat::apply_chat_template_tools_ex(
10933 Some(&tmpl),
10934 &turns,
10935 true,
10936 &[],
10937 &[],
10938 ThinkMode::Default,
10939 None,
10940 None,
10941 )
10942 .unwrap();
10943 let think = chat::apply_chat_template_tools_ex(
10944 Some(&tmpl),
10945 &turns,
10946 true,
10947 &[],
10948 &[],
10949 ThinkMode::Think,
10950 None,
10951 None,
10952 )
10953 .unwrap();
10954 assert_eq!(dflt, think);
10955 assert!(
10956 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
10957 "{dflt:?}"
10958 );
10959 let chat_mode = chat::apply_chat_template_tools_ex(
10960 Some(&tmpl),
10961 &turns,
10962 true,
10963 &[],
10964 &[],
10965 ThinkMode::NoThink,
10966 None,
10967 None,
10968 )
10969 .unwrap();
10970 assert!(
10971 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
10972 "{chat_mode:?}"
10973 );
10974 }
10975
10976 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
10981 let base = format!(
10982 "{}/../../research/dsv4-template-20260818",
10983 env!("CARGO_MANIFEST_DIR")
10984 );
10985 let refdir = std::path::Path::new(&base).join("ref");
10986 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
10987 .expect("load dsv4 tokenizer from ref dir");
10988 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
10989 let banked: serde_json::Value = serde_json::from_str(
10990 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
10991 .unwrap(),
10992 )
10993 .unwrap();
10994 let obj = banked.as_object().unwrap();
10995 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
10996 for (name, ids_v) in obj {
10997 let rendered =
10998 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
10999 let want: Vec<u32> = ids_v
11000 .as_array()
11001 .unwrap()
11002 .iter()
11003 .map(|v| v.as_u64().unwrap() as u32)
11004 .collect();
11005 let got = tok.encode(&rendered, true);
11006 assert_eq!(got, want, "tokenization diverged for {name}");
11007 }
11008 }
11009
11010 #[test]
11011 fn dsv4_tokenization_crosscheck_matches_official_ids() {
11012 dsv4_run_tokenization_crosscheck("fixtures");
11013 }
11014
11015 #[test]
11019 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
11020 dsv4_run_tokenization_crosscheck("fixtures-0731");
11021 }
11022
11023 #[test]
11024 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
11025 let base = format!(
11035 "{}/../../research/dsv4-template-20260818",
11036 env!("CARGO_MANIFEST_DIR")
11037 );
11038 let refdir = std::path::Path::new(&base).join("ref");
11039 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
11040 .expect("load dsv4 tokenizer from ref dir");
11041 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
11042 let tmpl = dsv4_sentinel();
11043 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
11044 {"type": "function", "function": {
11045 "name": "get_data",
11046 "description": "Fetch a blob",
11047 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
11048 "required": ["key"]}
11049 }}
11050 ])));
11051
11052 let cases: Vec<(&str, String)> = vec![
11053 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
11055 ("space-131k", " ".repeat(131_072)),
11056 ("digit-131k", "7".repeat(131_072)),
11057 (
11058 "mixed-runs",
11059 format!(
11060 "{}{}{}{}",
11061 "Z".repeat(65_536),
11062 " ".repeat(65_536),
11063 "7".repeat(65_536),
11064 "\n".repeat(65_536)
11065 ),
11066 ),
11067 ("cjk-64k", "中".repeat(65_536)),
11068 ("accented-letter-64k", "é".repeat(65_536)),
11069 ];
11070 for (name, blob) in &cases {
11071 let msgs = serde_json::json!([
11072 {"role": "system", "content": "You are a tool-using assistant."},
11073 {"role": "user", "content": "Fetch the blob."},
11074 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
11075 "tool_calls": [{"id": "call_001", "type": "function",
11076 "function": {"name": "get_data",
11077 "arguments": "{\"key\": \"blob\"}"}}]},
11078 {"role": "tool", "tool_call_id": "call_001", "content": blob}
11079 ]);
11080 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
11081 let rendered = chat::apply_chat_template_tools_ex(
11082 Some(&tmpl),
11083 &turns,
11084 true,
11085 &[],
11086 &req_tools,
11087 ThinkMode::Think,
11088 None,
11089 None,
11090 )
11091 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
11092 assert!(
11093 rendered.contains(blob.as_str()),
11094 "{name}: tool result missing from render"
11095 );
11096 let t0 = std::time::Instant::now();
11097 let ids = tok.encode(&rendered, true);
11098 let encode_dt = t0.elapsed();
11099 assert!(!ids.is_empty(), "{name}: empty encode");
11100 let back = tok.decode(&ids);
11101 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
11102 assert!(
11106 encode_dt < std::time::Duration::from_secs(60),
11107 "{name}: encode took {encode_dt:?}"
11108 );
11109 if *name == "ascii-letter-131k"
11113 && let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR")
11114 {
11115 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
11116 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
11117 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
11118 }
11119 }
11120 }
11121
11122 #[test]
11123 fn models_v1_entry_advertises_thinking_support() {
11124 let step_caps = ModelCaps {
11127 effort_levels: true,
11128 ..tool_caps()
11129 };
11130 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
11131 assert_eq!(entry["capabilities"]["reasoning"], true);
11132 assert_eq!(entry["capabilities"]["tools"], true);
11133
11134 let plain = ModelCaps {
11136 chat_ok: true,
11137 ..Default::default()
11138 };
11139 let entry = model_entry_v1("plain", Some(&plain), None);
11140 assert_eq!(entry["capabilities"]["reasoning"], false);
11141 assert_eq!(entry["capabilities"]["tools"], false);
11142 let entry = model_entry_v1("unknown", None, None);
11144 assert_eq!(entry["capabilities"]["reasoning"], false);
11145 assert_eq!(entry["capabilities"]["streaming"], true);
11146 }
11147
11148 #[test]
11149 fn chat_request_preserves_turns_and_openai_stop_forms() {
11150 let payload = serde_json::json!({
11151 "model": "plain_quant",
11152 "messages": [
11153 {"role": "system", "content": "rules"},
11154 {"role": "developer", "content": "dev rules"},
11155 {"role": "user", "content": "task"},
11156 {"role": "assistant", "content": "work"}
11157 ],
11158 "max_tokens": 64,
11159 "temperature": 0.0,
11160 "stop": "<stop>"
11161 });
11162 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11163 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11164 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
11165 let request = plan.request;
11166 assert!(
11167 plan.parser.is_none(),
11168 "no tools -> no parser (isolation contract)"
11169 );
11170 assert!(request.tools_json.is_empty());
11171 assert_eq!(request.think, ThinkMode::Default);
11172 assert_eq!(request.model, "plain_quant");
11173 assert_eq!(request.params.max_new, 64);
11174 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11176 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
11177 }))
11178 .unwrap();
11179 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11180 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
11181 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
11182 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11184 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
11185 "max_completion_tokens": 7
11186 }))
11187 .unwrap();
11188 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11189 assert_eq!(
11190 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
11191 .unwrap()
11192 .request
11193 .params
11194 .max_new,
11195 7
11196 );
11197 let req: CompletionReq = serde_json::from_value(serde_json::json!({
11199 "model": "plain_quant", "prompt": "task"
11200 }))
11201 .unwrap();
11202 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11203 assert_eq!(
11204 build_request(&req, tx, lanes::Lane::Interactive, None)
11205 .params
11206 .max_new,
11207 worker::MAX_NEW_CTX_BOUNDED
11208 );
11209 let turns: Vec<(String, String)> = request
11210 .chat_turns
11211 .iter()
11212 .map(|t| (t.role.clone(), t.content.clone()))
11213 .collect();
11214 assert_eq!(
11215 turns,
11216 vec![
11217 ("system".into(), "rules".into()),
11218 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
11220 ("assistant".into(), "work".into()),
11221 ]
11222 );
11223 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
11224 assert_eq!(request.stop_strings, vec!["<stop>"]);
11225
11226 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11227 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
11228 "stop": ["a", "b"]
11229 }))
11230 .unwrap();
11231 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
11232
11233 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11237 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
11238 "stop": ["", "real", ""]
11239 }))
11240 .unwrap();
11241 assert_eq!(req.stop.into_vec(), vec!["real"]);
11242 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11243 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
11244 "stop": ""
11245 }))
11246 .unwrap();
11247 assert!(req.stop.into_vec().is_empty());
11248
11249 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
11250 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
11251 "stop": null
11252 }))
11253 .unwrap();
11254 assert!(req.stop.into_vec().is_empty());
11255 }
11256
11257 #[tokio::test]
11258 async fn chat_response_has_openai_message_shape() {
11259 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
11260 tx.send(Event::Token {
11261 id: 1,
11262 text: "hello".into(),
11263 })
11264 .unwrap();
11265 tx.send(Event::Done {
11266 stop_reason: "Eos".into(),
11267 n_tokens: 1,
11268 n_prompt: 42,
11269 n_cached: 30,
11270 elapsed_s: 0.5,
11271 spec: None,
11272 })
11273 .unwrap();
11274 drop(tx);
11275 let response = blocking_response(
11276 rx,
11277 "plain_quant".into(),
11278 true,
11279 Vec::new(),
11280 None,
11281 Envelope::new(true),
11282 )
11283 .await;
11284 assert_eq!(response.status(), StatusCode::OK);
11285 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
11286 .await
11287 .unwrap();
11288 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
11289 assert_eq!(payload["object"], "chat.completion");
11290 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
11292 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
11293 assert!(
11294 payload["system_fingerprint"]
11295 .as_str()
11296 .unwrap()
11297 .starts_with("memra-")
11298 );
11299 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
11300 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
11301 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
11302 assert_eq!(payload["usage"]["prompt_tokens"], 42);
11304 assert_eq!(payload["usage"]["completion_tokens"], 1);
11305 assert_eq!(payload["usage"]["total_tokens"], 43);
11306 assert_eq!(
11307 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
11308 30
11309 );
11310 assert!(payload["usage"].get("spec").is_none());
11313 }
11314
11315 #[tokio::test]
11316 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
11317 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
11318 tx.send(Event::Token {
11320 id: 4,
11321 text: "hello".into(),
11322 })
11323 .unwrap();
11324 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
11325 tx.send(Event::Done {
11326 stop_reason: "MaxNew".into(),
11327 n_tokens: 4,
11328 n_prompt: 2,
11329 n_cached: 0,
11330 elapsed_s: 0.5,
11331 spec: None,
11332 })
11333 .unwrap();
11334 drop(tx);
11335
11336 let response = blocking_response(
11337 rx,
11338 "plain_quant".into(),
11339 false,
11340 Vec::new(),
11341 None,
11342 Envelope::new(false),
11343 )
11344 .await;
11345 assert_eq!(response.status(), StatusCode::OK);
11346 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
11347 .await
11348 .unwrap();
11349 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
11350 assert_eq!(payload["text"], "hello");
11351 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
11352 assert_eq!(payload["n_tokens"], 4);
11353 }
11354
11355 #[tokio::test]
11358 async fn chat_usage_carries_spec_acceptance_summary() {
11359 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
11360 tx.send(Event::Token {
11361 id: 1,
11362 text: "hello".into(),
11363 })
11364 .unwrap();
11365 tx.send(Event::Done {
11366 stop_reason: "Eos".into(),
11367 n_tokens: 1,
11368 n_prompt: 42,
11369 n_cached: 0,
11370 elapsed_s: 0.5,
11371 spec: Some(worker::SpecUsage {
11372 rounds: 10,
11373 drafted: 30,
11374 accepted: 21,
11375 }),
11376 })
11377 .unwrap();
11378 drop(tx);
11379 let response = blocking_response(
11380 rx,
11381 "plain_quant".into(),
11382 true,
11383 Vec::new(),
11384 None,
11385 Envelope::new(true),
11386 )
11387 .await;
11388 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
11389 .await
11390 .unwrap();
11391 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
11392 let sp = &payload["usage"]["spec"];
11393 assert_eq!(sp["rounds"], 10);
11394 assert_eq!(sp["drafted"], 30);
11395 assert_eq!(sp["accepted"], 21);
11396 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
11397 assert_eq!(payload["usage"]["total_tokens"], 43);
11399 }
11400
11401 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
11402 let mut payload = serde_json::json!({
11403 "model": "m",
11404 "messages": [{"role": "user", "content": "Weather in Paris?"}],
11405 "tools": [{"type": "function", "function": {
11406 "name": "get_weather",
11407 "description": "Get current weather",
11408 "parameters": {"type": "object",
11409 "properties": {"city": {"type": "string"},
11410 "days": {"type": "integer"}},
11411 "required": ["city"]}}}],
11412 });
11413 if let Some(obj) = extra.as_object() {
11414 for (k, v) in obj {
11415 payload[k] = v.clone();
11416 }
11417 }
11418 serde_json::from_value(payload).unwrap()
11419 }
11420
11421 #[test]
11425 fn glm5_vision_decode_is_deferred_and_grid_pinned() {
11426 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11427 let req: ChatCompletionReq = serde_json::from_value(json!({
11428 "model": "m", "messages": [{"role": "user", "content": "hi"}],
11429 }))
11430 .unwrap();
11431 let mut plan = build_chat_request(
11432 req,
11433 Some(&ModelCaps {
11434 chat_ok: true,
11435 ..Default::default()
11436 }),
11437 tx,
11438 lanes::Lane::Interactive,
11439 None,
11440 )
11441 .unwrap();
11442 let bmp = |w: u32, h: u32| -> Vec<u8> {
11445 let row = (w * 3).div_ceil(4) * 4;
11446 let size = 54 + row * h;
11447 let mut b = vec![0x42u8, 0x4d];
11448 b.extend_from_slice(&size.to_le_bytes());
11449 b.extend_from_slice(&[0; 4]);
11450 b.extend_from_slice(&54u32.to_le_bytes());
11451 b.extend_from_slice(&40u32.to_le_bytes());
11452 b.extend_from_slice(&w.to_le_bytes());
11453 b.extend_from_slice(&h.to_le_bytes());
11454 b.extend_from_slice(&1u16.to_le_bytes());
11455 b.extend_from_slice(&24u16.to_le_bytes());
11456 b.extend_from_slice(&[0u8; 24]);
11457 b.extend(std::iter::repeat_n(0x7fu8, (row * h) as usize));
11458 b
11459 };
11460 let bytes = bmp(112, 112);
11461 let (gh, gw) = memra_engine::vision_glm5::glm5_plan_image(&bytes).unwrap();
11462 assert_eq!((gh, gw), (8, 8), "identity resize grid");
11463 assert_eq!(memra_engine::vision_glm5::n_merged_for_grid(gh, gw), 16);
11464 plan.pending_glm5.push(PendingGlm5Image {
11465 bytes: bytes.clone(),
11466 gh,
11467 gw,
11468 });
11469 decode_pending_vision(&mut plan).unwrap();
11470 assert_eq!(plan.request.glm5_images.len(), 1);
11471 let unit = &plan.request.glm5_images[0];
11472 assert_eq!((unit.gh, unit.gw), (gh, gw));
11473 assert_eq!(
11474 unit.patches.len(),
11475 gh * gw * memra_engine::vision_glm5::G5V_PATCH_IN
11476 );
11477 plan.request.glm5_images.clear();
11479 plan.pending_glm5.push(PendingGlm5Image {
11480 bytes,
11481 gh: gh + 2,
11482 gw,
11483 });
11484 let err = decode_pending_vision(&mut plan).unwrap_err();
11485 assert!(err.contains("header-planned"), "got: {err}");
11486 }
11487
11488 #[test]
11489 fn vision_decode_is_deferred_and_grid_pinned() {
11490 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11495 let req: ChatCompletionReq = serde_json::from_value(json!({
11496 "model": "m", "messages": [{"role": "user", "content": "hi"}],
11497 }))
11498 .unwrap();
11499 let mut plan = build_chat_request(
11500 req,
11501 Some(&ModelCaps {
11502 chat_ok: true,
11503 ..Default::default()
11504 }),
11505 tx,
11506 lanes::Lane::Interactive,
11507 None,
11508 )
11509 .unwrap();
11510 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
11514 let mut b = Vec::new();
11515 b.extend_from_slice(b"BM");
11516 b.extend_from_slice(&54u32.to_le_bytes());
11517 b.extend_from_slice(&0u32.to_le_bytes());
11518 b.extend_from_slice(&54u32.to_le_bytes());
11519 b.extend_from_slice(&40u32.to_le_bytes());
11520 b.extend_from_slice(&w.to_le_bytes());
11521 b.extend_from_slice(&h.to_le_bytes());
11522 b.extend_from_slice(&1u16.to_le_bytes());
11523 b.extend_from_slice(&24u16.to_le_bytes());
11524 b.extend_from_slice(&[0u8; 24]);
11525 if with_pixels {
11526 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
11527 }
11528 b
11529 };
11530 let bytes = bmp(64, 64, true);
11531 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
11532 plan.pending_images.push(PendingVisionUnit::Still {
11533 bytes: bytes.clone(),
11534 gh,
11535 gw,
11536 });
11537 decode_pending_vision(&mut plan).unwrap();
11538 assert_eq!(plan.request.images.len(), 1);
11539 assert_eq!(
11540 (
11541 plan.request.images[0].prep.gh,
11542 plan.request.images[0].prep.gw
11543 ),
11544 (gh, gw),
11545 "decoded grid must equal the header-planned grid the pad run was rendered from"
11546 );
11547 plan.request.images.clear();
11549 plan.pending_images.push(PendingVisionUnit::Still {
11550 bytes,
11551 gh: gh + 2,
11552 gw,
11553 });
11554 let err = decode_pending_vision(&mut plan).unwrap_err();
11555 assert!(err.contains("header-planned"), "got: {err}");
11556 let bomb = bmp(16_000, 16_000, false);
11559 plan.pending_images.clear();
11560 plan.pending_images.push(PendingVisionUnit::Still {
11561 bytes: bomb,
11562 gh: 2,
11563 gw: 2,
11564 });
11565 let err = decode_pending_vision(&mut plan).unwrap_err();
11566 assert!(err.contains("exceeds the decode budget"), "got: {err}");
11567 }
11568
11569 #[test]
11570 fn tools_request_renders_client_key_order_and_arms_parser() {
11571 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11572 let plan = build_chat_request(
11573 weather_request(json!({})),
11574 Some(&tool_caps()),
11575 tx,
11576 lanes::Lane::Interactive,
11577 None,
11578 )
11579 .unwrap();
11580 assert!(plan.parser.is_some());
11581 assert_eq!(plan.request.tools_json.len(), 1);
11582 assert_eq!(
11584 plan.request.tools_json[0],
11585 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
11586 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
11587 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
11588 \"integer\"}}, \"required\": [\"city\"]}}}"
11589 );
11590 }
11591
11592 #[test]
11593 fn hy3_tools_and_reasoning_flow_through_the_real_chat_plan() {
11594 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11595 let plan = build_chat_request(
11596 weather_request(json!({"reasoning_effort": "high"})),
11597 Some(&hy3_tool_caps()),
11598 tx,
11599 lanes::Lane::Interactive,
11600 None,
11601 )
11602 .unwrap();
11603 assert_eq!(plan.request.think, ThinkMode::Think);
11604 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
11605 assert!(
11606 plan.request
11607 .stop_strings
11608 .iter()
11609 .any(|stop| stop == "</tool_calls:opensource>")
11610 );
11611 let rendered = chat::apply_chat_template_tools_ex(
11612 Some("... hy_User ... <tools> ..."),
11613 &plan.request.chat_turns,
11614 true,
11615 &plan.request.tools_json,
11616 &plan.request.tools_struct,
11617 plan.request.think,
11618 plan.request.reasoning_effort.as_deref(),
11619 None,
11620 )
11621 .unwrap();
11622 assert!(rendered.contains("<tool_calls:opensource>"));
11623 assert!(rendered.ends_with("<think:opensource>"));
11624
11625 let mut parser = plan.parser.expect("HY3 tools arm its native parser");
11626 let pieces = parser.push(concat!(
11627 "Need weather.</think:opensource>",
11628 "<tool_calls:opensource><tool_call:opensource>get_weather",
11629 "<tool_sep:opensource>\n<arg_key:opensource>city</arg_key:opensource>\n",
11630 "<arg_value:opensource>Paris</arg_value:opensource>\n",
11631 "</tool_call:opensource></tool_calls:opensource>",
11632 ));
11633 assert!(pieces.contains(&Piece::Reasoning("Need weather.".into())));
11634 assert!(pieces.iter().any(|piece| matches!(piece, Piece::Call(call)
11635 if call.name == "get_weather" && call.arguments == r#"{"city":"Paris"}"#)));
11636 }
11637
11638 #[test]
11639 fn tool_choice_none_strips_tools_and_parser() {
11640 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11641 let plan = build_chat_request(
11642 weather_request(json!({"tool_choice": "none"})),
11643 Some(&tool_caps()),
11644 tx,
11645 lanes::Lane::Interactive,
11646 None,
11647 )
11648 .unwrap();
11649 let mut p = plan
11652 .parser
11653 .expect("think-open chat arms the reasoning splitter");
11654 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
11655 assert_eq!(
11656 pieces,
11657 vec![
11658 Piece::Reasoning("x".into()),
11659 Piece::Content("<tool_call> stays prose".into()),
11660 ]
11661 );
11662 assert!(plan.request.tools_json.is_empty());
11663 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11665 assert!(
11666 build_chat_request(
11667 weather_request(json!({"tool_choice": "required"})),
11668 Some(&tool_caps()),
11669 tx,
11670 lanes::Lane::Interactive,
11671 None
11672 )
11673 .is_err()
11674 );
11675 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11676 assert!(
11677 build_chat_request(
11678 weather_request(json!({"tool_choice":
11679 {"type": "function", "function": {"name": "get_weather"}}})),
11680 Some(&tool_caps()),
11681 tx,
11682 lanes::Lane::Interactive,
11683 None
11684 )
11685 .is_err()
11686 );
11687 }
11688
11689 #[test]
11690 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
11691 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
11692 let _ = std::fs::remove_dir_all(&root);
11693
11694 let st = root.join("st_single");
11696 std::fs::create_dir_all(&st).unwrap();
11697 std::fs::write(st.join("config.json"), "{}").unwrap();
11698 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
11699 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
11700
11701 let sh = root.join("st_sharded");
11703 std::fs::create_dir_all(&sh).unwrap();
11704 std::fs::write(sh.join("config.json"), "{}").unwrap();
11705 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
11706 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
11707
11708 let rp = root.join("repack");
11710 std::fs::create_dir_all(&rp).unwrap();
11711 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
11712 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
11713
11714 let bogus = root.join("bogus");
11716 std::fs::create_dir_all(&bogus).unwrap();
11717 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
11718 assert!(
11719 err.contains("model.safetensors"),
11720 "error should say what is missing: {err}"
11721 );
11722 assert!(
11723 err.contains("manifest.json"),
11724 "error should mention the repack form: {err}"
11725 );
11726
11727 let nc = root.join("no_config");
11729 std::fs::create_dir_all(&nc).unwrap();
11730 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
11731 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
11732 assert!(
11733 err.contains("config.json"),
11734 "error should name config.json: {err}"
11735 );
11736
11737 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
11739 assert!(err.contains("does not exist"), "{err}");
11740
11741 let f = root.join("model.gguf");
11743 std::fs::write(&f, b"g").unwrap();
11744 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
11745
11746 let _ = std::fs::remove_dir_all(&root);
11747 }
11748
11749 #[test]
11750 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
11751 let caps = ModelCaps {
11754 tools_branch: false,
11755 qwen_think: false,
11756 think_switch: false,
11757 chat_ok: false,
11758 ..Default::default()
11759 };
11760 let payload = serde_json::json!({
11761 "model": "st_model",
11762 "messages": [{"role": "user", "content": "hello"}],
11763 });
11764 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11765 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11766 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
11767 Err(e) => e,
11768 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
11769 };
11770 assert!(
11771 err.contains("no chat template"),
11772 "message should name the cause: {err}"
11773 );
11774 assert!(
11775 err.contains("/v1/completions"),
11776 "message should point at the raw-prompt escape hatch: {err}"
11777 );
11778 }
11779
11780 #[test]
11781 fn tools_on_model_without_tools_branch_is_rejected() {
11782 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11783 let caps = ModelCaps {
11784 chat_ok: true,
11785 ..Default::default()
11786 };
11787 assert!(
11788 build_chat_request(
11789 weather_request(json!({})),
11790 Some(&caps),
11791 tx,
11792 lanes::Lane::Interactive,
11793 None
11794 )
11795 .is_err()
11796 );
11797 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11798 assert!(
11799 build_chat_request(
11800 weather_request(json!({})),
11801 None,
11802 tx,
11803 lanes::Lane::Interactive,
11804 None
11805 )
11806 .is_err()
11807 );
11808 }
11809
11810 #[test]
11811 fn reasoning_effort_maps_to_think_switch() {
11812 for (extra, want) in [
11818 (json!({}), ThinkMode::Default),
11819 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
11820 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
11821 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
11822 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
11823 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
11824 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
11825 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
11826 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
11827 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
11831 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
11832 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
11833 (
11837 json!({"reasoning": {"enabled": true, "effort": "none"}}),
11838 ThinkMode::Think,
11839 ),
11840 (
11841 json!({"reasoning": {"enabled": false, "effort": "high"}}),
11842 ThinkMode::NoThink,
11843 ),
11844 ] {
11845 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11846 let plan = build_chat_request(
11847 weather_request(extra.clone()),
11848 Some(&ladder_caps()),
11853 tx,
11854 lanes::Lane::Interactive,
11855 None,
11856 )
11857 .unwrap();
11858 assert_eq!(plan.request.think, want, "extra={extra}");
11859 }
11860 for extra in [
11864 json!({"reasoning_effort": "extreme"}),
11865 json!({"reasoning": {"effort": "banana"}}),
11866 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
11867 json!({"reasoning": {"enabled": true, "effort": ""}}),
11868 ] {
11869 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11870 assert!(
11871 build_chat_request(
11872 weather_request(extra.clone()),
11873 Some(&tool_caps()),
11874 tx,
11875 lanes::Lane::Interactive,
11876 None
11877 )
11878 .is_err(),
11879 "extra={extra} must be rejected by the one allowlist"
11880 );
11881 }
11882 for (raw, want) in [
11885 ("none", Some("none")),
11886 ("minimal", Some("minimal")),
11887 ("low", Some("low")),
11888 ("medium", Some("medium")),
11889 ("high", Some("high")),
11890 ("xhigh", Some("high")),
11891 ("max", Some("high")),
11892 ("ultra", Some("high")),
11893 ("banana", None),
11894 ("", None),
11895 ("HIGH", None),
11896 ] {
11897 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
11898 }
11899 for (raw, want) in [
11902 ("none", Some("none")),
11903 ("minimal", Some("minimal")),
11904 ("low", Some("low")),
11905 ("medium", Some("medium")),
11906 ("high", Some("high")),
11907 ("xhigh", Some("max")),
11908 ("max", Some("max")),
11909 ("ultra", Some("max")),
11910 ("banana", None),
11911 ("", None),
11912 ("MAX", None),
11913 ] {
11914 assert_eq!(
11915 canonical_effort_for(raw, true),
11916 want,
11917 "canonical_effort_for({raw:?}, dsv4)"
11918 );
11919 }
11920 }
11921
11922 #[test]
11923 fn dsv4_reasoning_effort_max_survives_canonicalization() {
11924 let dsv4_caps = ModelCaps {
11930 chat_ok: true,
11931 dsv4: true,
11932 ..Default::default()
11933 };
11934 let build = |caps: &ModelCaps, effort: &str| {
11935 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11936 let req: ChatCompletionReq = serde_json::from_value(json!({
11937 "model": "m",
11938 "messages": [{"role": "user", "content": "hi"}],
11939 "reasoning_effort": effort,
11940 }))
11941 .unwrap();
11942 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
11943 };
11944 for raw in ["max", "xhigh", "ultra"] {
11945 let plan = build(&dsv4_caps, raw).unwrap();
11946 assert_eq!(
11947 plan.request.reasoning_effort.as_deref(),
11948 Some("max"),
11949 "dsv4 {raw:?} must reach the renderer as the max rung"
11950 );
11951 assert_eq!(plan.request.think, chat::ThinkMode::Think);
11952 }
11953 let plan = build(&dsv4_caps, "high").unwrap();
11955 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
11956 let step_caps = ModelCaps {
11958 chat_ok: true,
11959 effort_levels: true,
11960 ..Default::default()
11961 };
11962 let plan = build(&step_caps, "max").unwrap();
11963 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
11964 }
11965
11966 #[test]
11967 fn default_reasoning_effort_flips_only_the_unset_request() {
11968 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
11975 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11976 build_chat_request_with_trace(
11977 weather_request(extra),
11978 Some(&ladder_caps()),
11979 tx,
11980 lanes::Lane::Interactive,
11981 None,
11982 None,
11983 default_effort,
11984 &ModelSamplingDefaults::default(),
11985 )
11986 .unwrap()
11987 };
11988 for (extra, want) in [
11989 (json!({}), ThinkMode::Think),
11991 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
11995 (json!({"include_reasoning": false}), ThinkMode::NoThink),
11996 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
11998 (json!({"include_reasoning": true}), ThinkMode::Think),
11999 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
12001 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
12002 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
12003 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
12005 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
12006 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
12007 ] {
12008 let plan = build(extra.clone(), Some("high"));
12009 assert_eq!(plan.request.think, want, "extra={extra}");
12010 }
12011 assert_eq!(
12013 build(json!({}), Some("none")).request.think,
12014 ThinkMode::NoThink
12015 );
12016 assert_eq!(
12017 build(json!({"reasoning_effort": "high"}), Some("none"))
12018 .request
12019 .think,
12020 ThinkMode::Think
12021 );
12022 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
12026 }
12027
12028 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
12034 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
12035 {%- else %}'<think>\\n'{%- endif %}";
12036
12037 #[test]
12038 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
12039 let build = |extra: serde_json::Value| {
12045 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12046 build_chat_request(
12047 weather_request(extra),
12048 Some(&tool_caps()),
12049 tx,
12050 lanes::Lane::Interactive,
12051 None,
12052 )
12053 };
12054 for (extra, want) in [
12055 (json!({"enable_thinking": false}), ThinkMode::NoThink),
12056 (json!({"enable_thinking": true}), ThinkMode::Think),
12057 (
12058 json!({"chat_template_kwargs": {"enable_thinking": false}}),
12059 ThinkMode::NoThink,
12060 ),
12061 (
12062 json!({"chat_template_kwargs": {"enable_thinking": true}}),
12063 ThinkMode::Think,
12064 ),
12065 (
12068 json!({"enable_thinking": false, "reasoning_effort": "high"}),
12069 ThinkMode::NoThink,
12070 ),
12071 (
12073 json!({"enable_thinking": false,
12074 "chat_template_kwargs": {"enable_thinking": false}}),
12075 ThinkMode::NoThink,
12076 ),
12077 ] {
12078 let plan = build(extra.clone()).unwrap_or_else(|e| {
12079 panic!("{extra} must be accepted and honored, got 400: {e}");
12080 });
12081 assert_eq!(
12082 plan.request.think, want,
12083 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
12084 );
12085 }
12086 let render = |extra: serde_json::Value| -> String {
12089 let plan = build(extra).unwrap();
12090 chat::apply_chat_template_tools_ex(
12091 Some(SWITCHED_QWEN_TMPL),
12092 &plan.request.chat_turns,
12093 true,
12094 &plan.request.tools_json,
12095 &plan.request.tools_struct,
12096 plan.request.think,
12097 plan.request.reasoning_effort.as_deref(),
12098 None,
12099 )
12100 .unwrap()
12101 };
12102 let off = render(json!({"enable_thinking": false}));
12103 assert!(
12104 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
12105 "enable_thinking:false must render the CLOSED think pair: {off:?}"
12106 );
12107 let on = render(json!({}));
12108 assert!(
12109 on.ends_with("<|im_start|>assistant\n<think>\n"),
12110 "an unset request must still render the template's OPEN think tail: {on:?}"
12111 );
12112 assert_eq!(
12113 off,
12114 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
12115 "both vLLM spellings must render byte-identically"
12116 );
12117 assert_eq!(
12118 off,
12119 render(json!({"reasoning_effort": "none"})),
12120 "the vLLM spelling must render byte-identically to the OpenAI spelling"
12121 );
12122 }
12123
12124 #[test]
12125 fn unknown_chat_template_kwarg_refuses_by_name() {
12126 let build = |extra: serde_json::Value| {
12129 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12130 build_chat_request(
12131 weather_request(extra),
12132 Some(&tool_caps()),
12133 tx,
12134 lanes::Lane::Interactive,
12135 None,
12136 )
12137 };
12138 let refusal = |extra: serde_json::Value, why: &str| -> String {
12139 build(extra).err().unwrap_or_else(|| panic!("{why}"))
12140 };
12141 let err = refusal(
12142 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
12143 "an unimplementable template kwarg must not be accepted",
12144 );
12145 assert!(
12146 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
12147 "the refusal must name the offending key AND the supported one: {err}"
12148 );
12149 let err = refusal(
12150 json!({"chat_template_kwargs": "enable_thinking=false"}),
12151 "a non-object chat_template_kwargs must not be accepted",
12152 );
12153 assert!(
12154 err.contains("must be an object"),
12155 "refusal must say what shape is expected: {err}"
12156 );
12157 let err = refusal(
12158 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
12159 "a stringly-typed switch must not be accepted",
12160 );
12161 assert!(
12162 err.contains("true or false"),
12163 "refusal must name the expected type: {err}"
12164 );
12165 let plan = build(json!({"chat_template_kwargs": null}))
12167 .expect("null chat_template_kwargs is the unset case");
12168 assert_eq!(plan.request.think, ThinkMode::Default);
12169 }
12170
12171 const Q38_TMPL: &str =
12186 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
12187
12188 fn render_with(
12191 tmpl: &str,
12192 caps: &ModelCaps,
12193 extra: serde_json::Value,
12194 default_effort: Option<&str>,
12195 ) -> Result<String, String> {
12196 let mut payload = serde_json::json!({
12197 "model": "m",
12198 "messages": [{"role": "user", "content": "hi"}],
12199 });
12200 if let Some(obj) = extra.as_object() {
12201 for (k, v) in obj {
12202 payload[k] = v.clone();
12203 }
12204 }
12205 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12206 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12207 let plan = build_chat_request_with_trace(
12208 req,
12209 Some(caps),
12210 tx,
12211 lanes::Lane::Interactive,
12212 None,
12213 None,
12214 default_effort,
12215 &ModelSamplingDefaults::default(),
12216 )?;
12217 Ok(chat::apply_chat_template_tools_ex(
12218 Some(tmpl),
12219 &plan.request.chat_turns,
12220 true,
12221 &plan.request.tools_json,
12222 &plan.request.tools_struct,
12223 plan.request.think,
12224 plan.request.reasoning_effort.as_deref(),
12225 None,
12226 )
12227 .unwrap())
12228 }
12229
12230 #[test]
12231 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
12232 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
12239 let xhigh = "Reasoning effort is set to xhigh.";
12240 let low = "Reasoning effort is set to low.";
12241 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
12243 assert!(
12244 r(json!({"reasoning_effort": "high"}))
12245 .unwrap()
12246 .contains(xhigh)
12247 );
12248 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
12251 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
12252 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
12254 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
12255 assert_ne!(low_p, high_p);
12256 assert_ne!(low_p, medium);
12257 assert_ne!(high_p, medium);
12258 for alias in ["xhigh", "max", "ultra"] {
12261 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
12262 }
12263 assert_eq!(r(json!({})).unwrap(), high_p);
12266 assert_eq!(
12269 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
12270 medium
12271 );
12272 let off = r(json!({"reasoning_effort": "none"})).unwrap();
12275 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
12276 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
12277 }
12278
12279 #[test]
12280 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
12281 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
12294focused, moving directly to the conclusion without unnecessary elaboration.";
12295 let expected = format!(
12296 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
12297 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
12298 );
12299 let after_fix = render_with(
12301 Q38_TMPL,
12302 &ladder_caps(),
12303 json!({"reasoning_effort": "low"}),
12304 None,
12305 )
12306 .unwrap();
12307 assert_eq!(
12308 after_fix, expected,
12309 "the shipped prompt for reasoning_effort:\"low\""
12310 );
12311 const ORNITH_TMPL: &str = include_str!(
12314 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
12315 );
12316 let on_deployed_binary = render_with(
12317 ORNITH_TMPL,
12318 &tool_caps(),
12319 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
12320 {"role": "user", "content": "hi"}]}),
12321 None,
12322 )
12323 .unwrap();
12324 assert_eq!(
12325 on_deployed_binary, expected,
12326 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
12327 level, or its reasoning-volume numbers do not describe the shipped prompt"
12328 );
12329 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
12332 assert!(
12333 !ladderless_unset.contains("Reasoning effort is set to"),
12334 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
12335 );
12336 assert_eq!(
12337 ladderless_unset,
12338 render_with(
12339 Q38_TMPL,
12340 &ladder_caps(),
12341 json!({"reasoning_effort": "medium"}),
12342 None
12343 )
12344 .unwrap(),
12345 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
12346 );
12347 }
12348
12349 #[test]
12350 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
12351 let off = render_with(
12358 Q38_TMPL,
12359 &ladder_caps(),
12360 json!({"reasoning_effort": "none"}),
12361 None,
12362 )
12363 .unwrap();
12364 for extra in [
12365 json!({"include_reasoning": false}),
12366 json!({"reasoning": {"exclude": true}}),
12367 ] {
12368 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
12369 assert!(
12370 got.ends_with("<think>\n\n</think>\n\n"),
12371 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
12372 );
12373 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
12374 }
12375 for extra in [
12380 json!({"enable_thinking": true, "include_reasoning": false}),
12381 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
12382 json!({"reasoning": {"enabled": true, "exclude": true}}),
12383 ] {
12384 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
12385 .err()
12386 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
12387 assert!(e.contains("contradictory"), "{extra}: {e}");
12388 assert!(
12389 e.contains("include_reasoning") || e.contains("exclude"),
12390 "{extra}: the refusal must name the suppression field the caller sent: {e}"
12391 );
12392 }
12393 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
12396 for extra in [
12397 json!({"include_reasoning": true}),
12398 json!({"reasoning": {"exclude": false}}),
12399 ] {
12400 assert_eq!(
12401 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
12402 dflt,
12403 "{extra} must not perturb the model's default"
12404 );
12405 }
12406 let switchless = ModelCaps {
12410 think_switch: false,
12411 ..tool_caps()
12412 };
12413 let err = render_with(
12414 Q38_TMPL,
12415 &switchless,
12416 json!({"include_reasoning": false}),
12417 None,
12418 )
12419 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
12420 assert!(err.contains("cannot disable reasoning"), "{err}");
12421 }
12422
12423 #[test]
12424 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
12425 let build = |extra: serde_json::Value| {
12426 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12427 build_chat_request(
12428 weather_request(extra),
12429 Some(&ladder_caps()),
12430 tx,
12431 lanes::Lane::Interactive,
12432 None,
12433 )
12434 };
12435 let err = |extra: serde_json::Value, why: &str| -> String {
12436 build(extra).err().unwrap_or_else(|| panic!("{why}"))
12437 };
12438 let e = err(
12442 json!({"reasoning": {"max_tokens": 1024}}),
12443 "reasoning.max_tokens must not be accepted-and-ignored",
12444 );
12445 assert!(e.contains("reasoning.max_tokens"), "{e}");
12446 assert!(e.contains("ONE output budget"), "{e}");
12447 for extra in [
12452 json!({"reasoning": {"max_tokens": null}}),
12453 json!({"reasoning": {"banana": null}}),
12454 ] {
12455 let e = err(
12456 extra.clone(),
12457 "a null-valued unhonourable key must still refuse",
12458 );
12459 assert!(
12460 e.contains("max_tokens") || e.contains("banana"),
12461 "{extra}: {e}"
12462 );
12463 }
12464 let e = err(
12466 json!({"reasoning": {"budget": 5}}),
12467 "an unknown reasoning key must not be accepted",
12468 );
12469 assert!(
12470 e.contains("reasoning.budget") && e.contains("enabled"),
12471 "{e}"
12472 );
12473 for (extra, want) in [
12477 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
12478 (json!({"reasoning": {"exclude": 1}}), "true or false"),
12479 (json!({"reasoning": {"effort": 3}}), "must be a string"),
12480 ] {
12481 let e = err(
12482 extra.clone(),
12483 "a wrong-typed reasoning key must not be ignored",
12484 );
12485 assert!(e.contains(want), "{extra}: {e}");
12486 }
12487 for extra in [
12492 json!({"reasoning": {"enabled": true}}),
12493 json!({"reasoning": {"effort": "low"}}),
12494 json!({"reasoning": {"exclude": false}}),
12495 json!({"reasoning": null}),
12496 json!({"reasoning": {"effort": null}}),
12497 json!({"reasoning": {"enabled": null, "exclude": null}}),
12498 ] {
12499 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
12500 }
12501 }
12502
12503 #[test]
12504 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
12505 const ORNITH_TMPL: &str = include_str!(
12513 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
12514 );
12515 let explicit_on = render_with(
12519 ORNITH_TMPL,
12520 &tool_caps(),
12521 json!({"reasoning": {"enabled": true}}),
12522 None,
12523 )
12524 .unwrap();
12525 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
12526 for extra in [
12527 json!({"reasoning_effort": "low"}),
12528 json!({"reasoning_effort": "medium"}),
12529 json!({"reasoning_effort": "high"}),
12530 json!({"reasoning_effort": "xhigh"}),
12532 json!({"reasoning": {"effort": "xhigh"}}),
12533 ] {
12534 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
12535 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
12536 assert_eq!(
12537 got, explicit_on,
12538 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
12539 documented translation, not a decorative accept"
12540 );
12541 }
12542 for extra in [
12544 json!({}),
12545 json!({"reasoning_effort": "none"}),
12546 json!({"reasoning_effort": "minimal"}),
12547 json!({"enable_thinking": false}),
12548 ] {
12549 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
12550 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
12551 }
12552 let minimal = render_with(
12555 ORNITH_TMPL,
12556 &tool_caps(),
12557 json!({"reasoning_effort": "minimal"}),
12558 None,
12559 )
12560 .unwrap();
12561 assert!(
12562 minimal.ends_with("<think>\n\n</think>\n\n"),
12563 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
12564 );
12565 let ladder_low = render_with(
12568 Q38_TMPL,
12569 &ladder_caps(),
12570 json!({"reasoning_effort": "low"}),
12571 None,
12572 )
12573 .unwrap();
12574 assert!(
12575 ladder_low.contains("Reasoning effort is set to low."),
12576 "{ladder_low:?}"
12577 );
12578 assert_ne!(
12579 ladder_low,
12580 render_with(
12581 Q38_TMPL,
12582 &ladder_caps(),
12583 json!({"reasoning_effort": "high"}),
12584 None
12585 )
12586 .unwrap(),
12587 "the ladder model's rungs stay distinct prompts"
12588 );
12589 }
12590
12591 #[test]
12592 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
12593 let render_chat = |body: serde_json::Value| -> Result<String, String> {
12604 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
12605 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12606 let plan = build_chat_request(
12607 req,
12608 Some(&ladder_caps()),
12609 tx,
12610 lanes::Lane::Interactive,
12611 None,
12612 )?;
12613 Ok(chat::apply_chat_template_tools_ex(
12614 Some(Q38_TMPL),
12615 &plan.request.chat_turns,
12616 true,
12617 &plan.request.tools_json,
12618 &plan.request.tools_struct,
12619 plan.request.think,
12620 plan.request.reasoning_effort.as_deref(),
12621 None,
12622 )
12623 .unwrap())
12624 };
12625 for (intent, chat_body, responses_body, messages_body) in [
12630 (
12631 "reasoning OFF",
12632 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
12633 "reasoning_effort": "none"}),
12634 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
12635 json!({"model": "m", "max_tokens": 16,
12636 "messages": [{"role": "user", "content": "hi"}],
12637 "thinking": {"type": "disabled"}}),
12638 ),
12639 (
12640 "reasoning ON at the top rung",
12641 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
12642 "reasoning_effort": "xhigh"}),
12643 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
12644 json!({"model": "m", "max_tokens": 16,
12645 "messages": [{"role": "user", "content": "hi"}],
12646 "output_config": {"effort": "xhigh"}}),
12647 ),
12648 (
12649 "reasoning ON at the bottom rung",
12650 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
12651 "reasoning_effort": "low"}),
12652 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
12653 json!({"model": "m", "max_tokens": 16,
12654 "messages": [{"role": "user", "content": "hi"}],
12655 "output_config": {"effort": "low"}}),
12656 ),
12657 (
12658 "the model's own default",
12659 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
12660 json!({"model": "m", "input": "hi"}),
12661 json!({"model": "m", "max_tokens": 16,
12662 "messages": [{"role": "user", "content": "hi"}]}),
12663 ),
12664 ] {
12665 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
12666 let via_responses = responses_api::translate(&responses_body)
12667 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
12668 let via_messages = anthropic::translate(&messages_body)
12669 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
12670 for (surface, translated) in [
12671 ("/v1/responses", via_responses),
12672 ("/v1/messages", via_messages),
12673 ] {
12674 let got = render_chat(translated)
12675 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
12676 assert_eq!(
12677 got, chat,
12678 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
12679 /v1/chat/completions — the parameter is honoured on one format and not \
12680 the other"
12681 );
12682 }
12683 }
12684 let switchless = ModelCaps {
12687 think_switch: false,
12688 ..ladder_caps()
12689 };
12690 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
12691 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
12692 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12693 let plan =
12694 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
12695 Ok(format!("{:?}", plan.request.think))
12696 };
12697 for (surface, body) in [
12698 (
12699 "/v1/responses",
12700 responses_api::translate(&json!({
12701 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
12702 .unwrap(),
12703 ),
12704 (
12705 "/v1/messages",
12706 anthropic::translate(&json!({
12707 "model": "m", "max_tokens": 16,
12708 "messages": [{"role": "user", "content": "hi"}],
12709 "thinking": {"type": "disabled"}}))
12710 .unwrap(),
12711 ),
12712 ] {
12713 let err = render_switchless(body)
12714 .err()
12715 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
12716 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
12717 }
12718 }
12719
12720 #[test]
12721 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
12722 let build = |extra: serde_json::Value| {
12730 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12731 build_chat_request(
12732 weather_request(extra),
12733 Some(&ladder_caps()),
12734 tx,
12735 lanes::Lane::Interactive,
12736 None,
12737 )
12738 };
12739 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
12740 .expect("preserve_thinking:true is the vendor default the renderer implements");
12741 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
12742 .err()
12743 .expect("preserve_thinking:false (the strip arm) must refuse");
12744 assert!(e.contains("preserve_thinking"), "{e}");
12745 assert!(e.contains("strip"), "{e}");
12746 assert_eq!(
12749 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
12750 .unwrap()
12751 .request
12752 .think,
12753 ThinkMode::NoThink
12754 );
12755 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
12757 .err()
12758 .expect("a stringly-typed preserve_thinking must not be accepted");
12759 assert!(e.contains("true or false"), "{e}");
12760 }
12761
12762 #[test]
12763 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
12764 let dsv4_caps = ModelCaps {
12770 qwen_think: true,
12771 think_switch: false,
12772 dsv4: true,
12773 ..tool_caps()
12774 };
12775 for extra in [
12776 json!({"reasoning_effort": "none"}),
12777 json!({"reasoning": {"enabled": false}}),
12778 json!({"enable_thinking": false}),
12779 json!({"include_reasoning": false}),
12780 ] {
12781 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12782 let plan = build_chat_request(
12783 weather_request(extra.clone()),
12784 Some(&dsv4_caps),
12785 tx,
12786 lanes::Lane::Interactive,
12787 None,
12788 )
12789 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
12790 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
12791 }
12792 }
12793
12794 #[test]
12795 fn contradictory_think_switches_refuse_instead_of_picking_one() {
12796 let build = |extra: serde_json::Value| {
12799 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12800 build_chat_request(
12801 weather_request(extra),
12802 Some(&tool_caps()),
12803 tx,
12804 lanes::Lane::Interactive,
12805 None,
12806 )
12807 };
12808 for extra in [
12809 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
12810 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
12811 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
12812 ] {
12813 match build(extra.clone()) {
12814 Err(err) => assert!(
12815 err.contains("contradictory"),
12816 "the refusal must say the switches contradict: {err}"
12817 ),
12818 Ok(plan) => panic!(
12819 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
12820 plan.request.think
12821 ),
12822 }
12823 }
12824 for extra in [
12826 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
12827 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
12828 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
12829 ] {
12830 build(extra.clone())
12831 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
12832 }
12833 }
12834
12835 #[test]
12836 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
12837 let switchless = ModelCaps {
12842 tools_branch: true,
12843 qwen_think: true,
12844 think_switch: false,
12845 chat_ok: true,
12846 ..Default::default()
12847 };
12848 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
12849 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12850 build_chat_request_with_trace(
12851 weather_request(extra),
12852 Some(caps),
12853 tx,
12854 lanes::Lane::Interactive,
12855 None,
12856 None,
12857 default_effort,
12858 &ModelSamplingDefaults::default(),
12859 )
12860 };
12861 for extra in [
12862 json!({"reasoning_effort": "none"}),
12863 json!({"reasoning_effort": "minimal"}),
12864 json!({"reasoning": {"enabled": false}}),
12865 json!({"enable_thinking": false}),
12866 json!({"chat_template_kwargs": {"enable_thinking": false}}),
12867 ] {
12868 let err = build(extra.clone(), &switchless, None)
12869 .err()
12870 .unwrap_or_else(|| {
12871 panic!(
12872 "{extra} on a switchless think template must not be accepted-and-ignored"
12873 )
12874 });
12875 assert!(
12876 err.contains("cannot disable reasoning"),
12877 "the refusal must say the model cannot disable reasoning: {err}"
12878 );
12879 }
12880 for (extra, default_effort) in [
12884 (json!({}), None),
12885 (json!({"reasoning_effort": "high"}), None),
12888 (json!({"reasoning": {"enabled": true}}), None),
12889 (json!({"enable_thinking": true}), None),
12890 (json!({}), Some("none")),
12891 (json!({}), Some("minimal")),
12892 (json!({}), Some("high")),
12893 ] {
12894 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
12895 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
12896 });
12897 }
12898 assert_eq!(
12901 build(json!({"enable_thinking": false}), &tool_caps(), None)
12902 .unwrap()
12903 .request
12904 .think,
12905 ThinkMode::NoThink
12906 );
12907 }
12908
12909 #[test]
12910 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
12911 let gemma_caps = ModelCaps {
12918 tools_branch: true,
12919 chat_ok: true,
12920 gemma_think: true,
12921 instruct_type: Some("gemma".into()),
12922 ..Default::default()
12923 };
12924 let render =
12925 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
12926 let mut payload = serde_json::json!({
12927 "model": "google/gemma-4-31b-it",
12928 "messages": [{"role": "user", "content": "Weather in Paris?"}],
12929 });
12930 if let Some(obj) = extra.as_object() {
12931 for (k, v) in obj {
12932 payload[k] = v.clone();
12933 }
12934 }
12935 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12936 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12937 let plan = build_chat_request_with_trace(
12938 req,
12939 Some(&gemma_caps),
12940 tx,
12941 lanes::Lane::Interactive,
12942 None,
12943 None,
12944 default_effort,
12945 &ModelSamplingDefaults::default(),
12946 )
12947 .unwrap();
12948 chat::apply_chat_template_tools_ex(
12949 Some(tmpl),
12950 &plan.request.chat_turns,
12951 true,
12952 &plan.request.tools_json,
12953 &plan.request.tools_struct,
12954 plan.request.think,
12955 plan.request.reasoning_effort.as_deref(),
12956 None, )
12958 .unwrap()
12959 };
12960 let official = gemma_template("official");
12961 let unset_with_knob = render(&official, json!({}), Some("high"));
12962 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
12963 assert_eq!(
12964 unset_with_knob, explicit_on,
12965 "knob render must be byte-identical to the explicit think-on render"
12966 );
12967 assert!(
12968 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
12969 "think-on injects the <|think|> system token: {unset_with_knob:?}"
12970 );
12971 assert!(
12972 unset_with_knob.ends_with("<|turn>model\n"),
12973 "think-on generation turn is OPEN: {unset_with_knob:?}"
12974 );
12975 let explicit_off_with_knob =
12979 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
12980 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
12981 assert_eq!(explicit_off_with_knob, explicit_off);
12982 assert!(
12983 !explicit_off_with_knob.contains("<|think|>")
12984 && explicit_off_with_knob.ends_with("<|turn>model\n"),
12985 "explicit off keeps the official template's thinking-off bytes: \
12986 {explicit_off_with_knob:?}"
12987 );
12988 let unset_no_knob = render(&official, json!({}), None);
12990 assert_eq!(
12991 unset_no_knob, explicit_off,
12992 "knobless unset stays the template's own thinking-off default"
12993 );
12994 assert_ne!(unset_no_knob, unset_with_knob);
12995 let qat = gemma_template("qat");
12998 assert!(
12999 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
13000 "QAT knobless unset keeps the closed-channel default"
13001 );
13002 assert_eq!(
13003 render(&qat, json!({}), Some("high")),
13004 render(&qat, json!({"reasoning_effort": "high"}), None),
13005 "QAT knob render must equal the explicit think-on render"
13006 );
13007 }
13008
13009 #[test]
13010 fn default_reasoning_effort_is_validated_at_metadata_load() {
13011 let parsed = OpenRouterMetadataFile::from_toml(
13013 r#"
13014[models.g]
13015default_reasoning_effort = "high"
13016"#,
13017 )
13018 .unwrap();
13019 assert_eq!(
13020 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
13021 Some("high")
13022 );
13023 let err = OpenRouterMetadataFile::from_toml(
13024 r#"
13025[models.g]
13026default_reasoning_effort = "always"
13027"#,
13028 )
13029 .unwrap_err();
13030 assert!(err.contains("default_reasoning_effort"), "{err}");
13031 }
13032
13033 #[test]
13034 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
13035 let effort_caps = ModelCaps {
13047 effort_levels: true,
13048 think_switch: false,
13049 ..tool_caps()
13050 };
13051 for (extra, want) in [
13052 (json!({}), None),
13053 (json!({"reasoning_effort": "low"}), Some("low")),
13054 (json!({"reasoning_effort": "medium"}), Some("medium")),
13055 (json!({"reasoning_effort": "high"}), Some("high")),
13056 (json!({"reasoning": {"effort": "high"}}), Some("high")),
13057 (json!({"reasoning_effort": "xhigh"}), Some("high")),
13059 (json!({"reasoning": {"effort": "max"}}), Some("high")),
13060 ] {
13061 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13062 let plan = build_chat_request(
13063 weather_request(extra.clone()),
13064 Some(&effort_caps),
13065 tx,
13066 lanes::Lane::Interactive,
13067 None,
13068 )
13069 .unwrap();
13070 assert_eq!(
13071 plan.request.reasoning_effort.as_deref(),
13072 want,
13073 "extra={extra}"
13074 );
13075 }
13076 for extra in [
13083 json!({"reasoning_effort": "none"}),
13084 json!({"reasoning_effort": "minimal"}),
13085 json!({"reasoning": {"enabled": false}}),
13086 json!({"enable_thinking": false}),
13087 json!({"include_reasoning": false}),
13088 ] {
13089 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13090 let err = build_chat_request(
13091 weather_request(extra.clone()),
13092 Some(&effort_caps),
13093 tx,
13094 lanes::Lane::Interactive,
13095 None,
13096 )
13097 .err()
13098 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
13099 assert!(
13100 err.contains("cannot disable reasoning"),
13101 "extra={extra}: {err}"
13102 );
13103 }
13104 for extra in [
13111 json!({"reasoning_effort": "high"}),
13112 json!({"reasoning": {"effort": "low"}}),
13113 ] {
13114 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13115 let plan = build_chat_request(
13116 weather_request(extra.clone()),
13117 Some(&tool_caps()),
13118 tx,
13119 lanes::Lane::Interactive,
13120 None,
13121 )
13122 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
13123 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
13124 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
13125 }
13126 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13128 let plan = build_chat_request(
13129 weather_request(json!({})),
13130 Some(&tool_caps()),
13131 tx,
13132 lanes::Lane::Interactive,
13133 None,
13134 )
13135 .unwrap();
13136 assert_eq!(plan.request.reasoning_effort, None);
13137 }
13138
13139 #[test]
13140 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
13141 let payload = serde_json::json!({
13142 "model": "m",
13143 "messages": [
13144 {"role": "user", "content": "Weather in Paris?"},
13145 {"role": "assistant", "content": null, "tool_calls": [
13146 {"id": "call_x", "type": "function", "function": {
13147 "name": "get_weather",
13148 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
13149 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
13150 ],
13151 });
13152 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
13153 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13154 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
13155 .unwrap();
13156 let turns = &plan.request.chat_turns;
13157 assert_eq!(turns[1].tool_calls.len(), 1);
13158 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
13159 assert_eq!(
13160 turns[1].tool_calls[0].params,
13161 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
13162 );
13163 assert_eq!(turns[2].role, "tool");
13164 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
13165 let mut p = plan
13168 .parser
13169 .expect("think-open chat arms the reasoning splitter");
13170 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
13171 assert_eq!(
13172 pieces,
13173 vec![
13174 Piece::Reasoning("thought".into()),
13175 Piece::Content("answer <tool_call> is prose here".into()),
13176 ]
13177 );
13178 }
13179
13180 #[tokio::test]
13181 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
13182 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13183 tx.send(Event::Token {
13184 id: 1,
13185 text: "plan</think>\n\n".into(),
13186 })
13187 .unwrap();
13188 tx.send(Event::Token {
13189 id: 2,
13190 text: "<tool_call>\n<function=get_weather>\n\
13191<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
13192 .into(),
13193 })
13194 .unwrap();
13195 tx.send(Event::Done {
13196 stop_reason: "Eos".into(),
13197 n_tokens: 2,
13198 n_prompt: 40,
13199 n_cached: 0,
13200 elapsed_s: 0.5,
13201 spec: None,
13202 })
13203 .unwrap();
13204 drop(tx);
13205 let parser = ToolStreamParser::new(HashMap::new(), true);
13206 let response = blocking_response(
13207 rx,
13208 "m".into(),
13209 true,
13210 Vec::new(),
13211 Some(parser),
13212 Envelope::new(true),
13213 )
13214 .await;
13215 assert_eq!(response.status(), StatusCode::OK);
13216 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
13217 .await
13218 .unwrap();
13219 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13220 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
13221 assert_eq!(
13224 payload["choices"][0]["message"]["content"],
13225 serde_json::Value::Null
13226 );
13227 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
13228 assert_eq!(
13229 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
13230 "plan"
13231 );
13232 let call = &payload["choices"][0]["message"]["tool_calls"][0];
13233 assert_eq!(call["type"], "function");
13234 assert_eq!(call["function"]["name"], "get_weather");
13235 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
13236 assert_eq!(payload["usage"]["prompt_tokens"], 40);
13239 assert_eq!(payload["usage"]["completion_tokens"], 2);
13240 assert_eq!(payload["usage"]["total_tokens"], 42);
13241 assert_eq!(
13242 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
13243 0
13244 );
13245 }
13246
13247 #[test]
13248 fn cache_salt_plumbs_to_the_worker_namespace() {
13249 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13251 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
13252 }))
13253 .unwrap();
13254 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13255 assert_eq!(
13256 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
13257 "tenant-a"
13258 );
13259
13260 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13261 "model": "m", "messages": [{"role": "user", "content": "task"}],
13262 "cache_salt": "tenant-b"
13263 }))
13264 .unwrap();
13265 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13266 assert_eq!(
13267 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13268 .unwrap()
13269 .request
13270 .cache_ns,
13271 "tenant-b"
13272 );
13273
13274 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13276 "model": "m", "prompt": "task"
13277 }))
13278 .unwrap();
13279 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13280 assert_eq!(
13281 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
13282 ""
13283 );
13284 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13285 "model": "m", "messages": [{"role": "user", "content": "task"}]
13286 }))
13287 .unwrap();
13288 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13289 assert_eq!(
13290 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13291 .unwrap()
13292 .request
13293 .cache_ns,
13294 ""
13295 );
13296 }
13297
13298 #[test]
13299 fn cache_salt_validation_rejects_oversized_value() {
13300 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
13301 assert_eq!(
13302 validate_cache_namespace(&salt, false),
13303 Err("cache_salt must be at most 64 bytes")
13304 );
13305 }
13306
13307 #[test]
13308 fn cache_salt_validation_rejects_reserved_open_namespace() {
13309 let salt = Some("t:acme\u{1f}private".to_string());
13310 assert_eq!(
13311 validate_cache_namespace(&salt, false),
13312 Err("cache_salt must not use the reserved t: prefix without a keyring")
13313 );
13314 }
13315
13316 #[test]
13317 fn cache_salt_validation_accepts_normal_value() {
13318 let raw = "tenant-A_7.c2VjcmV0LXNjb3Bl+/=";
13319 let salt = Some(raw.to_string());
13320 assert_eq!(validate_cache_namespace(&salt, false).unwrap(), raw);
13321 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
13322 let max_raw = "a".repeat(CACHE_SALT_MAX_BYTES);
13323 let max = Some(max_raw.clone());
13324 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max_raw);
13325 }
13326
13327 #[test]
13328 fn cache_salt_validation_rejects_unsupported_characters() {
13329 let salt = Some("tenant salt".to_string());
13330 assert_eq!(
13331 validate_cache_namespace(&salt, false),
13332 Err("cache_salt contains unsupported characters")
13333 );
13334 }
13335
13336 #[test]
13337 fn affinity_key_honors_both_client_conventions_in_priority_order() {
13338 use axum::http::HeaderMap;
13339 let hdr = |v: &str| {
13340 let mut h = HeaderMap::new();
13341 h.insert("x-session-id", v.parse().unwrap());
13342 h
13343 };
13344 let empty = HeaderMap::new();
13345 let s = |v: &str| Some(v.to_string());
13346 assert_eq!(affinity_key(&s("explicit"), &None, &empty), s("explicit"));
13348 assert_eq!(
13349 affinity_key(&None, &s("openai-user"), &empty),
13350 s("openai-user")
13351 );
13352 assert_eq!(affinity_key(&None, &None, &hdr("hdr-id")), s("hdr-id"));
13353 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")), s("a"));
13356 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")), s("b"));
13357 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")), None);
13360 assert_eq!(affinity_key(&s(""), &s("real"), &empty), s("real"));
13361 assert_eq!(affinity_key(&s(" padded "), &None, &empty), s("padded"));
13363 assert_eq!(affinity_key(&None, &None, &empty), None);
13365 }
13366
13367 #[test]
13368 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
13369 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13370 "model": "m", "prompt": "task", "session_id": "conv-1"
13371 }))
13372 .unwrap();
13373 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13374 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
13375 assert_eq!(
13376 build_request(&req, tx, lanes::Lane::Interactive, key)
13377 .affinity
13378 .as_deref(),
13379 Some("conv-1")
13380 );
13381 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13383 "model": "m", "messages": [{"role": "user", "content": "task"}],
13384 "user": "conv-2"
13385 }))
13386 .unwrap();
13387 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13388 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
13389 assert_eq!(
13390 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
13391 .unwrap()
13392 .request
13393 .affinity
13394 .as_deref(),
13395 Some("conv-2")
13396 );
13397 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13399 "model": "m", "prompt": "task"
13400 }))
13401 .unwrap();
13402 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13403 assert!(
13404 build_request(&req, tx, lanes::Lane::Interactive, None)
13405 .affinity
13406 .is_none()
13407 );
13408 }
13409
13410 async fn sse_data_lines(resp: Response) -> Vec<String> {
13412 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13413 .await
13414 .unwrap();
13415 String::from_utf8(bytes.to_vec())
13416 .unwrap()
13417 .lines()
13418 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
13419 .collect()
13420 }
13421
13422 #[tokio::test]
13423 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
13424 let feed = |think: bool| {
13430 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13431 let body = if think {
13432 "a plan</think>\n\nanswer"
13433 } else {
13434 "answer"
13435 };
13436 tx.send(Event::Token {
13437 id: 1,
13438 text: body.into(),
13439 })
13440 .unwrap();
13441 tx.send(Event::Done {
13442 stop_reason: "Eos".into(),
13443 n_tokens: 3,
13444 n_prompt: 10,
13445 n_cached: 0,
13446 elapsed_s: 0.1,
13447 spec: None,
13448 })
13449 .unwrap();
13450 drop(tx);
13451 rx
13452 };
13453 let resp = blocking_response(
13455 feed(true),
13456 "m".into(),
13457 true,
13458 Vec::new(),
13459 Some(ToolStreamParser::reasoning_only()),
13460 Envelope::new(true),
13461 )
13462 .await;
13463 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13464 .await
13465 .unwrap();
13466 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13467 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
13468 assert_eq!(
13469 v["choices"][0]["message"]["reasoning_details"][0]["text"],
13470 "a plan"
13471 );
13472 assert_eq!(v["choices"][0]["message"]["content"], "answer");
13473 let resp = blocking_response(
13476 feed(false),
13477 "m".into(),
13478 true,
13479 Vec::new(),
13480 None,
13481 Envelope::new(true),
13482 )
13483 .await;
13484 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13485 .await
13486 .unwrap();
13487 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13488 assert!(
13489 v["choices"][0]["message"].get("reasoning").is_none(),
13490 "a reasoning-off response must carry no reasoning field: {v}"
13491 );
13492 assert_eq!(v["choices"][0]["message"]["content"], "answer");
13493 let resp = sse_response(
13495 feed(true),
13496 "m".into(),
13497 true,
13498 Some(ToolStreamParser::reasoning_only()),
13499 Envelope::new(true),
13500 Vec::new(),
13501 None,
13502 )
13503 .into_response();
13504 let lines = sse_data_lines(resp).await;
13505 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
13506 .iter()
13507 .map(|l| serde_json::from_str(l).unwrap())
13508 .collect();
13509 let reasoning: String = chunks
13510 .iter()
13511 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
13512 .collect();
13513 assert_eq!(
13514 reasoning, "a plan",
13515 "think text must stream as delta.reasoning"
13516 );
13517 let content: String = chunks
13518 .iter()
13519 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
13520 .collect();
13521 assert_eq!(content, "answer", "content must exclude the think segment");
13522 let resp = sse_response(
13524 feed(false),
13525 "m".into(),
13526 true,
13527 None,
13528 Envelope::new(true),
13529 Vec::new(),
13530 None,
13531 )
13532 .into_response();
13533 let lines = sse_data_lines(resp).await;
13534 for l in &lines[..lines.len() - 1] {
13535 let c: serde_json::Value = serde_json::from_str(l).unwrap();
13536 assert!(
13537 c["choices"][0]["delta"].get("reasoning").is_none(),
13538 "a reasoning-off stream must carry no reasoning deltas: {c}"
13539 );
13540 }
13541 }
13542
13543 #[tokio::test]
13544 async fn stream_chunks_carry_envelope_and_first_delta_role() {
13545 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13546 tx.send(Event::Token {
13547 id: 1,
13548 text: "he".into(),
13549 })
13550 .unwrap();
13551 tx.send(Event::Token {
13552 id: 2,
13553 text: "llo".into(),
13554 })
13555 .unwrap();
13556 tx.send(Event::Done {
13557 stop_reason: "Eos".into(),
13558 n_tokens: 2,
13559 n_prompt: 10,
13560 n_cached: 0,
13561 elapsed_s: 0.1,
13562 spec: None,
13563 })
13564 .unwrap();
13565 drop(tx);
13566 let resp = sse_response(
13567 rx,
13568 "m".into(),
13569 true,
13570 None,
13571 Envelope::new(true),
13572 Vec::new(),
13573 None,
13574 )
13575 .into_response();
13576 let lines = sse_data_lines(resp).await;
13577 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
13578 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
13579 .iter()
13580 .map(|l| serde_json::from_str(l).unwrap())
13581 .collect();
13582 let id = chunks[0]["id"].as_str().unwrap().to_string();
13584 assert!(id.starts_with("chatcmpl-"));
13585 for c in &chunks {
13586 assert_eq!(c["id"], id.as_str());
13587 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
13588 assert!(
13589 c["system_fingerprint"]
13590 .as_str()
13591 .unwrap()
13592 .starts_with("memra-")
13593 );
13594 assert_eq!(c["object"], "chat.completion.chunk");
13595 }
13596 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
13598 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
13599 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
13600 let fin = chunks.last().unwrap();
13602 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
13603 assert_eq!(fin["usage"]["prompt_tokens"], 10);
13604 }
13605
13606 #[tokio::test]
13607 async fn stream_token_events_equal_usage_on_every_finish_path() {
13608 for (stop_reason, expected_finish) in [
13609 ("Eos", "stop"),
13610 ("Callback", "stop"),
13611 ("MaxNew", "length"),
13612 ("ContextFull", "length"),
13613 ] {
13614 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13615 tx.send(Event::Token {
13618 id: 248_046,
13619 text: String::new(),
13620 })
13621 .unwrap();
13622 tx.send(Event::Done {
13623 stop_reason: stop_reason.into(),
13624 n_tokens: 1,
13625 n_prompt: 8,
13626 n_cached: 8,
13627 elapsed_s: 0.1,
13628 spec: None,
13629 })
13630 .unwrap();
13631 drop(tx);
13632
13633 let resp = sse_response(
13634 rx,
13635 "m".into(),
13636 true,
13637 None,
13638 Envelope::new(true),
13639 Vec::new(),
13640 None,
13641 )
13642 .into_response();
13643 let lines = sse_data_lines(resp).await;
13644 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
13645 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
13646 .iter()
13647 .map(|line| serde_json::from_str(line).unwrap())
13648 .collect();
13649 let token_events = chunks
13650 .iter()
13651 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
13652 .count();
13653 let terminal = chunks.last().unwrap();
13654 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
13655 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
13656 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
13657 }
13658 }
13659
13660 #[tokio::test]
13661 async fn stream_excludes_stop_text_like_non_stream_does() {
13662 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13666 tx.send(Event::Token {
13667 id: 1,
13668 text: "answer\nPro".into(),
13669 })
13670 .unwrap();
13671 tx.send(Event::Token {
13672 id: 2,
13673 text: "blem: leaked prompt".into(),
13674 })
13675 .unwrap();
13676 tx.send(Event::Done {
13677 stop_reason: "Callback".into(),
13678 n_tokens: 2,
13679 n_prompt: 8,
13680 n_cached: 0,
13681 elapsed_s: 0.1,
13682 spec: None,
13683 })
13684 .unwrap();
13685 drop(tx);
13686 let resp = sse_response(
13687 rx,
13688 "m".into(),
13689 true,
13690 None,
13691 Envelope::new(true),
13692 vec!["Problem:".into()],
13693 None,
13694 )
13695 .into_response();
13696 let lines = sse_data_lines(resp).await;
13697 let content: String = lines
13698 .iter()
13699 .filter(|l| *l != "[DONE]")
13700 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
13701 .filter_map(|c| {
13702 c["choices"][0]["delta"]["content"]
13703 .as_str()
13704 .map(str::to_string)
13705 })
13706 .collect();
13707 assert_eq!(content, "answer\n");
13708
13709 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13711 tx.send(Event::Token {
13712 id: 1,
13713 text: "ends in Pro".into(),
13714 })
13715 .unwrap();
13716 tx.send(Event::Done {
13717 stop_reason: "Eos".into(),
13718 n_tokens: 1,
13719 n_prompt: 8,
13720 n_cached: 0,
13721 elapsed_s: 0.1,
13722 spec: None,
13723 })
13724 .unwrap();
13725 drop(tx);
13726 let resp = sse_response(
13727 rx,
13728 "m".into(),
13729 true,
13730 None,
13731 Envelope::new(true),
13732 vec!["Problem:".into()],
13733 None,
13734 )
13735 .into_response();
13736 let lines = sse_data_lines(resp).await;
13737 let content: String = lines
13738 .iter()
13739 .filter(|l| *l != "[DONE]")
13740 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
13741 .filter_map(|c| {
13742 c["choices"][0]["delta"]["content"]
13743 .as_str()
13744 .map(str::to_string)
13745 })
13746 .collect();
13747 assert_eq!(content, "ends in Pro");
13748 }
13749
13750 #[tokio::test]
13751 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
13752 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13753 tx.send(Event::Error(worker::EngineError::engine("boom")))
13754 .unwrap();
13755 drop(tx);
13756 let resp = sse_response(
13757 rx,
13758 "m".into(),
13759 true,
13760 None,
13761 Envelope::new(true),
13762 Vec::new(),
13763 None,
13764 )
13765 .into_response();
13766 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13767 .await
13768 .unwrap();
13769 let body = String::from_utf8(bytes.to_vec()).unwrap();
13770 assert!(
13772 !body.contains("event: error"),
13773 "named SSE event leaked: {body}"
13774 );
13775 let lines: Vec<&str> = body
13776 .lines()
13777 .filter_map(|l| l.strip_prefix("data: "))
13778 .collect();
13779 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
13780 assert_eq!(err["error"]["message"], "boom");
13781 assert_eq!(err["error"]["type"], "server_error");
13782 assert_eq!(err["error"]["code"], "engine_error");
13783 assert_eq!(lines.last(), Some(&"[DONE]"));
13784 }
13785
13786 #[test]
13787 fn ttft_sse_marker_ignores_keepalive_comments() {
13788 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
13789 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
13790 assert!(is_sse_data_frame(
13791 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
13792 ));
13793 }
13794
13795 #[tokio::test]
13796 async fn error_bodies_use_the_openai_object_shape() {
13797 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13798 tx.send(Event::Error(worker::EngineError::model_not_found(
13799 "unknown model \"x\"",
13800 )))
13801 .unwrap();
13802 drop(tx);
13803 let response =
13804 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
13805 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
13806 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
13807 .await
13808 .unwrap();
13809 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13810 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
13812 assert_eq!(payload["error"]["type"], "invalid_request_error");
13813 assert_eq!(payload["error"]["param"], "model");
13814 assert_eq!(payload["error"]["code"], "model_not_found");
13815 }
13816
13817 fn retry_after(resp: &Response) -> Option<String> {
13824 resp.headers()
13825 .get(axum::http::header::RETRY_AFTER)
13826 .and_then(|v| v.to_str().ok())
13827 .map(str::to_string)
13828 }
13829
13830 async fn body_value(resp: Response) -> serde_json::Value {
13833 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13834 .await
13835 .expect("body");
13836 serde_json::from_slice(&bytes).expect("json body")
13837 }
13838
13839 #[test]
13840 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
13841 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
13843 assert_eq!(
13844 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
13845 TIMEOUT_MS_DEFAULT
13846 );
13847 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
13851 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
13852 }
13853 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
13855 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
13856 assert!(err.contains("timeout_ms"), "{err}");
13857 assert!(
13858 err.contains(&TIMEOUT_MS_MIN.to_string())
13859 && err.contains(&TIMEOUT_MS_MAX.to_string()),
13860 "the message must state the range: {err}"
13861 );
13862 assert!(
13863 err.contains("stream"),
13864 "the message must point at streaming for longer work: {err}"
13865 );
13866 }
13867 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
13869 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
13870 assert!(
13871 err.contains("timeout_ms") && err.contains("stream"),
13872 "{err}"
13873 );
13874 }
13875 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
13877 }
13878
13879 #[tokio::test]
13882 #[allow(clippy::await_holding_lock)] async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
13884 let _l = DRAIN_LOCK.lock().unwrap();
13885 let st = fake_worker_state();
13886
13887 let comp = completions(
13888 State(st.clone()),
13889 HeaderMap::new(),
13890 None,
13891 Json(
13892 serde_json::from_value(json!({
13893 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
13894 .unwrap(),
13895 ),
13896 )
13897 .await;
13898 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
13899 let chat = chat_completions(
13900 State(st.clone()),
13901 HeaderMap::new(),
13902 None,
13903 Json(
13904 serde_json::from_value(json!({
13905 "model": "m", "messages": [{"role": "user", "content": "t"}],
13906 "timeout_ms": 90_001}))
13907 .unwrap(),
13908 ),
13909 )
13910 .await;
13911 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
13912 let resp_api = responses_api::responses(
13913 State(st.clone()),
13914 HeaderMap::new(),
13915 None,
13916 axum::body::Bytes::from(
13917 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
13918 ),
13919 )
13920 .await;
13921 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
13922 let msgs = anthropic::messages(
13923 State(st.clone()),
13924 HeaderMap::new(),
13925 None,
13926 axum::body::Bytes::from(
13927 json!({"model": "m", "max_tokens": 16,
13928 "messages": [{"role": "user", "content": "t"}],
13929 "timeout_ms": 90_001})
13930 .to_string(),
13931 ),
13932 )
13933 .await;
13934 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
13935
13936 for (surface, resp) in [
13938 ("/v1/completions", comp),
13939 ("/v1/chat/completions", chat),
13940 ("/v1/responses", resp_api),
13941 ] {
13942 let body = body_value(resp).await;
13943 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
13944 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
13945 let m = body["error"]["message"].as_str().unwrap();
13946 assert!(
13947 m.contains("90000") && m.contains("stream"),
13948 "{surface}: {m}"
13949 );
13950 }
13951 let body = body_value(msgs).await;
13953 assert_eq!(body["error"]["type"], "invalid_request_error");
13954 let m = body["error"]["message"].as_str().unwrap();
13955 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
13956 }
13957
13958 #[tokio::test]
13961 #[allow(clippy::await_holding_lock)] async fn a_non_integer_timeout_ms_is_a_named_400() {
13963 let _l = DRAIN_LOCK.lock().unwrap();
13964 let st = fake_worker_state();
13965 let resp = chat_completions(
13966 State(st),
13967 HeaderMap::new(),
13968 None,
13969 Json(
13970 serde_json::from_value(json!({
13971 "model": "m", "messages": [{"role": "user", "content": "t"}],
13972 "timeout_ms": "30s"}))
13973 .unwrap(),
13974 ),
13975 )
13976 .await;
13977 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13978 let body = body_value(resp).await;
13979 assert_eq!(body["error"]["param"], "timeout_ms");
13980 }
13981
13982 #[tokio::test]
13988 #[allow(clippy::await_holding_lock)] async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
13990 let _l = DRAIN_LOCK.lock().unwrap();
13991 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13996 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
13997 let worker_cancel = cancel_seen.clone();
13998 let health = health::WorkerHealth::new();
13999 let h = health.clone();
14000 std::thread::spawn(move || {
14001 h.mark_ready();
14002 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
14003 worker::release_pending_admit();
14004 worker::release_admission_reservation(req.lane);
14005 let _ = req.tx.send(Event::PromptUsage {
14006 n_prompt: 1,
14007 n_cached: 0,
14008 });
14009 let _ = req.tx.send(Event::Token {
14010 id: 1,
14011 text: "partial".into(),
14012 });
14013 for _ in 0..5_000 {
14016 if req.tx.is_closed() {
14017 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
14018 break;
14019 }
14020 std::thread::sleep(std::time::Duration::from_millis(1));
14021 }
14022 }
14023 });
14024 for _ in 0..2_000 {
14025 if health.live().is_ok() {
14026 break;
14027 }
14028 std::thread::sleep(std::time::Duration::from_millis(1));
14029 }
14030 let mut st = fake_worker_state();
14031 st.cmd_tx = cmd_tx;
14032 st.health = health;
14033 let mock = MockMetering::admit_all();
14034 st.metering = Some(mock.clone());
14035
14036 let resp = chat_completions(
14037 State(st),
14038 HeaderMap::new(),
14039 None,
14040 Json(
14041 serde_json::from_value(json!({
14042 "model": "m", "messages": [{"role": "user", "content": "t"}],
14043 "timeout_ms": 1_000}))
14044 .unwrap(),
14045 ),
14046 )
14047 .await;
14048
14049 assert_eq!(resp.status(), StatusCode::OK);
14054 let body = body_value(resp).await;
14055 assert!(
14056 body["choices"][0]["message"]["content"]
14057 .as_str()
14058 .unwrap()
14059 .contains("partial"),
14060 "the tokens generated before the cut must be delivered: {body}"
14061 );
14062 assert_eq!(body["choices"][0]["finish_reason"], "error");
14066 assert_eq!(
14067 body["choices"][0]["native_finish_reason"],
14068 "deadline_exceeded"
14069 );
14070 assert_eq!(body["error"]["code"], "deadline_exceeded");
14071 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
14072 let message = body["error"]["message"].as_str().unwrap();
14073 assert!(
14074 message.contains("1000") && message.contains("stream"),
14075 "the partial must name the deadline and the streaming alternative: {message}"
14076 );
14077 assert_eq!(body["usage"]["completion_tokens"], 1);
14078
14079 let mut cancelled = false;
14084 for _ in 0..500 {
14085 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
14086 cancelled = true;
14087 break;
14088 }
14089 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
14090 }
14091 assert!(
14092 cancelled,
14093 "the deadline must CANCEL generation (worker's event channel closed)"
14094 );
14095
14096 let events = mock.events();
14101 assert!(
14102 events.contains(&MeterEvent::DeadlinePartial {
14103 prompt: 1,
14104 cached: 0,
14105 completion: 1,
14106 }),
14107 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
14108 );
14109 assert!(
14110 !events
14111 .iter()
14112 .any(|e| matches!(e, MeterEvent::Complete { .. })),
14113 "a deadline cut must stay distinguishable from a full answer: {events:?}"
14114 );
14115 }
14116
14117 #[tokio::test]
14121 #[allow(clippy::await_holding_lock)] async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
14123 let _l = DRAIN_LOCK.lock().unwrap();
14124 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
14125 let health = health::WorkerHealth::new();
14126 let h = health.clone();
14127 std::thread::spawn(move || {
14128 h.mark_ready();
14129 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
14132 worker::release_pending_admit();
14133 worker::release_admission_reservation(req.lane);
14134 let _ = req.tx.send(Event::PromptUsage {
14135 n_prompt: 1,
14136 n_cached: 0,
14137 });
14138 for _ in 0..5_000 {
14139 if req.tx.is_closed() {
14140 break;
14141 }
14142 std::thread::sleep(std::time::Duration::from_millis(1));
14143 }
14144 }
14145 });
14146 for _ in 0..2_000 {
14147 if health.live().is_ok() {
14148 break;
14149 }
14150 std::thread::sleep(std::time::Duration::from_millis(1));
14151 }
14152 let mut st = fake_worker_state();
14153 st.cmd_tx = cmd_tx;
14154 st.health = health;
14155 let mock = MockMetering::admit_all();
14156 st.metering = Some(mock.clone());
14157 let resp = chat_completions(
14158 State(st),
14159 HeaderMap::new(),
14160 None,
14161 Json(
14162 serde_json::from_value(json!({
14163 "model": "m", "messages": [{"role": "user", "content": "t"}],
14164 "timeout_ms": 1_000}))
14165 .unwrap(),
14166 ),
14167 )
14168 .await;
14169 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
14170 assert!(resp.headers().get("x-should-retry").is_none());
14172 assert_eq!(retry_after(&resp), None);
14173 let body = body_value(resp).await;
14174 assert_eq!(body["error"]["code"], "deadline_exceeded");
14175 assert!(
14176 body["error"]["message"]
14177 .as_str()
14178 .unwrap()
14179 .contains("not billed"),
14180 "the zero-token 408 keeps the billing promise: {body}"
14181 );
14182 let events = mock.events();
14183 assert!(
14184 events.contains(&MeterEvent::Unbilled {
14185 outcome: "deadline_exceeded",
14186 status: 408,
14187 code: "deadline_exceeded".into(),
14188 }),
14189 "the named zero-debit census outcome, not the generic reject — every sibling \
14190 deadline path settles this one: {events:?}"
14191 );
14192 }
14193
14194 #[tokio::test]
14197 #[allow(clippy::await_holding_lock)] async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
14199 let _l = DRAIN_LOCK.lock().unwrap();
14200 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
14202 let health = health::WorkerHealth::new();
14203 let h = health.clone();
14204 std::thread::spawn(move || {
14205 h.mark_ready();
14206 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
14207 worker::release_pending_admit();
14208 worker::release_admission_reservation(req.lane);
14209 let _ = req.tx.send(Event::PromptUsage {
14210 n_prompt: 1,
14211 n_cached: 0,
14212 });
14213 while !req.tx.is_closed() {
14214 std::thread::sleep(std::time::Duration::from_millis(1));
14215 }
14216 }
14217 });
14218 for _ in 0..2_000 {
14219 if health.live().is_ok() {
14220 break;
14221 }
14222 std::thread::sleep(std::time::Duration::from_millis(1));
14223 }
14224 let mut st = fake_worker_state();
14225 st.cmd_tx = cmd_tx;
14226 st.health = health;
14227 let mock = MockMetering::admit_all();
14228 st.metering = Some(mock.clone());
14229
14230 let resp = chat_completions(
14231 State(st),
14232 HeaderMap::new(),
14233 None,
14234 Json(
14235 serde_json::from_value(json!({
14236 "model": "m", "messages": [{"role": "user", "content": "t"}],
14237 "stream": true, "timeout_ms": 1_000}))
14238 .unwrap(),
14239 ),
14240 )
14241 .await;
14242 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
14245 let body = body_value(resp).await;
14246 assert_eq!(body["error"]["code"], "deadline_exceeded");
14247 assert!(
14248 body["error"]["message"]
14249 .as_str()
14250 .unwrap()
14251 .contains("first token"),
14252 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
14253 );
14254 let events = mock.events();
14255 assert!(
14256 events.contains(&MeterEvent::Unbilled {
14257 outcome: "deadline_exceeded",
14258 status: 408,
14259 code: "deadline_exceeded".into(),
14260 }),
14261 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
14262 );
14263 }
14264
14265 #[tokio::test]
14269 #[allow(clippy::await_holding_lock)] async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
14271 let _l = DRAIN_LOCK.lock().unwrap();
14272 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
14275 let mock = MockMetering::admit_all();
14276 st.metering = Some(mock.clone());
14277 let resp = chat_completions(
14278 State(st),
14279 HeaderMap::new(),
14280 None,
14281 Json(
14282 serde_json::from_value(json!({
14283 "model": "m", "messages": [{"role": "user", "content": "t"}],
14284 "stream": true, "timeout_ms": 1_000}))
14285 .unwrap(),
14286 ),
14287 )
14288 .await;
14289 assert_eq!(
14290 resp.status(),
14291 StatusCode::OK,
14292 "TTFT was met — 200 is correct"
14293 );
14294 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14295 .await
14296 .expect("the stream must run to completion past the deadline");
14297 let text = String::from_utf8(bytes.to_vec()).unwrap();
14298 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
14299 let events = mock.events();
14300 assert!(
14301 events
14302 .iter()
14303 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
14304 "a stream past its deadline after first token still settles as COMPLETE with \
14305 all four tokens: {events:?}"
14306 );
14307 }
14308
14309 fn admission_counters_guard() -> std::sync::MutexGuard<'static, ()> {
14316 static COUNTERS: std::sync::Mutex<()> = std::sync::Mutex::new(());
14317 COUNTERS
14318 .lock()
14319 .unwrap_or_else(|poisoned| poisoned.into_inner())
14320 }
14321
14322 #[test]
14325 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
14326 let _counters = admission_counters_guard();
14327 let st = fake_worker_state();
14328 let lane = lanes::Lane::Interactive;
14329 let cap = lane_cap(lane);
14330 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
14331 let prev = counter.swap(max_queue_depth(cap), std::sync::atomic::Ordering::AcqRel);
14332 let rl = RateLimit {
14333 limit: cap,
14334 remaining: 0,
14335 reset_s: 1,
14336 };
14337 let (resp, outcome) = reserve_pending_admit(
14338 &st,
14339 lane,
14340 &rl,
14341 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
14342 )
14343 .map(|_| ())
14344 .expect_err("a backlog at the bound must shed");
14345 counter.store(prev, std::sync::atomic::Ordering::Release);
14346 assert_eq!(outcome, "shed_queue");
14347 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
14348 assert!(
14349 retry_after(&resp).is_some(),
14350 "a shed must carry Retry-After so the router's spill can act on it"
14351 );
14352 let stamped = rl.attach(resp);
14354 for h in [
14355 "x-ratelimit-limit",
14356 "x-ratelimit-remaining",
14357 "x-ratelimit-reset",
14358 ] {
14359 assert!(stamped.headers().get(h).is_some(), "missing {h}");
14360 }
14361 }
14362
14363 #[test]
14367 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
14368 let _counters = admission_counters_guard();
14369 let st = fake_worker_state();
14370 let lane = lanes::Lane::Interactive;
14371 let cap = lane_cap(lane);
14372 {
14373 let mut m = st.metrics.lock().unwrap();
14374 m.completed = 10;
14375 m.tokens_out = 1_000;
14376 m.step_p50_ms = 10.0; }
14378 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
14379 let prev = counter.swap(cap, std::sync::atomic::Ordering::AcqRel); let rl = RateLimit {
14381 limit: cap,
14382 remaining: 0,
14383 reset_s: 1,
14384 };
14385 let admitted = reserve_pending_admit(
14387 &st,
14388 lane,
14389 &rl,
14390 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
14391 );
14392 assert!(
14393 admitted.is_ok(),
14394 "a request whose deadline covers the estimate must be admitted"
14395 );
14396 drop(admitted); let (resp, outcome) = reserve_pending_admit(
14399 &st,
14400 lane,
14401 &rl,
14402 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
14403 )
14404 .map(|_| ())
14405 .expect_err("a deadline shorter than the estimated wait must shed");
14406 counter.store(prev, std::sync::atomic::Ordering::Release);
14407 assert_eq!(outcome, "shed_deadline");
14408 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
14409 assert!(retry_after(&resp).is_some());
14410 }
14411
14412 #[test]
14415 fn deadline_shed_is_interactive_only_and_silent_with_free_slots() {
14416 let _counters = admission_counters_guard();
14417 let st = fake_worker_state();
14418 let cap = lane_cap(lanes::Lane::Interactive);
14419 {
14420 let mut m = st.metrics.lock().unwrap();
14421 m.completed = 10;
14422 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
14424 }
14425 let free = RateLimit {
14427 limit: cap,
14428 remaining: 1,
14429 reset_s: 0,
14430 };
14431 let g = reserve_pending_admit(
14432 &st,
14433 lanes::Lane::Interactive,
14434 &free,
14435 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
14436 );
14437 assert!(
14438 g.is_ok(),
14439 "free capacity must admit regardless of the estimate"
14440 );
14441 drop(g);
14442 let full = RateLimit {
14445 limit: cap,
14446 remaining: 0,
14447 reset_s: 5,
14448 };
14449 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
14450 let counter = &worker::ADMISSION_RESERVATIONS[lane.idx()];
14451 let prev = counter.swap(1, std::sync::atomic::Ordering::AcqRel); let g = reserve_pending_admit(
14453 &st,
14454 lane,
14455 &full,
14456 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
14457 );
14458 assert!(
14459 g.is_ok(),
14460 "{lane:?} must not be deadline-shed by the interactive gate"
14461 );
14462 drop(g);
14463 counter.store(prev, std::sync::atomic::Ordering::Release);
14464 }
14465 }
14466
14467 #[test]
14468 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
14469 let _counters = admission_counters_guard();
14470 let st = fake_worker_state();
14471 let cap = lane_cap(lanes::Lane::Interactive);
14472 let bound = max_queue_depth(cap);
14473 assert!(bound > 0, "the queue bound must admit at least one request");
14474 let rl = RateLimit {
14475 limit: cap,
14476 remaining: 0,
14477 reset_s: 1,
14478 };
14479 let _ = worker::PENDING_ADMITS.fetch_update(
14480 std::sync::atomic::Ordering::AcqRel,
14481 std::sync::atomic::Ordering::Acquire,
14482 |_| Some(0),
14483 );
14484 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
14485 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
14486 let guard = reserve_pending_admit(
14487 &st,
14488 lanes::Lane::Interactive,
14489 &rl,
14490 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
14491 )
14492 .expect("the final queue slot should be reservable");
14493 assert_eq!(
14494 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
14495 1
14496 );
14497 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
14498 drop(guard);
14499 assert_eq!(
14500 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
14501 0
14502 );
14503 assert_eq!(
14504 counter.load(std::sync::atomic::Ordering::Acquire),
14505 bound - 1
14506 );
14507
14508 counter.store(bound, std::sync::atomic::Ordering::Release);
14509 let rejected = reserve_pending_admit(
14510 &st,
14511 lanes::Lane::Interactive,
14512 &rl,
14513 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
14514 );
14515 assert!(matches!(rejected, Err((_, "shed_queue"))));
14516 counter.store(0, std::sync::atomic::Ordering::Release);
14517 }
14518
14519 #[test]
14520 fn admission_reservations_are_lane_scoped() {
14521 let _counters = admission_counters_guard();
14522 let st = fake_worker_state();
14523 let harvest = lanes::Lane::Harvest;
14524 let interactive = lanes::Lane::Interactive;
14525 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
14526 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
14527 harvest_counter.store(
14528 max_queue_depth(lane_cap(harvest)),
14529 std::sync::atomic::Ordering::Release,
14530 );
14531 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
14532 let free = RateLimit {
14533 limit: lane_cap(interactive),
14534 remaining: 1,
14535 reset_s: 0,
14536 };
14537 let guard = reserve_pending_admit(
14538 &st,
14539 interactive,
14540 &free,
14541 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
14542 )
14543 .expect("a full harvest queue must not consume interactive capacity");
14544 drop(guard);
14545 let harvest_rl = RateLimit {
14546 limit: lane_cap(harvest),
14547 remaining: 0,
14548 reset_s: 1,
14549 };
14550 assert!(matches!(
14551 reserve_pending_admit(
14552 &st,
14553 harvest,
14554 &harvest_rl,
14555 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
14556 ),
14557 Err((_, "shed_queue"))
14558 ));
14559 harvest_counter.store(0, std::sync::atomic::Ordering::Release);
14560 }
14561
14562 #[test]
14563 fn taxonomy_maps_every_class_to_its_status_and_code() {
14564 use worker::{EngineError as E, ErrClass as C};
14565 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
14566 (
14567 E::invalid_param("bad json", "response_format"),
14568 StatusCode::BAD_REQUEST,
14569 "invalid_request_error",
14570 "",
14571 ),
14572 (
14573 E::context_length("prompt (9000 tok) >= context cap (8192)"),
14574 StatusCode::BAD_REQUEST,
14575 "invalid_request_error",
14576 "context_length_exceeded",
14577 ),
14578 (
14579 E::model_not_found("unknown model \"nope\""),
14580 StatusCode::BAD_REQUEST,
14581 "invalid_request_error",
14582 "model_not_found",
14583 ),
14584 (
14585 E::rate_limit("lane judge is at capacity, retry"),
14586 StatusCode::TOO_MANY_REQUESTS,
14587 "rate_limit_error",
14588 "rate_limit_exceeded",
14589 ),
14590 (
14591 E::overloaded("no VRAM for a new session"),
14592 StatusCode::SERVICE_UNAVAILABLE,
14593 "server_error",
14594 "overloaded",
14595 ),
14596 (
14597 E::engine("graph step failed: launch error"),
14598 StatusCode::INTERNAL_SERVER_ERROR,
14599 "server_error",
14600 "engine_error",
14601 ),
14602 ];
14603 for (err, want_status, want_type, want_code) in cases {
14604 let (status, etype, code) = class_http(err.class);
14605 assert_eq!(status, want_status, "{:?}", err);
14606 assert_eq!(etype, want_type, "{:?}", err);
14607 if !want_code.is_empty() {
14608 assert_eq!(code, Some(want_code), "{:?}", err);
14609 }
14610 let body = engine_error_body(&err);
14612 assert_eq!(body["error"]["message"], err.message);
14613 assert_eq!(body["error"]["type"], want_type);
14614 }
14615 for c in [
14617 C::InvalidRequest,
14618 C::ContextLength,
14619 C::ModelNotFound,
14620 C::RateLimit,
14621 C::Overloaded,
14622 C::Engine,
14623 ] {
14624 let (s, t, _) = class_http(c);
14625 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
14626 assert!(!t.is_empty());
14627 }
14628 }
14629
14630 #[test]
14631 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
14632 let e = worker::EngineError::engine(
14637 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
14638 );
14639 let resp = engine_error_response(&e);
14640 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
14641 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
14642 }
14643
14644 #[test]
14645 fn retry_headers_follow_the_sdk_contract() {
14646 for e in [
14650 worker::EngineError::rate_limit("shed"),
14651 worker::EngineError::overloaded("no VRAM"),
14652 ] {
14653 let resp = engine_error_response(&e);
14654 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
14655 let secs: u64 = ra
14656 .parse()
14657 .expect("Retry-After must be integer delay-seconds");
14658 assert!(
14659 secs > 0 && secs <= 60,
14660 "Retry-After {secs}s outside the honored window"
14661 );
14662 let ms = resp
14663 .headers()
14664 .get("retry-after-ms")
14665 .unwrap()
14666 .to_str()
14667 .unwrap();
14668 assert_eq!(
14669 ms.parse::<u64>().unwrap(),
14670 secs * 1000,
14671 "the two headers disagree"
14672 );
14673 assert!(
14674 resp.headers().get("x-should-retry").is_none(),
14675 "a retryable class must not say x-should-retry: false"
14676 );
14677 }
14678 }
14679
14680 #[tokio::test]
14687 async fn admit_predict_reject_matches_shed_contract() {
14688 let shed = retry_contract_response(
14690 (
14691 StatusCode::TOO_MANY_REQUESTS,
14692 Json(error_body(
14693 "interactive queue is at its bound",
14694 "rate_limit_error",
14695 None,
14696 Some("shed_queue"),
14697 )),
14698 )
14699 .into_response(),
14700 Some(7),
14701 );
14702 let predict = engine_error_response(&worker::EngineError::rate_limit_after(
14705 "predicted KV-to-completion exceeds the box budget; retry",
14706 7,
14707 ));
14708 assert_eq!(shed.status(), predict.status());
14709 for header in ["retry-after", "retry-after-ms"] {
14710 assert_eq!(
14711 shed.headers().get(header),
14712 predict.headers().get(header),
14713 "header {header} must be byte-identical to the shed contract"
14714 );
14715 }
14716 let shed_body: serde_json::Value = serde_json::from_slice(
14717 &axum::body::to_bytes(shed.into_body(), usize::MAX)
14718 .await
14719 .unwrap(),
14720 )
14721 .unwrap();
14722 let predict_body: serde_json::Value = serde_json::from_slice(
14723 &axum::body::to_bytes(predict.into_body(), usize::MAX)
14724 .await
14725 .unwrap(),
14726 )
14727 .unwrap();
14728 assert_eq!(shed_body["error"]["type"], predict_body["error"]["type"]);
14729 assert_eq!(predict_body["error"]["type"], "rate_limit_error");
14730 let shed_keys: Vec<&String> = shed_body["error"].as_object().unwrap().keys().collect();
14731 let predict_keys: Vec<&String> =
14732 predict_body["error"].as_object().unwrap().keys().collect();
14733 assert_eq!(shed_keys, predict_keys, "same body schema, key for key");
14734 assert_eq!(predict_body["error"]["code"], "rate_limit_exceeded");
14735
14736 let clamped = engine_error_response(&worker::EngineError::rate_limit_after("m", 400));
14738 assert_eq!(retry_after(&clamped).as_deref(), Some("60"));
14739 let plain = engine_error_response(&worker::EngineError::rate_limit("m"));
14741 assert_eq!(retry_after(&plain).as_deref(), Some("2"));
14742 }
14743
14744 #[tokio::test]
14745 #[allow(clippy::await_holding_lock)] async fn command_send_failure_obeys_the_retry_contract() {
14747 let _l = DRAIN_LOCK.lock().unwrap();
14748 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
14749 let mut st = fake_worker_state();
14750 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
14751 drop(cmd_rx);
14752 st.cmd_tx = cmd_tx;
14753
14754 let completion = completions(
14755 State(st.clone()),
14756 axum::http::HeaderMap::new(),
14757 None,
14758 Json(
14759 serde_json::from_value(serde_json::json!({
14760 "model": "m", "prompt": "test"
14761 }))
14762 .unwrap(),
14763 ),
14764 )
14765 .await;
14766 let chat = chat_completions(
14767 State(st),
14768 axum::http::HeaderMap::new(),
14769 None,
14770 Json(
14771 serde_json::from_value(serde_json::json!({
14772 "model": "m", "messages": [{"role": "user", "content": "test"}]
14773 }))
14774 .unwrap(),
14775 ),
14776 )
14777 .await;
14778
14779 for resp in [completion, chat] {
14780 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
14781 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
14782 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
14783 assert_ne!(
14784 resp.headers()
14785 .get("x-should-retry")
14786 .and_then(|v| v.to_str().ok()),
14787 Some("false")
14788 );
14789 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14790 .await
14791 .unwrap();
14792 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14793 assert_eq!(payload["error"]["type"], "server_error");
14794 assert_eq!(payload["error"]["code"], "overloaded");
14795 }
14796 }
14797
14798 #[test]
14799 fn unfixable_client_errors_say_x_should_retry_false() {
14800 for e in [
14803 worker::EngineError::model_not_found("unknown model \"x\""),
14804 worker::EngineError::context_length("prompt too long"),
14805 worker::EngineError::invalid_param("bad", "messages"),
14806 ] {
14807 let resp = engine_error_response(&e);
14808 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
14809 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
14810 assert!(
14811 retry_after(&resp).is_none(),
14812 "a 400 must not promise a retry window"
14813 );
14814 }
14815 }
14816
14817 #[tokio::test]
14818 async fn a_closed_worker_channel_is_503_not_500() {
14819 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
14823 drop(tx);
14824 let resp =
14825 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
14826 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
14827 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
14828 }
14829
14830 #[tokio::test]
14831 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
14832 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
14835 tx.send(Event::Error(worker::EngineError::rate_limit(
14836 "lane judge shed: interactive p99 over budget, retry",
14837 )))
14838 .unwrap();
14839 let (resp, error_code) = peek_admission(rx)
14840 .await
14841 .expect_err("a shed must not be forwarded into the stream");
14842 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
14843 assert_eq!(error_code, "rate_limit_exceeded");
14844 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
14845 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
14846 .await
14847 .unwrap();
14848 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
14849 assert!(
14850 payload["error"].is_object(),
14851 "bare-string error body: {payload}"
14852 );
14853 assert_eq!(payload["error"]["type"], "rate_limit_error");
14854 assert!(
14855 payload["error"]["message"]
14856 .as_str()
14857 .unwrap()
14858 .contains("shed")
14859 );
14860 }
14861
14862 #[tokio::test]
14863 async fn interactive_admission_error_is_a_preheader_429() {
14864 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
14867 tx.send(Event::Error(worker::EngineError::rate_limit(
14868 "KV capacity unavailable",
14869 )))
14870 .unwrap();
14871 let (resp, error_code) = peek_admission(rx)
14872 .await
14873 .expect_err("admission error must stay pre-header");
14874 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
14875 assert_eq!(error_code, "rate_limit_exceeded");
14876 }
14877
14878 #[tokio::test]
14879 async fn admission_peek_preserves_context_error_for_the_ledger() {
14880 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
14881 tx.send(Event::Error(worker::EngineError::context_length(
14882 "prompt exceeds configured model maximum",
14883 )))
14884 .unwrap();
14885 let (resp, error_code) = peek_admission(rx)
14886 .await
14887 .expect_err("context rejection must stay pre-header");
14888 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
14889 assert_eq!(error_code, "context_length_exceeded");
14890 }
14891
14892 #[tokio::test]
14893 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
14894 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
14895 tx.send(Event::PromptUsage {
14896 n_prompt: 262_143,
14897 n_cached: 0,
14898 })
14899 .unwrap();
14900 let mut replay = peek_admission(rx).await.expect("successful admission");
14901 assert!(matches!(
14902 replay.recv().await,
14903 Some(Event::PromptUsage {
14904 n_prompt: 262_143,
14905 n_cached: 0
14906 }),
14907 ));
14908 }
14909
14910 #[test]
14911 fn penalties_plumb_from_http_to_sampler_config() {
14912 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
14915 "model": "m", "messages": [{"role": "user", "content": "task"}],
14916 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
14917 }))
14918 .unwrap();
14919 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14920 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14921 .unwrap()
14922 .request
14923 .sampler_cfg;
14924 assert_eq!(cfg.penalty_freq, 0.5);
14925 assert_eq!(cfg.penalty_present, 0.25);
14926 assert_eq!(cfg.penalty_repeat, 1.1);
14927 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14928
14929 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14930 "model": "m", "prompt": "task", "frequency_penalty": 1.5
14931 }))
14932 .unwrap();
14933 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14934 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
14935 assert_eq!(cfg.penalty_freq, 1.5);
14936 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14937
14938 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14940 "model": "m", "prompt": "task"
14941 }))
14942 .unwrap();
14943 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14944 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
14945 assert_eq!(cfg.penalty_last_n, 0);
14946 assert_eq!(cfg.penalty_repeat, 1.0);
14947 }
14948
14949 #[test]
14950 fn omitted_temperature_is_openai_default_not_greedy() {
14951 let chat_temp = |body: serde_json::Value| {
14967 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14968 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14969 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14970 .unwrap()
14971 .request
14972 .sampler_cfg
14973 .temperature
14974 };
14975 let comp_temp = |body: serde_json::Value| {
14976 let req: CompletionReq = serde_json::from_value(body).unwrap();
14977 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14978 build_request(&req, tx, lanes::Lane::Interactive, None)
14979 .sampler_cfg
14980 .temperature
14981 };
14982
14983 assert_eq!(
14985 chat_temp(serde_json::json!({
14986 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
14987 1.0,
14988 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
14989 );
14990 assert_eq!(
14991 comp_temp(serde_json::json!({
14992 "model": "m", "prompt": "t"})),
14993 1.0,
14994 "omitted completions temperature must be the OpenAI 1.0 default"
14995 );
14996
14997 assert_eq!(
14999 chat_temp(serde_json::json!({
15000 "model": "m", "messages": [{"role": "user", "content": "t"}],
15001 "temperature": 0.0})),
15002 0.0,
15003 "explicit temperature 0 must stay greedy"
15004 );
15005 assert_eq!(
15006 comp_temp(serde_json::json!({
15007 "model": "m", "prompt": "t", "temperature": 0})),
15008 0.0,
15009 "explicit temperature 0 must stay greedy"
15010 );
15011 assert!(
15013 memra_engine::sampler::Sampler::new(sampler_config(
15014 0.0,
15015 0,
15016 1.0,
15017 0.0,
15018 0.0,
15019 0.0,
15020 1.0,
15021 Some(0)
15022 ))
15023 .is_greedy()
15024 );
15025 assert!(
15026 !memra_engine::sampler::Sampler::new(sampler_config(
15027 1.0,
15028 0,
15029 1.0,
15030 0.0,
15031 0.0,
15032 0.0,
15033 1.0,
15034 Some(0)
15035 ))
15036 .is_greedy()
15037 );
15038
15039 assert_eq!(
15041 chat_temp(serde_json::json!({
15042 "model": "m", "messages": [{"role": "user", "content": "t"}],
15043 "temperature": 0.7})),
15044 0.7
15045 );
15046
15047 let req: CompletionReq = serde_json::from_value(serde_json::json!({
15051 "model": "m", "prompt": "t"}))
15052 .unwrap();
15053 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15054 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
15055 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
15056 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
15057 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
15058 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
15059 assert!(
15064 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
15065 "the omitted-temperature default must ride sampled spec's pure-temp regime"
15066 );
15067 }
15068
15069 #[test]
15070 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
15071 let caps = ModelCaps {
15072 chat_temperature_default: Some(0.5),
15073 chat_top_p_default: Some(0.9),
15074 chat_ok: true,
15075 ..Default::default()
15076 };
15077 let cfg = |extra: serde_json::Value| {
15078 let mut body = serde_json::json!({
15079 "model": "step35",
15080 "messages": [{"role": "user", "content": "task"}]
15081 });
15082 body.as_object_mut()
15083 .unwrap()
15084 .extend(extra.as_object().unwrap().clone());
15085 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15086 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15087 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
15088 .unwrap()
15089 .request
15090 .sampler_cfg
15091 };
15092
15093 let omitted = cfg(serde_json::json!({}));
15094 assert_eq!(omitted.temperature, 0.5);
15095 assert_eq!(omitted.top_p, 0.9);
15096
15097 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
15098 assert_eq!(explicit_temp.temperature, 0.7);
15099 assert_eq!(
15100 explicit_temp.top_p, 0.9,
15101 "omitting top_p must retain StepFun's nucleus default"
15102 );
15103
15104 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
15105 assert_eq!(
15106 explicit.temperature, 0.0,
15107 "explicit greedy must remain authoritative"
15108 );
15109 assert_eq!(
15110 explicit.top_p, 1.0,
15111 "explicit untruncated sampling must remain authoritative"
15112 );
15113 }
15114
15115 fn qwen38_vendor_defaults() -> SamplingDefaults {
15119 SamplingDefaults {
15120 temperature: Some(1.0),
15121 top_p: Some(0.95),
15122 top_k: Some(20),
15123 min_p: Some(0.0),
15124 presence_penalty: Some(0.0),
15125 repetition_penalty: Some(1.0),
15126 frequency_penalty: None,
15127 }
15128 }
15129
15130 fn gemma4_vendor_defaults() -> SamplingDefaults {
15135 SamplingDefaults {
15136 temperature: Some(1.0),
15137 top_p: Some(0.95),
15138 top_k: Some(64),
15139 ..Default::default()
15140 }
15141 }
15142
15143 #[test]
15144 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
15145 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
15150 let chat = |extra: serde_json::Value| {
15151 let mut body = serde_json::json!({
15152 "model": "google/gemma-4-31b-it",
15153 "messages": [{"role": "user", "content": "task"}],
15154 "seed": 7
15156 });
15157 body.as_object_mut()
15158 .unwrap()
15159 .extend(extra.as_object().unwrap().clone());
15160 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15161 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15162 build_chat_request_with_trace(
15163 req,
15164 Some(&ModelCaps {
15165 chat_ok: true,
15166 ..Default::default()
15167 }),
15168 tx,
15169 lanes::Lane::Interactive,
15170 None,
15171 None,
15172 None,
15173 &d,
15174 )
15175 .unwrap()
15176 .request
15177 .sampler_cfg
15178 };
15179
15180 let omitted = chat(serde_json::json!({}));
15182 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
15183 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
15184 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
15185 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
15187 assert_eq!(omitted.penalty_repeat, 1.0);
15188 assert_eq!(omitted.penalty_freq, 0.0);
15189 assert_eq!(omitted.penalty_present, 0.0);
15190 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
15191 assert!(
15192 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
15193 "the vendor default must NOT be greedy — that is the whole point of the lane"
15194 );
15195
15196 let greedy = chat(serde_json::json!({"temperature": 0}));
15199 assert_eq!(
15200 greedy.temperature, 0.0,
15201 "explicit temperature 0 stays greedy"
15202 );
15203 assert!(
15204 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
15205 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
15206 spec/graph exactness arms"
15207 );
15208
15209 let one_field = chat(serde_json::json!({"top_k": 3}));
15211 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
15212 assert_eq!(
15213 one_field.temperature, 1.0,
15214 "omitting temperature still takes the vendor value"
15215 );
15216 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
15217
15218 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
15221 assert_eq!(
15222 disabled.top_k, 0,
15223 "an explicit top_k 0 means KEEP ALL, not 'unset'"
15224 );
15225 assert_eq!(
15226 disabled.top_p, 1.0,
15227 "an explicit top_p 1.0 means untruncated"
15228 );
15229
15230 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
15232 assert_eq!(penal.penalty_present, 1.5);
15233 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
15234 }
15235
15236 #[test]
15237 fn vendor_sampling_defaults_are_identical_on_every_surface() {
15238 let d = qwen38_vendor_defaults();
15250 let md = ModelSamplingDefaults::single(d);
15251 let comp = |extra: serde_json::Value| {
15252 let mut body = serde_json::json!({
15253 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
15254 body.as_object_mut()
15255 .unwrap()
15256 .extend(extra.as_object().unwrap().clone());
15257 let req: CompletionReq = serde_json::from_value(body).unwrap();
15258 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15259 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
15260 };
15261 let chat = |extra: serde_json::Value| {
15262 let mut body = serde_json::json!({
15263 "model": "qwen/qwen3.8-27b",
15264 "messages": [{"role": "user", "content": "task"}],
15265 "seed": 11 });
15266 body.as_object_mut()
15267 .unwrap()
15268 .extend(extra.as_object().unwrap().clone());
15269 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15270 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15271 build_chat_request_with_trace(
15272 req,
15273 Some(&ModelCaps {
15274 chat_ok: true,
15275 ..Default::default()
15276 }),
15277 tx,
15278 lanes::Lane::Interactive,
15279 None,
15280 None,
15281 None,
15282 &md,
15283 )
15284 .unwrap()
15285 .request
15286 .sampler_cfg
15287 };
15288
15289 for extra in [
15290 serde_json::json!({}),
15291 serde_json::json!({"temperature": 0}),
15292 serde_json::json!({"temperature": 0.0}),
15293 serde_json::json!({"temperature": 0.7}),
15294 serde_json::json!({"top_p": 1.0}),
15295 serde_json::json!({"top_k": 0}),
15296 serde_json::json!({"min_p": 0.05}),
15297 serde_json::json!({"repetition_penalty": 1.1}),
15298 serde_json::json!({"frequency_penalty": 0.5}),
15299 serde_json::json!({"presence_penalty": 1.5}),
15300 serde_json::json!({
15301 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
15302 "frequency_penalty": 0.1, "presence_penalty": 0.2,
15303 "repetition_penalty": 1.05 }),
15304 ] {
15305 let c = comp(extra.clone());
15306 let h = chat(extra.clone());
15307 assert_eq!(
15308 (
15309 c.temperature,
15310 c.top_p,
15311 c.top_k,
15312 c.min_p,
15313 c.penalty_repeat,
15314 c.penalty_freq,
15315 c.penalty_present,
15316 c.penalty_last_n,
15317 c.seed
15318 ),
15319 (
15320 h.temperature,
15321 h.top_p,
15322 h.top_k,
15323 h.min_p,
15324 h.penalty_repeat,
15325 h.penalty_freq,
15326 h.penalty_present,
15327 h.penalty_last_n,
15328 h.seed
15329 ),
15330 "/v1/completions and /v1/chat/completions disagree on {extra} — \
15331 standard-surface-law violation"
15332 );
15333 }
15334
15335 let omitted = comp(serde_json::json!({}));
15337 assert_eq!(
15338 omitted.temperature, 1.0,
15339 "qwen3.8 card thinking temperature"
15340 );
15341 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
15342 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
15343 assert!(
15345 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
15346 .is_greedy()
15347 );
15348 }
15349
15350 #[tokio::test]
15363 #[allow(clippy::await_holding_lock)] async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
15365 let _l = DRAIN_LOCK.lock().unwrap();
15366 let step_caps = ModelCaps {
15367 chat_ok: true,
15368 chat_temperature_default: Some(0.5),
15369 chat_top_p_default: Some(0.9),
15370 ..Default::default()
15371 };
15372 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
15373 let st = fake_worker_state_full(
15374 1,
15375 std::time::Duration::ZERO,
15376 HashMap::from([("m".to_string(), step_caps)]),
15377 Some(cfg_tx),
15378 );
15379 let fields = |saw: &WorkerSaw| {
15384 let c = &saw.sampler_cfg;
15385 (
15386 c.temperature,
15387 c.top_p,
15388 c.top_k,
15389 c.min_p,
15390 c.penalty_repeat,
15391 c.penalty_freq,
15392 c.penalty_present,
15393 c.penalty_last_n,
15394 )
15395 };
15396 let worker_saw = |surface: &str| {
15397 cfg_rx
15398 .recv_timeout(std::time::Duration::from_secs(10))
15399 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
15400 };
15401
15402 let resp = completions(
15403 State(st.clone()),
15404 axum::http::HeaderMap::new(),
15405 None,
15406 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
15407 )
15408 .await;
15409 assert_eq!(
15410 resp.status(),
15411 StatusCode::OK,
15412 "/v1/completions rejected the omitted-sampling request"
15413 );
15414 let comp = worker_saw("/v1/completions");
15415
15416 let resp = chat_completions(
15417 State(st.clone()),
15418 axum::http::HeaderMap::new(),
15419 None,
15420 Json(
15421 serde_json::from_value(serde_json::json!({
15422 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
15423 .unwrap(),
15424 ),
15425 )
15426 .await;
15427 assert_eq!(
15428 resp.status(),
15429 StatusCode::OK,
15430 "/v1/chat/completions rejected the omitted-sampling request"
15431 );
15432 let chat = worker_saw("/v1/chat/completions");
15433
15434 let resp = anthropic::messages(
15435 State(st.clone()),
15436 axum::http::HeaderMap::new(),
15437 None,
15438 axum::body::Bytes::from(
15439 serde_json::json!({
15440 "model": "m", "max_tokens": 16,
15441 "messages": [{"role": "user", "content": "t"}]})
15442 .to_string(),
15443 ),
15444 )
15445 .await;
15446 assert_eq!(
15447 resp.status(),
15448 StatusCode::OK,
15449 "/v1/messages rejected the omitted-sampling request"
15450 );
15451 let msg = worker_saw("/v1/messages");
15452
15453 let resp = responses_api::responses(
15454 State(st.clone()),
15455 axum::http::HeaderMap::new(),
15456 None,
15457 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
15458 )
15459 .await;
15460 assert_eq!(
15461 resp.status(),
15462 StatusCode::OK,
15463 "/v1/responses rejected the omitted-sampling request"
15464 );
15465 let rsp = worker_saw("/v1/responses");
15466
15467 for (surface, cfg) in [
15468 ("/v1/completions", &comp),
15469 ("/v1/messages", &msg),
15470 ("/v1/responses", &rsp),
15471 ] {
15472 assert_eq!(
15473 fields(cfg),
15474 fields(&chat),
15475 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
15476 for the same omitted-sampling request — standard-surface-law violation \
15477 (hermes d991b51699218285)"
15478 );
15479 }
15480 assert_eq!(
15483 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
15484 (0.5, 0.9),
15485 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
15486 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
15487 );
15488 }
15489
15490 #[tokio::test]
15500 #[allow(clippy::await_holding_lock)] async fn same_effort_value_resolves_identically_on_every_surface() {
15502 let _l = DRAIN_LOCK.lock().unwrap();
15503 let caps = ModelCaps {
15506 chat_ok: true,
15507 effort_levels: true,
15508 ..Default::default()
15509 };
15510 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
15511 let st = fake_worker_state_full(
15512 1,
15513 std::time::Duration::ZERO,
15514 HashMap::from([("m".to_string(), caps)]),
15515 Some(saw_tx),
15516 );
15517 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
15518 match surface {
15519 "/v1/chat/completions" => {
15520 chat_completions(
15521 State(st),
15522 axum::http::HeaderMap::new(),
15523 None,
15524 Json(
15525 serde_json::from_value(serde_json::json!({
15526 "model": "m", "max_tokens": 8,
15527 "reasoning_effort": effort,
15528 "messages": [{"role": "user", "content": "t"}]}))
15529 .unwrap(),
15530 ),
15531 )
15532 .await
15533 }
15534 "/v1/responses" => {
15535 responses_api::responses(
15536 State(st),
15537 axum::http::HeaderMap::new(),
15538 None,
15539 axum::body::Bytes::from(
15540 serde_json::json!({
15541 "model": "m", "max_output_tokens": 8, "input": "t",
15542 "reasoning": {"effort": effort}})
15543 .to_string(),
15544 ),
15545 )
15546 .await
15547 }
15548 "/v1/messages" => {
15549 anthropic::messages(
15550 State(st),
15551 axum::http::HeaderMap::new(),
15552 None,
15553 axum::body::Bytes::from(
15554 serde_json::json!({
15555 "model": "m", "max_tokens": 8,
15556 "messages": [{"role": "user", "content": "t"}],
15557 "output_config": {"effort": effort}})
15558 .to_string(),
15559 ),
15560 )
15561 .await
15562 }
15563 other => panic!("unknown surface {other}"),
15564 }
15565 };
15566 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
15567
15568 for (effort, want_think, want_level) in [
15571 ("none", ThinkMode::NoThink, Some("low")),
15572 ("minimal", ThinkMode::NoThink, Some("low")),
15573 ("low", ThinkMode::Think, Some("low")),
15574 ("medium", ThinkMode::Think, Some("medium")),
15575 ("high", ThinkMode::Think, Some("high")),
15576 ("xhigh", ThinkMode::Think, Some("high")),
15578 ] {
15579 for surface in SURFACES {
15580 let resp = send(st.clone(), surface, effort).await;
15581 assert_eq!(
15582 resp.status(),
15583 StatusCode::OK,
15584 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
15585 diverged again (issue #31)"
15586 );
15587 let saw = saw_rx
15588 .recv_timeout(std::time::Duration::from_secs(10))
15589 .unwrap_or_else(|_| {
15590 panic!("{surface}: effort {effort:?} request never reached the worker")
15591 });
15592 assert_eq!(
15593 (saw.think, saw.reasoning_effort.as_deref()),
15594 (want_think, want_level),
15595 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
15596 reasoning surface — the parameter was dropped or remapped before \
15597 parse_think (issue #31 regression)"
15598 );
15599 }
15600 }
15601
15602 for effort in ["bogus", "banana", ""] {
15605 for surface in SURFACES {
15606 let resp = send(st.clone(), surface, effort).await;
15607 assert_eq!(
15608 resp.status(),
15609 StatusCode::BAD_REQUEST,
15610 "{surface} accepted effort {effort:?} — silent-accept regression \
15611 (issue #31: the value never reached parse_think's allowlist)"
15612 );
15613 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
15615 .await
15616 .unwrap();
15617 let v: serde_json::Value = serde_json::from_slice(&body)
15618 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
15619 match surface {
15620 "/v1/messages" => {
15621 assert_eq!(v["type"], "error", "{surface} error envelope");
15622 assert_eq!(
15623 v["error"]["type"], "invalid_request_error",
15624 "{surface} error type"
15625 );
15626 }
15627 _ => {
15628 assert!(
15629 v["error"]["message"].is_string(),
15630 "{surface} OpenAI-shaped error body: {v}"
15631 );
15632 }
15633 }
15634 }
15635 }
15636
15637 let resp = anthropic::messages(
15641 State(st.clone()),
15642 axum::http::HeaderMap::new(),
15643 None,
15644 axum::body::Bytes::from(
15645 serde_json::json!({
15646 "model": "m", "max_tokens": 8,
15647 "messages": [{"role": "user", "content": "t"}],
15648 "thinking": {"type": "enabled"},
15649 "output_config": {"effort": "none"}})
15650 .to_string(),
15651 ),
15652 )
15653 .await;
15654 assert_eq!(resp.status(), StatusCode::OK);
15655 let saw = saw_rx
15656 .recv_timeout(std::time::Duration::from_secs(10))
15657 .expect("thinking+effort request never reached the worker");
15658 assert_eq!(
15659 saw.think,
15660 ThinkMode::Think,
15661 "thinking.type (the documented Anthropic lever) must win the switch over \
15662 output_config.effort"
15663 );
15664 let resp = anthropic::messages(
15665 State(st.clone()),
15666 axum::http::HeaderMap::new(),
15667 None,
15668 axum::body::Bytes::from(
15669 serde_json::json!({
15670 "model": "m", "max_tokens": 8,
15671 "messages": [{"role": "user", "content": "t"}],
15672 "thinking": {"type": "enabled"},
15673 "output_config": {"effort": "banana"}})
15674 .to_string(),
15675 ),
15676 )
15677 .await;
15678 assert_eq!(
15679 resp.status(),
15680 StatusCode::BAD_REQUEST,
15681 "an invalid effort must 400 even next to an explicit thinking.type — \
15682 precedence must not re-open the silent-accept hole"
15683 );
15684 }
15685
15686 #[test]
15687 fn vendor_sampling_defaults_are_boot_validated() {
15688 let parsed = OpenRouterMetadataFile::from_toml(
15691 r#"
15692[models.g]
15693default_temperature = 1.0
15694default_top_p = 0.95
15695default_top_k = 64
15696default_min_p = 0.0
15697default_presence_penalty = 0.0
15698default_frequency_penalty = 0.0
15699default_repetition_penalty = 1.0
15700"#,
15701 )
15702 .unwrap();
15703 let g = parsed.get("g").unwrap();
15704 assert_eq!(g.default_temperature, Some(1.0));
15705 assert_eq!(g.default_top_p, Some(0.95));
15706 assert_eq!(g.default_top_k, Some(64));
15707
15708 let err = OpenRouterMetadataFile::from_toml(
15712 r#"
15713[models.g]
15714default_temperature = 0.0
15715"#,
15716 )
15717 .unwrap_err();
15718 assert!(err.contains("default_temperature"), "{err}");
15719 assert!(
15720 err.contains("greedy"),
15721 "the refusal must say WHY a zero default is refused: {err}"
15722 );
15723
15724 for bad in [
15725 "default_temperature = 2.5",
15726 "default_temperature = -1.0",
15727 "default_top_p = 0.0",
15728 "default_top_p = 1.5",
15729 "default_min_p = 1.0",
15730 "default_min_p = -0.1",
15731 "default_presence_penalty = 3.0",
15732 "default_frequency_penalty = -2.5",
15733 "default_repetition_penalty = 0.0",
15734 ] {
15735 let err =
15736 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
15737 let key = bad.split(' ').next().unwrap();
15738 assert!(err.contains(key), "{bad} must be refused by name: {err}");
15739 }
15740
15741 let err = OpenRouterMetadataFile::from_toml(
15745 r#"
15746[models.g]
15747default_temperture = 1.0
15748"#,
15749 )
15750 .unwrap_err();
15751 assert!(
15752 err.contains("unknown field"),
15753 "an unknown key must be fatal, which is what makes binary-first ordering \
15754 mandatory: {err}"
15755 );
15756 }
15757
15758 #[test]
15759 fn non_thinking_sampling_arm_is_boot_validated() {
15760 let parsed = OpenRouterMetadataFile::from_toml(
15764 r#"
15765[models.q]
15766default_temperature = 1.0
15767default_top_p = 0.95
15768default_top_k = 20
15769
15770[models.q.non_thinking_sampling]
15771temperature = 0.7
15772top_p = 0.8
15773top_k = 20
15774presence_penalty = 1.5
15775"#,
15776 )
15777 .unwrap();
15778 let arm = parsed
15779 .get("q")
15780 .unwrap()
15781 .non_thinking_sampling
15782 .as_ref()
15783 .unwrap();
15784 assert_eq!(arm.temperature, Some(0.7));
15785 assert_eq!(arm.top_p, Some(0.8));
15786 assert_eq!(arm.top_k, Some(20));
15787 assert_eq!(arm.presence_penalty, Some(1.5));
15788 assert_eq!(
15789 arm.min_p, None,
15790 "undeclared arm fields stay undeclared, never invented"
15791 );
15792
15793 let err = OpenRouterMetadataFile::from_toml(
15797 r#"
15798[models.q]
15799[models.q.non_thinking_sampling]
15800temperature = 0.0
15801"#,
15802 )
15803 .unwrap_err();
15804 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
15805 assert!(err.contains("greedy"), "{err}");
15806
15807 let err = OpenRouterMetadataFile::from_toml(
15810 r#"
15811[models.q]
15812[models.q.non_thinking_sampling]
15813"#,
15814 )
15815 .unwrap_err();
15816 assert!(err.contains("non_thinking_sampling"), "{err}");
15817 assert!(err.contains("declare"), "{err}");
15818
15819 for bad in [
15821 "temperature = 2.5",
15822 "top_p = 0.0",
15823 "top_p = 1.5",
15824 "min_p = 1.0",
15825 "presence_penalty = 3.0",
15826 "frequency_penalty = -2.5",
15827 "repetition_penalty = 0.0",
15828 ] {
15829 let err = OpenRouterMetadataFile::from_toml(&format!(
15830 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
15831 ))
15832 .unwrap_err();
15833 let key = bad.split(' ').next().unwrap();
15834 assert!(
15835 err.contains(&format!("non_thinking_sampling.{key}")),
15836 "the refusal for {bad:?} must name the nested key: {err}"
15837 );
15838 }
15839
15840 let err = OpenRouterMetadataFile::from_toml(
15844 r#"
15845[models.q]
15846[models.q.non_thinking_sampling]
15847temperture = 0.7
15848"#,
15849 )
15850 .unwrap_err();
15851 assert!(err.contains("unknown field"), "{err}");
15852 }
15853
15854 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
15859 SamplingDefaults {
15860 temperature: Some(0.7),
15861 top_p: Some(0.8),
15862 top_k: Some(20),
15863 presence_penalty: Some(1.5),
15864 ..Default::default()
15865 }
15866 }
15867
15868 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
15869 ModelSamplingDefaults {
15870 thinking: qwen38_vendor_defaults(),
15871 non_thinking: Some(qwen38_non_thinking_defaults()),
15872 }
15873 }
15874
15875 fn qwen38_caps() -> ModelCaps {
15879 ModelCaps {
15880 chat_ok: true,
15881 qwen_think: true,
15882 think_switch: true,
15883 ..Default::default()
15884 }
15885 }
15886
15887 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
15890 (
15891 c.temperature,
15892 c.top_p,
15893 c.top_k,
15894 c.min_p,
15895 c.penalty_present,
15896 c.penalty_freq,
15897 c.penalty_repeat,
15898 c.penalty_last_n,
15899 c.seed,
15900 )
15901 }
15902
15903 fn build_with_arms(
15904 defaults: &ModelSamplingDefaults,
15905 caps: &ModelCaps,
15906 default_effort: Option<&str>,
15907 extra: serde_json::Value,
15908 ) -> Request {
15909 let mut body = serde_json::json!({
15910 "model": "m",
15911 "messages": [{"role": "user", "content": "task"}],
15912 "seed": 3
15914 });
15915 body.as_object_mut()
15916 .unwrap()
15917 .extend(extra.as_object().unwrap().clone());
15918 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15919 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15920 build_chat_request_with_trace(
15921 req,
15922 Some(caps),
15923 tx,
15924 lanes::Lane::Interactive,
15925 None,
15926 None,
15927 default_effort,
15928 defaults,
15929 )
15930 .unwrap()
15931 .request
15932 }
15933
15934 #[test]
15935 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
15936 let two_arm = qwen38_two_arm_defaults();
15941 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
15942 let caps = qwen38_caps();
15943
15944 let off_spellings = [
15946 serde_json::json!({"reasoning_effort": "none"}),
15947 serde_json::json!({"enable_thinking": false}),
15948 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
15949 serde_json::json!({"reasoning": {"enabled": false}}),
15950 ];
15951 for extra in &off_spellings {
15952 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
15953 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
15954 let c = &r.sampler_cfg;
15955 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
15956 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
15957 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
15958 assert_eq!(
15959 c.penalty_present, 1.5,
15960 "{extra}: non-thinking presence_penalty"
15961 );
15962 assert_eq!(
15963 c.penalty_last_n,
15964 memra_engine::spec::PEN_WINDOW_MAX,
15965 "{extra}: the arm's presence penalty uses the cross-path history window"
15966 );
15967 assert_eq!(
15968 c.min_p, 0.0,
15969 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
15970 );
15971
15972 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
15975 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
15976 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
15977 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
15978 assert_eq!(
15979 s.sampler_cfg.penalty_present, 0.0,
15980 "{extra}: single-arm model"
15981 );
15982 }
15983
15984 for extra in [
15987 serde_json::json!({}),
15988 serde_json::json!({"enable_thinking": true}),
15989 serde_json::json!({"reasoning_effort": "high"}),
15990 serde_json::json!({"reasoning": {"enabled": true}}),
15991 ] {
15992 for defaults in [&two_arm, &single_arm] {
15993 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
15994 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
15995 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
15996 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
15997 assert_eq!(
15998 c.penalty_present, 0.0,
15999 "{extra}: thinking arm has no presence"
16000 );
16001 }
16002 }
16003
16004 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
16007 assert_eq!(
16008 c.temperature, 0.7,
16009 "deployment-default off = non-thinking arm"
16010 );
16011 let c = build_with_arms(
16013 &two_arm,
16014 &caps,
16015 Some("none"),
16016 serde_json::json!({"enable_thinking": true}),
16017 )
16018 .sampler_cfg;
16019 assert_eq!(
16020 c.temperature, 1.0,
16021 "explicit ON beats the deployment default"
16022 );
16023
16024 let c = build_with_arms(
16026 &two_arm,
16027 &caps,
16028 None,
16029 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
16030 )
16031 .sampler_cfg;
16032 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
16033 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
16034 let c = build_with_arms(
16035 &two_arm,
16036 &caps,
16037 None,
16038 serde_json::json!({
16039 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
16040 )
16041 .sampler_cfg;
16042 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
16043 assert_eq!(
16044 c.penalty_present, 0.0,
16045 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
16046 is a value, not an absence"
16047 );
16048 assert_eq!(
16049 c.penalty_last_n, 0,
16050 "all penalties off => no history window"
16051 );
16052 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
16053
16054 let c = build_with_arms(
16057 &two_arm,
16058 &caps,
16059 None,
16060 serde_json::json!({"enable_thinking": false, "temperature": 0}),
16061 )
16062 .sampler_cfg;
16063 assert!(
16064 memra_engine::sampler::Sampler::new(c).is_greedy(),
16065 "explicit temperature 0 must stay greedy on the non-thinking arm"
16066 );
16067
16068 let c = build_with_arms(
16071 &single_arm,
16072 &caps,
16073 None,
16074 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
16075 )
16076 .sampler_cfg;
16077 assert_eq!(c.temperature, 0.55);
16078 assert_eq!(
16079 c.top_p, 0.95,
16080 "single-arm model: unset top_p takes its one arm"
16081 );
16082 }
16083
16084 #[test]
16085 fn sampling_arms_never_blend_field_by_field() {
16086 let parsed = OpenRouterMetadataFile::from_toml(
16091 r#"
16092[models.m]
16093default_temperature = 1.0
16094default_min_p = 0.05
16095
16096[models.m.non_thinking_sampling]
16097temperature = 0.6
16098"#,
16099 )
16100 .unwrap();
16101 let caps = ModelCaps {
16102 chat_temperature_default: Some(0.5),
16103 chat_top_p_default: Some(0.9),
16104 ..Default::default()
16105 };
16106 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
16107 let client = ClientSampling {
16108 seed: Some(1),
16109 ..Default::default()
16110 };
16111
16112 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
16113 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
16114 assert_eq!(
16115 off.min_p, 0.0,
16116 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
16117 );
16118 assert_eq!(
16119 off.top_p, 1.0,
16120 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
16121 );
16122
16123 for mode in [ThinkMode::Default, ThinkMode::Think] {
16125 let on = resolve_sampler_config(client, d.for_mode(mode));
16126 assert_eq!(on.temperature, 1.0);
16127 assert_eq!(on.min_p, 0.05);
16128 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
16129 }
16130 }
16131
16132 #[test]
16133 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
16134 let caps = qwen38_caps();
16143 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
16144 let two_arm = qwen38_two_arm_defaults();
16145
16146 let bodies = [
16147 serde_json::json!({}),
16148 serde_json::json!({"enable_thinking": true}),
16149 serde_json::json!({"reasoning_effort": "high"}),
16150 serde_json::json!({"reasoning_effort": "none"}),
16151 serde_json::json!({"enable_thinking": false}),
16152 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
16153 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
16154 serde_json::json!({"enable_thinking": false, "temperature": 0}),
16155 ];
16156 for extra in &bodies {
16157 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
16159 let mut client = ClientSampling {
16160 seed: Some(3),
16161 ..Default::default()
16162 };
16163 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
16164 client.temperature = Some(t as f32);
16165 }
16166 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
16167 client.top_p = Some(p as f32);
16168 }
16169 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
16170 assert_eq!(
16171 sampler_key(&r.sampler_cfg),
16172 sampler_key(&pre_arm),
16173 "{extra}: single-arm model diverged from the pre-arm resolution law"
16174 );
16175
16176 if r.think != ThinkMode::NoThink {
16179 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
16180 assert_eq!(t.think, r.think, "{extra}");
16181 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
16182 assert_eq!(
16183 sampler_key(&t.sampler_cfg),
16184 sampler_key(&r.sampler_cfg),
16185 "{extra}: a thinking-on request must not feel the non-thinking arm"
16186 );
16187 }
16188 }
16189 }
16190
16191 #[test]
16192 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
16193 let r = build_with_arms(
16199 &qwen38_two_arm_defaults(),
16200 &qwen38_caps(),
16201 None,
16202 serde_json::json!({"response_format": {"type": "json_object"}}),
16203 );
16204 assert_eq!(
16205 r.think,
16206 ThinkMode::NoThink,
16207 "constraint forces the switch off"
16208 );
16209 assert_eq!(
16210 r.sampler_cfg.temperature, 0.7,
16211 "and the arm follows the real mode"
16212 );
16213 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
16214 }
16215
16216 #[test]
16217 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
16218 let caps = ModelCaps {
16223 chat_temperature_default: Some(0.5),
16224 chat_top_p_default: Some(0.9),
16225 chat_ok: true,
16226 ..Default::default()
16227 };
16228 let metadata = OpenRouterModelMetadata {
16229 default_temperature: Some(1.0),
16230 default_top_p: Some(0.95),
16231 default_top_k: Some(64),
16232 ..Default::default()
16233 };
16234
16235 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
16236 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
16237 assert_eq!(caps_only.top_p, Some(0.9));
16238 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
16239
16240 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
16241 assert_eq!(
16242 both.temperature,
16243 Some(1.0),
16244 "metadata outranks the arch cap"
16245 );
16246 assert_eq!(both.top_p, Some(0.95));
16247 assert_eq!(both.top_k, Some(64));
16248
16249 let partial = SamplingDefaults::resolve(
16251 Some(&OpenRouterModelMetadata {
16252 default_temperature: Some(0.7),
16253 ..Default::default()
16254 }),
16255 Some(&caps),
16256 );
16257 assert_eq!(partial.temperature, Some(0.7));
16258 assert_eq!(
16259 partial.top_p,
16260 Some(0.9),
16261 "an undeclared metadata field must fall through to the cap, not to 1.0"
16262 );
16263
16264 assert_eq!(
16266 SamplingDefaults::resolve(None, None),
16267 SamplingDefaults::default()
16268 );
16269 }
16270
16271 #[test]
16272 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
16273 let resolved = |d: &SamplingDefaults| {
16286 resolve_sampler_config(
16287 ClientSampling {
16288 seed: Some(1),
16289 ..Default::default()
16290 },
16291 d,
16292 )
16293 };
16294
16295 assert!(
16297 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
16298 .is_spec_sampling(),
16299 "the API-standard default must stay in the fast pure-temp regime"
16300 );
16301
16302 for (name, d) in [
16303 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
16304 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
16305 ] {
16306 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
16307 assert!(
16308 !sampler.is_greedy(),
16309 "{name}: vendor default must not be greedy"
16310 );
16311 assert!(
16312 !sampler.is_spec_sampling(),
16313 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
16314 starts passing, either the vendor numbers changed or the in-graph draft \
16315 learned filters, and the perf note in docs/SERVING.md needs revisiting"
16316 );
16317 }
16318
16319 let opted_out = resolve_sampler_config(
16321 ClientSampling {
16322 top_p: Some(1.0),
16323 top_k: Some(0),
16324 seed: Some(1),
16325 ..Default::default()
16326 },
16327 &qwen38_vendor_defaults(),
16328 );
16329 assert!(
16330 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
16331 "explicitly disabling the filters must restore the pure-temp regime"
16332 );
16333 }
16334
16335 #[test]
16336 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
16337 let comp_seed = |body: serde_json::Value| {
16344 let req: CompletionReq = serde_json::from_value(body).unwrap();
16345 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
16346 build_request(&req, tx, lanes::Lane::Interactive, None)
16347 .sampler_cfg
16348 .seed
16349 };
16350 let chat_seed = |body: serde_json::Value| {
16351 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16352 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
16353 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16354 .unwrap()
16355 .request
16356 .sampler_cfg
16357 .seed
16358 };
16359
16360 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
16363 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
16364 let c = chat_seed(serde_json::json!({
16365 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
16366 assert_ne!(
16367 a, 0,
16368 "omitted seed must not be the pinned 0 that caused the loop"
16369 );
16370 assert_ne!(b, 0);
16371 assert_ne!(c, 0);
16372 assert_ne!(
16373 a, b,
16374 "two seed-omitting requests must get DIFFERENT streams"
16375 );
16376 assert_ne!(a, c);
16377
16378 assert_eq!(
16381 comp_seed(serde_json::json!({
16382 "model": "m", "prompt": "t", "seed": 0})),
16383 0,
16384 "explicit seed 0 must stay 0 — the determinism gates depend on it"
16385 );
16386 assert_eq!(
16387 comp_seed(serde_json::json!({
16388 "model": "m", "prompt": "t", "seed": 12345})),
16389 12345
16390 );
16391 assert_eq!(
16392 chat_seed(serde_json::json!({
16393 "model": "m", "messages": [{"role": "user", "content": "t"}],
16394 "seed": 777})),
16395 777
16396 );
16397 assert_eq!(
16399 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
16400 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
16401 );
16402
16403 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
16406 assert_eq!(
16407 seeds.len(),
16408 256,
16409 "fresh_seed must not collide across rapid calls"
16410 );
16411 assert!(!seeds.contains(&0));
16412 }
16413
16414 #[test]
16415 fn response_format_builds_grammar_only_when_present() {
16416 let mk = |rf: Option<serde_json::Value>| {
16420 let mut body = serde_json::json!({
16421 "model": "m", "messages": [{"role": "user", "content": "t"}]});
16422 if let Some(rf) = rf {
16423 body["response_format"] = rf;
16424 }
16425 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
16426 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
16427 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
16428 };
16429 assert!(mk(None).unwrap().request.grammar.is_none());
16430 assert!(
16431 mk(Some(serde_json::json!({"type": "text"})))
16432 .unwrap()
16433 .request
16434 .grammar
16435 .is_none()
16436 );
16437 assert!(matches!(
16438 mk(Some(serde_json::json!({"type": "json_object"})))
16439 .unwrap()
16440 .request
16441 .grammar,
16442 Some(constrained::GrammarSpec::JsonObject)
16443 ));
16444 assert!(matches!(
16445 mk(Some(serde_json::json!({"type": "json_schema",
16446 "json_schema": {"schema": {"type": "object"}}})))
16447 .unwrap()
16448 .request
16449 .grammar,
16450 Some(constrained::GrammarSpec::JsonSchema(_))
16451 ));
16452 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
16454 }
16455
16456 #[test]
16466 fn response_format_think_table_switch_postthink_refusal() {
16467 let mk = |caps: &ModelCaps| {
16468 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
16469 "model": "m", "messages": [{"role": "user", "content": "t"}],
16470 "response_format": {"type": "json_object"}}))
16471 .unwrap();
16472 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
16473 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
16474 };
16475 let switch = ModelCaps {
16477 chat_ok: true,
16478 qwen_think: true,
16479 think_switch: true,
16480 ..Default::default()
16481 };
16482 let plan = mk(&switch).unwrap();
16483 assert_eq!(
16484 plan.request.think,
16485 memra_tokenizer::chat::ThinkMode::NoThink,
16486 "switch-carrying template must keep the grammar-from-token-1 path"
16487 );
16488 assert!(plan.request.grammar.is_some());
16489
16490 let postthink = ModelCaps {
16492 chat_ok: true,
16493 qwen_think: true,
16494 think_switch: false,
16495 think_close: vec![128799],
16496 ..Default::default()
16497 };
16498 let plan = mk(&postthink).unwrap();
16499 assert_ne!(
16500 plan.request.think,
16501 memra_tokenizer::chat::ThinkMode::NoThink,
16502 "post-think constrained request must keep the think channel ON"
16503 );
16504 assert!(plan.request.grammar.is_some());
16505
16506 let no_contract = ModelCaps {
16508 chat_ok: true,
16509 qwen_think: true,
16510 think_switch: false,
16511 think_close: Vec::new(),
16512 ..Default::default()
16513 };
16514 let err = match mk(&no_contract) {
16515 Err(err) => err,
16516 Ok(_) => panic!("think-forced template with no close contract must refuse"),
16517 };
16518 assert!(
16519 err.contains("think-close"),
16520 "refusal must name the missing close contract: {err}"
16521 );
16522 }
16523
16524 #[test]
16525 fn unsupported_semantic_params_are_named_rejections() {
16526 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
16528 "model": "m", "messages": [{"role": "user", "content": "t"}],
16529 "response_format": {"type": "json_object"}
16530 }))
16531 .unwrap();
16532 assert!(req.response_format.is_some());
16533 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
16534 "model": "m", "messages": [{"role": "user", "content": "t"}],
16535 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
16536 "user": "u-1", "stream_options": {"include_usage": true}
16537 }))
16538 .unwrap();
16539 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
16541 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
16542 assert_eq!(req.n, Some(1));
16543 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
16545 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
16546 assert_eq!(param, "logit_bias");
16547 assert_eq!(msg, "logit_bias is not supported (why)");
16548 }
16549
16550 #[test]
16551 fn completions_accept_openai_stop_forms() {
16552 for (value, expected) in [
16553 (serde_json::json!("Problem:"), vec!["Problem:"]),
16554 (
16555 serde_json::json!(["Question:", "Problem:"]),
16556 vec!["Question:", "Problem:"],
16557 ),
16558 (serde_json::Value::Null, Vec::<&str>::new()),
16559 ] {
16560 let req: CompletionReq = serde_json::from_value(serde_json::json!({
16561 "model": "plain_quant", "prompt": "task", "stop": value
16562 }))
16563 .unwrap();
16564 assert_eq!(req.stop.into_vec(), expected);
16565 }
16566 }
16567
16568 fn fake_worker_state() -> AppState {
16575 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
16576 }
16577
16578 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
16579 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
16580 }
16581
16582 struct WorkerSaw {
16587 sampler_cfg: SamplerConfig,
16588 think: ThinkMode,
16589 reasoning_effort: Option<String>,
16590 }
16591
16592 fn fake_worker_state_full(
16599 steps: usize,
16600 step_delay: std::time::Duration,
16601 caps: HashMap<String, ModelCaps>,
16602 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
16603 ) -> AppState {
16604 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
16605 let health = health::WorkerHealth::new();
16606 let h = health.clone();
16607 std::thread::spawn(move || {
16608 h.mark_ready();
16609 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
16610 if let Some(tx) = &saw_tx {
16611 let _ = tx.send(WorkerSaw {
16612 sampler_cfg: req.sampler_cfg.clone(),
16613 think: req.think,
16614 reasoning_effort: req.reasoning_effort.clone(),
16615 });
16616 }
16617 worker::release_pending_admit();
16621 worker::release_admission_reservation(req.lane);
16622 h.beat_busy();
16623 if let Some(ready) = req.constraint_ready.take() {
16624 let _ = ready.send(Ok(()));
16625 }
16626 let _ = req.tx.send(Event::PromptUsage {
16627 n_prompt: 1,
16628 n_cached: 0,
16629 });
16630 for step in 0..steps {
16631 h.beat_busy();
16632 let text = if steps == 1 { "ok" } else { "x" };
16633 let _ = req.tx.send(Event::Token {
16634 id: step as u32 + 1,
16635 text: text.into(),
16636 });
16637 if !step_delay.is_zero() {
16638 std::thread::sleep(step_delay);
16639 }
16640 }
16641 let _ = req.tx.send(Event::Done {
16642 stop_reason: "Eos".into(),
16643 n_tokens: steps,
16644 n_prompt: 1,
16645 n_cached: 0,
16646 elapsed_s: 0.01,
16647 spec: None,
16648 });
16649 h.set_phase(health::PHASE_IDLE);
16650 }
16651 });
16652 for _ in 0..2000 {
16655 if health.live().is_ok() {
16656 break;
16657 }
16658 std::thread::sleep(std::time::Duration::from_millis(1));
16659 }
16660 AppState {
16661 cmd_tx,
16662 models: Arc::new(vec!["m".into()]),
16663 caps: Arc::new(caps),
16664 openrouter_metadata: Arc::new(HashMap::new()),
16665 provider_metadata: Arc::new(None),
16666 metering: None,
16667
16668 budget_tokenizers: None,
16669 api_auth: ApiAuth::default(),
16670 metrics_auth: MetricsAuth::default(),
16671 metrics: SharedMetrics::default(),
16672 inflight: Arc::new(Default::default()),
16673 tenant_inflight: Arc::new(Default::default()),
16674 health,
16675 bg: None,
16676 }
16677 }
16678
16679 #[tokio::test]
16680 #[allow(clippy::await_holding_lock)] async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
16682 let _l = DRAIN_LOCK.lock().unwrap();
16683 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
16684 let normal_state = st.clone();
16685 let normal = tokio::spawn(async move {
16686 chat_completions(
16687 State(normal_state),
16688 axum::http::HeaderMap::new(),
16689 None,
16690 Json(
16691 serde_json::from_value(serde_json::json!({
16692 "model": "m",
16693 "messages": [{"role": "user", "content": "keep decoding"}],
16694 }))
16695 .unwrap(),
16696 ),
16697 )
16698 .await
16699 });
16700 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
16701
16702 let mut deep = serde_json::json!({"type": "string"});
16703 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
16704 deep = serde_json::json!({"allOf": [deep]});
16705 }
16706 let bad = chat_completions(
16707 State(st.clone()),
16708 axum::http::HeaderMap::new(),
16709 None,
16710 Json(
16711 serde_json::from_value(serde_json::json!({
16712 "model": "m",
16713 "messages": [{"role": "user", "content": "bad schema"}],
16714 "response_format": {
16715 "type": "json_schema",
16716 "json_schema": {"schema": deep},
16717 },
16718 }))
16719 .unwrap(),
16720 ),
16721 )
16722 .await;
16723 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
16724 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
16725 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
16726 .await
16727 .unwrap();
16728 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16729 assert!(
16730 payload["error"]["message"]
16731 .as_str()
16732 .unwrap()
16733 .contains("maximum nesting depth")
16734 );
16735 assert!(
16736 !normal.is_finished(),
16737 "bad schema stalled or replaced the normal decode"
16738 );
16739
16740 let normal_response = normal.await.unwrap();
16741 assert_eq!(normal_response.status(), StatusCode::OK);
16742 let snapshot = st.health.snapshot();
16743 assert!(
16744 st.health.live().is_ok(),
16745 "normal decode left health stalled"
16746 );
16747 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
16748 }
16749
16750 #[tokio::test]
16751 #[allow(clippy::await_holding_lock)] async fn valid_response_format_preflight_preserves_generation() {
16753 let _l = DRAIN_LOCK.lock().unwrap();
16754 let response = chat_completions(
16755 State(fake_worker_state()),
16756 axum::http::HeaderMap::new(),
16757 None,
16758 Json(
16759 serde_json::from_value(serde_json::json!({
16760 "model": "m",
16761 "messages": [{"role": "user", "content": "valid schema"}],
16762 "response_format": {"type": "json_object"},
16763 }))
16764 .unwrap(),
16765 ),
16766 )
16767 .await;
16768 assert_eq!(response.status(), StatusCode::OK);
16769 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16770 .await
16771 .unwrap();
16772 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16773 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
16774 }
16775
16776 #[tokio::test]
16777 #[allow(clippy::await_holding_lock)] async fn unknown_model_refuses_model_not_found_before_admission() {
16779 let _l = DRAIN_LOCK.lock().unwrap();
16780 let response = chat_completions(
16785 State(fake_worker_state()),
16786 axum::http::HeaderMap::new(),
16787 None,
16788 Json(
16789 serde_json::from_value(serde_json::json!({
16790 "model": "qwen/qwen3.8-27b-typo",
16791 "messages": [{"role": "user", "content": "hi"}],
16792 }))
16793 .unwrap(),
16794 ),
16795 )
16796 .await;
16797 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
16798 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16799 .await
16800 .unwrap();
16801 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16802 assert_eq!(payload["error"]["code"], "model_not_found");
16803 assert_eq!(payload["error"]["type"], "invalid_request_error");
16804
16805 let response = completions(
16807 State(fake_worker_state()),
16808 axum::http::HeaderMap::new(),
16809 None,
16810 Json(
16811 serde_json::from_value(serde_json::json!({
16812 "model": "nope",
16813 "prompt": "hi",
16814 }))
16815 .unwrap(),
16816 ),
16817 )
16818 .await;
16819 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
16820 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16821 .await
16822 .unwrap();
16823 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16824 assert_eq!(payload["error"]["code"], "model_not_found");
16825 }
16826
16827 const METRICS_KEY_ACME: &str = "completion-acme-secret";
16828 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
16829
16830 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
16831 let spec = format!(
16832 "acme:{},blue:{}",
16833 auth::sha256_hex(METRICS_KEY_ACME),
16834 auth::sha256_hex(METRICS_KEY_BLUE),
16835 );
16836 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
16837 let mut st = fake_worker_state();
16838 st.api_auth.keyring = Some(keyring);
16839 st.metrics_auth = MetricsAuth::new(
16840 true,
16841 st.api_auth.configured(),
16842 metrics_token.map(str::to_string),
16843 );
16844 {
16845 let mut metrics = st.metrics.lock().unwrap();
16846 metrics.admitted = 17;
16847 metrics.prompt_tokens_in = 400;
16848 metrics.cached_tokens_in = 60;
16849 metrics.prefix_hits = 2;
16850 metrics.prefix_misses = 3;
16851 metrics.prefix_inserts = 5;
16852 metrics.prefix_evictions = 7;
16853 metrics.prefix_skips_budget = 9;
16854 metrics.prefix_skips_pinned = 10;
16855 metrics.prefix_hit_tokens = 11;
16856 metrics.lcp_hist[4] = 13;
16857 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
16858 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
16859 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
16860 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
16861 metrics.prefix_entries = 29;
16862 metrics.prefix_bytes = 31;
16863 metrics.active_sessions = 3;
16864 metrics.queued_requests = 5;
16865 metrics.admission_inflight.insert("m".into(), 4);
16866 metrics
16867 .admission_booked_bytes
16868 .insert("m".into(), 41_000_000);
16869 metrics.continuation_pool_entries = 7;
16870 metrics.spec_pool_entries = 11;
16871 metrics.cuda_driver_free_bytes = 13;
16872 metrics.cuda_pool_reserved_bytes = 17;
16873 metrics.cuda_pool_used_bytes = 19;
16874 metrics.cuda_pool_cached_bytes = 23;
16875 metrics.batch_size_last = 37;
16876 metrics.spec.insert(
16877 "m".into(),
16878 memra_engine::spec::SpecTelemetry {
16879 rounds: 2,
16880 drafted: 6,
16881 accepted: 4,
16882 ..Default::default()
16883 },
16884 );
16885 let mut spec_window = memra_engine::spec::SpecTelemetry {
16886 rounds: 4,
16887 drafted: 12,
16888 accepted: 6,
16889 ..Default::default()
16890 };
16891 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
16892 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
16893 metrics.spec_window.insert("m".into(), spec_window);
16894 metrics.constraint_compiler_fail_closed.insert(
16895 "m".into(),
16896 Arc::new(std::sync::atomic::AtomicBool::new(true)),
16897 );
16898 }
16899 st
16900 }
16901
16902 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
16903 let mut headers = HeaderMap::new();
16904 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
16905 let response = get_metrics(State(st), headers).await;
16906 assert_eq!(response.status(), StatusCode::OK);
16907 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16908 .await
16909 .unwrap();
16910 serde_json::from_slice(&bytes).unwrap()
16911 }
16912
16913 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
16914 let mut headers = HeaderMap::new();
16915 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
16916 let response = yield_metrics(State(st), headers).await;
16917 assert_eq!(response.status(), StatusCode::OK);
16918 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16919 .await
16920 .unwrap();
16921 serde_json::from_slice(&bytes).unwrap()
16922 }
16923
16924 #[test]
16925 fn exposed_open_bind_is_refused_before_server_start() {
16926 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
16927 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
16928
16929 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
16930 assert!(err.contains("refusing unauthenticated non-loopback bind"));
16931 assert!(err.contains("MEMRA_API_KEY"));
16932 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
16933 assert!(validate_bind_security("[::]:8000", false, false).is_err());
16934
16935 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
16936 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
16937 }
16938
16939 #[tokio::test]
16940 async fn keyed_metrics_require_and_accept_api_bearer() {
16941 let mut st = fake_worker_state();
16942 st.api_auth.single_key = Some(Arc::from("completion-secret"));
16943 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
16944
16945 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
16946 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
16947 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
16948 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
16949
16950 let mut headers = HeaderMap::new();
16951 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
16952 assert_eq!(
16953 get_metrics(State(st.clone()), headers.clone())
16954 .await
16955 .status(),
16956 StatusCode::OK,
16957 );
16958 let body = metrics_json(st.clone(), "completion-secret").await;
16959 assert!(
16960 body.get("admitted").is_some(),
16961 "the legacy single-key domain keeps cumulative counters",
16962 );
16963 assert!(
16964 body.get("active_sessions").is_none(),
16965 "a static completion key is not an operator metrics principal",
16966 );
16967 assert_eq!(
16968 yield_metrics(State(st), headers).await.status(),
16969 StatusCode::OK
16970 );
16971 }
16972
16973 #[tokio::test]
16974 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
16975 let st = multi_key_metrics_state(None);
16976 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
16977 assert_eq!(
16978 body.as_object().unwrap().len(),
16979 2,
16980 "completion metrics must contain only tenant-scoped rows",
16981 );
16982 let tenants = body["tenants"].as_object().unwrap();
16983 assert_eq!(tenants.len(), 1);
16984 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
16985 assert!(!tenants.contains_key("t:blue"));
16986 let adsd = body["adsd_suspect_total"].as_object().unwrap();
16987 assert_eq!(adsd.len(), 1);
16988 assert_eq!(adsd["t:acme"], 1);
16989 assert!(!adsd.contains_key("t:blue"));
16990
16991 let mut headers = HeaderMap::new();
16992 headers.insert(
16993 "authorization",
16994 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
16995 );
16996 assert_eq!(
16997 yield_metrics(State(st), headers).await.status(),
16998 StatusCode::FORBIDDEN,
16999 "the process-wide yield view requires an operator metrics token",
17000 );
17001 }
17002
17003 #[tokio::test]
17004 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
17005 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
17006 for operator_only in [
17007 "prefix_cache_entries",
17008 "prefix_cache_bytes",
17009 "prefix_cache_skips_budget",
17010 "prefix_cache_skips_pinned",
17011 "active_sessions",
17012 "queued_requests",
17013 "admission_inflight",
17014 "admission_booked_bytes",
17015 "continuation_pool_entries",
17016 "spec_pool_entries",
17017 "cuda_driver_free_bytes",
17018 "cuda_pool_reserved_bytes",
17019 "cuda_pool_used_bytes",
17020 "cuda_pool_cached_bytes",
17021 "constraint_compiler_fail_closed",
17022 "serve_idle_seconds",
17023 "spec",
17024 "spec_tau",
17025 "spec_accept_by_position",
17026 "dual_pp",
17027 "pp_wave",
17028 "peer_probe_bypassed",
17029 "peer_probe_boundary_copies",
17030 "peer_probe_runtime_reprobes",
17031 "peer_probe_runtime_failures",
17032 "peer_probe_deferred_total",
17033 "peer_probe_integrity_degraded",
17034 "peer_probe_degraded_to_host_bounce",
17035 ] {
17036 assert!(
17037 body.get(operator_only).is_none(),
17038 "tenant metrics must not expose operator field {operator_only}",
17039 );
17040 }
17041 }
17042
17043 #[test]
17044 fn populated_spec_acceptance_metrics_are_operator_only() {
17045 for scope in [
17046 MetricsScope::CompletionDomain,
17047 MetricsScope::Tenant("t:acme".into()),
17048 ] {
17049 let mut body = json!({});
17050 insert_spec_acceptance_metrics(&mut body, &scope, || {
17051 panic!("tenant scope evaluated the process-wide spec snapshot")
17052 });
17053 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
17054 assert!(
17055 body.get("spec_accept_by_position").is_none(),
17056 "{scope:?} leaked the accept histogram"
17057 );
17058 }
17059
17060 let mut telemetry = memra_engine::spec::SpecTelemetry {
17061 rounds: 4,
17062 drafted: 12,
17063 accepted: 6,
17064 ..Default::default()
17065 };
17066 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
17067 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
17068 let mut body = json!({});
17069 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
17070 HashMap::from([("model-a".to_string(), telemetry)])
17071 });
17072 assert_eq!(body["spec_tau"]["model-a"], 1.5);
17073 let histogram = &body["spec_accept_by_position"]["model-a"];
17074 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
17075 assert_eq!(histogram["rounds"], 4);
17076 assert_eq!(histogram["offered"], json!([4, 4, 4]));
17077 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
17078 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
17079 }
17080
17081 #[test]
17082 fn populated_dual_pp_metrics_are_operator_only() {
17083 let populated = DualPpMetricsSnapshot {
17084 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
17085 stage_samples: [1, 1, 1, 1],
17086 dropped_timing_samples: 0,
17087 overlaps: 17,
17088 slot_pairs: 19,
17089 slot_uses: [19, 19],
17090 slot_collisions: 0,
17091 };
17092 for scope in [
17093 MetricsScope::CompletionDomain,
17094 MetricsScope::Tenant("t:acme".into()),
17095 ] {
17096 let mut body = json!({});
17097 insert_dual_pp_metrics(&mut body, &scope, || populated);
17098 assert!(
17099 body.get("dual_pp").is_none(),
17100 "{scope:?} leaked dual PP topology"
17101 );
17102 }
17103
17104 let mut body = json!({});
17105 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
17106 assert_eq!(body["dual_pp"]["overlaps"], 17);
17107 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
17108 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
17109 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
17110 assert_eq!(
17111 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
17112 1.0
17113 );
17114 }
17115
17116 #[test]
17117 fn populated_pp_wave_metrics_are_operator_only() {
17118 let populated = PpWaveMetricsSnapshot {
17119 ticks: 11,
17120 cells: 96,
17121 overlaps: 37,
17122 };
17123 for scope in [
17124 MetricsScope::CompletionDomain,
17125 MetricsScope::Tenant("t:acme".into()),
17126 ] {
17127 let mut body = json!({});
17128 insert_pp_wave_metrics(&mut body, &scope, || populated);
17129 assert!(
17130 body.get("pp_wave").is_none(),
17131 "{scope:?} leaked PP wave topology"
17132 );
17133 }
17134
17135 let mut body = json!({});
17136 insert_pp_wave_metrics(&mut body, &MetricsScope::All, || populated);
17137 assert_eq!(body["pp_wave"]["ticks"], 11);
17138 assert_eq!(body["pp_wave"]["cells"], 96);
17139 assert_eq!(body["pp_wave"]["overlaps"], 37);
17140 }
17141
17142 #[test]
17143 fn peer_probe_metrics_are_operator_only() {
17144 let populated = memra_engine::pp::PeerProbeMetrics {
17145 bypassed: 1,
17146 boundary_copies: 8_192,
17147 runtime_probes: 1,
17148 runtime_failures: 0,
17149 deferred_total: 4,
17150 integrity_degraded: true,
17151 degraded_to_host_bounce: true,
17152 };
17153 for scope in [
17154 MetricsScope::CompletionDomain,
17155 MetricsScope::Tenant("t:acme".into()),
17156 ] {
17157 let mut body = json!({});
17158 insert_peer_probe_metrics(&mut body, &scope, || populated);
17159 assert!(body.get("peer_probe_bypassed").is_none());
17160 }
17161
17162 let mut body = json!({});
17163 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
17164 assert_eq!(body["peer_probe_bypassed"], 1);
17165 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
17166 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
17167 assert_eq!(body["peer_probe_runtime_failures"], 0);
17168 assert_eq!(body["peer_probe_deferred_total"], 4);
17169 assert_eq!(body["peer_probe_integrity_degraded"], true);
17170 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
17171 }
17172
17173 #[tokio::test]
17174 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
17175 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
17176 for operator_only in [
17177 "lcp_histogram",
17178 "cache_hit_token_ratio",
17179 "prefix_cache_hits",
17180 "prefix_cache_misses",
17181 "prefix_cache_inserts",
17182 "prefix_cache_evictions",
17183 "prefix_cache_skips_budget",
17184 "prefix_cache_skips_pinned",
17185 "prefix_cache_hit_tokens",
17186 ] {
17187 assert!(
17188 tenant_body.get(operator_only).is_none(),
17189 "tenant metrics must not expose global prefix field {operator_only}",
17190 );
17191 }
17192 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
17193 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
17194 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
17195 assert_eq!(
17196 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
17197 0.4
17198 );
17199
17200 let operator_body = metrics_json(
17201 multi_key_metrics_state(Some("scrape-secret")),
17202 "scrape-secret",
17203 )
17204 .await;
17205 assert_eq!(operator_body["prefix_cache_hits"], 2);
17206 assert_eq!(operator_body["prefix_cache_misses"], 3);
17207 assert_eq!(operator_body["prefix_cache_inserts"], 5);
17208 assert_eq!(operator_body["prefix_cache_evictions"], 7);
17209 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
17210 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
17211 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
17212 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
17213 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
17214 }
17215
17216 #[tokio::test]
17217 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
17218 let st = multi_key_metrics_state(Some("scrape-secret"));
17219 let mut completion_headers = HeaderMap::new();
17220 completion_headers.insert(
17221 "authorization",
17222 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
17223 );
17224 assert_eq!(
17225 get_metrics(State(st.clone()), completion_headers.clone())
17226 .await
17227 .status(),
17228 StatusCode::FORBIDDEN,
17229 );
17230 assert_eq!(
17231 yield_metrics(State(st.clone()), completion_headers)
17232 .await
17233 .status(),
17234 StatusCode::FORBIDDEN,
17235 );
17236
17237 let body = metrics_json(st.clone(), "scrape-secret").await;
17238 let tenants = body["tenants"].as_object().unwrap();
17239 assert_eq!(tenants.len(), 2);
17240 assert!(tenants.contains_key("t:acme"));
17241 assert!(tenants.contains_key("t:blue"));
17242 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
17243 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
17244 assert_eq!(body["active_sessions"], 3);
17245 assert_eq!(body["queued_requests"], 5);
17246 assert_eq!(body["admission_inflight"]["m"], 4);
17248 assert_eq!(body["admission_booked_bytes"]["m"], 41_000_000);
17249 assert_eq!(body["prefix_cache_bytes"], 31);
17250 assert_eq!(body["cuda_driver_free_bytes"], 13);
17251 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
17252 assert_eq!(body["spec"]["m"]["drafted"], 6);
17253 assert_eq!(body["spec_tau"]["m"], 1.5);
17254 assert_eq!(
17255 body["spec_accept_by_position"]["m"]["accepted"],
17256 json!([3, 2, 1])
17257 );
17258 let yield_body = yield_metrics_json(st, "scrape-secret").await;
17259 assert_eq!(yield_body["batch_size_last"], 37);
17260 }
17261
17262 #[tokio::test]
17263 async fn metrics_token_protects_public_override_without_api_keys() {
17264 let mut st = fake_worker_state();
17265 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
17266
17267 assert_eq!(
17268 get_metrics(State(st.clone()), HeaderMap::new())
17269 .await
17270 .status(),
17271 StatusCode::UNAUTHORIZED,
17272 );
17273 let mut headers = HeaderMap::new();
17274 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
17275 assert_eq!(
17276 get_metrics(State(st.clone()), headers.clone())
17277 .await
17278 .status(),
17279 StatusCode::OK,
17280 );
17281 assert_eq!(
17282 yield_metrics(State(st), headers).await.status(),
17283 StatusCode::OK
17284 );
17285 }
17286
17287 #[tokio::test]
17288 async fn no_key_loopback_metrics_remain_open_for_development() {
17289 let mut st = fake_worker_state();
17290 st.metrics_auth = MetricsAuth::new(true, false, None);
17291 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
17292 assert_eq!(response.status(), StatusCode::OK);
17293 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
17294 .await
17295 .unwrap();
17296 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17297 assert!(
17298 body.get("active_sessions").is_some(),
17299 "no-key loopback development keeps full operator visibility",
17300 );
17301 assert_eq!(
17302 yield_metrics(State(st), HeaderMap::new()).await.status(),
17303 StatusCode::OK,
17304 );
17305 }
17306
17307 #[test]
17308 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
17309 let metrics = SharedMetrics::default();
17310 let rl = RateLimit::compute(4, 1, &metrics);
17312 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
17313 let rl = RateLimit::compute(4, 3, &metrics);
17314 assert_eq!(rl.remaining, 1);
17315 let rl = RateLimit::compute(4, 4, &metrics);
17317 assert_eq!(rl.remaining, 0);
17318 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
17319 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
17321 let m = worker::Metrics {
17323 completed: 2,
17324 tokens_out: 200,
17325 step_p50_ms: 20.0,
17326 ..Default::default()
17327 };
17328 assert_eq!(reset_estimate_s(&m), 2); }
17330
17331 #[test]
17332 fn inflight_guard_counts_up_and_frees_on_drop() {
17333 let counts: InflightCounts = Arc::new(Default::default());
17334 let tenants: TenantGauge = Arc::new(Default::default());
17335 let (g1, n1, t1) = InflightGuard::try_acquire(
17336 counts.clone(),
17337 lanes::Lane::Interactive,
17338 tenants.clone(),
17339 "acme",
17340 None,
17341 )
17342 .unwrap();
17343 let (g2, n2, t2) = InflightGuard::try_acquire(
17344 counts.clone(),
17345 lanes::Lane::Interactive,
17346 tenants.clone(),
17347 "acme",
17348 None,
17349 )
17350 .unwrap();
17351 assert_eq!((n1, n2), (1, 2));
17352 assert_eq!((t1, t2), (1, 2));
17354 let (gj, nj, tj) = InflightGuard::try_acquire(
17356 counts.clone(),
17357 lanes::Lane::Judge,
17358 tenants.clone(),
17359 "blue",
17360 None,
17361 )
17362 .unwrap();
17363 assert_eq!((nj, tj), (1, 1));
17364 drop(g1);
17365 drop(gj);
17366 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
17367 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
17368 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
17369 assert!(tenants.lock().unwrap().get("blue").is_none());
17371 drop(g2);
17372 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
17373 assert!(tenants.lock().unwrap().is_empty());
17374 }
17375
17376 #[test]
17377 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
17378 let counts: InflightCounts = Arc::new(Default::default());
17379 let tenants: TenantGauge = Arc::new(Default::default());
17380 let start = Arc::new(std::sync::Barrier::new(3));
17381 let attempted = Arc::new(std::sync::Barrier::new(3));
17382 let mut joins = Vec::new();
17383 for _ in 0..2 {
17384 let counts = counts.clone();
17385 let tenants = tenants.clone();
17386 let start = start.clone();
17387 let attempted = attempted.clone();
17388 joins.push(std::thread::spawn(move || {
17389 start.wait();
17390 let result = InflightGuard::try_acquire(
17391 counts,
17392 lanes::Lane::Interactive,
17393 tenants,
17394 "preview_001",
17395 Some(1),
17396 );
17397 let won = result.is_ok();
17398 attempted.wait(); drop(result);
17400 won
17401 }));
17402 }
17403 start.wait();
17404 attempted.wait();
17405 let wins = joins
17406 .into_iter()
17407 .map(|join| join.join().unwrap())
17408 .filter(|won| *won)
17409 .count();
17410 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
17411 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
17412 assert!(tenants.lock().unwrap().is_empty());
17413 }
17414
17415 #[tokio::test]
17416 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
17417 let st = fake_worker_state();
17418 let tenant = auth::TenantCtx {
17419 tenant: "preview_001".into(),
17420 lane_class: auth::LaneClass::Interactive,
17421 rate_limit: Some(1),
17422 key_prefix: None,
17423 };
17424 let first_env = Envelope::new(true);
17425 let (guard, first_rl) =
17426 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
17427 Ok(slot) => slot,
17428 Err(_) => panic!("the first request must acquire the tenant slot"),
17429 };
17430 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
17431
17432 let second_env = Envelope::new(true);
17433 let response =
17434 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
17435 Err(response) => response,
17436 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
17437 };
17438 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
17439 assert_eq!(response.headers()["retry-after"], "2");
17440 assert_eq!(response.headers()["retry-after-ms"], "2000");
17441 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
17442 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
17443 assert_eq!(response.headers()["x-request-id"], second_env.id);
17444 assert_eq!(
17445 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
17446 1,
17447 "rejected request must not consume a lane slot"
17448 );
17449 assert_eq!(
17450 st.tenant_inflight
17451 .lock()
17452 .unwrap()
17453 .get("preview_001")
17454 .copied(),
17455 Some(1),
17456 "rejected request must not increment the tenant gauge"
17457 );
17458 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
17459 .await
17460 .unwrap();
17461 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17462 assert_eq!(payload["error"]["type"], "rate_limit_error");
17463 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
17464 assert!(
17465 payload["error"]["message"]
17466 .as_str()
17467 .unwrap()
17468 .contains("concurrent request limit")
17469 );
17470
17471 drop(guard);
17472 let _ = InflightGuard::try_acquire(
17473 st.inflight.clone(),
17474 lanes::Lane::Interactive,
17475 st.tenant_inflight.clone(),
17476 "preview_001",
17477 Some(1),
17478 )
17479 .expect("slot must reopen after the in-flight request completes");
17480 }
17481
17482 #[test]
17483 fn tenant_rate_limit_override_is_min_with_global_cap() {
17484 let metrics = SharedMetrics::default();
17485 let unlimited = auth::TenantCtx::default_tenant();
17486 let capped = auth::TenantCtx {
17487 tenant: "acme".into(),
17488 lane_class: auth::LaneClass::Interactive,
17489 rate_limit: Some(2),
17490 key_prefix: None,
17491 };
17492 let global = lane_cap(lanes::Lane::Interactive);
17493 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
17495 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
17496 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
17498 assert_eq!((rl.limit, rl.remaining), (2, 1));
17499 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
17500 assert_eq!(rl.remaining, 0);
17501 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
17502 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
17506 assert_eq!(rl.remaining, 0);
17507 let wide = auth::TenantCtx {
17508 rate_limit: Some(global + 100),
17509 ..capped.clone()
17510 };
17511 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
17512 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
17513 }
17514
17515 #[test]
17516 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
17517 let batch = auth::TenantCtx {
17518 tenant: "bulk".into(),
17519 lane_class: auth::LaneClass::Batch,
17520 rate_limit: None,
17521 key_prefix: None,
17522 };
17523 let interactive = auth::TenantCtx::default_tenant();
17524 let hdr = |v: Option<&str>| {
17525 let mut h = axum::http::HeaderMap::new();
17526 if let Some(v) = v {
17527 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
17528 }
17529 h
17530 };
17531 assert_eq!(
17533 lane_for_tenant(&hdr(None), &interactive).unwrap(),
17534 lanes::Lane::Interactive
17535 );
17536 assert_eq!(
17537 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
17538 lanes::Lane::Judge
17539 );
17540 assert_eq!(
17542 lane_for_tenant(&hdr(None), &batch).unwrap(),
17543 lanes::Lane::Harvest
17544 );
17545 assert_eq!(
17546 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
17547 lanes::Lane::Judge
17548 );
17549 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
17550 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
17551 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
17553 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
17554 }
17555
17556 #[tokio::test]
17557 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
17558 let hdr = |v: &str| {
17563 let mut h = axum::http::HeaderMap::new();
17564 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
17565 h
17566 };
17567 let body = |resp: Response| async move {
17568 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17569 .await
17570 .unwrap();
17571 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
17572 };
17573
17574 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
17575 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
17576 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
17577 let payload = body(resp).await;
17578 assert!(
17579 payload["error"].is_object(),
17580 "bare-string error body: {payload}"
17581 );
17582 assert_eq!(payload["error"]["type"], "invalid_request_error");
17583 assert_eq!(payload["error"]["param"], "x-lane");
17584 assert_eq!(payload["error"]["code"], "invalid_lane");
17585
17586 let batch = auth::TenantCtx {
17587 tenant: "bulk".into(),
17588 lane_class: auth::LaneClass::Batch,
17589 rate_limit: None,
17590 key_prefix: None,
17591 };
17592 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
17593 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
17594 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
17595 let payload = body(resp).await;
17596 assert_eq!(payload["error"]["type"], "authentication_error");
17597 assert_eq!(payload["error"]["param"], "x-lane");
17598 }
17599
17600 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
17603
17604 #[tokio::test]
17605 #[allow(clippy::await_holding_lock)] async fn responses_carry_rate_limit_headers_and_slot_frees() {
17607 let _l = DRAIN_LOCK.lock().unwrap();
17608 let st = fake_worker_state();
17609 let resp = chat_completions(
17612 State(st.clone()),
17613 axum::http::HeaderMap::new(),
17614 None,
17615 Json(
17616 serde_json::from_value(serde_json::json!({
17617 "model": "m", "messages": [{"role": "user", "content": "t"}]
17618 }))
17619 .unwrap(),
17620 ),
17621 )
17622 .await;
17623 assert_eq!(resp.status(), StatusCode::OK);
17624 let h = resp.headers();
17625 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
17626 let remaining: usize = h["x-ratelimit-remaining"]
17627 .to_str()
17628 .unwrap()
17629 .parse()
17630 .unwrap();
17631 assert_eq!(remaining, limit - 1);
17632 assert_eq!(h["x-ratelimit-reset"], "0");
17633 assert_eq!(
17634 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
17635 0,
17636 "slot must free at completion"
17637 );
17638 let resp = completions(
17641 State(st.clone()),
17642 axum::http::HeaderMap::new(),
17643 None,
17644 Json(
17645 serde_json::from_value(serde_json::json!({
17646 "model": "m", "prompt": "t", "stream": true
17647 }))
17648 .unwrap(),
17649 ),
17650 )
17651 .await;
17652 assert_eq!(resp.status(), StatusCode::OK);
17653 assert!(resp.headers().contains_key("x-ratelimit-limit"));
17654 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
17655 assert!(resp.headers().contains_key("x-ratelimit-reset"));
17656 assert_eq!(
17657 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
17658 1,
17659 "stream in flight holds the slot"
17660 );
17661 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
17662 .await
17663 .unwrap();
17664 assert_eq!(
17665 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
17666 0,
17667 "slot must free when the stream completes"
17668 );
17669 }
17670
17671 #[tokio::test]
17672 #[allow(clippy::await_holding_lock)] async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
17674 let _l = DRAIN_LOCK.lock().unwrap();
17675 let mut st = fake_worker_state();
17676 let mock = MockMetering::admit_all();
17677 st.metering = Some(mock.clone());
17678
17679 let nonstream = chat_completions(
17680 State(st.clone()),
17681 HeaderMap::new(),
17682 None,
17683 Json(
17684 serde_json::from_value(json!({
17685 "model": "m",
17686 "messages": [{"role": "user", "content": "t"}],
17687 }))
17688 .unwrap(),
17689 ),
17690 )
17691 .await;
17692 assert_eq!(nonstream.status(), StatusCode::OK);
17693 let nonstream_id = nonstream.headers()["x-request-id"]
17694 .to_str()
17695 .unwrap()
17696 .to_string();
17697
17698 let stream = completions(
17699 State(st),
17700 HeaderMap::new(),
17701 None,
17702 Json(
17703 serde_json::from_value(json!({
17704 "model": "m",
17705 "prompt": "t",
17706 "stream": true,
17707 }))
17708 .unwrap(),
17709 ),
17710 )
17711 .await;
17712 assert_eq!(stream.status(), StatusCode::OK);
17713 let stream_id = stream.headers()["x-request-id"]
17714 .to_str()
17715 .unwrap()
17716 .to_string();
17717 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
17718 .await
17719 .unwrap();
17720
17721 let events = mock.events();
17725 let opened: Vec<&str> = events
17726 .iter()
17727 .filter_map(|e| match e {
17728 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
17729 _ => None,
17730 })
17731 .collect();
17732 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
17733 let completes = events
17734 .iter()
17735 .filter(|e| {
17736 matches!(
17737 e,
17738 MeterEvent::Complete {
17739 prompt: 1,
17740 cached: 0,
17741 completion: 1,
17742 }
17743 )
17744 })
17745 .count();
17746 assert_eq!(
17747 completes, 2,
17748 "both surfaces settle complete with worker-truth usage: {events:?}"
17749 );
17750 }
17751
17752 #[tokio::test]
17753 #[allow(clippy::await_holding_lock)] async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
17755 let _l = DRAIN_LOCK.lock().unwrap();
17756 let mock = MockMetering::with_limits(vec![
17762 ReserveScript::Insufficient,
17763 ReserveScript::Admit { with_permit: false },
17764 ReserveScript::Blocked,
17765 ReserveScript::Admit { with_permit: true },
17766 ]);
17767 let mut st = fake_worker_state();
17768 st.metering = Some(mock.clone());
17769
17770 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
17772 assert_eq!(metrics.status(), StatusCode::OK);
17773 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
17774 .await
17775 .unwrap();
17776 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
17777 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
17778 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
17779 assert_eq!(metrics_body["budget_source_available"], true);
17780
17781 let request = || {
17782 Json(
17783 serde_json::from_value::<CompletionReq>(json!({
17784 "model": "m",
17785 "prompt_ids": [1],
17786 "max_tokens": 1,
17787 }))
17788 .unwrap(),
17789 )
17790 };
17791
17792 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
17793 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
17794 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
17795 .await
17796 .unwrap();
17797 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
17798 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
17799 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
17800
17801 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
17802 assert_eq!(included.status(), StatusCode::OK);
17803
17804 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
17807 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
17808
17809 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
17810 assert_eq!(admitted.status(), StatusCode::OK);
17811
17812 let events = mock.events();
17813 let terminal: Vec<&MeterEvent> = events
17814 .iter()
17815 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
17816 .collect();
17817 assert_eq!(
17818 terminal.len(),
17819 4,
17820 "four requests, four terminal settles: {events:?}"
17821 );
17822 assert!(matches!(
17823 terminal[0],
17824 MeterEvent::Reject { status: 402, .. }
17825 ));
17826 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
17827 assert!(matches!(
17828 terminal[2],
17829 MeterEvent::Reject { status: 402, .. }
17830 ));
17831 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
17832 let permits: Vec<bool> = events
17834 .iter()
17835 .filter_map(|e| match e {
17836 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
17837 _ => None,
17838 })
17839 .collect();
17840 assert_eq!(
17841 permits,
17842 vec![false, false, false, true],
17843 "the permit rides the receipt exactly when reserve minted one: {events:?}"
17844 );
17845 }
17846
17847 #[tokio::test]
17851 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
17852 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
17853 let mut st = fake_worker_state();
17854 st.metering = Some(mock.clone());
17855 let tenant = auth::TenantCtx {
17856 tenant: "acme".into(),
17857 lane_class: auth::LaneClass::Interactive,
17858 rate_limit: None,
17859 key_prefix: Some("mk-acme-testprefix00".into()),
17860 };
17861 let mut request = gate_request(1, 1);
17862 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
17863 .expect_err("a capped key must be refused at admission");
17864 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
17865 let (response, outcome) = rejection.into_response();
17866 assert_eq!(outcome, "key_spend_cap_reached");
17867 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
17868 let body = body_value(response).await;
17869 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
17870 assert!(
17871 body["error"]["message"].as_str().unwrap().contains("cap"),
17872 "the 402 must point at the KEY's cap, not tenant credit: {body}"
17873 );
17874 let events = mock.events();
17875 assert!(
17876 events.contains(&MeterEvent::Reserve {
17877 tenant: "acme".into(),
17878 principal: Some("mk-acme-testprefix00".into()),
17879 model: "qwen/qwen3.8-27b".into(),
17880 }),
17881 "the key prefix must reach reserve: {events:?}"
17882 );
17883 }
17884
17885 #[tokio::test]
17886 #[allow(clippy::await_holding_lock)] async fn streaming_client_disconnect_records_partial_usage_and_cost() {
17888 let _l = DRAIN_LOCK.lock().unwrap();
17889 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
17890 let mock = MockMetering::admit_all();
17891 st.metering = Some(mock.clone());
17892
17893 let response = completions(
17894 State(st),
17895 HeaderMap::new(),
17896 None,
17897 Json(
17898 serde_json::from_value(json!({
17899 "model": "m",
17900 "prompt": "disconnect after one delta",
17901 "stream": true,
17902 }))
17903 .unwrap(),
17904 ),
17905 )
17906 .await;
17907 assert_eq!(response.status(), StatusCode::OK);
17908 let request_id = response.headers()["x-request-id"]
17909 .to_str()
17910 .unwrap()
17911 .to_string();
17912 let mut body = Box::pin(response.into_body().into_data_stream());
17913 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
17914 .await
17915 .expect("stream ended before first delta")
17916 .expect("stream body failed");
17917 assert!(
17918 is_sse_data_frame(&first),
17919 "first frame was not SSE data: {first:?}"
17920 );
17921 drop(body);
17922
17923 let mut dropped = None;
17926 for _ in 0..500 {
17927 if let Some(event) = mock
17928 .events()
17929 .into_iter()
17930 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
17931 {
17932 dropped = Some(event);
17933 break;
17934 }
17935 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
17936 }
17937 let events = mock.events();
17938 assert!(
17939 events
17940 .iter()
17941 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
17942 "the receipt was opened under the caller-visible request id: {events:?}"
17943 );
17944 assert_eq!(
17945 dropped,
17946 Some(MeterEvent::Dropped {
17947 prompt: 1,
17948 cached: 0,
17949 completion: 1,
17950 }),
17951 "a client disconnect must leave the partial counts on the dropped receipt \
17952 (the implementation prices that drop): {events:?}"
17953 );
17954 }
17955
17956 #[tokio::test]
17957 #[allow(clippy::await_holding_lock)] async fn draining_rejects_new_requests_with_503_and_retry_after() {
17959 let _l = DRAIN_LOCK.lock().unwrap();
17960 let st = fake_worker_state();
17961 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
17962 let resp = chat_completions(
17964 State(st.clone()),
17965 axum::http::HeaderMap::new(),
17966 None,
17967 Json(
17968 serde_json::from_value(serde_json::json!({
17969 "model": "m", "messages": [{"role": "user", "content": "t"}]
17970 }))
17971 .unwrap(),
17972 ),
17973 )
17974 .await;
17975 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17976 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
17981 let ra_s: u64 = ra
17982 .parse()
17983 .expect("Retry-After must be integer delay-seconds");
17984 assert!(
17985 ra_s > 0 && ra_s <= 60,
17986 "Retry-After {ra_s}s is outside the honored window"
17987 );
17988 let ra_ms: u64 = resp.headers()["retry-after-ms"]
17989 .to_str()
17990 .unwrap()
17991 .parse()
17992 .unwrap();
17993 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
17994 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17995 .await
17996 .unwrap();
17997 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17998 assert!(
17999 payload["error"]["message"]
18000 .as_str()
18001 .unwrap()
18002 .contains("draining")
18003 );
18004 assert_eq!(payload["error"]["type"], "server_error");
18005 assert_eq!(payload["error"]["code"], "draining");
18006 let resp = completions(
18007 State(st.clone()),
18008 axum::http::HeaderMap::new(),
18009 None,
18010 Json(
18011 serde_json::from_value(serde_json::json!({
18012 "model": "m", "prompt": "t"
18013 }))
18014 .unwrap(),
18015 ),
18016 )
18017 .await;
18018 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
18019 assert!(resp.headers().contains_key("retry-after"));
18020 assert_eq!(
18021 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
18022 0,
18023 "rejected requests must not hold slots"
18024 );
18025 let resp = health_live(State(st.clone())).await.into_response();
18028 assert_eq!(
18029 resp.status(),
18030 StatusCode::OK,
18031 "a drain must not look like a liveness fault"
18032 );
18033 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18034 .await
18035 .unwrap();
18036 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18037 assert_eq!(payload["status"], "draining");
18038 let resp = health_ready(State(st.clone())).await.into_response();
18040 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
18041 let retry_s = drain_deadline_s().clamp(1, 60);
18042 let retry_s_text = retry_s.to_string();
18043 let retry_ms_text = (retry_s * 1000).to_string();
18044 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
18045 assert_eq!(
18046 resp.headers().get("retry-after-ms").unwrap(),
18047 retry_ms_text.as_str()
18048 );
18049 assert_ne!(
18050 resp.headers()
18051 .get("x-should-retry")
18052 .and_then(|v| v.to_str().ok()),
18053 Some("false")
18054 );
18055 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18056 .await
18057 .unwrap();
18058 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18059 assert_eq!(payload["status"], "not_ready");
18060 assert!(payload["detail"].as_str().unwrap().contains("draining"));
18061 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18062 let resp = chat_completions(
18064 State(st.clone()),
18065 axum::http::HeaderMap::new(),
18066 None,
18067 Json(
18068 serde_json::from_value(serde_json::json!({
18069 "model": "m", "messages": [{"role": "user", "content": "t"}]
18070 }))
18071 .unwrap(),
18072 ),
18073 )
18074 .await;
18075 assert_eq!(resp.status(), StatusCode::OK);
18076 }
18077
18078 #[tokio::test]
18081 #[allow(clippy::await_holding_lock)] async fn health_is_green_only_while_the_worker_is_alive() {
18083 let _l = DRAIN_LOCK.lock().unwrap();
18086 let st = fake_worker_state();
18087 let resp = health_live(State(st.clone())).await.into_response();
18090 assert_eq!(resp.status(), StatusCode::OK);
18091 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18092 .await
18093 .unwrap();
18094 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18095 assert_eq!(payload["status"], "ok");
18096 assert_eq!(payload["worker"]["phase"], "idle");
18097 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
18098 let ready = health_ready(State(st.clone())).await.into_response();
18099 assert_eq!(ready.status(), StatusCode::OK);
18100
18101 st.health.mark_dead("worker thread panicked: test-injected");
18105 let resp = health_live(State(st.clone())).await.into_response();
18106 assert_eq!(
18107 resp.status(),
18108 StatusCode::SERVICE_UNAVAILABLE,
18109 "a dead worker MUST NOT report a healthy liveness"
18110 );
18111 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18112 .await
18113 .unwrap();
18114 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18115 assert_eq!(payload["status"], "unhealthy");
18116 assert!(
18118 payload["detail"]
18119 .as_str()
18120 .unwrap()
18121 .contains("test-injected"),
18122 "cause not surfaced: {payload}"
18123 );
18124 let ready = health_ready(State(st.clone())).await.into_response();
18125 assert_eq!(
18126 ready.status(),
18127 StatusCode::SERVICE_UNAVAILABLE,
18128 "dead is also not ready"
18129 );
18130
18131 st.health.mark_ready();
18134 assert_eq!(
18135 health_live(State(st.clone()))
18136 .await
18137 .into_response()
18138 .status(),
18139 StatusCode::OK,
18140 "mark_ready must clear the latch (a successful respawn)"
18141 );
18142 }
18143
18144 #[tokio::test]
18145 #[allow(clippy::await_holding_lock)] async fn readyz_peer_probe_integrity_is_present_and_advisory() {
18147 let _l = DRAIN_LOCK.lock().unwrap();
18148 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18149 let st = fake_worker_state();
18150
18151 let ready = health_ready(State(st.clone())).await.into_response();
18152 assert_eq!(ready.status(), StatusCode::OK);
18153 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
18154 .await
18155 .unwrap();
18156 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18157 assert_eq!(payload["peer_probe_integrity"], "ok");
18158
18159 st.health.note_peer_probe_deferral(2, false);
18160 let deferred = health_ready(State(st.clone())).await.into_response();
18161 assert_eq!(deferred.status(), StatusCode::OK);
18162 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
18163 .await
18164 .unwrap();
18165 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18166 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
18167
18168 st.health.note_peer_probe_deferral(4, true);
18169 let degraded = health_ready(State(st.clone())).await.into_response();
18170 assert_eq!(
18171 degraded.status(),
18172 StatusCode::OK,
18173 "peer degradation is advisory while plain serving remains healthy"
18174 );
18175 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
18176 .await
18177 .unwrap();
18178 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18179 assert_eq!(payload["peer_probe_integrity"], "degraded");
18180
18181 st.health.mark_dead("test-injected worker failure");
18182 let unready = health_ready(State(st)).await.into_response();
18183 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
18184 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
18185 .await
18186 .unwrap();
18187 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18188 assert_eq!(
18189 payload["peer_probe_integrity"], "degraded",
18190 "the advisory field must also survive an unrelated readiness failure"
18191 );
18192 }
18193
18194 #[tokio::test]
18195 #[allow(clippy::await_holding_lock)] async fn liveness_failure_obeys_the_retry_contract() {
18197 let _l = DRAIN_LOCK.lock().unwrap();
18201 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18202 let st = fake_worker_state();
18203 st.health
18204 .mark_dead("worker thread panicked: retry-contract-test");
18205
18206 let resp = health_live(State(st)).await.into_response();
18207 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
18208 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
18209 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
18210 assert_ne!(
18211 resp.headers()
18212 .get("x-should-retry")
18213 .and_then(|v| v.to_str().ok()),
18214 Some("false")
18215 );
18216 }
18217
18218 #[tokio::test]
18219 #[allow(clippy::await_holding_lock)] async fn readiness_failure_obeys_the_retry_contract() {
18221 let _l = DRAIN_LOCK.lock().unwrap();
18222 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18223 let st = fake_worker_state();
18224 st.health
18225 .mark_dead("worker thread panicked: retry-contract-test");
18226
18227 let resp = health_ready(State(st)).await.into_response();
18228 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
18229 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
18230 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
18231 assert_ne!(
18232 resp.headers()
18233 .get("x-should-retry")
18234 .and_then(|v| v.to_str().ok()),
18235 Some("false")
18236 );
18237 }
18238
18239 #[tokio::test]
18240 #[allow(clippy::await_holding_lock)] async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
18242 let _l = DRAIN_LOCK.lock().unwrap();
18252 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
18253 let st = fake_worker_state();
18254 assert_eq!(
18255 health_live(State(st.clone()))
18256 .await
18257 .into_response()
18258 .status(),
18259 StatusCode::OK
18260 );
18261 st.health
18262 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
18263 let resp = health_live(State(st.clone())).await.into_response();
18264 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
18265 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
18266 .await
18267 .unwrap();
18268 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18269 assert!(
18270 payload["detail"]
18271 .as_str()
18272 .unwrap()
18273 .contains("probe exceeded")
18274 );
18275 st.health.mark_ready();
18278 assert_eq!(
18279 health_live(State(st.clone()))
18280 .await
18281 .into_response()
18282 .status(),
18283 StatusCode::SERVICE_UNAVAILABLE,
18284 "a GPU fault must not be cleared by an in-process respawn"
18285 );
18286 }
18287
18288 #[test]
18289 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
18290 let caps = ModelCaps {
18292 tools_branch: true,
18293 hy3: false,
18294 qwen_think: true,
18295 think_switch: true,
18296 chat_ok: true,
18297 context_length: 262144,
18298 tokenizer: "qwen2".into(),
18299 instruct_type: Some("chatml".into()),
18300 effort_levels: false,
18301 qwen_effort: false,
18302 gemma_think: false,
18303 dsv4: false,
18304 glm5: false,
18305 chat_temperature_default: None,
18306 chat_top_p_default: None,
18307 n_vocab: 151_936,
18308 think_close: Vec::new(),
18309 };
18310 let e = model_entry_v1("main", Some(&caps), None);
18311 assert_eq!(e["id"], "main");
18312 assert_eq!(e["name"], "main");
18313 assert_eq!(e["object"], "model");
18314 assert_eq!(e["context_length"], 262144);
18315 assert!(e["pricing"]["input"].is_null());
18317 assert!(e["pricing"]["output"].is_null());
18318
18319 let meta = OpenRouterModelMetadata {
18322 pricing: OpenRouterPricing {
18323 prompt: Some("0.00000038".into()),
18324 cached_prompt: Some("0.0000002".into()),
18325 completion: Some("0.0000026".into()),
18326 ..Default::default()
18327 },
18328 input_modalities: vec!["image".into(), "video".into()],
18329 max_output_length: Some(32768),
18330 ..Default::default()
18331 };
18332 let e = model_entry_v1("main", Some(&caps), Some(&meta));
18333 assert_eq!(e["pricing"]["currency"], "USD");
18336 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
18337 assert_eq!(e["pricing"]["input"], "0.38");
18338 assert_eq!(e["pricing"]["output"], "2.60");
18339 assert_eq!(e["pricing"]["cached_input"], "0.20");
18340 assert!(e["pricing"]["cache_write"].is_null());
18341 assert_eq!(e["pricing"]["minimum_request"], "0");
18342 assert_eq!(e["owned_by"], "main");
18343 assert_eq!(e["type"], "chat");
18344 assert_eq!(e["max_output_tokens"], 32768);
18345 assert_eq!(e["endpoints"], json!(["chat/completions"]));
18346 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
18347 assert_eq!(e["output_modalities"], json!(["text"]));
18348 assert_eq!(e["capabilities"]["streaming"], true);
18349 assert_eq!(e["capabilities"]["tools"], true);
18350 assert_eq!(e["lifecycle"]["status"], "active");
18351 assert!(e["lifecycle"]["deprecation_at"].is_null());
18352 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
18353 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
18354 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
18358 keys.sort_unstable();
18359 assert_eq!(
18360 keys,
18361 [
18362 "capabilities",
18363 "context_length",
18364 "endpoints",
18365 "id",
18366 "input_modalities",
18367 "lifecycle",
18368 "max_output_tokens",
18369 "name",
18370 "object",
18371 "output_modalities",
18372 "owned_by",
18373 "pricing",
18374 "reliability",
18375 "type",
18376 ],
18377 "unexpected /v1/models entry keys"
18378 );
18379 let mut price_keys: Vec<&str> = e["pricing"]
18380 .as_object()
18381 .unwrap()
18382 .keys()
18383 .map(String::as_str)
18384 .collect();
18385 price_keys.sort_unstable();
18386 assert_eq!(
18387 price_keys,
18388 [
18389 "cache_write",
18390 "cached_input",
18391 "currency",
18392 "input",
18393 "minimum_request",
18394 "output",
18395 "unit",
18396 ],
18397 "unexpected /v1/models pricing keys"
18398 );
18399
18400 let e = model_entry_v1("m", None, None);
18402 assert!(e["context_length"].is_null());
18403 assert!(e["max_output_tokens"].is_null());
18404 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
18406 assert!(e["context_length"].is_null());
18407 }
18408
18409 #[test]
18415 fn catalog_row_follows_the_declared_surface() {
18416 let caps = ModelCaps {
18417 tools_branch: true,
18418 ..Default::default()
18419 };
18420
18421 let embed = OpenRouterModelMetadata {
18422 surface: Some("embedding".into()),
18423 max_output_length: Some(1),
18424 ..Default::default()
18425 };
18426 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
18427 assert_eq!(e["type"], "embedding");
18428 assert_eq!(e["endpoints"], json!(["embeddings"]));
18429 assert_eq!(e["output_modalities"], json!(["embeddings"]));
18430 assert_eq!(e["capabilities"]["streaming"], false);
18431 assert_eq!(
18432 e["capabilities"]["tools"], false,
18433 "an embedder has no tools"
18434 );
18435 assert_eq!(e["capabilities"]["reasoning"], false);
18436 assert_eq!(e["capabilities"]["structured_output"], false);
18437 assert_eq!(e["capabilities"]["prompt_caching"], false);
18438 assert!(
18439 e["max_output_tokens"].is_null(),
18440 "a surface that emits no completion tokens must not advertise a ceiling"
18441 );
18442
18443 let rerank = OpenRouterModelMetadata {
18444 surface: Some("rerank".into()),
18445 ..Default::default()
18446 };
18447 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
18448 assert_eq!(r["type"], "rerank");
18449 assert_eq!(r["endpoints"], json!(["rerank"]));
18450 assert_eq!(r["output_modalities"], json!(["rerank"]));
18451 assert_eq!(r["capabilities"]["tools"], false);
18452 assert_eq!(r["capabilities"]["reasoning"], false);
18453
18454 let chat = OpenRouterModelMetadata {
18457 max_output_length: Some(32768),
18458 ..Default::default()
18459 };
18460 let c = model_entry_v1("main", Some(&caps), Some(&chat));
18461 assert_eq!(c["type"], "chat");
18462 assert_eq!(c["endpoints"], json!(["chat/completions"]));
18463 assert_eq!(c["output_modalities"], json!(["text"]));
18464 assert_eq!(c["capabilities"]["tools"], true);
18465 assert_eq!(c["max_output_tokens"], 32768);
18466 }
18467
18468 #[test]
18471 fn unknown_surface_is_rejected_at_config_load() {
18472 let bad = OpenRouterModelMetadata {
18473 surface: Some("embeddings".into()), ..Default::default()
18475 };
18476 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
18477 .expect_err("an unknown surface must not load");
18478 assert!(err.contains("surface"), "{err}");
18479
18480 for good in ["chat", "embedding", "rerank"] {
18481 let ok = OpenRouterModelMetadata {
18482 surface: Some(good.into()),
18483 ..Default::default()
18484 };
18485 assert!(
18486 validate_openrouter_metadata("m", &ok).is_ok(),
18487 "{good} must load"
18488 );
18489 }
18490 }
18491
18492 #[test]
18493 fn per_million_price_is_exact_decimal_shift() {
18494 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
18496 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
18497 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
18498 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
18499 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
18500 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
18501 assert_eq!(per_million_price("not-a-price"), None);
18502 assert_eq!(per_million_price(""), None);
18503 }
18504
18505 #[test]
18506 fn metadata_provider_block_parses_and_validates() {
18507 let (_, provider) = OpenRouterMetadataFile::parse(
18508 r#"
18509 [provider]
18510 id = "tiyuvta"
18511 status_url = "https://status.tiyuvta.ai"
18512 support_contact = "mailto:support@tiyuvta.ai"
18513 incident_contact = "mailto:incidents@tiyuvta.ai"
18514 regions = ["eu-central"]
18515 "#,
18516 )
18517 .unwrap();
18518 let provider = provider.unwrap();
18519 assert_eq!(provider.id, "tiyuvta");
18520 assert_eq!(provider.regions, vec!["eu-central"]);
18521 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
18523 assert!(err.contains("provider.id"), "{err}");
18524 let err = OpenRouterMetadataFile::parse(
18526 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
18527 )
18528 .unwrap_err();
18529 assert!(err.contains("must be a URI"), "{err}");
18530 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
18532 assert!(provider.is_none());
18533 }
18534
18535 #[test]
18536 fn models_openai_default_body_stays_byte_identical() {
18537 let body = models_openai_body(&["main".into(), "judge".into()]);
18538 let bytes = serde_json::to_vec(&body).unwrap();
18539 assert_eq!(
18540 bytes,
18541 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
18542 );
18543 }
18544
18545 #[test]
18546 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
18547 let loaded = vec![
18549 "qwen/qwen3.6-27b".to_string(),
18550 "qwen/qwen3.6-35b-a3b".to_string(),
18551 ];
18552 assert_eq!(
18553 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
18554 Some("qwen/qwen3.6-35b-a3b"),
18555 );
18556 assert_eq!(
18557 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
18558 Some("qwen/qwen3.6-27b"),
18559 );
18560 assert_eq!(
18562 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
18563 Some("qwen/qwen3.6-35b-a3b"),
18564 );
18565 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
18567 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
18568 assert_eq!(canonical_model_id(&loaded, ""), None);
18569 }
18570
18571 #[test]
18572 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
18573 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
18576 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
18577 assert_eq!(
18579 canonical_model_id(&loaded, "a/shared-name").as_deref(),
18580 Some("a/shared-name")
18581 );
18582 assert_eq!(
18583 canonical_model_id(&loaded, "b/shared-name").as_deref(),
18584 Some("b/shared-name")
18585 );
18586 let bare = vec!["solo".to_string()];
18588 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
18589 }
18590
18591 #[test]
18592 fn openrouter_models_entry_serializes_complete_metadata() {
18593 let metadata = OpenRouterMetadataFile::from_toml(
18594 r#"
18595[models.main]
18596hugging_face_id = "Qwen/Qwen3.6-27B"
18597created = 1786032000
18598quantization = "nvfp4"
18599description = "Qwen3.6 27B served by memra."
18600max_prompt_length = 245760
18601max_output_length = 16384
18602default_output_length = 4096
18603is_ready = true
18604is_free = false
18605discount_to_user = 0.1
18606openrouter_slug = "qwen/qwen3.6-27b"
18607datacenters = [{ country_code = "US", region = "us-east-1" }]
18608zdr = true
18609hipaa = false
18610
18611[models.main.pricing]
18612prompt = "0.000000234"
18613cached_prompt = "0.0000000585"
18614cache_write = "0.000000234"
18615completion = "0.000001872"
18616internal_reasoning = "0.000001872"
18617request = "0.01"
18618
18619[models.main.capacity]
18620prompt_tpm = 1000000
18621cached_prompt_tpm = 2000000
18622completion_tpm = 500000
18623request_rpm = 1000
18624concurrency = 64
18625"#,
18626 )
18627 .unwrap();
18628 let caps = ModelCaps {
18629 tools_branch: true,
18630 qwen_think: true,
18631 think_switch: true,
18632 chat_ok: true,
18633 context_length: 262144,
18634 tokenizer: "qwen2".into(),
18635 instruct_type: Some("chatml".into()),
18636 ..Default::default()
18637 };
18638 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
18639
18640 assert_eq!(entry["schema_version"], "2.4");
18641 assert_eq!(entry["id"], "main");
18642 assert_eq!(entry["name"], "main");
18643 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
18644 assert_eq!(entry["created"], 1786032000u64);
18645 assert_eq!(entry["quantization"], "nvfp4");
18646 assert_eq!(entry["tokenizer"], "qwen2");
18647 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
18648 assert!(
18649 entry.get("object").is_none(),
18650 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
18651 );
18652
18653 let input = &entry["input_modalities"][0];
18654 assert_eq!(input["type"], "text");
18655 assert_eq!(
18656 input["supported_inputs"]["max_context_length"]["value"],
18657 262144
18658 );
18659 assert_eq!(
18660 input["supported_inputs"]["max_prompt_length"]["value"],
18661 245760
18662 );
18663 let input_prices = input["pricing"].as_array().unwrap();
18664 let input_price = |kind: &str| {
18665 input_prices
18666 .iter()
18667 .find(|price| price["type"] == kind)
18668 .unwrap()
18669 };
18670 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
18671 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
18672 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
18673 assert_eq!(input["capacity"][0]["value"], 1000000);
18674 assert_eq!(input["capacity"][1]["value"], 2000000);
18675
18676 let output = &entry["output_modalities"][0];
18677 assert_eq!(output["type"], "text");
18678 assert_eq!(output["max_length"]["value"], 16384);
18679 assert_eq!(output["streaming"], true);
18680 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
18681 assert_eq!(
18682 output["supported_parameters"]["structured_outputs"]["type"],
18683 "boolean"
18684 );
18685 assert_eq!(
18686 output["supported_parameters"]["reasoning"]["type"],
18687 "boolean"
18688 );
18689 assert_eq!(output["pricing"][0]["type"], "completion");
18690 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
18691 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
18692 assert_eq!(output["capacity"][0]["value"], 500000);
18693 assert_eq!(output["capacity"][1]["type"], "concurrency");
18694 assert_eq!(output["capacity"][1]["value"], 64);
18695
18696 assert_eq!(entry["pricing"][0]["type"], "request");
18697 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
18698 assert_eq!(entry["capacity"][0]["value"], 1000);
18699 assert_eq!(entry["is_ready"], true);
18700 assert_eq!(entry["is_free"], false);
18701 assert_eq!(entry["discount_to_user"], 0.1);
18702 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
18703 assert_eq!(entry["datacenters"][0]["country_code"], "US");
18704 assert_eq!(entry["compliance"]["zdr"], true);
18705 assert_eq!(entry["compliance"]["hipaa"], false);
18706 }
18707
18708 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
18713[models."qwen/qwen3.6-35b-a3b"]
18714hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
18715created = 1777260255
18716quantization = "int4"
18717description = "Qwen3.6 35B-A3B fixture entry."
18718max_prompt_length = 262144
18719max_output_length = 262144
18720default_output_length = 8192
18721is_ready = true
18722is_free = false
18723discount_to_user = 0.0
18724openrouter_slug = "qwen/qwen3.6-35b-a3b"
18725zdr = false
18726hipaa = false
18727
18728[[models."qwen/qwen3.6-35b-a3b".datacenters]]
18729country_code = "CA"
18730region = "Ontario"
18731
18732[models."qwen/qwen3.6-35b-a3b".pricing]
18733prompt = "0.0000000931"
18734cached_prompt = "0.0000000652"
18735completion = "0.0000009025"
18736
18737[models."qwen/qwen3.6-35b-a3b".capacity]
18738prompt_tpm = 780000
18739cached_prompt_tpm = 310000
18740completion_tpm = 9600
18741request_rpm = 160
18742concurrency = 16
18743
18744[planned_models."qwen/qwen3.8-27b"]
18745description = "Planned fixture entry; must never be emitted."
18746max_prompt_length = 262144
18747max_output_length = 262144
18748default_output_length = 8192
18749is_ready = false
18750is_free = false
18751discount_to_user = 0.0
18752openrouter_slug = "qwen/qwen3.8-27b"
18753zdr = false
18754hipaa = false
18755
18756[planned_models."qwen/qwen3.8-27b".pricing]
18757prompt = "0.0000002745"
18758cached_prompt = "0.0000001922"
18759completion = "0.0000022800"
18760
18761[planned_models."google/gemma-4-26b-a4b-it"]
18762hugging_face_id = "google/gemma-4-26B-A4B-it"
18763created = 1775227989
18764quantization = "int4"
18765description = "Planned fixture entry; must never be emitted."
18766max_prompt_length = 262144
18767max_output_length = 262144
18768default_output_length = 8192
18769is_ready = false
18770is_free = false
18771discount_to_user = 0.0
18772openrouter_slug = "google/gemma-4-26b-a4b-it"
18773zdr = false
18774hipaa = false
18775
18776[planned_models."google/gemma-4-26b-a4b-it".pricing]
18777prompt = "0.0000000665"
18778cached_prompt = "0.0000000466"
18779completion = "0.0000003230"
18780"#;
18781
18782 #[test]
18783 fn gateway_registry_generates_the_staged_active_shape() {
18784 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
18785 let caps = ModelCaps {
18786 tools_branch: true,
18787 qwen_think: true,
18788 think_switch: true,
18789 chat_ok: true,
18790 context_length: 262144,
18791 tokenizer: "qwen2".into(),
18792 instruct_type: Some("chatml".into()),
18793 ..Default::default()
18794 };
18795 let q35_entry = model_entry_openrouter(
18796 "qwen/qwen3.6-35b-a3b",
18797 Some(&caps),
18798 metadata.get("qwen/qwen3.6-35b-a3b"),
18799 );
18800 assert_eq!(q35_entry["created"], 1777260255u64);
18801 assert_eq!(q35_entry["quantization"], "int4");
18802 assert_eq!(q35_entry["is_ready"], true);
18803 assert_eq!(
18804 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
18805 262144
18806 );
18807 assert_eq!(
18808 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
18809 262144
18810 );
18811 assert_eq!(
18812 q35_entry["output_modalities"][0]["max_length"]["value"],
18813 262144
18814 );
18815 let prices = q35_entry["input_modalities"][0]["pricing"]
18816 .as_array()
18817 .unwrap();
18818 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
18819 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
18820 assert_eq!(
18824 q35_entry["input_modalities"][0]["capacity"][0]["value"],
18825 780000
18826 );
18827 assert_eq!(
18828 q35_entry["input_modalities"][0]["capacity"][1]["value"],
18829 310000
18830 );
18831 assert_eq!(
18832 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
18833 262144
18834 );
18835 assert_eq!(
18836 q35_entry["output_modalities"][0]["capacity"][0]["value"],
18837 9600
18838 );
18839 assert_eq!(
18840 q35_entry["output_modalities"][0]["capacity"][1]["value"],
18841 16
18842 );
18843 assert_eq!(
18844 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
18845 "0.0000009025"
18846 );
18847 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
18849
18850 assert_eq!(
18851 metadata.len(),
18852 1,
18853 "planned models must never enter the active map"
18854 );
18855 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
18856 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
18857 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
18858
18859 let openmodels = model_entry_openmodels(
18860 "qwen/qwen3.6-35b-a3b",
18861 Some(&caps),
18862 metadata.get("qwen/qwen3.6-35b-a3b"),
18863 )
18864 .unwrap();
18865 assert_eq!(openmodels["currency"], "USD");
18866 assert_eq!(openmodels["max_output_length"], 262144);
18867 assert_eq!(openmodels["is_ready"], true);
18868 assert_eq!(openmodels["is_free"], false);
18869 assert_eq!(openmodels["discount_to_user"], 0.0);
18870 }
18871
18872 #[test]
18873 fn gateway_registry_limits_are_live_request_limits() {
18874 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
18875 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
18876 let caps = ModelCaps {
18877 context_length: 262_144,
18878 ..Default::default()
18879 };
18880 let build = |value: serde_json::Value| {
18881 let req: CompletionReq = serde_json::from_value(value).unwrap();
18882 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
18883 build_request(&req, tx, lanes::Lane::Interactive, None)
18884 };
18885
18886 let mut omitted = build(json!({
18887 "model": "qwen/qwen3.6-35b-a3b",
18888 "prompt_ids": [1, 2, 3]
18889 }));
18890 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
18891 assert_eq!(omitted.params.max_new, 8_192);
18892 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
18893
18894 let mut field_top = build(json!({
18895 "model": "qwen/qwen3.6-35b-a3b",
18896 "prompt_ids": [1],
18897 "max_tokens": 262144
18898 }));
18899 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
18900 assert_eq!(field_top.params.max_new, 262_144);
18901 assert_eq!(
18902 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
18903 262_044,
18904 "the field-top output request is accepted but bounded by remaining trained context",
18905 );
18906
18907 let mut too_much_output = build(json!({
18908 "model": "qwen/qwen3.6-35b-a3b",
18909 "prompt_ids": [1],
18910 "max_tokens": 262145
18911 }));
18912 let (message, param) =
18913 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
18914 .unwrap_err();
18915 assert_eq!(param, "max_tokens");
18916 assert!(message.contains("262145"));
18917
18918 let mut oversized_allocation = build(json!({
18919 "model": "qwen/qwen3.6-35b-a3b",
18920 "prompt_ids": [1],
18921 "max_tokens": 1,
18922 "max_ctx": 262145
18923 }));
18924 let (_, param) =
18925 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
18926 .unwrap_err();
18927 assert_eq!(param, "max_ctx");
18928 }
18929
18930 #[test]
18931 fn planned_registry_entries_are_validated_but_never_activated() {
18932 let parsed = OpenRouterMetadataFile::from_toml(
18933 r#"
18934[planned_models.future]
18935max_output_length = 262144
18936default_output_length = 8192
18937
18938[planned_models.future.pricing]
18939prompt = "0.0000001"
18940"#,
18941 )
18942 .unwrap();
18943 assert!(parsed.is_empty());
18944
18945 let error = OpenRouterMetadataFile::from_toml(
18946 r#"
18947[planned_models.future]
18948default_output_length = 8192
18949"#,
18950 )
18951 .unwrap_err();
18952 assert!(error.contains("requires max_output_length"));
18953 }
18954
18955 #[test]
18960 fn every_catalog_feed_honours_the_declared_surface() {
18961 let metadata = OpenRouterMetadataFile::from_toml(
18962 r#"
18963[models."qwen/qwen3-embedding-8b"]
18964surface = "embedding"
18965created = 1787961600
18966max_output_length = 1
18967is_ready = true
18968is_free = false
18969discount_to_user = 0.0
18970
18971[models."qwen/qwen3-embedding-8b".pricing]
18972prompt = "0.00000001"
18973cached_prompt = "0.0"
18974completion = "0.0"
18975
18976[models."main"]
18977created = 1787443200
18978max_output_length = 32768
18979is_ready = true
18980is_free = false
18981discount_to_user = 0.0
18982
18983[models."main".pricing]
18984prompt = "0.00000025"
18985cached_prompt = "0.00000009"
18986completion = "0.0000012"
18987"#,
18988 )
18989 .unwrap();
18990 let caps = ModelCaps {
18991 tools_branch: true,
18992 qwen_think: true,
18993 think_switch: true,
18998 chat_ok: true,
18999 context_length: 32768,
19000 ..Default::default()
19001 };
19002 let embed = metadata.get("qwen/qwen3-embedding-8b");
19003 let chat = metadata.get("main");
19004
19005 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
19007 let out = &or["output_modalities"][0];
19008 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
19009 assert!(
19010 out.get("streaming").is_none(),
19011 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
19012 );
19013 let params = &out["supported_parameters"];
19018 assert_eq!(
19019 params.as_object().map(|o| o.len()),
19020 Some(0),
19021 "no completion parameter belongs on an embedder row: {params}"
19022 );
19023 for field in [
19024 "tools",
19025 "tool_choice",
19026 "reasoning",
19027 "max_tokens",
19028 "json_mode",
19029 "structured_outputs",
19030 "stop",
19031 "temperature",
19032 "seed",
19033 ] {
19034 assert!(params[field].is_null(), "{field} leaked onto an embedder");
19035 }
19036 assert!(
19037 out["max_length"].is_null(),
19038 "a surface emitting no completion tokens advertises no ceiling: {out}"
19039 );
19040
19041 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
19043 .expect("openmodels entry builds");
19044 assert_eq!(om["output_modalities"], json!(["embeddings"]));
19045 let features = om["supported_features"].as_array().unwrap();
19046 assert!(
19047 !features
19048 .iter()
19049 .any(|f| f == "tool_calling" || f == "reasoning"),
19050 "chat-only features leaked onto an embedder: {features:?}"
19051 );
19052
19053 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
19055 assert_eq!(v1["type"], "embedding");
19056 assert_eq!(v1["capabilities"]["tools"], false);
19057
19058 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
19060 let out_chat = &or_chat["output_modalities"][0];
19061 assert_eq!(out_chat["type"], "text");
19062 assert_eq!(out_chat["streaming"], true);
19063 assert!(!out_chat["supported_parameters"]["tools"].is_null());
19064 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
19065 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
19066 assert_eq!(out_chat["max_length"]["value"], 32768u64);
19067 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
19068 assert_eq!(om_chat["output_modalities"], json!(["text"]));
19069 assert!(
19070 om_chat["supported_features"]
19071 .as_array()
19072 .unwrap()
19073 .iter()
19074 .any(|f| f == "tool_calling")
19075 );
19076 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
19077 }
19078
19079 #[test]
19086 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
19087 let raw = std::fs::read_to_string(concat!(
19088 env!("CARGO_MANIFEST_DIR"),
19089 "/../../research/gateway-20260812/raw/sources/",
19090 "openrouter-provider-schema-v2.4-20260812.json"
19091 ))
19092 .expect("vendored Provider Monitor 2.4 schema is in-tree");
19093 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
19094 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
19095 .as_array()
19096 .expect("OutputModality is a oneOf");
19097
19098 let metadata = OpenRouterMetadataFile::from_toml(
19099 r#"
19100[models."embed"]
19101surface = "embedding"
19102created = 1787961600
19103max_output_length = 1
19104is_ready = true
19105is_free = false
19106discount_to_user = 0.0
19107
19108[models."embed".pricing]
19109prompt = "0.00000001"
19110cached_prompt = "0.0"
19111completion = "0.0"
19112
19113[models."rr"]
19114surface = "rerank"
19115created = 1787961600
19116max_output_length = 1
19117is_ready = true
19118is_free = false
19119discount_to_user = 0.0
19120
19121[models."rr".pricing]
19122prompt = "0.00000003"
19123cached_prompt = "0.0"
19124completion = "0.0"
19125
19126[models."chatty"]
19127created = 1787443200
19128max_output_length = 32768
19129is_ready = true
19130is_free = false
19131discount_to_user = 0.0
19132
19133[models."chatty".pricing]
19134prompt = "0.00000025"
19135cached_prompt = "0.00000009"
19136completion = "0.0000012"
19137"#,
19138 )
19139 .unwrap();
19140 let caps = ModelCaps {
19141 tools_branch: true,
19142 qwen_think: true,
19143 chat_ok: true,
19144 context_length: 32768,
19145 ..Default::default()
19146 };
19147
19148 for (alias, want_type) in [
19149 ("embed", "embeddings"),
19150 ("rr", "rerank"),
19151 ("chatty", "text"),
19152 ] {
19153 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
19154 let modality = &row["output_modalities"][0];
19155 assert_eq!(modality["type"], want_type, "{alias}: {row}");
19156
19157 let branch = branches
19159 .iter()
19160 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
19161 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
19162 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
19163 .as_object()
19164 .expect("branch properties")
19165 .keys()
19166 .map(String::as_str)
19167 .collect();
19168 for key in modality.as_object().expect("modality object").keys() {
19169 assert!(
19170 allowed.contains(key.as_str()),
19171 "{alias}: {key:?} is not a property of the {want_type:?} branch \
19172 (additionalProperties:false); allowed = {allowed:?}"
19173 );
19174 }
19175 for req in branch["required"].as_array().into_iter().flatten() {
19176 let req = req.as_str().expect("required entry is a string");
19177 assert!(
19178 modality.get(req).is_some(),
19179 "{alias}: required property {req:?} missing from the {want_type:?} branch"
19180 );
19181 }
19182 }
19183 }
19184
19185 #[test]
19186 fn openrouter_models_entry_omits_undeclared_optional_fields() {
19187 let entry = model_entry_openrouter("minimal", None, None);
19188 let object = entry.as_object().unwrap();
19189 for field in [
19190 "hugging_face_id",
19191 "created",
19192 "quantization",
19193 "tokenizer",
19194 "description",
19195 "pricing",
19196 "capacity",
19197 "is_ready",
19198 "is_free",
19199 "discount_to_user",
19200 "openrouter",
19201 "datacenters",
19202 "compliance",
19203 ] {
19204 assert!(
19205 !object.contains_key(field),
19206 "optional field {field} must be absent, not null"
19207 );
19208 }
19209 assert_eq!(entry["schema_version"], "2.4");
19210 assert_eq!(entry["input_modalities"][0]["type"], "text");
19211 assert!(
19212 entry["input_modalities"][0]
19213 .get("supported_inputs")
19214 .is_none()
19215 );
19216 assert!(entry["input_modalities"][0].get("pricing").is_none());
19217 assert!(entry["input_modalities"][0].get("capacity").is_none());
19218 assert_eq!(entry["output_modalities"][0]["type"], "text");
19219 assert_eq!(entry["output_modalities"][0]["streaming"], true);
19220 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
19221 assert!(entry["output_modalities"][0].get("max_length").is_none());
19222 assert!(entry["output_modalities"][0].get("pricing").is_none());
19223 assert!(entry["output_modalities"][0].get("capacity").is_none());
19224 }
19225
19226 #[test]
19227 fn openmodels_entry_serializes_standard_provider_shape() {
19228 let metadata = OpenRouterMetadataFile::from_toml(
19229 r#"
19230[models."qwen/qwen3.6-27b"]
19231created = 1786032000
19232max_output_length = 16384
19233is_ready = true
19234is_free = false
19235discount_to_user = 0.05
19236
19237[models."qwen/qwen3.6-27b".pricing]
19238prompt = "0.000000291"
19239cached_prompt = "0.000000291"
19240completion = "0.000002763"
19241request = "0"
19242"#,
19243 )
19244 .unwrap();
19245 let caps = ModelCaps {
19246 tools_branch: true,
19247 qwen_think: true,
19248 chat_ok: true,
19249 context_length: 262144,
19250 ..Default::default()
19251 };
19252 let entry = model_entry_openmodels(
19253 "qwen/qwen3.6-27b",
19254 Some(&caps),
19255 metadata.get("qwen/qwen3.6-27b"),
19256 )
19257 .unwrap();
19258
19259 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
19260 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
19261 assert_eq!(entry["created"], 1786032000u64);
19262 assert_eq!(entry["input_modalities"], json!(["text"]));
19263 assert_eq!(entry["output_modalities"], json!(["text"]));
19264 assert_eq!(entry["context_length"], 262144u64);
19265 assert_eq!(entry["max_output_length"], 16384u64);
19266 assert_eq!(entry["currency"], "USD");
19267 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
19268 assert_eq!(entry["pricing"]["completion"], "0.000002763");
19269 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
19270 assert_eq!(entry["pricing"]["request"], "0");
19271 assert_eq!(
19272 entry["supported_features"],
19273 json!(["tool_calling", "reasoning"])
19274 );
19275 assert_eq!(entry["is_ready"], true);
19276 assert_eq!(entry["is_free"], false);
19277 assert_eq!(entry["discount_to_user"], 0.05);
19278 assert!(entry.get("schema_version").is_none());
19279 assert!(entry.get("quantization").is_none());
19280 }
19281
19282 #[test]
19283 fn openmodels_entry_rejects_missing_operator_metadata() {
19284 let caps = ModelCaps {
19285 context_length: 262144,
19286 ..Default::default()
19287 };
19288 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
19289 assert_eq!(
19290 error,
19291 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
19292 );
19293 }
19294
19295 #[tokio::test]
19296 async fn blocking_response_excludes_stop_text_across_token_events() {
19297 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
19298 tx.send(Event::Token {
19299 id: 1,
19300 text: "answer\nPro".into(),
19301 })
19302 .unwrap();
19303 tx.send(Event::Token {
19304 id: 2,
19305 text: "blem: leaked prompt".into(),
19306 })
19307 .unwrap();
19308 tx.send(Event::Done {
19309 stop_reason: "Callback".into(),
19310 n_tokens: 2,
19311 n_prompt: 8,
19312 n_cached: 0,
19313 elapsed_s: 0.5,
19314 spec: None,
19315 })
19316 .unwrap();
19317 drop(tx);
19318 let response = blocking_response(
19319 rx,
19320 "plain_quant".into(),
19321 false,
19322 vec!["Problem:".into()],
19323 None,
19324 Envelope::new(false),
19325 )
19326 .await;
19327 assert_eq!(response.status(), StatusCode::OK);
19328 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
19329 .await
19330 .unwrap();
19331 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
19332 assert_eq!(payload["text"], "answer\n");
19333 assert_eq!(payload["stop_reason"], "Callback");
19334 }
19335
19336 #[test]
19340 fn step_walker_expansion_and_separator_law() {
19341 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
19343 let uri = format!("data:image/png;base64,{PNG64}");
19344 let content = serde_json::json!([
19345 {"type": "text", "text": "look at"},
19346 {"type": "text", "text": "this:"},
19347 {"type": "image_url", "image_url": {"url": uri}},
19348 {"type": "text", "text": "what is it?"},
19349 ]);
19350 let mut pending: Vec<PendingStepImage> = Vec::new();
19351 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
19352 let mut expansion = String::from("<im_start>");
19353 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
19354 expansion.push_str("<im_patch>");
19355 }
19356 expansion.push_str("<im_end>");
19357 assert_eq!(out, format!("look at this:{expansion}what is it?"));
19360 assert_eq!(pending.len(), 1);
19361 assert_eq!(pending[0].plan.n_tiles, 0);
19362 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
19363
19364 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
19366 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
19367 let http = serde_json::json!([
19368 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
19369 ]);
19370 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
19371 }
19372}