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 anthropic;
80mod dsv4_serve;
81mod embed_api;
82pub mod metering;
89mod responses_api;
90mod surfaces;
91mod toolcall;
92mod ttft;
93mod worker;
94
95use std::collections::HashMap;
96use std::net::{SocketAddr, ToSocketAddrs};
97use std::sync::Arc;
98use std::sync::mpsc::Sender;
99
100use axum::{
101 Extension, Json, Router,
102 body::Body,
103 extract::{DefaultBodyLimit, Query, Request as AxumRequest, State},
104 http::{
105 HeaderMap, StatusCode,
106 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
107 },
108 middleware::{self, Next},
109 response::{
110 IntoResponse, Response,
111 sse::{Event as SseEvent, Sse},
112 },
113 routing::{get, post},
114};
115use futures_core::Stream as _;
116use serde::{Deserialize, Serialize};
117use serde_json::json;
118
119use memra_engine::decode::GenParams;
120use memra_engine::sampler::SamplerConfig;
121use memra_tokenizer::{
122 Tokenizer,
123 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
124};
125use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
126use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
127
128const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
145const MAX_BODY_ADMISSIONS: usize = 4;
146const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
147const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
151const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
152const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
153const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
154const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
155const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
156
157fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
158 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
159 SEMAPHORE
160 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
161 .clone()
162}
163
164fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
165 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
166 SEMAPHORE
167 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
168 .clone()
169}
170
171fn declared_body_length(req: &AxumRequest) -> Option<usize> {
172 req.headers()
173 .get(CONTENT_LENGTH)
174 .and_then(|value| value.to_str().ok())
175 .and_then(|value| value.parse().ok())
176}
177
178fn body_requires_admission(req: &AxumRequest) -> bool {
179 if req.headers().contains_key(TRANSFER_ENCODING) {
183 return true;
184 }
185 declared_body_length(req).map_or(true, |length| length > BODY_ADMISSION_BYPASS_BYTES)
186}
187
188fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
192 let Some(length) = declared_body_length(req) else {
193 return BODY_READ_TIMEOUT;
194 };
195 let bytes = length as u64;
196 let extra_seconds =
197 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
198 let seconds = BODY_READ_TIMEOUT
199 .as_secs()
200 .saturating_add(extra_seconds)
201 .min(BODY_READ_TIMEOUT_MAX.as_secs());
202 std::time::Duration::from_secs(seconds)
203}
204
205async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
210 let resp = next.run(req).await;
211 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
212 return resp;
213 }
214 error_response_coded(
215 StatusCode::PAYLOAD_TOO_LARGE,
216 &format!(
217 "request body exceeds the {} MiB limit",
218 MAX_BODY_BYTES / (1024 * 1024)
219 ),
220 "invalid_request_error",
221 None,
222 Some("request_too_large"),
223 )
224}
225
226fn apply_body_limit(app: Router) -> Router {
229 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
230 .layer(middleware::from_fn(shape_payload_too_large))
231}
232
233fn protected_inference_path(path: &str) -> bool {
234 matches!(
235 path,
236 "/v1/auth/check"
237 | "/v1/completions"
238 | "/v1/chat/completions"
239 | "/v1/messages"
240 | "/v1/responses"
241 | "/v1/embeddings"
242 | "/v1/rerank"
243 )
244}
245
246async fn shape_inference_early_response(path: &str, response: Response) -> Response {
250 let request_id = Envelope::new(path != "/v1/completions");
251 if path == "/v1/messages" {
252 anthropic::with_anthropic_request_id(
253 &request_id.id,
254 anthropic::reshape_error(response, &request_id.id).await,
255 )
256 } else {
257 with_request_id(&request_id.id, response)
258 }
259}
260
261async fn authenticate_inference_before_body(
266 State(st): State<AppState>,
267 mut req: AxumRequest,
268 next: Next,
269) -> Response {
270 if !protected_inference_path(req.uri().path()) {
271 return next.run(req).await;
272 }
273 let path = req.uri().path().to_string();
274 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
278 return shape_inference_early_response(
279 &path,
280 error_response_coded(
281 StatusCode::PAYLOAD_TOO_LARGE,
282 &format!(
283 "request body exceeds the {} MiB limit",
284 MAX_BODY_BYTES / (1024 * 1024)
285 ),
286 "invalid_request_error",
287 None,
288 Some("request_too_large"),
289 ),
290 )
291 .await;
292 }
293 let headers = req.headers();
294 let bearer = bearer_token(headers);
295 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
296 let api_key = headers
297 .get("x-api-key")
298 .and_then(|value| value.to_str().ok());
299 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
300 } else {
301 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
302 };
303 if let Err(why) = auth {
304 return shape_inference_early_response(&path, authentication_error(why)).await;
305 }
306 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
314 let body_admission = if body_requires_admission(&req) {
315 body_admission_semaphore()
316 } else {
317 small_body_admission_semaphore()
318 };
319 let body_permit = match body_admission.try_acquire_owned() {
320 Ok(permit) => Some(permit),
321 Err(tokio::sync::TryAcquireError::Closed) => {
322 let response = retry_contract_response(
323 error_response_coded(
324 StatusCode::SERVICE_UNAVAILABLE,
325 "request body admission is unavailable",
326 "server_error",
327 None,
328 Some("body_admission_unavailable"),
329 ),
330 Some(BODY_ADMISSION_RETRY_AFTER_S),
331 );
332 return shape_inference_early_response(&path, response).await;
333 }
334 Err(tokio::sync::TryAcquireError::NoPermits) => {
335 let response = retry_contract_response(
336 error_response_coded(
337 StatusCode::TOO_MANY_REQUESTS,
338 "request body admission is busy",
339 "rate_limit_error",
340 None,
341 Some("body_admission_busy"),
342 ),
343 Some(BODY_ADMISSION_RETRY_AFTER_S),
344 );
345 return shape_inference_early_response(&path, response).await;
346 }
347 };
348 let body = std::mem::replace(req.body_mut(), Body::empty());
352 let mut body = Box::pin(body.into_data_stream());
353 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
354 let body_timed_out_flag = body_timed_out.clone();
355 let guarded_body = async_stream::stream! {
356 loop {
357 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
358 if remaining.is_zero() {
359 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
360 yield Err(std::io::Error::new(
361 std::io::ErrorKind::TimedOut,
362 "request body read deadline exceeded",
363 ));
364 break;
365 }
366 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
367 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
368 Ok(frame) => frame,
369 Err(_) => {
370 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
371 yield Err(std::io::Error::new(
372 std::io::ErrorKind::TimedOut,
373 "request body idle timeout exceeded",
374 ));
375 break;
376 }
377 };
378 match frame {
379 Some(Ok(bytes)) => yield Ok(bytes),
380 Some(Err(error)) => {
381 yield Err(std::io::Error::other(error.to_string()));
382 break;
383 }
384 None => break,
385 }
386 }
387 drop(body_permit);
388 };
389 *req.body_mut() = Body::from_stream(guarded_body);
390 let response = next.run(req).await;
391 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
392 let request_id = Envelope::new(path != "/v1/completions");
393 let timeout = error_response_coded(
394 StatusCode::REQUEST_TIMEOUT,
395 "request body read timed out",
396 "invalid_request_error",
397 None,
398 Some("request_body_timeout"),
399 );
400 return if path == "/v1/messages" {
401 anthropic::with_anthropic_request_id(
402 &request_id.id,
403 anthropic::reshape_error(timeout, &request_id.id).await,
404 )
405 } else {
406 with_request_id(&request_id.id, timeout)
407 };
408 }
409 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
410 let request_id = Envelope::new(true);
411 return anthropic::with_anthropic_request_id(
412 &request_id.id,
413 anthropic::reshape_error(response, &request_id.id).await,
414 );
415 }
416 response
417}
418
419#[cfg(test)]
420mod body_limit_tests {
421 use super::*;
422 use tower::ServiceExt as _;
423
424 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
425
426 fn test_app() -> Router {
430 let app = Router::new()
431 .route(
432 "/bytes",
433 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
434 )
435 .route(
436 "/json",
437 post(|Json(v): Json<serde_json::Value>| async move {
438 v["pad"].as_str().unwrap_or("").len().to_string()
439 }),
440 );
441 apply_body_limit(app)
442 }
443
444 fn streamed_body(chunks: usize) -> Body {
445 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
448 Body::from_stream(async_stream::stream! {
449 for _ in 0..chunks {
450 yield Ok::<_, std::io::Error>(chunk.clone());
451 }
452 })
453 }
454
455 #[tokio::test]
456 async fn bodies_past_the_old_2mib_default_are_accepted() {
457 for (path, body) in [
460 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
461 (
462 "/json",
463 Body::from(
464 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
465 ),
466 ),
467 ] {
468 let resp = test_app()
469 .oneshot(
470 axum::http::Request::post(path)
471 .header(CONTENT_TYPE, "application/json")
472 .body(body)
473 .unwrap(),
474 )
475 .await
476 .unwrap();
477 assert_eq!(resp.status(), StatusCode::OK, "{path}");
478 }
479 }
480
481 #[tokio::test]
482 async fn body_at_exactly_the_limit_is_accepted() {
483 let resp = test_app()
484 .oneshot(
485 axum::http::Request::post("/bytes")
486 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
487 .unwrap(),
488 )
489 .await
490 .unwrap();
491 assert_eq!(resp.status(), StatusCode::OK);
492 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
493 .await
494 .unwrap();
495 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
496 }
497
498 #[tokio::test]
499 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
500 for path in ["/bytes", "/json"] {
504 let resp = test_app()
505 .oneshot(
506 axum::http::Request::post(path)
507 .header(CONTENT_TYPE, "application/json")
508 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
509 .unwrap(),
510 )
511 .await
512 .unwrap();
513 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
514 assert_eq!(
515 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
516 Some(b"false".as_ref()),
517 "{path}: retrying identical bytes cannot fix a 413"
518 );
519 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
520 .await
521 .unwrap();
522 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
523 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
524 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
525 assert!(
526 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
527 "{path}: message names the limit"
528 );
529 }
530 }
531
532 #[tokio::test]
533 async fn authenticated_body_admission_is_finite() {
534 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
535 let semaphore = body_admission_semaphore();
536 let mut permits = Vec::new();
537 for _ in 0..MAX_BODY_ADMISSIONS {
538 permits.push(semaphore.clone().acquire_owned().await.unwrap());
539 }
540 assert!(
541 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
542 .await
543 .is_err(),
544 "body parser admission must not be unbounded"
545 );
546 drop(permits);
547 assert!(semaphore.acquire().await.is_ok());
548 }
549
550 #[tokio::test]
551 async fn small_body_admission_is_finite_and_separate() {
552 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
553 let large = body_admission_semaphore();
554 let small = small_body_admission_semaphore();
555 let mut small_permits = Vec::new();
556 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
557 small_permits.push(small.clone().acquire_owned().await.unwrap());
558 }
559 assert!(
560 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
561 .await
562 .is_err(),
563 "small body parser admission must be bounded"
564 );
565 assert!(
566 large.clone().try_acquire().is_ok(),
567 "small uploads must not consume large-upload permits"
568 );
569 drop(small_permits);
570 assert!(small.acquire().await.is_ok());
571 }
572
573 #[test]
574 fn small_declared_bodies_bypass_large_upload_admission() {
575 let request = axum::http::Request::post("/v1/chat/completions")
576 .header(CONTENT_LENGTH, "2048")
577 .body(Body::empty())
578 .unwrap();
579 assert!(!body_requires_admission(&request));
580
581 let request = axum::http::Request::post("/v1/chat/completions")
582 .header(
583 CONTENT_LENGTH,
584 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
585 )
586 .body(Body::empty())
587 .unwrap();
588 assert!(body_requires_admission(&request));
589
590 let request = axum::http::Request::post("/v1/chat/completions")
591 .header(CONTENT_LENGTH, "2048")
592 .header(TRANSFER_ENCODING, "chunked")
593 .body(Body::empty())
594 .unwrap();
595 assert!(body_requires_admission(&request));
596 }
597
598 #[test]
599 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
600 let unknown = axum::http::Request::post("/v1/chat/completions")
601 .body(Body::empty())
602 .unwrap();
603 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
604
605 let large = axum::http::Request::post("/v1/chat/completions")
606 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
607 .body(Body::empty())
608 .unwrap();
609 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
610 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
611
612 let absurd = axum::http::Request::post("/v1/chat/completions")
613 .header(CONTENT_LENGTH, u64::MAX.to_string())
614 .body(Body::empty())
615 .unwrap();
616 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
617 }
618
619 #[tokio::test]
620 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
621 let too_large = shape_inference_early_response(
622 "/v1/messages",
623 error_response_coded(
624 StatusCode::PAYLOAD_TOO_LARGE,
625 "request body exceeds the 192 MiB limit",
626 "invalid_request_error",
627 None,
628 Some("request_too_large"),
629 ),
630 )
631 .await;
632 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
633 let house_id = too_large.headers()["x-request-id"].clone();
634 assert_eq!(too_large.headers()["request-id"], house_id);
635 assert_eq!(too_large.headers()["x-should-retry"], "false");
636 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
637 .await
638 .unwrap();
639 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
640 assert_eq!(payload["type"], "error");
641 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
642
643 let busy = shape_inference_early_response(
644 "/v1/chat/completions",
645 retry_contract_response(
646 error_response_coded(
647 StatusCode::TOO_MANY_REQUESTS,
648 "request body admission is busy",
649 "rate_limit_error",
650 None,
651 Some("body_admission_busy"),
652 ),
653 Some(BODY_ADMISSION_RETRY_AFTER_S),
654 ),
655 )
656 .await;
657 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
658 assert!(!busy.headers()["x-request-id"].is_empty());
659 assert_eq!(busy.headers()["retry-after"], "1");
660 assert_eq!(busy.headers()["retry-after-ms"], "1000");
661 assert!(busy.headers().get("x-should-retry").is_none());
662 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
663 .await
664 .unwrap();
665 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
666 assert_eq!(payload["error"]["code"], "body_admission_busy");
667 }
668}
669
670#[derive(Clone, Default)]
671struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
672
673fn is_sse_data_frame(bytes: &[u8]) -> bool {
674 bytes
675 .windows(b"data:".len())
676 .any(|window| window == b"data:")
677}
678
679async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
680 let trace = ttft::start(req.uri().path());
681 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
682 let response = next.run(req).await;
683 let Some(trace) = trace else {
684 return response;
685 };
686 let is_sse = response
687 .headers()
688 .get(CONTENT_TYPE)
689 .and_then(|value| value.to_str().ok())
690 .is_some_and(|value| value.starts_with("text/event-stream"));
691 if !is_sse {
692 return response;
693 }
694
695 let (parts, body) = response.into_parts();
698 let mut body = Box::pin(body.into_data_stream());
699 let stream = async_stream::stream! {
700 while let Some(frame) =
701 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
702 {
703 if frame
704 .as_ref()
705 .is_ok_and(|bytes| is_sse_data_frame(bytes))
706 {
707 trace.mark_first_sse_byte();
708 }
709 yield frame;
710 }
711 };
712 Response::from_parts(parts, Body::from_stream(stream))
713}
714
715const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
716const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
717
718#[derive(Debug, Clone, Default, Deserialize)]
719#[serde(deny_unknown_fields)]
720struct OpenRouterMetadataFile {
721 #[serde(default)]
722 models: HashMap<String, OpenRouterModelMetadata>,
723 #[serde(default)]
726 planned_models: HashMap<String, OpenRouterModelMetadata>,
727 #[serde(default)]
730 provider: Option<ProviderMetadata>,
731}
732
733#[derive(Debug, Clone, Deserialize)]
737#[serde(deny_unknown_fields)]
738struct ProviderMetadata {
739 id: String,
740 #[serde(default)]
741 status_url: Option<String>,
742 #[serde(default)]
743 support_contact: Option<String>,
744 #[serde(default)]
745 incident_contact: Option<String>,
746 #[serde(default)]
747 regions: Vec<String>,
748}
749
750#[derive(Debug, Clone, Default, Deserialize)]
752#[serde(deny_unknown_fields)]
753struct LifecycleMetadata {
754 #[serde(default)]
755 status: Option<String>,
756 #[serde(default)]
757 deprecation_at: Option<String>,
758 #[serde(default)]
759 retirement_at: Option<String>,
760 #[serde(default)]
761 replacement_model_id: Option<String>,
762}
763
764#[derive(Debug, Clone, Default, Deserialize)]
766#[serde(deny_unknown_fields)]
767struct ReliabilityMetadata {
768 #[serde(default)]
769 first_token_timeout_seconds: Option<u64>,
770 #[serde(default)]
771 completion_timeout_seconds: Option<u64>,
772 #[serde(default)]
773 stream_idle_timeout_seconds: Option<u64>,
774 #[serde(default)]
775 capacity_scope: Option<String>,
776}
777
778#[derive(Debug, Clone, Default, Deserialize)]
779#[serde(deny_unknown_fields)]
780struct OpenRouterModelMetadata {
781 #[serde(default)]
783 owned_by: Option<String>,
784 #[serde(default)]
785 lifecycle: Option<LifecycleMetadata>,
786 #[serde(default)]
787 reliability: Option<ReliabilityMetadata>,
788 #[serde(default)]
789 hugging_face_id: Option<String>,
790 #[serde(default)]
791 created: Option<u64>,
792 #[serde(default)]
793 quantization: Option<String>,
794 #[serde(default)]
795 description: Option<String>,
796 #[serde(default)]
797 max_prompt_length: Option<u64>,
798 #[serde(default)]
799 max_output_length: Option<u64>,
800 #[serde(default)]
803 default_output_length: Option<u64>,
804 #[serde(default)]
805 pricing: OpenRouterPricing,
806 #[serde(default)]
807 capacity: OpenRouterCapacity,
808 #[serde(default)]
809 is_ready: Option<bool>,
810 #[serde(default)]
811 is_free: Option<bool>,
812 #[serde(default)]
813 discount_to_user: Option<f64>,
814 #[serde(default)]
815 openrouter_slug: Option<String>,
816 #[serde(default)]
817 datacenters: Vec<OpenRouterDatacenter>,
818 #[serde(default)]
822 input_modalities: Vec<String>,
823 #[serde(default)]
838 surface: Option<String>,
839 #[serde(default)]
840 zdr: Option<bool>,
841 #[serde(default)]
842 hipaa: Option<bool>,
843 #[serde(default)]
853 default_reasoning_effort: Option<String>,
854 #[serde(default)]
872 default_temperature: Option<f32>,
873 #[serde(default)]
874 default_top_p: Option<f32>,
875 #[serde(default)]
877 default_top_k: Option<usize>,
878 #[serde(default)]
879 default_min_p: Option<f32>,
880 #[serde(default)]
881 default_presence_penalty: Option<f32>,
882 #[serde(default)]
883 default_frequency_penalty: Option<f32>,
884 #[serde(default)]
886 default_repetition_penalty: Option<f32>,
887 #[serde(default)]
908 non_thinking_sampling: Option<SamplingArmMetadata>,
909}
910
911#[derive(Debug, Clone, Default, Deserialize)]
917#[serde(deny_unknown_fields)]
918struct SamplingArmMetadata {
919 #[serde(default)]
920 temperature: Option<f32>,
921 #[serde(default)]
922 top_p: Option<f32>,
923 #[serde(default)]
924 top_k: Option<usize>,
925 #[serde(default)]
926 min_p: Option<f32>,
927 #[serde(default)]
928 presence_penalty: Option<f32>,
929 #[serde(default)]
930 frequency_penalty: Option<f32>,
931 #[serde(default)]
932 repetition_penalty: Option<f32>,
933}
934
935impl SamplingArmMetadata {
936 fn is_empty(&self) -> bool {
937 self.temperature.is_none()
938 && self.top_p.is_none()
939 && self.top_k.is_none()
940 && self.min_p.is_none()
941 && self.presence_penalty.is_none()
942 && self.frequency_penalty.is_none()
943 && self.repetition_penalty.is_none()
944 }
945}
946
947#[derive(Debug, Clone, Default, Deserialize)]
948#[serde(deny_unknown_fields)]
949struct OpenRouterPricing {
950 #[serde(default)]
951 prompt: Option<String>,
952 #[serde(default)]
953 cached_prompt: Option<String>,
954 #[serde(default)]
955 cache_write: Option<String>,
956 #[serde(default)]
957 completion: Option<String>,
958 #[serde(default)]
959 internal_reasoning: Option<String>,
960 #[serde(default)]
961 request: Option<String>,
962}
963
964#[derive(Debug, Clone, Default, Deserialize)]
965#[serde(deny_unknown_fields)]
966struct OpenRouterCapacity {
967 #[serde(default)]
968 prompt_tpm: Option<u64>,
969 #[serde(default)]
970 cached_prompt_tpm: Option<u64>,
971 #[serde(default)]
972 completion_tpm: Option<u64>,
973 #[serde(default)]
974 request_rpm: Option<u64>,
975 #[serde(default)]
976 concurrency: Option<u64>,
977}
978
979#[derive(Debug, Clone, Deserialize, Serialize)]
980#[serde(deny_unknown_fields)]
981struct OpenRouterDatacenter {
982 country_code: String,
983 #[serde(default, skip_serializing_if = "Option::is_none")]
984 region: Option<String>,
985}
986
987impl OpenRouterMetadataFile {
988 fn parse(
989 text: &str,
990 ) -> Result<
991 (
992 HashMap<String, OpenRouterModelMetadata>,
993 Option<ProviderMetadata>,
994 ),
995 String,
996 > {
997 let file: Self =
998 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
999 for (alias, metadata) in &file.models {
1000 validate_openrouter_metadata(alias, metadata)?;
1001 }
1002 for (alias, metadata) in &file.planned_models {
1003 validate_openrouter_metadata(alias, metadata)?;
1004 if file.models.contains_key(alias) {
1005 return Err(format!(
1006 "model alias {alias:?} appears in both models and planned_models"
1007 ));
1008 }
1009 }
1010 if let Some(provider) = &file.provider {
1011 if provider.id.is_empty() {
1012 return Err("provider.id must be a non-empty slug".into());
1013 }
1014 for (field, value) in [
1016 ("provider.support_contact", &provider.support_contact),
1017 ("provider.incident_contact", &provider.incident_contact),
1018 ] {
1019 if let Some(value) = value {
1020 if !value.contains(':') {
1021 return Err(format!(
1022 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1023 ));
1024 }
1025 }
1026 }
1027 }
1028 Ok((file.models, file.provider))
1029 }
1030
1031 #[cfg(test)]
1032 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1033 Self::parse(text).map(|(models, _)| models)
1034 }
1035}
1036
1037fn per_million_price(per_token: &str) -> Option<String> {
1041 if !valid_price_string(per_token) {
1042 return None;
1043 }
1044 let (whole, frac) = match per_token.split_once('.') {
1045 Some((whole, frac)) => (whole, frac),
1046 None => (per_token, ""),
1047 };
1048 let mut digits = format!("{whole}{frac}");
1049 let point = whole.len() + 6;
1050 while digits.len() < point {
1051 digits.push('0');
1052 }
1053 let (int_part, frac_part) = digits.split_at(point);
1054 let int_part = int_part.trim_start_matches('0');
1055 let int_part = if int_part.is_empty() { "0" } else { int_part };
1056 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1057 while frac_out.len() < 2 {
1058 frac_out.push('0');
1059 }
1060 Some(format!("{int_part}.{frac_out}"))
1061}
1062
1063fn valid_price_string(value: &str) -> bool {
1064 let mut parts = value.split('.');
1065 let whole = parts.next().unwrap_or_default();
1066 let fraction = parts.next();
1067 !whole.is_empty()
1068 && whole.bytes().all(|b| b.is_ascii_digit())
1069 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1070 && parts.next().is_none()
1071}
1072
1073fn validate_openrouter_metadata(
1074 alias: &str,
1075 metadata: &OpenRouterModelMetadata,
1076) -> Result<(), String> {
1077 if alias.is_empty() {
1078 return Err("models metadata contains an empty model alias".into());
1079 }
1080 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1083 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1084 {
1085 return Err(format!(
1086 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1087 reasoning_effort level (none|minimal|low|medium|high)"
1088 ));
1089 }
1090 validate_sampling_defaults(alias, metadata)?;
1091 for m in &metadata.input_modalities {
1092 if m != "image" && m != "video" {
1093 return Err(format!(
1094 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1095 ));
1096 }
1097 }
1098 if let Some(sfc) = metadata.surface.as_deref()
1099 && !matches!(sfc, "chat" | "embedding" | "rerank")
1100 {
1101 return Err(format!(
1102 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1103 ));
1104 }
1105 if let Some(q) = metadata.quantization.as_deref()
1106 && !matches!(
1107 q,
1108 "int4"
1109 | "int8"
1110 | "fp4"
1111 | "mxfp4"
1112 | "nvfp4"
1113 | "fp6"
1114 | "fp8"
1115 | "mxfp8"
1116 | "fp16"
1117 | "bf16"
1118 | "fp32"
1119 )
1120 {
1121 return Err(format!(
1122 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1123 ));
1124 }
1125 for (field, value) in [
1126 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1127 (
1128 "pricing.cached_prompt",
1129 metadata.pricing.cached_prompt.as_deref(),
1130 ),
1131 (
1132 "pricing.cache_write",
1133 metadata.pricing.cache_write.as_deref(),
1134 ),
1135 ("pricing.completion", metadata.pricing.completion.as_deref()),
1136 (
1137 "pricing.internal_reasoning",
1138 metadata.pricing.internal_reasoning.as_deref(),
1139 ),
1140 ("pricing.request", metadata.pricing.request.as_deref()),
1141 ] {
1142 if let Some(value) = value
1143 && !valid_price_string(value)
1144 {
1145 return Err(format!(
1146 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1147 ));
1148 }
1149 }
1150 for (field, value) in [
1151 ("created", metadata.created),
1152 ("max_prompt_length", metadata.max_prompt_length),
1153 ("max_output_length", metadata.max_output_length),
1154 ("default_output_length", metadata.default_output_length),
1155 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1156 (
1157 "capacity.cached_prompt_tpm",
1158 metadata.capacity.cached_prompt_tpm,
1159 ),
1160 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1161 ("capacity.request_rpm", metadata.capacity.request_rpm),
1162 ("capacity.concurrency", metadata.capacity.concurrency),
1163 ] {
1164 if let Some(value) = value
1165 && value > JSON_SAFE_INTEGER_MAX
1166 {
1167 return Err(format!(
1168 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1169 ));
1170 }
1171 }
1172 for (field, value) in [
1173 ("max_prompt_length", metadata.max_prompt_length),
1174 ("max_output_length", metadata.max_output_length),
1175 ("default_output_length", metadata.default_output_length),
1176 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1177 (
1178 "capacity.cached_prompt_tpm",
1179 metadata.capacity.cached_prompt_tpm,
1180 ),
1181 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1182 ("capacity.request_rpm", metadata.capacity.request_rpm),
1183 ("capacity.concurrency", metadata.capacity.concurrency),
1184 ] {
1185 if value == Some(0) {
1186 return Err(format!(
1187 "model {alias:?}: {field} must be greater than zero when declared"
1188 ));
1189 }
1190 }
1191 if let (Some(default), Some(maximum)) =
1192 (metadata.default_output_length, metadata.max_output_length)
1193 && default > maximum
1194 {
1195 return Err(format!(
1196 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1197 ));
1198 }
1199 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1200 return Err(format!(
1201 "model {alias:?}: default_output_length requires max_output_length"
1202 ));
1203 }
1204 if let Some(discount) = metadata.discount_to_user
1205 && (!discount.is_finite() || discount >= 1.0)
1206 {
1207 return Err(format!(
1208 "model {alias:?}: discount_to_user must be finite and less than 1"
1209 ));
1210 }
1211 if metadata
1212 .openrouter_slug
1213 .as_deref()
1214 .is_some_and(str::is_empty)
1215 {
1216 return Err(format!(
1217 "model {alias:?}: openrouter_slug must not be empty when declared"
1218 ));
1219 }
1220 for dc in &metadata.datacenters {
1221 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1222 return Err(format!(
1223 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1224 dc.country_code
1225 ));
1226 }
1227 }
1228 Ok(())
1229}
1230
1231fn validate_sampling_defaults(
1246 alias: &str,
1247 metadata: &OpenRouterModelMetadata,
1248) -> Result<(), String> {
1249 validate_sampling_arm(
1250 alias,
1251 &[
1252 "default_temperature",
1253 "default_top_p",
1254 "default_min_p",
1255 "default_presence_penalty",
1256 "default_frequency_penalty",
1257 "default_repetition_penalty",
1258 ],
1259 metadata.default_temperature,
1260 metadata.default_top_p,
1261 metadata.default_min_p,
1262 metadata.default_presence_penalty,
1263 metadata.default_frequency_penalty,
1264 metadata.default_repetition_penalty,
1265 )?;
1266 if let Some(arm) = &metadata.non_thinking_sampling {
1267 if arm.is_empty() {
1271 return Err(format!(
1272 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1273 least one vendor recommendation or delete the table"
1274 ));
1275 }
1276 validate_sampling_arm(
1277 alias,
1278 &[
1279 "non_thinking_sampling.temperature",
1280 "non_thinking_sampling.top_p",
1281 "non_thinking_sampling.min_p",
1282 "non_thinking_sampling.presence_penalty",
1283 "non_thinking_sampling.frequency_penalty",
1284 "non_thinking_sampling.repetition_penalty",
1285 ],
1286 arm.temperature,
1287 arm.top_p,
1288 arm.min_p,
1289 arm.presence_penalty,
1290 arm.frequency_penalty,
1291 arm.repetition_penalty,
1292 )?;
1293 }
1294 Ok(())
1295}
1296
1297#[allow(clippy::too_many_arguments)]
1303fn validate_sampling_arm(
1304 alias: &str,
1305 keys: &[&str; 6],
1306 temperature: Option<f32>,
1307 top_p: Option<f32>,
1308 min_p: Option<f32>,
1309 presence_penalty: Option<f32>,
1310 frequency_penalty: Option<f32>,
1311 repetition_penalty: Option<f32>,
1312) -> Result<(), String> {
1313 if let Some(t) = temperature {
1314 if !t.is_finite() || t <= 0.0 || t > 2.0 {
1315 return Err(format!(
1316 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1317 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1318 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1319 vendor recommendation, not greedy); clients reach greedy by sending an \
1320 explicit temperature 0.",
1321 keys[0]
1322 ));
1323 }
1324 }
1325 if let Some(p) = top_p
1326 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1327 {
1328 return Err(format!(
1329 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1330 keys[1]
1331 ));
1332 }
1333 if let Some(m) = min_p
1334 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1335 {
1336 return Err(format!(
1337 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1338 keys[2]
1339 ));
1340 }
1341 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1342 if let Some(v) = value
1343 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1344 {
1345 return Err(format!(
1346 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1347 ));
1348 }
1349 }
1350 if let Some(r) = repetition_penalty
1351 && (!r.is_finite() || r <= 0.0)
1352 {
1353 return Err(format!(
1354 "model {alias:?}: {} {r} must be finite and \
1355 greater than zero (1.0 = off)",
1356 keys[5]
1357 ));
1358 }
1359 Ok(())
1360}
1361
1362fn load_openrouter_metadata(
1363 models: &[(String, String, Option<String>)],
1364) -> Result<
1365 (
1366 HashMap<String, OpenRouterModelMetadata>,
1367 Option<ProviderMetadata>,
1368 ),
1369 String,
1370> {
1371 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1372 Ok(path) => path,
1373 Err(_) => return Ok((HashMap::new(), None)),
1374 };
1375 let p = std::path::Path::new(&path);
1376 if !p.is_file() {
1377 return Err(format!(
1378 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1379 ));
1380 }
1381 let text =
1382 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1383 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1384 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1385 for alias in metadata.keys() {
1386 if !models.iter().any(|(name, _, _)| name == alias) {
1387 return Err(format!(
1388 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1389 ));
1390 }
1391 }
1392 eprintln!(
1393 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1394 metadata.len()
1395 );
1396 Ok((metadata, provider))
1397}
1398
1399#[derive(Clone)]
1400struct AppState {
1401 cmd_tx: Sender<Cmd>,
1402 models: Arc<Vec<String>>,
1403 caps: Arc<HashMap<String, ModelCaps>>,
1404 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1405 provider_metadata: Arc<Option<ProviderMetadata>>,
1407 metering: Option<Arc<dyn metering::Metering>>,
1413 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1416 api_auth: ApiAuth,
1419 metrics_auth: MetricsAuth,
1421 metrics: SharedMetrics,
1422 started: u64,
1425 inflight: InflightCounts,
1429 tenant_inflight: TenantGauge,
1432 health: health::SharedHealth,
1436 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1440}
1441
1442impl AppState {
1443 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1457 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1458 }
1459}
1460
1461#[derive(Clone, Default)]
1462struct ApiAuth {
1463 keyring: Option<&'static auth::KeyStore>,
1464 single_key: Option<Arc<str>>,
1465}
1466
1467impl ApiAuth {
1468 fn from_env() -> Result<ApiAuth, String> {
1469 let single_key = match std::env::var("MEMRA_API_KEY") {
1470 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1471 Ok(key) => Some(Arc::from(key)),
1472 Err(std::env::VarError::NotPresent) => None,
1473 Err(std::env::VarError::NotUnicode(_)) => {
1474 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1475 }
1476 };
1477 Ok(ApiAuth {
1478 keyring: auth::global(),
1479 single_key,
1480 })
1481 }
1482
1483 fn configured(&self) -> bool {
1484 self.keyring.is_some() || self.single_key.is_some()
1485 }
1486}
1487
1488#[derive(Clone, Default)]
1489struct MetricsAuth {
1490 required: bool,
1491 token: Option<Arc<str>>,
1492}
1493
1494impl MetricsAuth {
1495 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1496 let token = token.map(Arc::from);
1497 MetricsAuth {
1498 required: !bind_loopback || api_auth_configured || token.is_some(),
1499 token,
1500 }
1501 }
1502}
1503
1504fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1505 let mut resolved = addr
1506 .to_socket_addrs()
1507 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1508 let first = resolved
1509 .next()
1510 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1511 let mut loopback = first.ip().to_canonical().is_loopback();
1512 for socket in resolved {
1513 loopback &= socket.ip().to_canonical().is_loopback();
1514 }
1515 Ok((first, loopback))
1516}
1517
1518fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1519 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1520}
1521
1522fn validate_bind_security(
1523 addr: &str,
1524 api_auth_configured: bool,
1525 allow_open_bind: bool,
1526) -> Result<bool, String> {
1527 let loopback = bind_is_loopback(addr)?;
1528 if !loopback && !api_auth_configured && !allow_open_bind {
1529 return Err(format!(
1530 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1531 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1532 ));
1533 }
1534 Ok(loopback)
1535}
1536
1537type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1555
1556type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1559
1560struct InflightGuard {
1564 counts: InflightCounts,
1565 idx: usize,
1566 tenants: TenantGauge,
1567 tenant: String,
1568}
1569
1570impl InflightGuard {
1571 fn try_acquire(
1575 counts: InflightCounts,
1576 lane: lanes::Lane,
1577 tenants: TenantGauge,
1578 tenant: &str,
1579 tenant_cap: Option<usize>,
1580 ) -> Result<(Self, usize, usize), usize> {
1581 let idx = lane.idx();
1582 let nt = {
1583 let mut m = tenants.lock().unwrap();
1584 let e = m.entry(tenant.to_string()).or_insert(0);
1585 if tenant_cap.is_some_and(|cap| *e >= cap) {
1586 return Err(*e);
1587 }
1588 *e += 1;
1589 *e
1590 };
1591 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1592 Ok((
1593 InflightGuard {
1594 counts,
1595 idx,
1596 tenants,
1597 tenant: tenant.to_string(),
1598 },
1599 n,
1600 nt,
1601 ))
1602 }
1603}
1604
1605impl Drop for InflightGuard {
1606 fn drop(&mut self) {
1607 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1608 let mut m = self.tenants.lock().unwrap();
1609 if let Some(e) = m.get_mut(&self.tenant) {
1610 *e -= 1;
1611 if *e == 0 {
1612 m.remove(&self.tenant);
1613 }
1614 }
1615 }
1616}
1617
1618fn lane_cap(lane: lanes::Lane) -> usize {
1622 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1623 CAPS.get_or_init(|| {
1624 let batching = std::env::var("MEMRA_SERVE_BATCH")
1625 .map(|v| v != "0")
1626 .unwrap_or(true);
1627 let interactive = if batching {
1628 std::env::var("MEMRA_MAX_SESSIONS")
1629 .ok()
1630 .and_then(|v| v.parse().ok())
1631 .unwrap_or(64)
1632 } else {
1633 worker::MAX_ACTIVE
1634 };
1635 let p = lanes::LanePolicy::from_env();
1636 [interactive, p.max_sessions[1], p.max_sessions[2]]
1637 })[lane.idx()]
1638}
1639
1640fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1643 if m.completed > 0 && m.step_p50_ms > 0.0 {
1644 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1645 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1646 }
1647 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1648 *D.get_or_init(|| {
1649 std::env::var("MEMRA_RL_RESET_S")
1650 .ok()
1651 .and_then(|v| v.parse().ok())
1652 .unwrap_or(2)
1653 })
1654}
1655
1656pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1673pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1674pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1675
1676pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1680 let Some(v) = v.filter(|v| !v.is_null()) else {
1681 return Ok(TIMEOUT_MS_DEFAULT);
1682 };
1683 let Some(ms) = v.as_u64() else {
1684 return Err(format!(
1685 "timeout_ms must be an integer number of milliseconds in \
1686 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX}, got {v}; for work longer than \
1687 {TIMEOUT_MS_MAX} ms use \"stream\": true — the deadline then bounds only the \
1688 time to first token and the stream may run as long as it needs"
1689 ));
1690 };
1691 if !(TIMEOUT_MS_MIN..=TIMEOUT_MS_MAX).contains(&ms) {
1692 return Err(format!(
1693 "timeout_ms {ms} is outside the accepted range \
1694 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX} (milliseconds). {TIMEOUT_MS_MAX} is a \
1695 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1696 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1697 lie. For work longer than {TIMEOUT_MS_MAX} ms use \"stream\": true — the \
1698 deadline then bounds only the time to first token and the stream may run as \
1699 long as it needs"
1700 ));
1701 }
1702 Ok(ms)
1703}
1704
1705#[derive(Clone, Copy)]
1708pub(crate) struct RequestDeadline {
1709 pub(crate) at: tokio::time::Instant,
1710 pub(crate) ms: u64,
1711}
1712
1713impl RequestDeadline {
1714 pub(crate) fn starting_now(ms: u64) -> Self {
1715 Self {
1716 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
1717 ms,
1718 }
1719 }
1720
1721 pub(crate) fn remaining(&self) -> std::time::Duration {
1722 self.at
1723 .saturating_duration_since(tokio::time::Instant::now())
1724 }
1725}
1726
1727pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
1733 let what = if stream {
1734 "the first token was produced"
1735 } else {
1736 "the response completed"
1737 };
1738 let msg = format!(
1739 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
1740 {what}; generation was cancelled and this request is not billed"
1741 );
1742 error_response_coded(
1743 StatusCode::REQUEST_TIMEOUT,
1744 &msg,
1745 "timeout",
1746 Some("timeout_ms"),
1747 Some("deadline_exceeded"),
1748 )
1749}
1750
1751pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
1791
1792pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
1797
1798pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
1802
1803fn env_flag_on(name: &'static str, default_on: bool) -> bool {
1809 match std::env::var(name) {
1810 Ok(v) => !matches!(
1811 v.trim().to_ascii_lowercase().as_str(),
1812 "0" | "off" | "false"
1813 ),
1814 Err(_) => default_on,
1815 }
1816}
1817
1818fn env_u64(name: &'static str, default: u64) -> u64 {
1821 std::env::var(name)
1822 .ok()
1823 .and_then(|v| v.parse::<u64>().ok())
1824 .filter(|v| *v > 0)
1825 .unwrap_or(default)
1826}
1827
1828const CHARS_PER_TOKEN_FLOOR: usize = 6;
1839
1840pub(crate) fn prompt_tokens_estimate(
1841 request: &worker::Request,
1842 tokenizer: Option<&Tokenizer>,
1843) -> u64 {
1844 if !request.prompt_ids.is_empty() {
1845 return request.prompt_ids.len() as u64;
1846 }
1847 let mut text = String::new();
1848 text.push_str(&request.prompt_text);
1849 for turn in &request.chat_turns {
1850 text.push_str(&turn.content);
1851 }
1852 for tool in &request.tools_json {
1853 text.push_str(tool);
1854 }
1855 if let Some(tokenizer) = tokenizer {
1856 return tokenizer.encode(text.as_str(), false).len() as u64;
1857 }
1858 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
1859}
1860
1861pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
1865 let prefill_ms = prompt_tokens
1866 .saturating_mul(1_000)
1867 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
1868 .unwrap_or(u64::MAX);
1869 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
1870 if decode_ms == 0 {
1871 return None;
1872 }
1873 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
1874}
1875
1876pub(crate) fn nonstream_deadline_gate(
1886 request: &worker::Request,
1887 stream: bool,
1888 deadline: RequestDeadline,
1889 caller_declared_max_tokens: bool,
1890 tokenizer: Option<&Tokenizer>,
1891) -> Result<(), String> {
1892 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
1893 return Ok(());
1894 }
1895 let max_new = request.params.max_new as u64;
1896 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
1903 {
1904 return Ok(());
1905 }
1906 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
1907 let remaining_ms = deadline.remaining().as_millis() as u64;
1908 let prefill_ms = prompt_tokens.saturating_mul(1_000)
1909 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
1910 let decode_ms = max_new.saturating_mul(1_000)
1911 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
1912 let est_ms = prefill_ms.saturating_add(decode_ms);
1913 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
1914 if est_ms <= bound_ms {
1915 return Ok(());
1916 }
1917 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
1918 let advice = match fits {
1919 Some(fits) if fits > 0 => format!(
1920 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
1921 stream's deadline bounds only the time to first token, so it may run as long \
1922 as it needs"
1923 ),
1924 _ => format!(
1925 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
1926 token, so no max_tokens fits: set \"stream\": true"
1927 ),
1928 };
1929 Err(format!(
1930 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
1931 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
1932 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
1933 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
1934 rather than after the deadline: {advice}",
1935 est_ms / 1_000,
1936 ))
1937}
1938
1939fn max_queue_depth(cap: usize) -> usize {
1943 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1944 D.get_or_init(|| {
1945 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
1946 .ok()
1947 .and_then(|v| v.parse().ok())
1948 })
1949 .unwrap_or(cap.saturating_mul(4))
1950}
1951
1952pub(crate) fn admission_backpressure(
1965 st: &AppState,
1966 lane: lanes::Lane,
1967 rl: &RateLimit,
1968 deadline: RequestDeadline,
1969) -> Result<(), (Response, &'static str)> {
1970 if lane != lanes::Lane::Interactive || rl.remaining > 0 {
1971 return Ok(());
1972 }
1973 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
1974 let backlog = m.queued_requests as usize
1975 + worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire);
1976 let cap = lane_cap(lane).max(1);
1977 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
1978 let bound = max_queue_depth(cap);
1979 if backlog >= bound {
1980 let msg = format!(
1981 "interactive queue is at its bound ({backlog} queued, bound {bound}); this \
1982 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
1983 coarse estimate, not a promise)"
1984 );
1985 let resp = retry_contract_response(
1986 (
1987 StatusCode::TOO_MANY_REQUESTS,
1988 Json(error_body(
1989 &msg,
1990 "rate_limit_error",
1991 None,
1992 Some("shed_queue"),
1993 )),
1994 )
1995 .into_response(),
1996 Some(est_wait_s),
1997 );
1998 return Err((resp, "shed_queue"));
1999 }
2000 let remaining_ms = deadline.remaining().as_millis() as u64;
2001 if est_wait_s.saturating_mul(1_000) > remaining_ms {
2002 let msg = format!(
2003 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2004 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2005 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2006 estimate, not a promise)"
2007 );
2008 let resp = retry_contract_response(
2009 (
2010 StatusCode::TOO_MANY_REQUESTS,
2011 Json(error_body(
2012 &msg,
2013 "rate_limit_error",
2014 None,
2015 Some("shed_deadline"),
2016 )),
2017 )
2018 .into_response(),
2019 Some(est_wait_s),
2020 );
2021 return Err((resp, "shed_deadline"));
2022 }
2023 Ok(())
2024}
2025
2026pub(crate) struct PendingAdmissionGuard {
2031 reserved: bool,
2032 lane: lanes::Lane,
2033}
2034
2035impl PendingAdmissionGuard {
2036 pub(crate) fn commit(mut self) {
2040 self.reserved = false;
2041 std::mem::forget(self);
2042 }
2043}
2044
2045impl Drop for PendingAdmissionGuard {
2046 fn drop(&mut self) {
2047 if self.reserved {
2048 worker::release_pending_admit();
2049 worker::release_admission_reservation(self.lane);
2050 }
2051 }
2052}
2053
2054pub(crate) fn reserve_pending_admit(
2055 st: &AppState,
2056 lane: lanes::Lane,
2057 rl: &RateLimit,
2058 deadline: RequestDeadline,
2059) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2060 let cap = lane_cap(lane).max(1);
2065 let bound = max_queue_depth(cap);
2066 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2067 loop {
2068 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2069 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2070 let backlog = reservations;
2074 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2075 if backlog >= bound {
2076 let msg = format!(
2077 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2078 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2079 coarse estimate, not a promise)",
2080 lane.as_str()
2081 );
2082 let resp = retry_contract_response(
2083 (
2084 StatusCode::TOO_MANY_REQUESTS,
2085 Json(error_body(
2086 &msg,
2087 "rate_limit_error",
2088 None,
2089 Some("shed_queue"),
2090 )),
2091 )
2092 .into_response(),
2093 Some(est_wait_s),
2094 );
2095 return Err((resp, "shed_queue"));
2096 }
2097 let remaining_ms = deadline.remaining().as_millis() as u64;
2098 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2102 if lane == lanes::Lane::Interactive
2103 && waits_for_capacity
2104 && est_wait_s.saturating_mul(1_000) > remaining_ms
2105 {
2106 let msg = format!(
2107 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2108 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2109 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2110 estimate, not a promise)"
2111 );
2112 let resp = retry_contract_response(
2113 (
2114 StatusCode::TOO_MANY_REQUESTS,
2115 Json(error_body(
2116 &msg,
2117 "rate_limit_error",
2118 None,
2119 Some("shed_deadline"),
2120 )),
2121 )
2122 .into_response(),
2123 Some(est_wait_s),
2124 );
2125 return Err((resp, "shed_deadline"));
2126 }
2127 if reservations_for_lane
2128 .compare_exchange(
2129 reservations,
2130 reservations.saturating_add(1),
2131 std::sync::atomic::Ordering::AcqRel,
2132 std::sync::atomic::Ordering::Acquire,
2133 )
2134 .is_ok()
2135 {
2136 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2140 return Ok(PendingAdmissionGuard {
2141 reserved: true,
2142 lane,
2143 });
2144 }
2145 }
2146}
2147
2148static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2159
2160fn draining() -> bool {
2161 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2162}
2163
2164fn drain_deadline_s() -> u64 {
2166 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2167 *D.get_or_init(|| {
2168 std::env::var("MEMRA_DRAIN_S")
2169 .ok()
2170 .and_then(|v| v.parse().ok())
2171 .unwrap_or(30)
2172 })
2173}
2174
2175fn drain_response() -> Response {
2185 let resp = (
2186 StatusCode::SERVICE_UNAVAILABLE,
2187 Json(error_body(
2188 "server is draining (shutdown in progress); retry",
2189 "server_error",
2190 None,
2191 Some("draining"),
2192 )),
2193 )
2194 .into_response();
2195 retry_contract_response(resp, Some(drain_deadline_s()))
2196}
2197
2198struct RateLimit {
2200 limit: usize,
2201 remaining: usize,
2202 reset_s: u64,
2203}
2204
2205impl RateLimit {
2206 fn at_admit(
2211 lane: lanes::Lane,
2212 n_inflight: usize,
2213 metrics: &SharedMetrics,
2214 tenant: &auth::TenantCtx,
2215 n_tenant: usize,
2216 ) -> Self {
2217 let global = lane_cap(lane);
2218 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2219 return Self::compute(global, n_inflight, metrics);
2220 };
2221 let headroom = t
2222 .saturating_sub(n_tenant)
2223 .min(global.saturating_sub(n_inflight));
2224 Self::compute(t, t - headroom, metrics)
2226 }
2227
2228 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2229 let remaining = limit.saturating_sub(n_inflight);
2230 let reset_s = if remaining > 0 {
2231 0
2232 } else {
2233 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2234 reset_estimate_s(&m)
2235 };
2236 RateLimit {
2237 limit,
2238 remaining,
2239 reset_s,
2240 }
2241 }
2242
2243 fn attach(&self, mut resp: Response) -> Response {
2245 let h = resp.headers_mut();
2246 for (k, v) in [
2247 ("x-ratelimit-limit", self.limit as u64),
2248 ("x-ratelimit-remaining", self.remaining as u64),
2249 ("x-ratelimit-reset", self.reset_s),
2250 ] {
2251 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2252 h.insert(axum::http::HeaderName::from_static(k), v);
2253 }
2254 }
2255 resp
2256 }
2257}
2258
2259fn acquire_request_slot(
2263 st: &AppState,
2264 lane: lanes::Lane,
2265 tenant: &auth::TenantCtx,
2266 env: &Envelope,
2267) -> Result<(InflightGuard, RateLimit), Response> {
2268 let global = lane_cap(lane);
2269 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2270 match InflightGuard::try_acquire(
2271 st.inflight.clone(),
2272 lane,
2273 st.tenant_inflight.clone(),
2274 &tenant.tenant,
2275 tenant_cap,
2276 ) {
2277 Ok((guard, n_inflight, n_tenant)) => {
2278 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2279 Ok((guard, rl))
2280 }
2281 Err(n_tenant) => {
2282 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2283 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2284 let error =
2285 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2286 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2287 }
2288 }
2289}
2290
2291#[derive(Deserialize)]
2293struct CompletionReq {
2294 model: String,
2295 #[serde(default)]
2296 prompt: String,
2297 #[serde(default)]
2299 prompt_ids: Vec<u32>,
2300 #[serde(default)]
2303 max_tokens: Option<usize>,
2304 #[serde(default)]
2317 temperature: Option<f32>,
2318 #[serde(default)]
2319 top_p: Option<f32>,
2320 #[serde(default)]
2322 top_k: Option<usize>,
2323 #[serde(default)]
2325 min_p: Option<f32>,
2326 #[serde(default)]
2328 frequency_penalty: Option<f32>,
2329 #[serde(default)]
2330 presence_penalty: Option<f32>,
2331 #[serde(default)]
2333 repetition_penalty: Option<f32>,
2334 #[serde(default)]
2340 seed: Option<u64>,
2341 #[serde(default)]
2342 stop: StopSequences,
2343 #[serde(default)]
2346 logit_bias: Option<serde_json::Value>,
2347 #[serde(default)]
2348 logprobs: Option<serde_json::Value>,
2349 #[serde(default)]
2350 n: Option<usize>,
2351 #[serde(default)]
2352 best_of: Option<usize>,
2353 #[serde(default)]
2355 chat: bool,
2356 #[serde(default)]
2358 stream: bool,
2359 #[serde(default)]
2361 max_ctx: Option<usize>,
2362 #[serde(default)]
2364 trace_id: Option<String>,
2365 #[serde(default)]
2369 cache_salt: Option<String>,
2370 #[serde(default)]
2374 session_id: Option<String>,
2375 #[serde(default)]
2376 user: Option<String>,
2377 #[serde(default)]
2381 timeout_ms: Option<serde_json::Value>,
2382}
2383
2384#[derive(Deserialize)]
2385struct ChatMessage {
2386 role: String,
2387 #[serde(default)]
2389 content: serde_json::Value,
2390 #[serde(default)]
2392 tool_calls: Vec<ReqToolCall>,
2393 #[serde(default)]
2396 tool_call_id: Option<String>,
2397 #[serde(default)]
2400 name: Option<String>,
2401 #[serde(default, alias = "reasoning_content")]
2411 reasoning: Option<String>,
2412}
2413
2414#[derive(Deserialize)]
2415struct ReqToolCall {
2416 #[serde(default)]
2417 #[allow(dead_code)]
2418 id: Option<String>,
2419 function: ReqToolFunction,
2420}
2421
2422#[derive(Deserialize)]
2423struct ReqToolFunction {
2424 name: String,
2425 #[serde(default)]
2427 arguments: serde_json::Value,
2428}
2429
2430#[derive(Clone, Default, Deserialize)]
2431#[serde(untagged)]
2432enum StopSequences {
2433 One(String),
2434 Many(Vec<String>),
2435 #[default]
2436 None,
2437}
2438
2439impl StopSequences {
2440 fn into_vec(self) -> Vec<String> {
2446 let stops = match self {
2447 Self::One(stop) => vec![stop],
2448 Self::Many(stops) => stops,
2449 Self::None => Vec::new(),
2450 };
2451 stops.into_iter().filter(|s| !s.is_empty()).collect()
2452 }
2453}
2454
2455#[derive(Deserialize)]
2461struct ChatCompletionReq {
2462 model: String,
2463 messages: Vec<ChatMessage>,
2464 #[serde(default, alias = "max_completion_tokens")]
2467 max_tokens: Option<usize>,
2468 #[serde(default)]
2471 temperature: Option<f32>,
2472 #[serde(default)]
2473 top_p: Option<f32>,
2474 #[serde(default)]
2478 top_k: Option<usize>,
2479 #[serde(default)]
2481 min_p: Option<f32>,
2482 #[serde(default)]
2484 frequency_penalty: Option<f32>,
2485 #[serde(default)]
2486 presence_penalty: Option<f32>,
2487 #[serde(default)]
2489 repetition_penalty: Option<f32>,
2490 #[serde(default)]
2492 seed: Option<u64>,
2493 #[serde(default)]
2494 stop: StopSequences,
2495 #[serde(default)]
2496 stream: bool,
2497 #[serde(default)]
2498 max_ctx: Option<usize>,
2499 #[serde(default)]
2504 response_format: Option<serde_json::Value>,
2505 #[serde(default)]
2506 logit_bias: Option<serde_json::Value>,
2507 #[serde(default)]
2508 logprobs: Option<serde_json::Value>,
2509 #[serde(default)]
2510 top_logprobs: Option<usize>,
2511 #[serde(default)]
2512 n: Option<usize>,
2513 #[serde(default)]
2515 tools: Vec<serde_json::Value>,
2516 #[serde(default)]
2518 tool_choice: Option<serde_json::Value>,
2519 #[serde(default)]
2525 reasoning_effort: Option<String>,
2526 #[serde(default)]
2535 reasoning: Option<serde_json::Value>,
2536 #[serde(default)]
2551 include_reasoning: Option<bool>,
2552 #[serde(default)]
2561 enable_thinking: Option<bool>,
2562 #[serde(default)]
2568 chat_template_kwargs: Option<serde_json::Value>,
2569 #[serde(default)]
2573 cache_salt: Option<String>,
2574 #[serde(default)]
2576 session_id: Option<String>,
2577 #[serde(default)]
2578 user: Option<String>,
2579 #[serde(default)]
2584 timeout_ms: Option<serde_json::Value>,
2585}
2586fn one() -> f32 {
2587 1.0
2588}
2589fn default_temperature() -> f32 {
2595 1.0
2596}
2597
2598#[derive(Debug, Clone, Copy, Default, PartialEq)]
2618struct SamplingDefaults {
2619 temperature: Option<f32>,
2620 top_p: Option<f32>,
2621 top_k: Option<usize>,
2622 min_p: Option<f32>,
2623 frequency_penalty: Option<f32>,
2624 presence_penalty: Option<f32>,
2625 repetition_penalty: Option<f32>,
2626}
2627
2628impl SamplingDefaults {
2629 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2632 SamplingDefaults {
2633 temperature: metadata
2634 .and_then(|m| m.default_temperature)
2635 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2636 top_p: metadata
2637 .and_then(|m| m.default_top_p)
2638 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2639 top_k: metadata.and_then(|m| m.default_top_k),
2640 min_p: metadata.and_then(|m| m.default_min_p),
2641 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2642 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2643 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2644 }
2645 }
2646}
2647
2648#[derive(Debug, Clone, Copy, Default, PartialEq)]
2661struct ModelSamplingDefaults {
2662 thinking: SamplingDefaults,
2663 non_thinking: Option<SamplingDefaults>,
2664}
2665
2666impl ModelSamplingDefaults {
2667 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2668 ModelSamplingDefaults {
2669 thinking: SamplingDefaults::resolve(metadata, caps),
2670 non_thinking: metadata
2675 .and_then(|m| m.non_thinking_sampling.as_ref())
2676 .map(|arm| SamplingDefaults {
2677 temperature: arm.temperature,
2678 top_p: arm.top_p,
2679 top_k: arm.top_k,
2680 min_p: arm.min_p,
2681 frequency_penalty: arm.frequency_penalty,
2682 presence_penalty: arm.presence_penalty,
2683 repetition_penalty: arm.repetition_penalty,
2684 }),
2685 }
2686 }
2687
2688 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
2701 match (think, &self.non_thinking) {
2702 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
2703 _ => &self.thinking,
2704 }
2705 }
2706
2707 fn single(thinking: SamplingDefaults) -> Self {
2710 ModelSamplingDefaults {
2711 thinking,
2712 non_thinking: None,
2713 }
2714 }
2715}
2716
2717#[derive(Debug, Clone, Copy, Default)]
2723struct ClientSampling {
2724 temperature: Option<f32>,
2725 top_p: Option<f32>,
2726 top_k: Option<usize>,
2727 min_p: Option<f32>,
2728 frequency_penalty: Option<f32>,
2729 presence_penalty: Option<f32>,
2730 repetition_penalty: Option<f32>,
2731 seed: Option<u64>,
2732}
2733
2734impl From<&CompletionReq> for ClientSampling {
2735 fn from(r: &CompletionReq) -> Self {
2736 ClientSampling {
2737 temperature: r.temperature,
2738 top_p: r.top_p,
2739 top_k: r.top_k,
2740 min_p: r.min_p,
2741 frequency_penalty: r.frequency_penalty,
2742 presence_penalty: r.presence_penalty,
2743 repetition_penalty: r.repetition_penalty,
2744 seed: r.seed,
2745 }
2746 }
2747}
2748
2749impl From<&ChatCompletionReq> for ClientSampling {
2750 fn from(r: &ChatCompletionReq) -> Self {
2751 ClientSampling {
2752 temperature: r.temperature,
2753 top_p: r.top_p,
2754 top_k: r.top_k,
2755 min_p: r.min_p,
2756 frequency_penalty: r.frequency_penalty,
2757 presence_penalty: r.presence_penalty,
2758 repetition_penalty: r.repetition_penalty,
2759 seed: r.seed,
2760 }
2761 }
2762}
2763
2764fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
2770 sampler_config(
2771 client
2772 .temperature
2773 .or(defaults.temperature)
2774 .unwrap_or_else(default_temperature),
2775 client.top_k.or(defaults.top_k).unwrap_or(0),
2776 client.top_p.or(defaults.top_p).unwrap_or_else(one),
2777 client.min_p.or(defaults.min_p).unwrap_or(0.0),
2778 client
2779 .frequency_penalty
2780 .or(defaults.frequency_penalty)
2781 .unwrap_or(0.0),
2782 client
2783 .presence_penalty
2784 .or(defaults.presence_penalty)
2785 .unwrap_or(0.0),
2786 client
2787 .repetition_penalty
2788 .or(defaults.repetition_penalty)
2789 .unwrap_or_else(one),
2790 client.seed,
2791 )
2792}
2793
2794#[derive(Serialize)]
2795struct CompletionResp {
2796 model: String,
2797 text: String,
2798 tokens: Vec<u32>,
2799 stop_reason: String,
2803 #[serde(default, skip_serializing_if = "Option::is_none")]
2808 error: Option<serde_json::Value>,
2809 n_tokens: usize,
2810 prompt_tokens: usize,
2813 cached_tokens: usize,
2814 elapsed_s: f64,
2815}
2816
2817fn usage_json(
2825 n_prompt: usize,
2826 n_tokens: usize,
2827 n_cached: usize,
2828 elapsed_s: f64,
2829 spec: Option<worker::SpecUsage>,
2830) -> serde_json::Value {
2831 let mut u = json!({
2832 "prompt_tokens": n_prompt,
2833 "completion_tokens": n_tokens,
2834 "total_tokens": n_prompt + n_tokens,
2835 "prompt_tokens_details": { "cached_tokens": n_cached },
2836 "elapsed_s": elapsed_s,
2837 });
2838 if let Some(sp) = spec {
2839 u["spec"] = json!({
2840 "rounds": sp.rounds,
2841 "drafted": sp.drafted,
2842 "accepted": sp.accepted,
2843 "acceptance_rate": if sp.drafted > 0 {
2844 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
2845 });
2846 }
2847 u
2848}
2849
2850const SYSTEM_FINGERPRINT: &str = concat!("memra-", env!("MEMRA_BUILD_SHA"));
2864
2865fn gen_hex128() -> String {
2868 use std::hash::{BuildHasher, Hasher};
2869 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2870 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2871 let t = std::time::SystemTime::now()
2872 .duration_since(std::time::UNIX_EPOCH)
2873 .map(|d| d.as_nanos() as u64)
2874 .unwrap_or(0);
2875 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
2876 h1.write_u64(n);
2877 h1.write_u64(t);
2878 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
2879 h2.write_u64(t.rotate_left(17));
2880 h2.write_u64(n);
2881 format!("{:016x}{:016x}", h1.finish(), h2.finish())
2882}
2883
2884#[derive(Clone)]
2887struct Envelope {
2888 id: String,
2889 created: u64,
2890}
2891
2892impl Envelope {
2893 fn new(chat: bool) -> Self {
2894 Envelope {
2895 id: format!(
2896 "{}-{}",
2897 if chat { "chatcmpl" } else { "cmpl" },
2898 gen_hex128()
2899 ),
2900 created: std::time::SystemTime::now()
2901 .duration_since(std::time::UNIX_EPOCH)
2902 .map(|d| d.as_secs())
2903 .unwrap_or(0),
2904 }
2905 }
2906
2907 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
2909 v["id"] = json!(self.id);
2910 v["created"] = json!(self.created);
2911 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
2912 v
2913 }
2914}
2915
2916fn with_request_id(id: &str, mut resp: Response) -> Response {
2918 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
2919 resp.headers_mut()
2920 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
2921 }
2922 resp
2923}
2924
2925fn openai_compat() -> bool {
2933 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2934 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
2935 Ok("openai") => true,
2936 Ok(_) => false,
2937 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
2938 })
2939}
2940
2941fn cache_namespace(cache_salt: &Option<String>) -> String {
2953 cache_salt.clone().unwrap_or_default()
2954}
2955
2956const CACHE_SALT_MAX_BYTES: usize = 64;
2957
2958fn validate_cache_namespace(
2959 cache_salt: &Option<String>,
2960 keyring_configured: bool,
2961) -> Result<String, &'static str> {
2962 let raw = cache_namespace(cache_salt);
2963 if raw.len() > CACHE_SALT_MAX_BYTES {
2964 return Err("cache_salt must be at most 64 bytes");
2965 }
2966 if !keyring_configured && raw.starts_with("t:") {
2967 return Err("cache_salt must not use the reserved t: prefix without a keyring");
2968 }
2969 if !raw
2970 .bytes()
2971 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
2972 {
2973 return Err("cache_salt contains unsupported characters");
2974 }
2975 Ok(raw)
2976}
2977
2978fn affinity_key(
2994 session_id: &Option<String>,
2995 user: &Option<String>,
2996 headers: &axum::http::HeaderMap,
2997) -> Option<String> {
2998 let clean = |s: &str| -> Option<String> {
2999 let t = s.trim();
3000 if t.is_empty() {
3001 None
3002 } else {
3003 Some(t.to_string())
3004 }
3005 };
3006 session_id
3007 .as_deref()
3008 .and_then(clean)
3009 .or_else(|| user.as_deref().and_then(clean))
3010 .or_else(|| {
3011 headers
3012 .get("x-session-id")
3013 .and_then(|v| v.to_str().ok())
3014 .and_then(clean)
3015 })
3016}
3017
3018fn error_body(
3023 message: &str,
3024 etype: &str,
3025 param: Option<&str>,
3026 code: Option<&str>,
3027) -> serde_json::Value {
3028 json!({ "error": {
3029 "message": message,
3030 "type": etype,
3031 "param": param,
3032 "code": code,
3033 } })
3034}
3035
3036fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3037 error_response_coded(status, message, etype, param, None)
3038}
3039
3040fn error_response_coded(
3045 status: StatusCode,
3046 message: &str,
3047 etype: &str,
3048 param: Option<&str>,
3049 code: Option<&str>,
3050) -> Response {
3051 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3052 if status.is_client_error()
3053 && status != StatusCode::TOO_MANY_REQUESTS
3054 && status != StatusCode::REQUEST_TIMEOUT
3055 && status != StatusCode::CONFLICT
3056 {
3057 resp.headers_mut().insert(
3058 "x-should-retry",
3059 axum::http::HeaderValue::from_static("false"),
3060 );
3061 }
3062 resp
3063}
3064
3065fn bad_request(message: &str, param: Option<&str>) -> Response {
3066 error_response(
3067 StatusCode::BAD_REQUEST,
3068 message,
3069 "invalid_request_error",
3070 param,
3071 )
3072}
3073
3074const 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>) {
3103 use worker::ErrClass as C;
3104 match class {
3105 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3106 C::ContextLength => (
3107 StatusCode::BAD_REQUEST,
3108 "invalid_request_error",
3109 Some("context_length_exceeded"),
3110 ),
3111 C::ModelNotFound => (
3112 StatusCode::BAD_REQUEST,
3113 "invalid_request_error",
3114 Some("model_not_found"),
3115 ),
3116 C::RateLimit => (
3117 StatusCode::TOO_MANY_REQUESTS,
3118 "rate_limit_error",
3119 Some("rate_limit_exceeded"),
3120 ),
3121 C::Overloaded => (
3122 StatusCode::SERVICE_UNAVAILABLE,
3123 "server_error",
3124 Some("overloaded"),
3125 ),
3126 C::Engine => (
3127 StatusCode::INTERNAL_SERVER_ERROR,
3128 "server_error",
3129 Some("engine_error"),
3130 ),
3131 }
3132}
3133
3134fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3136 use worker::ErrClass as C;
3137 match class {
3138 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3139 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3140 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3144 }
3145}
3146
3147fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3150 let (_, etype, code) = class_http(e.class);
3151 error_body(&e.message, etype, e.param, code)
3152}
3153
3154fn engine_error_response(e: &worker::EngineError) -> Response {
3160 engine_error_response_with_retry_after(
3161 e,
3162 e.retry_after_s.or_else(|| class_retry_after_s(e.class)),
3163 )
3164}
3165
3166fn engine_error_response_with_retry_after(
3167 e: &worker::EngineError,
3168 retry_after_s: Option<u64>,
3169) -> Response {
3170 let (status, _, _) = class_http(e.class);
3171 let resp = (status, Json(engine_error_body(e))).into_response();
3172 retry_contract_response(resp, retry_after_s)
3173}
3174
3175fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3177 let status = resp.status();
3178 let h = resp.headers_mut();
3179 match retry_after_s {
3180 Some(secs) => {
3181 let secs = secs.clamp(1, 60);
3183 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3184 h.insert(axum::http::header::RETRY_AFTER, v);
3185 }
3186 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3187 h.insert("retry-after-ms", v);
3188 }
3189 }
3190 None if status.is_client_error() => {
3191 h.insert(
3194 "x-should-retry",
3195 axum::http::HeaderValue::from_static("false"),
3196 );
3197 }
3198 None => {}
3199 }
3200 resp
3201}
3202
3203fn worker_unavailable_response() -> Response {
3204 engine_error_response_with_retry_after(
3205 &worker::EngineError::overloaded("worker unavailable"),
3206 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3207 )
3208}
3209
3210fn stop_reason_to_finish(r: &str) -> &'static str {
3211 match r {
3212 "Eos" | "Callback" => "stop",
3213 "MaxNew" | "ContextFull" => "length",
3214 _ => "stop",
3215 }
3216}
3217
3218fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3222 match v {
3223 serde_json::Value::Null => Ok(String::new()),
3224 serde_json::Value::String(s) => Ok(s.clone()),
3225 serde_json::Value::Array(parts) => {
3226 let mut out = String::new();
3227 for p in parts {
3228 match p.get("type").and_then(|t| t.as_str()) {
3229 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3230 Some(t) => out.push_str(t),
3231 None => return Err("content part has no text field".into()),
3232 },
3233 Some(other) => {
3234 return Err(format!(
3235 "unsupported content part type {other:?} (text only)"
3236 ));
3237 }
3238 }
3239 }
3240 Ok(out)
3241 }
3242 _ => Err("content must be a string, null, or an array of text parts".into()),
3243 }
3244}
3245
3246fn vision_enabled() -> bool {
3249 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3250 *ON.get_or_init(|| {
3251 std::env::var("MEMRA_VISION_DIR").is_ok()
3252 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3253 })
3254}
3255
3256fn gemma_vision_enabled() -> bool {
3261 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3262 *ON.get_or_init(|| {
3263 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3264 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3265 })
3266}
3267
3268fn step_vision_enabled() -> bool {
3274 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3275 *ON.get_or_init(|| {
3276 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3277 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3278 })
3279}
3280
3281const VISION_MAX_IMAGES: usize = 8;
3283
3284pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3290 std::sync::atomic::AtomicUsize::new(0);
3291pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3295 tokio::sync::Semaphore::const_new(1);
3296
3297pub(crate) struct VisionMemoryPermit {
3298 bytes: usize,
3299}
3300
3301#[derive(Debug)]
3302pub(crate) enum VisionMemoryError {
3303 Request(String),
3304 Capacity(String),
3305}
3306
3307impl Drop for VisionMemoryPermit {
3308 fn drop(&mut self) {
3309 if self.bytes != 0 {
3310 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3311 }
3312 }
3313}
3314
3315fn try_reserve_vision_memory(
3316 bytes: usize,
3317) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3318 if bytes == 0 {
3319 return Ok(None);
3320 }
3321 if bytes > MAX_VISION_PATCH_BYTES {
3322 return Err(VisionMemoryError::Request(format!(
3323 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3324 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3325 )));
3326 }
3327 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3328 loop {
3329 let Some(next) = in_use.checked_add(bytes) else {
3330 return Err(VisionMemoryError::Capacity(
3331 "vision patch memory reservation overflowed".into(),
3332 ));
3333 };
3334 if next > MAX_VISION_PATCH_BYTES {
3335 return Err(VisionMemoryError::Capacity(format!(
3336 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3337 in_use / (1024 * 1024),
3338 bytes / (1024 * 1024)
3339 )));
3340 }
3341 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3342 in_use,
3343 next,
3344 std::sync::atomic::Ordering::AcqRel,
3345 std::sync::atomic::Ordering::Acquire,
3346 ) {
3347 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3348 Err(actual) => in_use = actual,
3349 }
3350 }
3351}
3352
3353pub(crate) fn vision_memory_error_response(
3354 error: VisionMemoryError,
3355 param: Option<&str>,
3356) -> Response {
3357 match error {
3358 VisionMemoryError::Request(message) => bad_request(&message, param),
3359 VisionMemoryError::Capacity(message) => retry_contract_response(
3360 error_response_coded(
3361 StatusCode::SERVICE_UNAVAILABLE,
3362 &message,
3363 "server_error",
3364 None,
3365 Some("vision_memory_busy"),
3366 ),
3367 Some(RETRY_AFTER_S_OVERLOADED),
3368 ),
3369 }
3370}
3371
3372enum PendingVisionUnit {
3378 Still {
3379 bytes: Vec<u8>,
3380 gh: usize,
3381 gw: usize,
3382 },
3383 Video {
3384 bytes: Vec<u8>,
3385 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3386 video: usize,
3387 },
3388}
3389
3390struct PendingGemmaImage {
3392 bytes: Vec<u8>,
3393 gw: usize,
3394 gh: usize,
3395}
3396
3397struct PendingStepImage {
3400 bytes: Vec<u8>,
3401 plan: memra_engine::vision_step::StepImagePlan,
3402}
3403
3404fn content_to_text_vision_step(
3414 v: &serde_json::Value,
3415 step_images: &mut Vec<PendingStepImage>,
3416) -> Result<String, String> {
3417 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
3418 let parts = match v {
3419 serde_json::Value::Array(parts) => parts,
3420 _ => return content_to_text(v),
3421 };
3422 let mut out = String::new();
3423 let mut needs_sep = false;
3424 for p in parts {
3425 match p.get("type").and_then(|t| t.as_str()) {
3426 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3427 Some(t) => {
3428 if needs_sep {
3429 out.push(' ');
3430 }
3431 out.push_str(t);
3432 needs_sep = true;
3433 }
3434 None => return Err("content part has no text field".into()),
3435 },
3436 Some("image_url") => {
3437 let url = p
3438 .get("image_url")
3439 .and_then(|u| {
3440 if u.is_string() {
3441 u.as_str()
3442 } else {
3443 u.get("url").and_then(|x| x.as_str())
3444 }
3445 })
3446 .ok_or("image_url part has no url")?;
3447 if !url.starts_with("data:") {
3448 return Err(
3449 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3450 );
3451 }
3452 if step_images.len() >= VISION_MAX_IMAGES {
3453 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3454 }
3455 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3459 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3460 let plan = memra_engine::vision_step::step_plan_image(&bytes)
3461 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3462 for i in 0..plan.n_tiles {
3463 out.push_str("<patch_start>");
3464 for _ in 0..SV_TILE_ROWS {
3465 out.push_str("<im_patch>");
3466 }
3467 out.push_str("<patch_end>");
3468 if plan.newline_mask[i] {
3469 out.push_str("<patch_newline>");
3470 }
3471 }
3472 out.push_str("<im_start>");
3473 for _ in 0..SV_MAIN_ROWS {
3474 out.push_str("<im_patch>");
3475 }
3476 out.push_str("<im_end>");
3477 step_images.push(PendingStepImage { bytes, plan });
3478 needs_sep = false;
3479 }
3480 Some("video_url") => {
3481 return Err("step37 has no video input (image-only processor)".into());
3482 }
3483 Some(other) => {
3484 return Err(format!("unsupported content part type {other:?}"));
3485 }
3486 }
3487 }
3488 Ok(out)
3489}
3490
3491fn content_to_text_vision(
3500 v: &serde_json::Value,
3501 images: &mut Vec<PendingVisionUnit>,
3502 gemma_images: &mut Vec<PendingGemmaImage>,
3503 step_images: &mut Vec<PendingStepImage>,
3504 next_video: &mut usize,
3505) -> Result<String, String> {
3506 if step_vision_enabled() {
3510 return content_to_text_vision_step(v, step_images);
3511 }
3512 let parts = match v {
3513 serde_json::Value::Array(parts) => parts,
3514 _ => return content_to_text(v),
3515 };
3516 let mut out = String::new();
3517 for p in parts {
3518 match p.get("type").and_then(|t| t.as_str()) {
3519 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3520 Some(t) => out.push_str(t),
3521 None => return Err("content part has no text field".into()),
3522 },
3523 Some("image_url") if gemma_vision_enabled() => {
3524 let url = p
3525 .get("image_url")
3526 .and_then(|u| {
3527 if u.is_string() {
3528 u.as_str()
3529 } else {
3530 u.get("url").and_then(|x| x.as_str())
3531 }
3532 })
3533 .ok_or("image_url part has no url")?;
3534 if !url.starts_with("data:") {
3535 return Err(
3536 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3537 );
3538 }
3539 if gemma_images.len() >= VISION_MAX_IMAGES {
3540 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3541 }
3542 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
3546 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3547 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
3548 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3549 out.push_str("<|image>");
3551 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
3552 out.push_str("<|image|>");
3553 }
3554 out.push_str("<image|>");
3555 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
3556 }
3557 Some("image_url") => {
3558 if !vision_enabled() {
3559 return Err("image input is not enabled on this deployment".into());
3560 }
3561 let url = p
3562 .get("image_url")
3563 .and_then(|u| {
3564 if u.is_string() {
3565 u.as_str()
3566 } else {
3567 u.get("url").and_then(|x| x.as_str())
3568 }
3569 })
3570 .ok_or("image_url part has no url")?;
3571 if !url.starts_with("data:") {
3572 return Err(
3573 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3574 );
3575 }
3576 if images
3577 .iter()
3578 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
3579 .count()
3580 >= VISION_MAX_IMAGES
3581 {
3582 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3583 }
3584 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3589 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3590 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
3591 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3592 out.push_str("<|vision_start|>");
3593 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
3594 out.push_str("<|image_pad|>");
3595 }
3596 out.push_str("<|vision_end|>");
3597 images.push(PendingVisionUnit::Still { bytes, gh, gw });
3598 }
3599 Some("video_url") if gemma_vision_enabled() => {
3600 return Err("gemma-4 has no video input (image-only projector)".into());
3601 }
3602 Some("video_url") => {
3603 if !vision_enabled() {
3604 return Err("video input is not enabled on this deployment".into());
3605 }
3606 let url = p
3607 .get("video_url")
3608 .and_then(|u| {
3609 if u.is_string() {
3610 u.as_str()
3611 } else {
3612 u.get("url").and_then(|x| x.as_str())
3613 }
3614 })
3615 .ok_or("video_url part has no url")?;
3616 if !url.starts_with("data:") {
3617 return Err(
3618 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3619 );
3620 }
3621 if *next_video >= 2 {
3622 return Err("too many videos (max 2)".into());
3623 }
3624 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
3627 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
3628 .map_err(|e| format!("video: {e}"))?;
3629 let vidx = *next_video;
3630 *next_video += 1;
3631 for group in &vid.groups {
3633 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
3634 out.push_str("<|vision_start|>");
3635 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
3636 out.push_str("<|video_pad|>");
3637 }
3638 out.push_str("<|vision_end|>");
3639 }
3640 images.push(PendingVisionUnit::Video {
3643 bytes,
3644 groups: vid.groups,
3645 video: vidx,
3646 });
3647 }
3648 Some(other) => {
3649 return Err(format!("unsupported content part type {other:?}"));
3650 }
3651 }
3652 }
3653 Ok(out)
3654}
3655
3656fn pyjson(v: &serde_json::Value, out: &mut String) {
3660 match v {
3661 serde_json::Value::Object(m) => {
3662 out.push('{');
3663 for (i, (k, val)) in m.iter().enumerate() {
3664 if i > 0 {
3665 out.push_str(", ");
3666 }
3667 out.push_str(&serde_json::Value::String(k.clone()).to_string());
3668 out.push_str(": ");
3669 pyjson(val, out);
3670 }
3671 out.push('}');
3672 }
3673 serde_json::Value::Array(a) => {
3674 out.push('[');
3675 for (i, val) in a.iter().enumerate() {
3676 if i > 0 {
3677 out.push_str(", ");
3678 }
3679 pyjson(val, out);
3680 }
3681 out.push(']');
3682 }
3683 scalar => out.push_str(&scalar.to_string()),
3684 }
3685}
3686
3687fn pyjson_str(v: &serde_json::Value) -> String {
3688 let mut s = String::new();
3689 pyjson(v, &mut s);
3690 s
3691}
3692
3693fn sampler_config(
3700 temperature: f32,
3701 top_k: usize,
3702 top_p: f32,
3703 min_p: f32,
3704 frequency_penalty: f32,
3705 presence_penalty: f32,
3706 repetition_penalty: f32,
3707 seed: Option<u64>,
3708) -> SamplerConfig {
3709 let penalties_on =
3710 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
3711 SamplerConfig {
3712 temperature,
3713 top_k,
3714 top_p,
3715 min_p,
3716 penalty_last_n: if penalties_on {
3717 memra_engine::spec::PEN_WINDOW_MAX
3718 } else {
3719 0
3720 },
3721 penalty_repeat: repetition_penalty,
3722 penalty_freq: frequency_penalty,
3723 penalty_present: presence_penalty,
3724 seed: seed.unwrap_or_else(fresh_seed),
3727 }
3728}
3729
3730fn fresh_seed() -> u64 {
3735 use std::sync::atomic::{AtomicU64, Ordering};
3736 static COUNTER: AtomicU64 = AtomicU64::new(0);
3737 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
3738 let nanos = std::time::SystemTime::now()
3739 .duration_since(std::time::UNIX_EPOCH)
3740 .map(|d| d.as_nanos() as u64)
3741 .unwrap_or(0);
3742 let mut z = nanos
3743 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
3744 .wrapping_add(0x9E3779B97F4A7C15);
3745 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
3746 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
3747 z ^= z >> 31;
3748 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
3751}
3752
3753fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
3758 for (param, present, why) in fields {
3759 if *present {
3760 return Err((format!("{param} is not supported{why}"), param.to_string()));
3761 }
3762 }
3763 Ok(())
3764}
3765
3766#[derive(PartialEq)]
3767enum ToolChoice {
3768 Auto,
3769 None,
3770}
3771
3772fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
3773 match v {
3774 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
3775 Some(serde_json::Value::String(s)) => match s.as_str() {
3776 "auto" => Ok(ToolChoice::Auto),
3777 "none" => Ok(ToolChoice::None),
3778 "required" => Err("tool_choice \"required\" is not supported (no constrained \
3779 decoding); use \"auto\""
3780 .into()),
3781 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
3782 },
3783 Some(serde_json::Value::Object(_)) => {
3784 Err("named-function tool_choice is not supported; use \"auto\"".into())
3785 }
3786 Some(other) => Err(format!("bad tool_choice: {other}")),
3787 }
3788}
3789
3790fn parse_think(
3840 reasoning_effort: &Option<String>,
3841 reasoning: &Option<serde_json::Value>,
3842 vllm_switch: Option<bool>,
3843 suppress_switch: Option<bool>,
3844 default_effort: Option<&str>,
3845 dsv4: bool,
3846) -> Result<(ThinkMode, Option<String>, bool), String> {
3847 let mut effort = reasoning_effort.clone();
3848 let ReasoningObject {
3849 mut enabled,
3850 effort: object_effort,
3851 exclude,
3852 } = parse_reasoning_object(reasoning)?;
3853 if let Some(e) = object_effort {
3854 effort = Some(e);
3855 }
3856 match (enabled, vllm_switch) {
3861 (Some(a), Some(b)) if a != b => {
3862 return Err(format!(
3863 "contradictory reasoning switches: reasoning.enabled={a} and \
3864 enable_thinking={b} — send one"
3865 ));
3866 }
3867 (None, Some(b)) => enabled = Some(b),
3868 _ => {}
3869 }
3870 let suppress = match (exclude, suppress_switch) {
3882 (Some(true), _) | (_, Some(false)) => Some(false),
3883 _ => None,
3884 };
3885 match (enabled, suppress) {
3886 (Some(true), Some(false)) => {
3887 return Err(
3888 "contradictory reasoning switches: reasoning is enabled but \
3889 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
3890 on this server not delivering reasoning means not generating it, so send one"
3891 .into(),
3892 );
3893 }
3894 (None, Some(b)) => enabled = Some(b),
3895 _ => {}
3896 }
3897 let client_explicit = effort.is_some() || enabled.is_some();
3901 if effort.is_none() && enabled.is_none() {
3906 effort = default_effort.map(str::to_string);
3907 }
3908 let effort_arm = match effort.as_deref() {
3913 None => None,
3914 Some(raw) => {
3915 let level = canonical_effort_for(raw, dsv4).ok_or_else(|| {
3916 format!(
3917 "bad reasoning_effort {raw:?} \
3918 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
3919 highest level this model's template distinguishes)"
3920 )
3921 })?;
3922 Some(match level {
3923 "none" | "minimal" => (ThinkMode::NoThink, "low"),
3924 "low" => (ThinkMode::Think, "low"),
3925 "medium" => (ThinkMode::Think, "medium"),
3926 "max" => (ThinkMode::Think, "max"),
3927 _ => (ThinkMode::Think, "high"),
3928 })
3929 }
3930 };
3931 let (think, level) = match (enabled, effort_arm) {
3932 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
3935 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
3936 (None, Some((think, level))) => (think, Some(level.to_string())),
3937 (None, None) => (ThinkMode::Default, None),
3938 };
3939 Ok((think, level, client_explicit))
3940}
3941
3942struct ReasoningObject {
3944 enabled: Option<bool>,
3945 effort: Option<String>,
3946 exclude: Option<bool>,
3947}
3948
3949fn parse_reasoning_object(
3968 reasoning: &Option<serde_json::Value>,
3969) -> Result<ReasoningObject, String> {
3970 let mut out = ReasoningObject {
3971 enabled: None,
3972 effort: None,
3973 exclude: None,
3974 };
3975 let Some(v) = reasoning else { return Ok(out) };
3976 let obj = match v {
3977 serde_json::Value::Null => return Ok(out),
3978 serde_json::Value::Object(obj) => obj,
3979 _ => return Err("reasoning must be an object".into()),
3980 };
3981 for (key, value) in obj {
3982 match key.as_str() {
3990 "enabled" => {
3991 if !value.is_null() {
3992 out.enabled = Some(
3993 value
3994 .as_bool()
3995 .ok_or("reasoning.enabled must be true or false")?,
3996 );
3997 }
3998 }
3999 "exclude" => {
4000 if !value.is_null() {
4001 out.exclude = Some(
4002 value
4003 .as_bool()
4004 .ok_or("reasoning.exclude must be true or false")?,
4005 );
4006 }
4007 }
4008 "effort" => {
4009 if !value.is_null() {
4010 out.effort = Some(
4011 value
4012 .as_str()
4013 .ok_or("reasoning.effort must be a string")?
4014 .to_string(),
4015 );
4016 }
4017 }
4018 "max_tokens" => {
4019 return Err(
4020 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4021 are output tokens here, and max_tokens is the ONE output budget covering \
4022 reasoning and content together — there is no separate reasoning budget to \
4023 spend against, so honouring this field is impossible rather than merely \
4024 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4025 reasoning.enabled:false) to spend less of it on reasoning"
4026 .into(),
4027 );
4028 }
4029 other => {
4030 return Err(format!(
4031 "reasoning.{other} is not a field this server implements (it would change \
4032 nothing about the request); the supported keys are enabled, effort and \
4033 exclude"
4034 ));
4035 }
4036 }
4037 }
4038 Ok(out)
4039}
4040
4041fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4062 let Some(v) = kwargs else { return Ok(None) };
4063 let obj = match v {
4064 serde_json::Value::Null => return Ok(None),
4065 serde_json::Value::Object(obj) => obj,
4066 _ => return Err("chat_template_kwargs must be an object".into()),
4067 };
4068 let mut switch = None;
4069 for (key, value) in obj {
4070 match key.as_str() {
4071 "enable_thinking" => {
4072 switch = Some(
4073 value
4074 .as_bool()
4075 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4076 );
4077 }
4078 "preserve_thinking" => {
4079 let preserve = value
4080 .as_bool()
4081 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4082 if !preserve {
4083 return Err(
4084 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4085 server: the renderer implements the vendor DEFAULT (replay every prior \
4086 assistant turn's <think> block, empty when no reasoning was sent) but \
4087 not the strip arm — serving replay bytes under a strip request would \
4088 misdescribe the prompt. Omit the flag or send true"
4089 .into(),
4090 );
4091 }
4092 }
4094 other => {
4095 return Err(format!(
4096 "chat_template_kwargs.{other} is not supported by this server's \
4097 template renderer (it would change nothing about the prompt); the only \
4098 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4099 refuses in both directions — see its own message)"
4100 ));
4101 }
4102 }
4103 }
4104 Ok(switch)
4105}
4106
4107fn resolve_vllm_think_switch(
4111 enable_thinking: Option<bool>,
4112 kwargs: &Option<serde_json::Value>,
4113) -> Result<Option<bool>, String> {
4114 let from_kwargs = parse_template_kwargs(kwargs)?;
4115 match (enable_thinking, from_kwargs) {
4116 (Some(a), Some(b)) if a != b => Err(format!(
4117 "contradictory reasoning switches: enable_thinking={a} and \
4118 chat_template_kwargs.enable_thinking={b} — send one"
4119 )),
4120 (Some(a), _) => Ok(Some(a)),
4121 (None, b) => Ok(b),
4122 }
4123}
4124
4125pub(crate) fn canonical_effort_for(value: &str, dsv4_max: bool) -> Option<&'static str> {
4145 match value {
4146 "none" => Some("none"),
4147 "minimal" => Some("minimal"),
4148 "low" => Some("low"),
4149 "medium" => Some("medium"),
4150 "high" => Some("high"),
4151 "xhigh" | "max" | "ultra" => Some(if dsv4_max { "max" } else { "high" }),
4152 _ => None,
4153 }
4154}
4155
4156pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4159 canonical_effort_for(value, false)
4160}
4161
4162fn json_to_val(v: &serde_json::Value) -> chat::Val {
4165 match v {
4166 serde_json::Value::Null => chat::Val::Null,
4167 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4168 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4169 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4170 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4171 serde_json::Value::Object(o) => chat::Val::Obj(
4174 o.iter()
4175 .map(|(k, val)| (k.clone(), json_to_val(val)))
4176 .collect(),
4177 ),
4178 }
4179}
4180
4181#[allow(clippy::type_complexity)]
4185fn prepare_tools(
4186 tools: &[serde_json::Value],
4187) -> Result<
4188 (
4189 Vec<String>,
4190 Vec<chat::Val>,
4191 HashMap<String, HashMap<String, String>>,
4192 ),
4193 String,
4194> {
4195 let mut tools_json = Vec::with_capacity(tools.len());
4196 let mut tools_struct = Vec::with_capacity(tools.len());
4197 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4198 for t in tools {
4199 let f = t
4200 .get("function")
4201 .ok_or("each tool needs a function object")?;
4202 let name = f
4203 .get("name")
4204 .and_then(|n| n.as_str())
4205 .ok_or("each tool needs function.name")?;
4206 let mut params: HashMap<String, String> = HashMap::new();
4207 if let Some(props) = f
4208 .get("parameters")
4209 .and_then(|p| p.get("properties"))
4210 .and_then(|p| p.as_object())
4211 {
4212 for (p, def) in props {
4213 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4214 params.insert(p.clone(), ty.to_string());
4215 }
4216 }
4217 }
4218 schemas.insert(name.to_string(), params);
4219 tools_json.push(pyjson_str(t));
4220 tools_struct.push(json_to_val(f));
4222 }
4223 Ok((tools_json, tools_struct, schemas))
4224}
4225
4226fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4231 let parsed: serde_json::Value = match &tc.function.arguments {
4232 serde_json::Value::Null => json!({}),
4233 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4234 serde_json::Value::String(s) => serde_json::from_str(s)
4235 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4236 v @ serde_json::Value::Object(_) => v.clone(),
4237 _ => return Err("tool_calls arguments must be a JSON object".into()),
4238 };
4239 let obj = parsed
4240 .as_object()
4241 .ok_or("tool_calls arguments must decode to a JSON object")?;
4242 let params = obj
4243 .iter()
4244 .map(|(k, v)| {
4245 let rendered = match v {
4246 serde_json::Value::String(s) => s.clone(),
4247 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4248 scalar => scalar.to_string(),
4249 };
4250 (k.clone(), rendered)
4251 })
4252 .collect();
4253 let args = obj
4256 .iter()
4257 .map(|(k, v)| (k.clone(), json_to_val(v)))
4258 .collect();
4259 Ok(TmplToolCall {
4260 name: tc.function.name.clone(),
4261 params,
4262 args,
4263 id: tc.id.clone(),
4264 })
4265}
4266
4267fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4269 json!({ "id": c.id, "type": "function",
4270 "function": { "name": c.name, "arguments": c.arguments } })
4271}
4272
4273#[tokio::main]
4277pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4278 serve_with(ServerWiring::stock()).await
4279}
4280
4281enum MeteringWiring {
4283 Stock,
4287 Custom(metering::MeteringFactory),
4292}
4293
4294pub struct ServerWiring {
4298 metering: MeteringWiring,
4299 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
4302 claimed_env: Vec<&'static str>,
4306}
4307
4308impl ServerWiring {
4309 pub fn stock() -> Self {
4311 ServerWiring {
4312 metering: MeteringWiring::Stock,
4313 on_ready: None,
4314 claimed_env: Vec::new(),
4315 }
4316 }
4317
4318 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
4321 ServerWiring {
4322 metering: MeteringWiring::Custom(factory),
4323 on_ready: None,
4324 claimed_env: Vec::new(),
4325 }
4326 }
4327
4328 pub fn claiming(mut self, var: &'static str) -> Self {
4332 self.claimed_env.push(var);
4333 self
4334 }
4335
4336 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
4337 self.on_ready = Some(Box::new(hook));
4338 self
4339 }
4340}
4341
4342pub struct RuntimeHandles {
4345 pub trim: TrimHandle,
4346 pub purge: PurgeHandle,
4349 pub shutdown: tokio::sync::watch::Receiver<bool>,
4354}
4355
4356#[derive(Clone)]
4359pub struct TrimHandle {
4360 cmd_tx: Sender<Cmd>,
4361}
4362
4363impl TrimHandle {
4364 pub async fn trim(&self) -> Result<serde_json::Value, String> {
4366 let (tx, rx) = tokio::sync::oneshot::channel();
4367 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
4368 return Err("worker is down".into());
4369 }
4370 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4371 Ok(Ok(report)) => Ok(json!(report)),
4372 _ => Err("worker did not answer the trim within 30s".into()),
4373 }
4374 }
4375}
4376
4377#[derive(Clone)]
4386pub struct PurgeHandle {
4387 cmd_tx: Sender<Cmd>,
4388}
4389
4390impl PurgeHandle {
4391 pub async fn purge_tenant(&self, tenant: &str) -> Result<serde_json::Value, String> {
4393 let (tx, rx) = tokio::sync::oneshot::channel();
4394 let cmd = Cmd::PurgeTenantHost {
4395 tenant: tenant.to_string(),
4396 tx,
4397 };
4398 if self.cmd_tx.send(cmd).is_err() {
4399 return Err("worker is down".into());
4400 }
4401 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4402 Ok(Ok(report)) => Ok(json!(report)),
4403 _ => Err("worker did not answer the purge within 30s".into()),
4404 }
4405 }
4406}
4407
4408pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
4409 let args: Vec<String> = std::env::args().skip(1).collect();
4412 if let Some(code) = auth::run_cli(&args) {
4413 std::process::exit(code);
4414 }
4415 auth::init_from_env();
4418 let api_auth = match ApiAuth::from_env() {
4419 Ok(auth) => auth,
4420 Err(err) => {
4421 eprintln!("[server] FATAL: {err}");
4422 std::process::exit(1);
4423 }
4424 };
4425 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
4426 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
4427 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
4428 Ok(resolved) => resolved,
4429 Err(err) => {
4430 eprintln!("[server] FATAL: {err}");
4431 std::process::exit(1);
4432 }
4433 };
4434 if !bind_loopback && !api_auth.configured() && !allow_open_bind {
4435 let message = format!(
4436 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
4437 );
4438 eprintln!("[server] FATAL: {message}");
4439 std::process::exit(1);
4440 }
4441 if !bind_loopback && !api_auth.configured() {
4442 eprintln!(
4443 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
4444 metrics remain bearer-protected"
4445 );
4446 }
4447 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
4448 Ok(token) if token.is_empty() => {
4449 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
4450 std::process::exit(1);
4451 }
4452 Ok(token) => Some(token),
4453 Err(std::env::VarError::NotPresent) => None,
4454 Err(std::env::VarError::NotUnicode(_)) => {
4455 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
4456 std::process::exit(1);
4457 }
4458 };
4459 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
4460
4461 let models = parse_models_config();
4462 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
4463 Ok(loaded) => loaded,
4464 Err(err) => {
4465 eprintln!("[server] FATAL: {err}");
4466 std::process::exit(1);
4467 }
4468 };
4469 let metering_obj: Option<Arc<dyn metering::Metering>> = {
4475 let factory = match wiring.metering {
4476 MeteringWiring::Stock => None,
4477 MeteringWiring::Custom(factory) => Some(factory),
4478 };
4479 for deployment_only in [
4480 "MEMRA_REQUEST_LEDGER",
4481 "MEMRA_TENANT_BUDGETS",
4482 "MEMRA_ADMIN_ADDR",
4483 "MEMRA_ADMIN_TOKEN_FILE",
4484 "MEMRA_CAPTURE_DIR",
4485 ] {
4486 if std::env::var_os(deployment_only).is_some()
4487 && !wiring.claimed_env.contains(&deployment_only)
4488 {
4489 eprintln!(
4490 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
4491 build ships no accounting/admin/capture. Wire a Metering implementation \
4492 through ServerWiring and claim the vars it consumes."
4493 );
4494 std::process::exit(1);
4495 }
4496 }
4497 match factory {
4498 None => None,
4499 Some(factory) => {
4500 let model_ids: Vec<String> =
4501 models.iter().map(|(name, _, _)| name.clone()).collect();
4502 match factory(&metering::MeteringInit { models: &model_ids }) {
4503 Ok(metering_obj) => metering_obj,
4504 Err(err) => {
4505 eprintln!("[server] FATAL: metering wiring: {err}");
4506 std::process::exit(1);
4507 }
4508 }
4509 }
4510 }
4511 };
4512 let budget_tokenizers = if metering_obj
4513 .as_ref()
4514 .is_some_and(|manager| manager.enforces_limits())
4515 {
4516 match load_budget_tokenizers(&models) {
4517 Ok(tokenizers) => Some(tokenizers),
4518 Err(err) => {
4519 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
4520 std::process::exit(1);
4521 }
4522 }
4523 } else {
4524 None
4525 };
4526 eprintln!("[server] starting; models config = {models:?}");
4527
4528 let health_state = health::WorkerHealth::new();
4533 health::spawn_gpu_watch(health_state.clone());
4537 health::spawn_sd_watchdog(health_state.clone());
4538
4539 let (cmd_tx, model_names, caps, metrics, worker_thread) =
4541 match worker::spawn(models, health_state.clone()) {
4542 Ok(v) => v,
4543 Err(err) => {
4544 eprintln!("[server] FATAL: worker init failed: {err}");
4545 health_state.mark_dead(format!("worker init failed: {err}"));
4546 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
4547 std::process::exit(1);
4548 }
4549 };
4550 eprintln!("[server] worker ready; serving models: {model_names:?}");
4551
4552 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
4559 if let Some(on_ready) = wiring.on_ready {
4560 on_ready(RuntimeHandles {
4561 trim: TrimHandle {
4562 cmd_tx: cmd_tx.clone(),
4563 },
4564 purge: PurgeHandle {
4565 cmd_tx: cmd_tx.clone(),
4566 },
4567 shutdown: drain_shutdown_rx.clone(),
4568 });
4569 }
4570
4571 let bg_handle = darklane::spawn_from_env(health_state.clone());
4574 let bg_state = bg_handle.as_ref().map(|h| {
4575 let mode = darklane::BgConfig::from_env()
4576 .map(|c| c.yield_mode.as_str())
4577 .unwrap_or("stop");
4578 (h.state.clone(), mode)
4579 });
4580
4581 let state = AppState {
4582 cmd_tx,
4583 models: model_names,
4584 caps,
4585 openrouter_metadata: Arc::new(openrouter_metadata),
4586 provider_metadata: Arc::new(provider_metadata),
4587 metering: metering_obj,
4588 budget_tokenizers,
4589 api_auth,
4590 metrics_auth,
4591 metrics,
4592 started: std::time::SystemTime::now()
4593 .duration_since(std::time::UNIX_EPOCH)
4594 .map(|d| d.as_secs())
4595 .unwrap_or(0),
4596 inflight: Arc::new(Default::default()),
4597 tenant_inflight: Arc::new(Default::default()),
4598 health: health_state.clone(),
4599 bg: bg_state,
4600 };
4601 let inflight_handle = state.inflight.clone();
4602 let drain_metering = state.metering.clone();
4605 worker::register_http_inflight(state.inflight.clone());
4610 let app = Router::new()
4611 .route("/health", get(health_live))
4616 .route("/livez", get(health_live))
4617 .route("/readyz", get(health_ready))
4618 .route("/models", get(list_models))
4619 .route("/v1/models", get(list_models_v1))
4620 .route("/v1/auth/check", get(auth_check))
4621 .route("/v1/completions", post(completions))
4622 .route("/v1/embeddings", post(embed_api::embeddings))
4623 .route("/v1/rerank", post(embed_api::rerank))
4624 .route("/v1/chat/completions", post(chat_completions))
4625 .route("/v1/messages", post(anthropic::messages))
4629 .route("/v1/responses", post(responses_api::responses))
4630 .route("/metrics", get(get_metrics))
4631 .route("/yield/metrics", get(yield_metrics))
4632 .with_state(state.clone());
4633 let app = apply_body_limit(app);
4636 let app = app.layer(middleware::from_fn_with_state(
4639 state,
4640 authenticate_inference_before_body,
4641 ));
4642 let app = if ttft::enabled() {
4643 app.layer(middleware::from_fn(ttft_request_start))
4644 } else {
4645 app
4646 };
4647
4648 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
4649 eprintln!("[server] listening on http://{bind_addr}");
4650 drop(drain_shutdown_rx);
4651 health::sd_notify("READY=1\nSTATUS=serving");
4655 let inflight = inflight_handle;
4662 let signal_admin_shutdown = drain_shutdown_tx.clone();
4663 let serve_result = axum::serve(listener, app)
4664 .with_graceful_shutdown(async move {
4665 let mut sigterm =
4666 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
4667 Ok(s) => s,
4668 Err(err) => {
4669 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
4670 std::future::pending::<()>().await;
4671 unreachable!()
4672 }
4673 };
4674 sigterm.recv().await;
4675 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
4676 let _ = signal_admin_shutdown.send(true);
4677 health::sd_notify(&format!(
4681 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
4682 (drain_deadline_s() + 5) * 1_000_000
4683 ));
4684 let n: usize = inflight
4685 .iter()
4686 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4687 .sum();
4688 eprintln!(
4689 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
4690 drain_deadline_s()
4691 );
4692 let deadline = std::time::Duration::from_secs(drain_deadline_s());
4693 let t0 = std::time::Instant::now();
4694 loop {
4695 let n: usize = inflight
4696 .iter()
4697 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4698 .sum();
4699 if n == 0 {
4700 eprintln!(
4701 "[server] drain complete in {:.1}s; exiting",
4702 t0.elapsed().as_secs_f64()
4703 );
4704 break;
4705 }
4706 if t0.elapsed() >= deadline {
4707 eprintln!(
4708 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
4709 drain_deadline_s()
4710 );
4711 if let Some(metering) = drain_metering.as_ref() {
4718 metering.drain_kill();
4719 }
4720 break;
4721 }
4722 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4723 }
4724 })
4725 .await;
4726 let _ = drain_shutdown_tx.send(true);
4729 serve_result?;
4730 if let Some(h) = bg_handle {
4734 h.shutdown();
4735 }
4736 worker_thread.join().map_err(|_| {
4745 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
4746 })?;
4747 eprintln!("[server] GPU worker shutdown complete");
4748 Ok(())
4749}
4750
4751fn validate_model_path(path: &str) -> Result<(), String> {
4757 let p = std::path::Path::new(path);
4758 if !p.exists() {
4759 return Err(format!("model path {path:?} does not exist"));
4760 }
4761 if p.is_file() {
4762 return Ok(()); }
4764 if p.join("manifest.json").exists() {
4765 return Ok(()); }
4767 let has_st =
4768 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
4769 if !has_st {
4770 return Err(format!(
4771 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
4772 model.safetensors.index.json + config.json (HF safetensors dir), or \
4773 manifest.json (memra repack dir)"
4774 ));
4775 }
4776 if !p.join("config.json").exists() {
4777 return Err(format!(
4778 "model dir {path:?} has safetensors weights but no config.json"
4779 ));
4780 }
4781 Ok(())
4782}
4783
4784fn parse_models_config() -> Vec<(String, String, Option<String>)> {
4792 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
4793 let mut out = Vec::new();
4794 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
4795 if let Some((name, path)) = entry.split_once('=') {
4796 let (mpath, dpath) = match path.trim().split_once('+') {
4799 Some((m, d)) => (m.trim(), Some(d.trim())),
4800 None => (path.trim(), None),
4801 };
4802 let resolve = |p: &str| {
4803 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
4804 eprintln!("[server] FATAL: model {name:?}: {err}");
4805 std::process::exit(1);
4806 })
4807 };
4808 let mpath = resolve(mpath);
4809 if let Err(err) = validate_model_path(&mpath) {
4810 eprintln!("[server] FATAL: model {name:?}: {err}");
4811 std::process::exit(1);
4812 }
4813 let dpath = dpath.map(|d| {
4823 let d = resolve(d);
4824 let p = std::path::Path::new(&d);
4825 if !p.exists() {
4826 eprintln!(
4827 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
4828 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
4829 rather than serving plain decode under a config that asked \
4830 for speculative decoding."
4831 );
4832 std::process::exit(1);
4833 }
4834 if !p.is_file() {
4835 eprintln!(
4836 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
4837 file — a '+draft' attach must be a NextN/MTP GGUF file."
4838 );
4839 std::process::exit(1);
4840 }
4841 d
4842 });
4843 out.push((name.trim().to_string(), mpath, dpath));
4844 } else {
4845 eprintln!(
4846 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
4847 );
4848 }
4849 }
4850 if !out.is_empty() {
4851 return out;
4852 }
4853 }
4854 vec![
4856 (
4857 "main".into(),
4858 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
4859 None,
4860 ),
4861 (
4862 "judge".into(),
4863 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
4864 None,
4865 ),
4866 ]
4867}
4868
4869fn load_budget_tokenizers(
4870 models: &[(String, String, Option<String>)],
4871) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
4872 let mut tokenizers = HashMap::new();
4873 for (alias, path, _) in models {
4874 let path = std::path::Path::new(path);
4875 let tokenizer = if path.is_dir() {
4876 let tokenizer_dir = if path.join("manifest.json").exists() {
4877 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
4878 format!("model {alias:?}: open repack tokenizer source: {err}")
4879 })?;
4880 repack
4881 .source_dir()
4882 .filter(|source| source.join("tokenizer.json").exists())
4883 .unwrap_or(path)
4884 .to_path_buf()
4885 } else {
4886 path.to_path_buf()
4887 };
4888 Tokenizer::from_hf_dir(&tokenizer_dir)
4889 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4890 } else {
4891 let gguf = memra_gguf::GgufFile::open(path)
4892 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
4893 Tokenizer::from_gguf(&gguf)
4894 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4895 };
4896 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
4897 }
4898 Ok(Arc::new(tokenizers))
4899}
4900
4901fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4903 let s = st.health.snapshot();
4904 let mut v = json!({
4905 "status": status,
4906 "models": *st.models,
4907 "worker": {
4908 "phase": health::phase_name(s.phase),
4909 "beat_age_ms": s.beat_age_ms,
4910 "tick_max_ms": s.tick_max_ms,
4911 "stall_threshold_ms": s.stall_threshold_ms,
4912 "generation": s.generation,
4913 "xid_warnings": s.xid_warns,
4914 },
4915 });
4916 if let Some(d) = detail {
4917 v["detail"] = json!(d);
4918 }
4919 v
4920}
4921
4922fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4926 let mut v = health_payload(st, status, detail);
4927 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
4928 v
4929}
4930
4931async fn auth_check() -> impl IntoResponse {
4935 StatusCode::NO_CONTENT
4936}
4937
4938async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
4954 if draining() {
4955 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
4958 }
4959 match st.health.live() {
4960 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
4961 Err(why) => retry_contract_response(
4962 (
4963 StatusCode::SERVICE_UNAVAILABLE,
4964 Json(health_payload(&st, "unhealthy", Some(&why))),
4965 )
4966 .into_response(),
4967 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
4968 ),
4969 }
4970}
4971
4972async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
4983 let is_draining = draining();
4984 match st.health.ready(is_draining) {
4985 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
4986 Err(why) => retry_contract_response(
4987 (
4988 StatusCode::SERVICE_UNAVAILABLE,
4989 Json(readiness_payload(&st, "not_ready", Some(&why))),
4990 )
4991 .into_response(),
4992 Some(if is_draining {
4993 drain_deadline_s()
4994 } else {
4995 worker::WORKER_RESPAWN_BACKOFF_BASE_S
4996 }),
4997 ),
4998 }
4999}
5000
5001#[derive(Clone, Copy)]
5002struct DualPpMetricsSnapshot {
5003 stage_ns: [u64; 4],
5004 stage_samples: [usize; 4],
5005 dropped_timing_samples: usize,
5006 overlaps: usize,
5007 slot_pairs: usize,
5008 slot_uses: [usize; 2],
5009 slot_collisions: usize,
5010}
5011
5012impl DualPpMetricsSnapshot {
5013 fn current() -> Self {
5014 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
5015 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
5016 Self {
5017 stage_ns,
5018 stage_samples,
5019 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
5020 overlaps: memra_engine::pp::dual_pp_overlaps(),
5021 slot_pairs,
5022 slot_uses,
5023 slot_collisions,
5024 }
5025 }
5026
5027 fn populated(self) -> bool {
5028 self.stage_samples.iter().any(|&n| n > 0)
5029 || self.dropped_timing_samples > 0
5030 || self.slot_pairs > 0
5031 || self.slot_collisions > 0
5032 }
5033}
5034
5035fn insert_dual_pp_metrics(
5036 body: &mut serde_json::Value,
5037 metrics_scope: &MetricsScope,
5038 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
5039) {
5040 if !metrics_scope.operator() {
5043 return;
5044 }
5045 let snapshot = snapshot();
5046 if !snapshot.populated() {
5047 return;
5048 }
5049 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
5050 .iter()
5051 .enumerate()
5052 .map(|(i, name)| {
5053 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
5054 (
5055 name.to_string(),
5056 json!({
5057 "samples": snapshot.stage_samples[i],
5058 "total_ms": total_ms,
5059 "mean_ms": if snapshot.stage_samples[i] > 0 {
5060 total_ms / snapshot.stage_samples[i] as f64
5061 } else { 0.0 },
5062 }),
5063 )
5064 })
5065 .collect();
5066 body["dual_pp"] = json!({
5067 "overlaps": snapshot.overlaps,
5068 "slot_pairs": snapshot.slot_pairs,
5069 "slot_uses": snapshot.slot_uses,
5070 "slot_collisions": snapshot.slot_collisions,
5071 "cuda_event_spans": timings,
5072 "dropped_timing_samples": snapshot.dropped_timing_samples,
5073 });
5074}
5075
5076fn insert_spec_acceptance_metrics(
5077 body: &mut serde_json::Value,
5078 metrics_scope: &MetricsScope,
5079 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
5080) {
5081 if !metrics_scope.operator() {
5084 return;
5085 }
5086 let snapshot = snapshot();
5087 if snapshot.is_empty() {
5088 return;
5089 }
5090
5091 let mut tau = serde_json::Map::new();
5092 let mut by_position = serde_json::Map::new();
5093 for (model, telemetry) in snapshot {
5094 if telemetry.rounds == 0 {
5095 continue;
5096 }
5097 let n_pos = telemetry
5098 .pos_drafted
5099 .iter()
5100 .rposition(|&n| n > 0)
5101 .map_or(0, |position| position + 1);
5102 tau.insert(model.clone(), json!(telemetry.tau()));
5103 by_position.insert(
5104 model,
5105 json!({
5106 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
5107 "rounds": telemetry.rounds,
5108 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
5109 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
5110 "accept_rate": (0..n_pos).map(|position| {
5111 let offered = telemetry.pos_drafted[position];
5112 if offered > 0 {
5113 telemetry.pos_accepted[position] as f64 / offered as f64
5114 } else {
5115 0.0
5116 }
5117 }).collect::<Vec<f64>>(),
5118 }),
5119 );
5120 }
5121 if !tau.is_empty() {
5122 body["spec_tau"] = serde_json::Value::Object(tau);
5123 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
5124 }
5125}
5126
5127fn insert_peer_probe_metrics(
5128 body: &mut serde_json::Value,
5129 metrics_scope: &MetricsScope,
5130 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
5131) {
5132 if !metrics_scope.operator() {
5135 return;
5136 }
5137 let snapshot = snapshot();
5138 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
5139 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
5140 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
5141 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
5142 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
5143 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
5144 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
5145}
5146
5147async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5149 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5150 Ok(scope) => scope,
5151 Err(response) => return response,
5152 };
5153 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5154 let mut body = if metrics_scope.process_wide() {
5158 json!({
5159 "admitted": m.admitted,
5160 "completed": m.completed,
5161 "tokens_out": m.tokens_out,
5162 "step_p50_ms": m.step_p50_ms,
5163 "step_p99_ms": m.step_p99_ms,
5164 "prompt_tokens_in": m.prompt_tokens_in,
5166 "cached_tokens_in": m.cached_tokens_in,
5167 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5170 "admission_session_defers": m.admission_session_defers,
5173 "admission_vram_defers": m.admission_vram_defers,
5174 "step_oom_parks": m.step_oom_parks,
5175 "continuation_pool_hits": m.continuation_pool_hits,
5176 "continuation_pool_evictions": m.continuation_pool_evictions,
5177 "plain_affinity_rewinds": m.plain_affinity_rewinds,
5178 "served_dspark": m.served_dspark,
5179 "served_spec": m.served_spec,
5180 "served_plain": m.served_plain,
5181 "spec_pool_hits": m.spec_pool_hits,
5182 "spec_pool_misses": m.spec_pool_misses,
5183 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
5184 "spec_pool_evictions": m.spec_pool_evictions,
5185 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
5188 })
5189 } else {
5190 json!({})
5191 };
5192 if metrics_scope.operator() {
5196 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
5197 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
5198 body["budget_source_reload_consecutive"] =
5199 json!(budget_health.source_reload_consecutive);
5200 body["budget_source_available"] = json!(budget_health.source_available);
5201 }
5202 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
5204 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
5205 } else {
5206 0.0
5207 });
5208 body["prefix_cache_hits"] = json!(m.prefix_hits);
5209 body["prefix_cache_misses"] = json!(m.prefix_misses);
5210 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
5211 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
5212 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
5213 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
5214 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
5215 body["prefix_host_entries"] = json!(m.prefix_host_entries);
5219 body["prefix_host_bytes"] = json!(m.prefix_host_bytes);
5220 body["prefix_host_demotions"] = json!(m.prefix_host_demotions);
5221 body["prefix_host_promotions"] = json!(m.prefix_host_promotions);
5222 body["prefix_host_demote_ms"] = json!(m.prefix_host_demote_ms);
5223 body["prefix_host_promote_ms"] = json!(m.prefix_host_promote_ms);
5224 body["prefix_host_rejected_allocs"] = json!(m.prefix_host_rejected_allocs);
5225 body["prefix_host_purges"] = json!(m.prefix_host_purges);
5226 body["prefix_host_purged_entries"] = json!(m.prefix_host_purged_entries);
5227 body["prefix_host_purged_bytes"] = json!(m.prefix_host_purged_bytes);
5228 body["prefix_host_tenant_rejects"] = json!(m.prefix_host_tenant_rejects);
5229 body["prefix_host_pause_demotes"] = json!(m.prefix_host_pause_demotes);
5233 body["prefix_host_pause_cancels"] = json!(m.prefix_host_pause_cancels);
5234 body["kv_flex_borrowed_bytes"] = json!(m.kv_flex_borrowed_bytes);
5239 body["kv_flex_sheds"] = json!(m.kv_flex_sheds);
5240 body["kv_flex_shed_ms"] = json!(m.kv_flex_shed_ms);
5241 body["lcp_histogram"] = json!({
5244 "edges": worker::LCP_HIST_EDGES.to_vec(),
5245 "counts": m.lcp_hist.to_vec(),
5246 });
5247 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
5251 body["prefix_cache_entries"] = json!(m.prefix_entries);
5252 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
5253 body["active_sessions"] = json!(m.active_sessions);
5254 body["queued_requests"] = json!(m.queued_requests);
5255 body["admission_inflight"] = json!(m.admission_inflight);
5259 body["admission_booked_bytes"] = json!(m.admission_booked_bytes);
5260 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
5261 body["spec_pool_entries"] = json!(m.spec_pool_entries);
5262 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
5263 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
5264 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
5265 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
5266 if !m.constraint_compiler_fail_closed.is_empty() {
5267 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
5268 m.constraint_compiler_fail_closed
5269 .iter()
5270 .map(|(model, gauge)| {
5271 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
5272 (model.clone(), json!(value))
5273 })
5274 .collect(),
5275 );
5276 }
5277 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
5278 }
5279 if !m.ns_tokens.is_empty() {
5284 let tenants: serde_json::Map<String, serde_json::Value> = m
5285 .ns_tokens
5286 .iter()
5287 .filter(|(ns, _)| metrics_scope.includes(ns))
5288 .map(|(ns, [p, c])| {
5289 (
5290 ns.clone(),
5291 json!({
5292 "prompt_tokens_in": p,
5293 "cached_tokens_in": c,
5294 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
5295 }),
5296 )
5297 })
5298 .collect();
5299 if !tenants.is_empty() {
5300 body["tenants"] = serde_json::Value::Object(tenants);
5301 }
5302 }
5303 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
5304 .adsd_suspect_total
5305 .iter()
5306 .filter(|(tenant, _)| metrics_scope.includes(tenant))
5307 .map(|(tenant, total)| (tenant.clone(), json!(total)))
5308 .collect();
5309 if !adsd_suspect_total.is_empty() {
5310 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
5311 }
5312 if metrics_scope.operator() {
5314 if let Some((bg, mode)) = &st.bg {
5315 body["bg"] = bg.to_json(mode);
5316 }
5317 }
5318 if metrics_scope.operator() {
5325 let spec: serde_json::Map<String, serde_json::Value> = m
5326 .spec
5327 .iter()
5328 .map(|(model, t)| {
5329 let n_pos = t
5330 .pos_drafted
5331 .iter()
5332 .rposition(|&d| d > 0)
5333 .map_or(0, |p| p + 1);
5334 (
5335 model.clone(),
5336 json!({
5337 "rounds": t.rounds,
5338 "drafted": t.drafted,
5339 "accepted": t.accepted,
5340 "acceptance_rate": if t.drafted > 0 {
5341 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
5342 "tokens_per_round": if t.rounds > 0 {
5343 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
5344 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
5345 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
5346 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
5347 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
5348 .collect::<Vec<f64>>(),
5349 }),
5350 )
5351 })
5352 .collect();
5353 if !spec.is_empty() {
5354 body["spec"] = serde_json::Value::Object(spec);
5355 }
5356 }
5357 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
5358 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
5359 insert_peer_probe_metrics(
5360 &mut body,
5361 &metrics_scope,
5362 memra_engine::pp::peer_probe_metrics,
5363 );
5364 Json(body).into_response()
5365}
5366
5367#[derive(Debug, Default, Deserialize)]
5368struct ModelsQuery {
5369 #[serde(default)]
5370 schema: Option<String>,
5371}
5372
5373fn models_openai_body(models: &[String]) -> serde_json::Value {
5374 let data: Vec<_> = models
5375 .iter()
5376 .map(|m| json!({ "id": m, "object": "model" }))
5377 .collect();
5378 json!({ "object": "list", "data": data })
5379}
5380
5381fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
5386 match metadata.and_then(|m| m.surface.as_deref()) {
5387 Some("embedding") => "embedding",
5388 Some("rerank") => "rerank",
5389 _ => "chat",
5390 }
5391}
5392
5393fn openrouter_supported_parameters(
5394 caps: Option<&ModelCaps>,
5395 max_output_length: Option<u64>,
5396 is_chat: bool,
5397) -> serde_json::Value {
5398 let mut parameters = serde_json::Map::new();
5399 if !is_chat {
5407 return serde_json::Value::Object(parameters);
5408 }
5409 for name in [
5410 "temperature",
5411 "top_p",
5412 "min_p",
5413 "frequency_penalty",
5414 "presence_penalty",
5415 "repetition_penalty",
5416 "stop",
5417 ] {
5418 parameters.insert(name.into(), json!({ "type": "unknown" }));
5419 }
5420 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
5421 parameters.insert(
5422 "seed".into(),
5423 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
5424 );
5425 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
5426 if let Some(max) = max_output_length {
5427 max_tokens["max"] = json!(max);
5428 }
5429 parameters.insert("max_tokens".into(), max_tokens);
5430 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
5431 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
5432 if is_chat && caps.is_some_and(|c| c.tools_branch) {
5433 parameters.insert("tools".into(), json!({ "type": "boolean" }));
5434 parameters.insert(
5435 "tool_choice".into(),
5436 json!({ "type": "enum", "values": ["auto", "none"] }),
5437 );
5438 }
5439 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5440 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
5441 }
5442 serde_json::Value::Object(parameters)
5443}
5444
5445fn model_entry_openrouter(
5446 name: &str,
5447 caps: Option<&ModelCaps>,
5448 metadata: Option<&OpenRouterModelMetadata>,
5449) -> serde_json::Value {
5450 let empty = OpenRouterModelMetadata::default();
5451 let metadata = metadata.unwrap_or(&empty);
5452 let context_length = caps
5453 .map(|c| c.context_length as u64)
5454 .filter(|&v| v > 0 && v <= JSON_SAFE_INTEGER_MAX);
5455 let tokenizer = caps
5456 .map(|c| c.tokenizer.as_str())
5457 .filter(|tokenizer| !tokenizer.is_empty());
5458
5459 let mut input = serde_json::Map::new();
5460 input.insert("type".into(), json!("text"));
5461 let mut supported_inputs = serde_json::Map::new();
5462 if let Some(value) = context_length {
5463 supported_inputs.insert(
5464 "max_context_length".into(),
5465 json!({ "value": value, "unit": "token" }),
5466 );
5467 }
5468 if let Some(value) = metadata.max_prompt_length {
5469 supported_inputs.insert(
5470 "max_prompt_length".into(),
5471 json!({ "value": value, "unit": "token" }),
5472 );
5473 }
5474 if !supported_inputs.is_empty() {
5475 input.insert(
5476 "supported_inputs".into(),
5477 serde_json::Value::Object(supported_inputs),
5478 );
5479 }
5480 let mut input_pricing = Vec::new();
5481 for (kind, cost) in [
5482 ("prompt", metadata.pricing.prompt.as_deref()),
5483 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
5484 ("cache_write", metadata.pricing.cache_write.as_deref()),
5485 ] {
5486 if let Some(cost) = cost {
5487 input_pricing.push(json!({
5488 "type": kind,
5489 "unit": "token",
5490 "cost_usd": cost,
5491 }));
5492 }
5493 }
5494 if !input_pricing.is_empty() {
5495 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
5496 }
5497 let mut input_capacity = Vec::new();
5498 for (kind, value) in [
5499 ("prompt", metadata.capacity.prompt_tpm),
5500 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
5501 ] {
5502 if let Some(value) = value {
5503 input_capacity.push(json!({
5504 "type": kind,
5505 "unit": "token",
5506 "per": "minute",
5507 "value": value,
5508 }));
5509 }
5510 }
5511 if !input_capacity.is_empty() {
5512 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
5513 }
5514
5515 let or_surface = declared_surface(Some(metadata));
5516 let or_is_chat = or_surface == "chat";
5517 let mut output = serde_json::Map::new();
5518 output.insert(
5526 "type".into(),
5527 json!(match or_surface {
5528 "embedding" => "embeddings",
5529 "rerank" => "rerank",
5530 _ => "text",
5531 }),
5532 );
5533 output.insert(
5534 "supported_parameters".into(),
5535 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
5536 );
5537 if or_is_chat {
5541 output.insert("streaming".into(), json!(true));
5542 }
5543 if let Some(value) = metadata.max_output_length
5546 && or_is_chat
5547 {
5548 output.insert(
5549 "max_length".into(),
5550 json!({ "value": value, "unit": "token" }),
5551 );
5552 }
5553 let mut output_pricing = Vec::new();
5554 for (kind, cost) in [
5555 ("completion", metadata.pricing.completion.as_deref()),
5556 (
5557 "internal_reasoning",
5558 metadata.pricing.internal_reasoning.as_deref(),
5559 ),
5560 ] {
5561 if let Some(cost) = cost {
5562 output_pricing.push(json!({
5563 "type": kind,
5564 "unit": "token",
5565 "cost_usd": cost,
5566 }));
5567 }
5568 }
5569 if !output_pricing.is_empty() {
5570 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
5571 }
5572 let mut output_capacity = Vec::new();
5573 if let Some(value) = metadata.capacity.completion_tpm {
5574 output_capacity.push(json!({
5575 "type": "completion",
5576 "unit": "token",
5577 "per": "minute",
5578 "value": value,
5579 }));
5580 }
5581 if let Some(value) = metadata.capacity.concurrency {
5582 output_capacity.push(json!({
5583 "type": "concurrency",
5584 "unit": "request",
5585 "value": value,
5586 }));
5587 }
5588 if !output_capacity.is_empty() {
5589 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
5590 }
5591
5592 let mut entry = serde_json::Map::new();
5593 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
5594 entry.insert("id".into(), json!(name));
5595 entry.insert("name".into(), json!(name));
5596 if let Some(value) = metadata.hugging_face_id.as_deref() {
5597 entry.insert("hugging_face_id".into(), json!(value));
5598 }
5599 if let Some(value) = metadata.created {
5600 entry.insert("created".into(), json!(value));
5601 }
5602 if let Some(value) = metadata.quantization.as_deref() {
5603 entry.insert("quantization".into(), json!(value));
5604 }
5605 if let Some(value) = tokenizer {
5606 entry.insert("tokenizer".into(), json!(value));
5607 }
5608 if let Some(value) = metadata.description.as_deref() {
5609 entry.insert("description".into(), json!(value));
5610 }
5611 let mut input_modalities = vec![serde_json::Value::Object(input)];
5612 for m in &metadata.input_modalities {
5613 let mut extra = serde_json::Map::new();
5614 extra.insert("type".into(), json!(m));
5615 if let Some(cost) = metadata.pricing.prompt.as_deref() {
5616 extra.insert(
5618 "pricing".into(),
5619 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
5620 );
5621 }
5622 input_modalities.push(serde_json::Value::Object(extra));
5623 }
5624 entry.insert(
5625 "input_modalities".into(),
5626 serde_json::Value::Array(input_modalities),
5627 );
5628 entry.insert(
5629 "output_modalities".into(),
5630 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
5631 );
5632 if let Some(cost) = metadata.pricing.request.as_deref() {
5633 entry.insert(
5634 "pricing".into(),
5635 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
5636 );
5637 }
5638 if let Some(value) = metadata.capacity.request_rpm {
5639 entry.insert(
5640 "capacity".into(),
5641 json!([{
5642 "type": "request",
5643 "unit": "request",
5644 "per": "minute",
5645 "value": value,
5646 }]),
5647 );
5648 }
5649 if let Some(value) = metadata.is_ready {
5650 entry.insert("is_ready".into(), json!(value));
5651 }
5652 if let Some(value) = metadata.is_free {
5653 entry.insert("is_free".into(), json!(value));
5654 }
5655 if let Some(value) = metadata.discount_to_user {
5656 entry.insert("discount_to_user".into(), json!(value));
5657 }
5658 if let Some(value) = metadata.openrouter_slug.as_deref() {
5659 entry.insert("openrouter".into(), json!({ "slug": value }));
5660 }
5661 if !metadata.datacenters.is_empty() {
5662 entry.insert("datacenters".into(), json!(metadata.datacenters));
5663 }
5664 let mut compliance = serde_json::Map::new();
5665 if let Some(value) = metadata.zdr {
5666 compliance.insert("zdr".into(), json!(value));
5667 }
5668 if let Some(value) = metadata.hipaa {
5669 compliance.insert("hipaa".into(), json!(value));
5670 }
5671 if !compliance.is_empty() {
5672 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
5673 }
5674 serde_json::Value::Object(entry)
5675}
5676
5677fn models_openrouter_body(st: &AppState) -> serde_json::Value {
5678 let data: Vec<_> = st
5679 .models
5680 .iter()
5681 .map(|model| {
5682 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
5683 })
5684 .collect();
5685 json!({ "data": data })
5686}
5687
5688fn model_entry_openmodels(
5689 name: &str,
5690 caps: Option<&ModelCaps>,
5691 metadata: Option<&OpenRouterModelMetadata>,
5692) -> Result<serde_json::Value, String> {
5693 let metadata = metadata.ok_or_else(|| {
5694 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
5695 })?;
5696 let context_length = caps
5697 .map(|c| c.context_length as u64)
5698 .filter(|&value| value > 0 && value <= JSON_SAFE_INTEGER_MAX)
5699 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
5700 let created = metadata
5701 .created
5702 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
5703 let max_output_length = metadata
5704 .max_output_length
5705 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
5706 let prompt = metadata
5707 .pricing
5708 .prompt
5709 .as_deref()
5710 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
5711 let completion =
5712 metadata.pricing.completion.as_deref().ok_or_else(|| {
5713 format!("OpenModels feed requires pricing.completion for model {name:?}")
5714 })?;
5715 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
5716 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
5717 })?;
5718 let is_ready = metadata
5719 .is_ready
5720 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
5721 let is_free = metadata
5722 .is_free
5723 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
5724 let discount_to_user = metadata
5725 .discount_to_user
5726 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
5727
5728 let mut pricing = serde_json::Map::new();
5729 pricing.insert("prompt".into(), json!(prompt));
5730 pricing.insert("completion".into(), json!(completion));
5731 pricing.insert("input_cache_read".into(), json!(input_cache_read));
5732 if let Some(value) = metadata.pricing.request.as_deref() {
5733 pricing.insert("request".into(), json!(value));
5734 }
5735
5736 let om_surface = declared_surface(Some(metadata));
5737 let om_is_chat = om_surface == "chat";
5738 let mut supported_features = Vec::new();
5739 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
5740 supported_features.push("tool_calling");
5741 }
5742 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5743 supported_features.push("reasoning");
5744 }
5745
5746 let mut entry = serde_json::Map::new();
5747 entry.insert("id".into(), json!(name));
5748 entry.insert("name".into(), json!(name));
5749 entry.insert("created".into(), json!(created));
5750 entry.insert("input_modalities".into(), json!(["text"]));
5751 entry.insert(
5752 "output_modalities".into(),
5753 json!(match om_surface {
5754 "embedding" => ["embeddings"],
5755 "rerank" => ["rerank"],
5756 _ => ["text"],
5757 }),
5758 );
5759 entry.insert("context_length".into(), json!(context_length));
5760 entry.insert("max_output_length".into(), json!(max_output_length));
5761 entry.insert("currency".into(), json!("USD"));
5764 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
5765 entry.insert("supported_features".into(), json!(supported_features));
5766 entry.insert("is_ready".into(), json!(is_ready));
5767 entry.insert("is_free".into(), json!(is_free));
5768 entry.insert("discount_to_user".into(), json!(discount_to_user));
5769 Ok(serde_json::Value::Object(entry))
5770}
5771
5772fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
5773 let data: Result<Vec<_>, _> = st
5774 .models
5775 .iter()
5776 .map(|model| {
5777 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
5778 })
5779 .collect();
5780 Ok(json!({ "data": data? }))
5781}
5782
5783async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
5784 match query.schema.as_deref() {
5785 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
5786 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
5787 Some("openmodels") => match models_openmodels_body(&st) {
5788 Ok(body) => Json(body).into_response(),
5789 Err(error) => bad_request(&error, Some("schema")),
5790 },
5791 Some(schema) => bad_request(
5792 &format!(
5793 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
5794 ),
5795 Some("schema"),
5796 ),
5797 }
5798}
5799
5800fn model_entry_v1(
5808 name: &str,
5809 caps: Option<&ModelCaps>,
5810 metadata: Option<&OpenRouterModelMetadata>,
5811) -> serde_json::Value {
5812 let ctx = caps.map(|c| c.context_length).filter(|&c| c > 0);
5813 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
5817 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
5820 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
5821 Some(p) => json!(p),
5822 None => serde_json::Value::Null,
5823 };
5824 let owned_by = metadata
5825 .and_then(|m| m.owned_by.as_deref())
5826 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
5827 let mut input_modalities = vec!["text"];
5828 if let Some(meta) = metadata {
5829 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
5830 }
5831 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
5832 let reliability = metadata.and_then(|m| m.reliability.as_ref());
5833 let surface = declared_surface(metadata);
5839 let (model_type, endpoints, output_modalities) = match surface {
5840 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
5844 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
5845 _ => ("chat", vec!["chat/completions"], vec!["text"]),
5846 };
5847 let is_chat = surface == "chat";
5848 json!({
5849 "id": name,
5850 "name": name,
5851 "object": "model",
5852 "owned_by": owned_by,
5853 "type": model_type,
5854 "context_length": ctx,
5855 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
5858 "endpoints": endpoints,
5859 "input_modalities": input_modalities,
5860 "output_modalities": output_modalities,
5861 "capabilities": {
5862 "streaming": is_chat,
5865 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
5866 "structured_output": is_chat && !is_dsv4,
5867 "reasoning": is_chat && thinking,
5868 "prompt_caching": is_chat && !is_dsv4,
5869 },
5870 "pricing": {
5871 "currency": "USD",
5872 "unit": "per_1m_tokens",
5873 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
5874 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
5875 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
5876 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
5877 "minimum_request": metadata
5879 .and_then(|m| m.pricing.request.as_deref())
5880 .unwrap_or("0"),
5881 },
5882 "lifecycle": {
5883 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
5884 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
5885 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
5886 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
5887 },
5888 "reliability": {
5889 "first_token_timeout_seconds":
5890 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
5891 "completion_timeout_seconds":
5892 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
5893 "stream_idle_timeout_seconds":
5894 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
5895 "capacity_scope":
5896 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
5897 },
5898 })
5899}
5900
5901async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
5904 let data: Vec<_> = st
5905 .models
5906 .iter()
5907 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
5908 .collect();
5909 let mut body = json!({
5910 "object": "list",
5911 "contract_version": "2.0",
5912 "data": data,
5913 });
5914 if let Some(provider) = st.provider_metadata.as_ref() {
5919 body["provider"] = json!({
5920 "id": provider.id,
5921 "status_url": provider.status_url,
5922 "support_contact": provider.support_contact,
5923 "incident_contact": provider.incident_contact,
5924 "regions": provider.regions,
5925 "request_id_header": "x-request-id",
5926 "error_contract": {
5927 "rate_limit_status": 429,
5928 "overload_status": 503,
5929 "retry_after_header": "Retry-After",
5930 "account_quota_error_codes": ["insufficient_balance"],
5931 },
5932 });
5933 }
5934 Json(body)
5935}
5936
5937async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5940 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5941 Ok(scope) => scope,
5942 Err(response) => return response,
5943 };
5944 if !metrics_scope.process_wide() {
5945 return error_response(
5946 StatusCode::FORBIDDEN,
5947 "completion api keys do not authorize process-wide yield metrics; configure \
5948 MEMRA_METRICS_TOKEN",
5949 "authentication_error",
5950 None,
5951 );
5952 }
5953 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5954 let lane = |i: usize| {
5955 json!({
5956 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
5957 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
5958 })
5959 };
5960 let mut body = json!({
5961 "lanes": {
5962 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
5963 },
5964 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
5965 });
5966 if metrics_scope.operator() {
5967 body["batch_size_last"] = json!(m.batch_size_last);
5968 }
5969 Json(body).into_response()
5970}
5971
5972async fn peek_admission(
5985 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5986) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, (Response, &'static str)> {
5987 match rx.recv().await {
5988 Some(Event::Error(e)) => {
5993 let error_code = engine_error_code(e.class);
5994 Err((engine_error_response(&e), error_code))
5995 }
5996 first => {
5997 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
5998 if let Some(ev) = first {
5999 let _ = tx2.send(ev);
6000 }
6001 tokio::spawn(forward_events(rx, tx2));
6002 Ok(rx2)
6003 }
6004 }
6005}
6006
6007async fn forward_events(
6018 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
6019 tx2: tokio::sync::mpsc::UnboundedSender<Event>,
6020) {
6021 loop {
6022 tokio::select! {
6023 biased;
6024 () = tx2.closed() => break,
6025 ev = rx.recv() => match ev {
6026 Some(ev) => {
6027 if tx2.send(ev).is_err() {
6028 break;
6029 }
6030 }
6031 None => break,
6032 },
6033 }
6034 }
6035}
6036
6037async fn peek_first_token(
6052 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
6053 deadline: RequestDeadline,
6054) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, ()> {
6055 let mut buffered: Vec<Event> = Vec::new();
6056 loop {
6057 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
6058 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
6061 let first_delivery = matches!(
6062 ev,
6063 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
6064 );
6065 buffered.push(ev);
6066 if first_delivery {
6067 break;
6068 }
6069 }
6070 }
6071 }
6072 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
6073 for ev in buffered {
6074 let _ = tx2.send(ev);
6075 }
6076 tokio::spawn(forward_events(rx, tx2));
6077 Ok(rx2)
6078}
6079
6080#[cfg(test)]
6082fn build_request(
6086 req: &CompletionReq,
6087 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6088 lane: lanes::Lane,
6089 affinity: Option<String>,
6090) -> Request {
6091 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
6092}
6093
6094fn build_request_with_trace(
6095 req: &CompletionReq,
6096 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6097 lane: lanes::Lane,
6098 affinity: Option<String>,
6099 ttft: Option<Arc<ttft::Trace>>,
6100 sampling_defaults: &SamplingDefaults,
6101) -> Request {
6102 let params = GenParams {
6103 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6104 max_ctx: req.max_ctx,
6105 eos: Vec::new(), };
6107 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
6112 Request {
6113 model: req.model.clone(),
6114 prompt_ids: req.prompt_ids.clone(),
6115 prompt_text: req.prompt.clone(),
6116 chat: req.chat,
6117 chat_turns: Vec::new(),
6118 tools_json: Vec::new(),
6119 tools_struct: Vec::new(),
6120 think: ThinkMode::Default,
6121 reasoning_effort: None, params,
6123 sampler_cfg,
6124 stop_strings: req.stop.clone().into_vec(),
6125 trace_id: req.trace_id.clone(),
6126 request_id: String::new(),
6129 admit_predict_logged: false,
6130 max_prompt_tokens: None,
6131 cache_ns: cache_namespace(&req.cache_salt),
6132 affinity,
6133 lane,
6134 grammar: None, prepared_constraint: None,
6136 constraint_ready: None,
6137 oom_retries: 0, spec_k_replay: None,
6139 prepared_prompt: None,
6140 capture: None, images: Vec::new(), gemma_images: Vec::new(),
6143 step_images: Vec::new(),
6144 vision_memory: None,
6145 ttft,
6146 tx,
6147 }
6148}
6149
6150struct ChatPlan {
6153 request: Request,
6154 parser: Option<ToolStreamParser>,
6157 pending_images: Vec<PendingVisionUnit>,
6160 pending_gemma: Vec<PendingGemmaImage>,
6161 pending_step: Vec<PendingStepImage>,
6162 vision_memory: Option<VisionMemoryPermit>,
6166}
6167
6168pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
6169 req.messages.iter().any(|message| {
6170 message.content.as_array().is_some_and(|parts| {
6171 parts.iter().any(|part| {
6172 matches!(
6173 part.get("type").and_then(serde_json::Value::as_str),
6174 Some("image_url" | "video_url")
6175 )
6176 })
6177 })
6178 })
6179}
6180
6181fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
6182 let mut total = 0usize;
6183 let mut add = |bytes: usize| {
6184 total = total.checked_add(bytes).ok_or_else(|| {
6185 "vision patch memory reservation overflowed while planning".to_string()
6186 })?;
6187 Ok::<(), String>(())
6188 };
6189 for unit in &plan.pending_images {
6190 let bytes = match unit {
6191 PendingVisionUnit::Still { gh, gw, .. } => gh
6192 .checked_mul(*gw)
6193 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6194 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6195 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
6196 PendingVisionUnit::Video { groups, .. } => {
6197 groups.iter().try_fold(0usize, |total, group| {
6198 let bytes = group
6199 .gh
6200 .checked_mul(group.gw)
6201 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6202 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6203 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6204 total.checked_add(bytes).ok_or_else(|| {
6205 "vision patch memory reservation overflowed while planning".to_string()
6206 })
6207 })?
6208 }
6209 };
6210 add(bytes)?;
6211 }
6212 for unit in &plan.pending_gemma {
6213 let bytes = unit
6214 .gw
6215 .checked_mul(unit.gh)
6216 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
6217 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6218 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6219 add(bytes)?;
6220 }
6221 for unit in &plan.pending_step {
6222 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
6223 let patches = unit
6225 .plan
6226 .n_tiles
6227 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
6228 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
6229 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6230 let bytes = patches
6231 .checked_mul(SV_PATCH_IN)
6232 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6233 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6234 add(bytes)?;
6235 }
6236 Ok(total)
6237}
6238
6239pub(crate) fn reserve_vision_memory(
6240 plan: &ChatPlan,
6241) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
6242 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
6243 try_reserve_vision_memory(bytes)
6244}
6245
6246#[cfg(test)]
6247fn build_chat_request(
6248 req: ChatCompletionReq,
6249 caps: Option<&ModelCaps>,
6250 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6251 lane: lanes::Lane,
6252 affinity: Option<String>,
6253) -> Result<ChatPlan, String> {
6254 let defaults = ModelSamplingDefaults::resolve(None, caps);
6257 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
6258}
6259
6260#[allow(clippy::too_many_arguments)]
6274fn build_chat_request_with_trace(
6275 req: ChatCompletionReq,
6276 caps: Option<&ModelCaps>,
6277 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6278 lane: lanes::Lane,
6279 affinity: Option<String>,
6280 ttft: Option<Arc<ttft::Trace>>,
6281 default_effort: Option<&str>,
6282 sampling_defaults: &ModelSamplingDefaults,
6283) -> Result<ChatPlan, String> {
6284 let client_sampling: ClientSampling = (&req).into();
6287 let tool_choice = parse_tool_choice(&req.tool_choice)?;
6288 if let Some(c) = caps {
6293 if !c.chat_ok {
6294 return Err(format!(
6295 "model {:?} has no chat template (checkpoint carries neither \
6296 tokenizer_config.json chat_template nor chat_template.jinja) — \
6297 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
6298 req.model
6299 ));
6300 }
6301 }
6302 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
6303 let (mut think, effort_level, think_client_explicit) = parse_think(
6304 &req.reasoning_effort,
6305 &req.reasoning,
6306 vllm_switch,
6307 req.include_reasoning,
6308 default_effort,
6309 caps.is_some_and(|c| c.dsv4),
6310 )?;
6311 let level_template = caps
6315 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort)
6316 .unwrap_or(false);
6317 if think_client_explicit && think == ThinkMode::NoThink {
6338 if let Some(c) = caps {
6339 if c.qwen_think && !c.think_switch && !c.dsv4 {
6340 return Err(format!(
6341 "model {:?} cannot disable reasoning: its chat template opens a think \
6342 tail unconditionally and carries no enable_thinking switch, so \
6343 reasoning_effort/enable_thinking cannot turn it off on this model",
6344 req.model
6345 ));
6346 }
6347 }
6348 }
6349 let reasoning_effort = if level_template { effort_level } else { None };
6382 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
6385 if grammar.is_some() {
6399 if let Some(c) = caps {
6400 if c.qwen_think && think != ThinkMode::NoThink {
6401 if c.think_switch {
6402 think = ThinkMode::NoThink;
6403 } else if c.think_close.is_empty() {
6404 return Err(
6405 "response_format requires the model's think channel to close \
6406 before the grammar can engage, but this chat template has \
6407 neither an enable_thinking switch nor a recognizable \
6408 think-close token sequence"
6409 .into(),
6410 );
6411 }
6412 }
6416 }
6417 }
6418
6419 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
6426
6427 let (tools_json, tools_struct, schemas) =
6430 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
6431 prepare_tools(&req.tools)?
6432 } else {
6433 (Vec::new(), Vec::new(), HashMap::new())
6434 };
6435
6436 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
6437 let mut images: Vec<PendingVisionUnit> = Vec::new();
6438 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
6439 let mut step_images: Vec<PendingStepImage> = Vec::new();
6440 let mut next_video = 0usize;
6441 for msg in &req.messages {
6442 let content = content_to_text_vision(
6443 &msg.content,
6444 &mut images,
6445 &mut gemma_images,
6446 &mut step_images,
6447 &mut next_video,
6448 )
6449 .map_err(|e| format!("{} message: {e}", msg.role))?;
6450 let tool_calls = msg
6451 .tool_calls
6452 .iter()
6453 .map(render_req_tool_call)
6454 .collect::<Result<Vec<_>, _>>()?;
6455 if !tool_calls.is_empty() && msg.role != "assistant" {
6456 return Err("tool_calls are only valid on assistant messages".into());
6457 }
6458 let role = if msg.role == "developer" {
6461 "system".to_string()
6462 } else {
6463 msg.role.clone()
6464 };
6465 turns.push(TmplTurn {
6466 role,
6467 content,
6468 tool_calls,
6469 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
6471 tool_call_id: msg.tool_call_id.clone(),
6472 tool_name: msg.name.clone(),
6473 tool_responses: Vec::new(),
6474 task: None,
6478 tools: Vec::new(),
6479 });
6480 }
6481
6482 let has_tool_features = !tools_json.is_empty()
6485 || turns
6486 .iter()
6487 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
6488 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
6489 return Err(format!(
6490 "model {:?} chat template has no tools branch",
6491 req.model
6492 ));
6493 }
6494
6495 let think_open = caps
6498 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
6499 .unwrap_or(false);
6500 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
6518 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
6524 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
6525 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
6526 let parser = if is_dsv4 && (dsv4_tools || dsv4_think_open) {
6527 Some(ToolStreamParser::dsv4(dsv4_think_open))
6528 } else if gemma_tools {
6529 Some(ToolStreamParser::gemma_tools())
6530 } else if !tools_json.is_empty() {
6531 Some(ToolStreamParser::new(schemas, think_open))
6532 } else if think_open {
6533 Some(ToolStreamParser::reasoning_only())
6534 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
6535 Some(ToolStreamParser::gemma_thought())
6543 } else {
6544 None
6545 };
6546
6547 Ok(ChatPlan {
6548 request: Request {
6549 model: req.model,
6550 prompt_ids: Vec::new(),
6551 prompt_text: String::new(),
6552 chat: false,
6553 chat_turns: turns,
6554 tools_json,
6555 tools_struct,
6556 think,
6557 reasoning_effort,
6558 params: GenParams {
6559 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6560 max_ctx: req.max_ctx,
6561 eos: Vec::new(),
6562 },
6563 sampler_cfg,
6564 stop_strings: {
6565 let mut stops = req.stop.into_vec();
6570 if gemma_tools {
6571 stops.push("<tool_call|>".to_string());
6572 }
6573 if dsv4_tools {
6578 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
6579 }
6580 stops
6581 },
6582 trace_id: None,
6583 request_id: String::new(),
6586 admit_predict_logged: false,
6587 max_prompt_tokens: None,
6588 cache_ns: cache_namespace(&req.cache_salt),
6589 affinity,
6590 lane,
6591 grammar,
6592 prepared_constraint: None,
6593 constraint_ready: None,
6594 oom_retries: 0, spec_k_replay: None,
6596 prepared_prompt: None,
6597 images: Vec::new(),
6602 gemma_images: Vec::new(),
6603 step_images: Vec::new(),
6604 capture: None, vision_memory: None,
6606 ttft,
6607 tx,
6608 },
6609 parser,
6610 pending_images: images,
6611 pending_gemma: gemma_images,
6612 pending_step: step_images,
6613 vision_memory: None,
6614 })
6615}
6616
6617fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
6623 for (i, unit) in plan.pending_images.drain(..).enumerate() {
6624 match unit {
6625 PendingVisionUnit::Still { bytes, gh, gw } => {
6626 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
6627 .map_err(|e| format!("image {}: {e}", i + 1))?;
6628 if (prep.gh, prep.gw) != (gh, gw) {
6629 return Err(format!(
6630 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
6631 i + 1,
6632 prep.gh,
6633 prep.gw
6634 ));
6635 }
6636 plan.request
6637 .images
6638 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
6639 }
6640 PendingVisionUnit::Video {
6641 bytes,
6642 groups,
6643 video,
6644 } => {
6645 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
6646 .map_err(|e| format!("video {}: {e}", i + 1))?;
6647 if prepared.groups.len() != groups.len() {
6648 return Err(format!(
6649 "video {}: decoded {} groups differ from its header-planned {} groups",
6650 i + 1,
6651 prepared.groups.len(),
6652 groups.len()
6653 ));
6654 }
6655 for ((group, prep), timestamp) in
6656 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
6657 {
6658 if (prep.gh, prep.gw) != (group.gh, group.gw) {
6659 return Err(format!(
6660 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
6661 i + 1,
6662 prep.gh,
6663 prep.gw,
6664 group.gh,
6665 group.gw
6666 ));
6667 }
6668 if (timestamp - group.timestamp).abs() > 0.001 {
6669 return Err(format!(
6670 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
6671 i + 1,
6672 group.timestamp
6673 ));
6674 }
6675 plan.request
6676 .images
6677 .push(memra_engine::vision_pre::VisionUnit {
6678 prep,
6679 video: Some(video),
6680 });
6681 }
6682 }
6683 }
6684 }
6685 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
6686 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
6687 .map_err(|e| format!("image {}: {e}", i + 1))?;
6688 if (gw, gh) != (unit.gw, unit.gh) {
6689 return Err(format!(
6690 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
6691 i + 1,
6692 unit.gw,
6693 unit.gh
6694 ));
6695 }
6696 plan.request
6697 .gemma_images
6698 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
6699 }
6700 for (i, unit) in plan.pending_step.drain(..).enumerate() {
6701 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
6702 .map_err(|e| format!("image {}: {e}", i + 1))?;
6703 if prepped.tiles.len() != unit.plan.n_tiles
6704 || prepped.newline_mask != unit.plan.newline_mask
6705 {
6706 return Err(format!(
6707 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
6708 ({} tiles) — refusing (pad runs already rendered)",
6709 i + 1,
6710 prepped.tiles.len(),
6711 unit.plan.n_tiles
6712 ));
6713 }
6714 plan.request.step_images.push(prepped);
6715 }
6716 Ok(())
6717}
6718
6719fn bearer_token(headers: &HeaderMap) -> Option<&str> {
6727 headers
6728 .get("authorization")
6729 .and_then(|value| value.to_str().ok())
6730 .and_then(|value| value.strip_prefix("Bearer "))
6731}
6732
6733fn authentication_error(why: auth::AuthDenied) -> Response {
6734 match why {
6735 auth::AuthDenied::Unknown => error_response(
6736 StatusCode::UNAUTHORIZED,
6737 "invalid api key",
6738 "authentication_error",
6739 None,
6740 ),
6741 auth::AuthDenied::Disabled => error_response(
6742 StatusCode::FORBIDDEN,
6743 "api key is disabled",
6744 "authentication_error",
6745 None,
6746 ),
6747 }
6748}
6749
6750fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
6751 auth::authenticate_with(
6752 api_auth.keyring,
6753 api_auth.single_key.as_deref(),
6754 bearer_token(headers),
6755 )
6756 .map_err(authentication_error)
6757}
6758
6759#[derive(Debug, Clone, PartialEq, Eq)]
6760enum MetricsScope {
6761 All,
6762 CompletionDomain,
6763 Tenant(String),
6764}
6765
6766impl MetricsScope {
6767 fn operator(&self) -> bool {
6768 matches!(self, MetricsScope::All)
6769 }
6770
6771 fn process_wide(&self) -> bool {
6772 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
6773 }
6774
6775 fn includes(&self, tenant_row: &str) -> bool {
6776 match self {
6777 MetricsScope::All | MetricsScope::CompletionDomain => true,
6778 MetricsScope::Tenant(tenant) => tenant == tenant_row,
6779 }
6780 }
6781}
6782
6783fn authorize_metrics(
6784 api_auth: &ApiAuth,
6785 metrics_auth: &MetricsAuth,
6786 headers: &HeaderMap,
6787) -> Result<MetricsScope, Response> {
6788 if !metrics_auth.required {
6789 return Ok(MetricsScope::All);
6790 }
6791 let Some(candidate) = bearer_token(headers) else {
6792 return Err(authentication_error(auth::AuthDenied::Unknown));
6793 };
6794 if let Some(token) = metrics_auth.token.as_deref() {
6795 if auth::constant_time_secret_eq(token, candidate) {
6796 return Ok(MetricsScope::All);
6797 }
6798 if api_auth.configured() {
6799 return match auth::authenticate_with(
6800 api_auth.keyring,
6801 api_auth.single_key.as_deref(),
6802 Some(candidate),
6803 ) {
6804 Ok(_) => Err(error_response(
6805 StatusCode::FORBIDDEN,
6806 "completion api keys do not authorize metrics while \
6807 MEMRA_METRICS_TOKEN is configured",
6808 "authentication_error",
6809 None,
6810 )),
6811 Err(why) => Err(authentication_error(why)),
6812 };
6813 }
6814 return Err(authentication_error(auth::AuthDenied::Unknown));
6815 }
6816 if api_auth.configured() {
6817 let tenant = authenticate(api_auth, headers)?;
6818 return Ok(if api_auth.keyring.is_some() {
6819 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
6820 } else {
6821 MetricsScope::CompletionDomain
6825 });
6826 }
6827 Err(authentication_error(auth::AuthDenied::Unknown))
6828}
6829
6830fn lane_for_tenant(
6836 headers: &axum::http::HeaderMap,
6837 tenant: &auth::TenantCtx,
6838) -> Result<lanes::Lane, Response> {
6839 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
6840 None => None,
6841 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
6846 error_response_coded(
6847 StatusCode::BAD_REQUEST,
6848 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
6849 "invalid_request_error",
6850 Some("x-lane"),
6851 Some("invalid_lane"),
6852 )
6853 })?),
6854 };
6855 match tenant.lane_class {
6856 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
6857 auth::LaneClass::Batch => match requested {
6858 None => Ok(lanes::Lane::Harvest),
6859 Some(lanes::Lane::Interactive) => Err(error_response(
6860 StatusCode::FORBIDDEN,
6861 "this api key is batch-class: x-lane interactive is not permitted \
6862 (use judge or harvest)",
6863 "authentication_error",
6864 Some("x-lane"),
6865 )),
6866 Some(l) => Ok(l),
6867 },
6868 }
6869}
6870
6871fn tenant_namespace(
6875 tenant: &auth::TenantCtx,
6876 cache_salt: &Option<String>,
6877) -> Result<String, &'static str> {
6878 let keyring_configured = auth::global().is_some();
6879 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
6880 if keyring_configured {
6881 Ok(auth::scope_namespace(&tenant.tenant, &raw))
6882 } else {
6883 Ok(raw)
6884 }
6885}
6886
6887fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
6892 eprintln!(
6893 "[meter] admit id={} tenant={} lane={} model={:?}",
6894 env.id,
6895 tenant.tenant,
6896 lane.as_str(),
6897 model
6898 );
6899}
6900
6901fn apply_model_request_limits(
6902 request: &mut Request,
6903 metadata: Option<&OpenRouterModelMetadata>,
6904 caps: Option<&ModelCaps>,
6905) -> Result<(), (String, &'static str)> {
6906 let Some(metadata) = metadata else {
6907 return Ok(());
6908 };
6909 let max_prompt = metadata
6910 .max_prompt_length
6911 .map(usize::try_from)
6912 .transpose()
6913 .map_err(|_| {
6914 (
6915 "configured model prompt limit does not fit this platform".into(),
6916 "model",
6917 )
6918 })?;
6919 let max_output = metadata
6920 .max_output_length
6921 .map(usize::try_from)
6922 .transpose()
6923 .map_err(|_| {
6924 (
6925 "configured model output limit does not fit this platform".into(),
6926 "model",
6927 )
6928 })?;
6929
6930 request.max_prompt_tokens = max_prompt;
6931 if let Some(max_output) = max_output {
6932 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
6933 request.params.max_new = metadata
6934 .default_output_length
6935 .map(usize::try_from)
6936 .transpose()
6937 .map_err(|_| {
6938 (
6939 "configured default output length does not fit this platform".into(),
6940 "model",
6941 )
6942 })?
6943 .unwrap_or(max_output);
6944 } else if request.params.max_new > max_output {
6945 return Err((
6946 format!(
6947 "max_tokens {} exceeds configured model maximum {max_output}",
6948 request.params.max_new
6949 ),
6950 "max_tokens",
6951 ));
6952 }
6953 }
6954
6955 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
6959 (max_prompt, max_output, request.params.max_ctx)
6960 {
6961 let operational_ctx = max_prompt
6962 .checked_add(max_output)
6963 .and_then(|value| value.checked_add(8))
6964 .ok_or_else(|| {
6965 (
6966 "configured model context envelope overflowed".into(),
6967 "model",
6968 )
6969 })?;
6970 let operational_ctx = caps
6971 .map(|caps| caps.context_length)
6972 .filter(|&context| context > 0)
6973 .map_or(operational_ctx, |context| operational_ctx.min(context));
6974 if requested_ctx > operational_ctx {
6975 return Err((
6976 format!(
6977 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
6978 ),
6979 "max_ctx",
6980 ));
6981 }
6982 }
6983 Ok(())
6984}
6985
6986fn effective_max_tokens(request: &worker::Request) -> Option<u64> {
6990 (request.params.max_new != worker::MAX_NEW_CTX_BOUNDED).then_some(request.params.max_new as u64)
6991}
6992
6993#[allow(clippy::too_many_arguments)]
6994fn start_request_receipt(
6995 st: &AppState,
6996 env: &Envelope,
6997 tenant: &auth::TenantCtx,
6998 model: &str,
6999 route: &'static str,
7000 lane: lanes::Lane,
7001 stream: bool,
7002 max_tokens: Option<u64>,
7003 reserved_ctx: Option<u64>,
7004 budget_permit: Option<metering::Permit>,
7005) -> Option<Box<dyn metering::Receipt>> {
7006 st.metering.as_ref().map(|accounting| {
7007 accounting.open(
7008 &metering::RequestMeta {
7009 request_id: &env.id,
7010 tenant: &tenant.tenant,
7011 principal: tenant.key_prefix.as_deref(),
7012 model,
7013 route,
7014 lane: lane.as_str(),
7015 stream,
7016 max_tokens,
7017 reserved_ctx,
7018 },
7019 budget_permit,
7020 )
7021 })
7022}
7023
7024fn arm_capture(
7030 mut receipt: Option<Box<dyn metering::Receipt>>,
7031 prompt: impl FnOnce() -> serde_json::Value,
7032) -> Option<Box<dyn metering::Receipt>> {
7033 if let Some(receipt) = receipt.as_mut()
7034 && receipt.wants_capture()
7035 {
7036 receipt.arm_capture(prompt());
7037 }
7038 receipt
7039}
7040
7041fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
7045 serde_json::Value::Array(
7046 messages
7047 .iter()
7048 .map(|message| {
7049 let mut row = json!({ "role": message.role, "content": message.content });
7050 if !message.tool_calls.is_empty() {
7051 row["tool_calls"] = serde_json::Value::Array(
7052 message
7053 .tool_calls
7054 .iter()
7055 .map(|call| {
7056 json!({
7057 "id": call.id,
7058 "function": {
7059 "name": call.function.name,
7060 "arguments": call.function.arguments,
7061 },
7062 })
7063 })
7064 .collect(),
7065 );
7066 }
7067 row
7068 })
7069 .collect(),
7070 )
7071}
7072
7073enum BudgetRejection {
7074 Invalid(String),
7075 Insufficient,
7076 Unenrolled,
7077 PrincipalCapped,
7080 Unavailable(String),
7081}
7082
7083impl BudgetRejection {
7084 fn into_response(self) -> (Response, &'static str) {
7085 match self {
7086 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
7087 Self::Insufficient => (
7088 error_response_coded(
7089 StatusCode::PAYMENT_REQUIRED,
7090 "tenant prepaid balance is insufficient for this request",
7091 "insufficient_balance",
7092 None,
7093 Some("insufficient_balance"),
7094 ),
7095 "insufficient_balance",
7096 ),
7097 Self::Unenrolled => (
7098 error_response_coded(
7099 StatusCode::PAYMENT_REQUIRED,
7100 "tenant is not enrolled for prepaid billing",
7101 "tenant_not_enrolled",
7102 None,
7103 Some("tenant_not_enrolled"),
7104 ),
7105 "tenant_not_enrolled",
7106 ),
7107 Self::PrincipalCapped => (
7108 error_response_coded(
7109 StatusCode::PAYMENT_REQUIRED,
7110 "this API key's spend cap is reached; raise or clear the key's cap to continue",
7111 "key_spend_cap_reached",
7112 None,
7113 Some("key_spend_cap_reached"),
7114 ),
7115 "key_spend_cap_reached",
7116 ),
7117 Self::Unavailable(err) => {
7118 eprintln!("[budget] ERROR: admission unavailable: {err}");
7119 (
7120 error_response_coded(
7121 StatusCode::SERVICE_UNAVAILABLE,
7122 "tenant budget accounting is unavailable",
7123 "server_error",
7124 None,
7125 Some("tenant_budget_unavailable"),
7126 ),
7127 "tenant_budget_unavailable",
7128 )
7129 }
7130 }
7131 }
7132}
7133
7134fn prepare_budget_prompt(
7135 request: &mut Request,
7136 tokenizer: Option<&Tokenizer>,
7137) -> Result<usize, String> {
7138 if let Some(error) = worker::prompt_source_limit_error(request) {
7139 return Err(error);
7140 }
7141 if request.prepared_prompt.is_none() {
7142 if let Some(trace) = request.ttft.as_ref() {
7143 trace.mark_tokenize_start();
7144 }
7145 let prompt = if !request.prompt_ids.is_empty() {
7146 request.prompt_ids.clone()
7147 } else if !request.chat_turns.is_empty() {
7148 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7149 let plain = worker::plain_chat_render_path(
7156 &request.tools_json,
7157 &request.think,
7158 request.reasoning_effort.as_deref(),
7159 &request.chat_turns,
7160 tokenizer.has_qwen_effort_ladder(),
7161 );
7162 let rendered = if plain {
7163 let messages: Vec<_> = request
7164 .chat_turns
7165 .iter()
7166 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
7167 .collect();
7168 tokenizer.apply_chat_template(&messages, true)
7169 } else {
7170 tokenizer
7171 .apply_chat_template_tools_ex(
7172 &request.chat_turns,
7173 true,
7174 &request.tools_json,
7175 &request.tools_struct,
7176 request.think,
7177 request.reasoning_effort.as_deref(),
7178 )
7179 .map_err(|err| format!("chat template: {err}"))?
7180 };
7181 tokenizer.encode(&rendered, true)
7182 } else if request.chat {
7183 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7184 let rendered =
7185 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
7186 tokenizer.encode(&rendered, true)
7187 } else {
7188 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7189 tokenizer.encode(&request.prompt_text, true)
7190 };
7191 if prompt.is_empty() {
7192 return Err("empty prompt after tokenization".into());
7193 }
7194 if let Some(trace) = request.ttft.as_ref() {
7195 trace.mark_tokenize_end(prompt.len());
7196 }
7197 request.prepared_prompt = Some(prompt);
7198 }
7199 let prompt_tokens = request
7200 .prepared_prompt
7201 .as_ref()
7202 .expect("budget prompt was prepared")
7203 .len();
7204 if let Some(limit) = request.max_prompt_tokens
7205 && prompt_tokens > limit
7206 {
7207 return Err(format!(
7208 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
7209 ));
7210 }
7211 Ok(prompt_tokens)
7212}
7213
7214fn budget_completion_bound(
7215 request: &Request,
7216 prompt_tokens: usize,
7217 caps: Option<&ModelCaps>,
7218) -> Result<usize, String> {
7219 let max_new = request.params.max_new;
7220 let requested_ctx = match (request.params.max_ctx, max_new) {
7221 (Some(cap), _) => cap,
7222 (None, worker::MAX_NEW_CTX_BOUNDED) => {
7223 let server_ctx = std::env::var("MEMRA_CTX")
7224 .ok()
7225 .and_then(|value| value.parse().ok())
7226 .unwrap_or(8192usize);
7227 let mut cap = server_ctx;
7228 if prompt_tokens.saturating_add(16) > cap {
7229 cap = prompt_tokens.saturating_add(server_ctx);
7230 }
7231 cap
7232 }
7233 (None, max_new) => prompt_tokens
7234 .checked_add(max_new)
7235 .and_then(|value| value.checked_add(8))
7236 .ok_or_else(|| "request context bound overflowed".to_string())?,
7237 };
7238 let ctx_cap = caps
7239 .map(|caps| caps.context_length)
7240 .filter(|&context| context > 0)
7241 .map_or(requested_ctx, |context| requested_ctx.min(context));
7242 if prompt_tokens >= ctx_cap {
7243 return Err(format!(
7244 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
7245 ));
7246 }
7247 Ok(max_new.min(ctx_cap - prompt_tokens))
7248}
7249
7250struct BudgetAdmission {
7255 permit: Option<metering::Permit>,
7256 reserved_ctx: Option<u64>,
7257}
7258
7259impl std::fmt::Debug for BudgetAdmission {
7261 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
7262 f.debug_struct("BudgetAdmission")
7263 .field("permit", &self.permit.is_some())
7264 .field("reserved_ctx", &self.reserved_ctx)
7265 .finish()
7266 }
7267}
7268
7269fn admit_tenant_budget(
7270 st: &AppState,
7271 tenant: &auth::TenantCtx,
7272 request: &mut Request,
7273) -> Result<BudgetAdmission, BudgetRejection> {
7274 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
7275 return Ok(BudgetAdmission {
7276 permit: None,
7277 reserved_ctx: None,
7278 });
7279 };
7280 match accounting.is_limited(&tenant.tenant) {
7281 Ok(false) => return Err(BudgetRejection::Unenrolled),
7282 Ok(true) => {}
7283 Err(metering::AdmitError::Unavailable(err)) => {
7284 return Err(BudgetRejection::Unavailable(err));
7285 }
7286 Err(other) => {
7287 return Err(BudgetRejection::Unavailable(format!(
7288 "unexpected budget enrollment result: {other:?}"
7289 )));
7290 }
7291 }
7292 let tokenizer = st
7293 .budget_tokenizers
7294 .as_ref()
7295 .and_then(|tokenizers| tokenizers.get(&request.model))
7296 .map(Arc::as_ref);
7297 if request.prompt_ids.is_empty() && tokenizer.is_none() {
7298 return Err(BudgetRejection::Unavailable(format!(
7299 "no reservation tokenizer for model {:?}",
7300 request.model
7301 )));
7302 }
7303 let prompt_tokens =
7304 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
7305 let completion_tokens =
7306 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
7307 .map_err(BudgetRejection::Invalid)?;
7308 let prompt_tokens = u64::try_from(prompt_tokens)
7309 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
7310 let completion_tokens = u64::try_from(completion_tokens)
7311 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
7312 match accounting.reserve(
7313 &tenant.tenant,
7314 tenant.key_prefix.as_deref(),
7315 &request.model,
7316 prompt_tokens,
7317 completion_tokens,
7318 ) {
7319 Ok(permit) => Ok(BudgetAdmission {
7320 permit,
7321 reserved_ctx: Some(prompt_tokens.saturating_add(completion_tokens)),
7322 }),
7323 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
7324 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
7325 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
7329 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
7330 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
7331 }
7332}
7333
7334fn request_ledger_error_response() -> Response {
7335 error_response_coded(
7336 StatusCode::INTERNAL_SERVER_ERROR,
7337 "request completion could not be committed to the billing ledger",
7338 "server_error",
7339 None,
7340 Some("request_ledger_unavailable"),
7341 )
7342}
7343
7344fn request_ledger_error_body() -> serde_json::Value {
7345 error_body(
7346 "request completion could not be committed to the billing ledger",
7347 "server_error",
7348 None,
7349 Some("request_ledger_unavailable"),
7350 )
7351}
7352
7353fn ledger_rejected(
7354 mut receipt: Option<Box<dyn metering::Receipt>>,
7355 response: Response,
7356 error_code: &str,
7357 request_id: &str,
7358) -> Response {
7359 let status = response.status().as_u16();
7360 if let Some(receipt) = receipt.as_mut()
7361 && let Err(err) = receipt.reject(status, error_code)
7362 {
7363 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
7364 return with_request_id(request_id, request_ledger_error_response());
7365 }
7366 with_request_id(request_id, response)
7367}
7368
7369fn ledger_unbilled(
7374 mut receipt: Option<Box<dyn metering::Receipt>>,
7375 response: Response,
7376 outcome: &'static str,
7377 error_code: &str,
7378 request_id: &str,
7379) -> Response {
7380 let status = response.status().as_u16();
7381 if let Some(receipt) = receipt.as_mut()
7382 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
7383 {
7384 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
7385 return with_request_id(request_id, request_ledger_error_response());
7386 }
7387 with_request_id(request_id, response)
7388}
7389
7390fn engine_error_code(class: worker::ErrClass) -> &'static str {
7391 use worker::ErrClass as C;
7392 match class {
7393 C::InvalidRequest => "invalid_request",
7394 C::ContextLength => "context_length_exceeded",
7395 C::ModelNotFound => "model_not_found",
7396 C::RateLimit => "rate_limit_exceeded",
7397 C::Overloaded => "overloaded",
7398 C::Engine => "engine_error",
7399 }
7400}
7401
7402fn model_not_found_response(models: &[String], requested: &str) -> Response {
7422 error_response_coded(
7423 StatusCode::BAD_REQUEST,
7424 &format!("unknown model {requested:?}; loaded: {models:?}"),
7425 "invalid_request_error",
7426 Some("model"),
7427 Some("model_not_found"),
7428 )
7429}
7430
7431fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
7440 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
7441 return Ok(());
7442 };
7443 if let Some((pos, &id)) = ids
7444 .iter()
7445 .enumerate()
7446 .find(|&(_, &id)| id as usize >= n_vocab)
7447 {
7448 return Err(format!(
7449 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
7450 ));
7451 }
7452 Ok(())
7453}
7454
7455#[cfg(test)]
7456mod prompt_ids_tests {
7457 use super::*;
7458
7459 #[test]
7460 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
7461 let caps = ModelCaps {
7462 n_vocab: 8,
7463 ..Default::default()
7464 };
7465 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
7467 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
7468 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
7470 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
7471 assert!(err.contains("vocab size 8"), "{err}");
7472 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
7473 assert!(err.contains("4294967295"), "{err}");
7474 let unknown = ModelCaps::default();
7476 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
7477 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
7478 }
7479}
7480
7481fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
7482 if models.iter().any(|m| m == requested) {
7483 return Some(requested.to_string());
7484 }
7485 if requested.is_empty() || requested.contains('/') {
7486 return None;
7487 }
7488 let mut matches = models.iter().filter(|m| {
7489 m.rsplit('/')
7490 .next()
7491 .is_some_and(|suffix| suffix == requested)
7492 });
7493 match (matches.next(), matches.next()) {
7494 (Some(only), None) => Some(only.clone()),
7495 _ => None,
7496 }
7497}
7498
7499async fn completions(
7500 State(st): State<AppState>,
7501 headers: axum::http::HeaderMap,
7502 trace: Option<Extension<TtftRequestTrace>>,
7503 Json(mut req): Json<CompletionReq>,
7504) -> Response {
7505 let env = Envelope::new(false);
7506 match canonical_model_id(&st.models, &req.model) {
7507 Some(canonical) => req.model = canonical,
7508 None => {
7509 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7510 }
7511 }
7512 let ttft = trace.and_then(|Extension(trace)| trace.0);
7515 if let Some(trace) = ttft.as_ref() {
7516 trace.mark_parsed();
7517 trace.bind_request(&env.id, &req.model);
7518 }
7519 let tenant = match authenticate(&st.api_auth, &headers) {
7520 Ok(t) => t,
7521 Err(resp) => return with_request_id(&env.id, resp),
7522 };
7523 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7524 Ok(ns) => ns,
7525 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7526 };
7527 if let Err((msg, param)) = reject_unsupported(&[
7529 (
7530 "logit_bias",
7531 req.logit_bias.is_some(),
7532 " (device-side sampling has no bias hook yet)",
7533 ),
7534 ("logprobs", req.logprobs.is_some(), ""),
7535 (
7536 "n",
7537 req.n.is_some_and(|n| n != 1),
7538 " for n != 1 (single choice only)",
7539 ),
7540 (
7541 "best_of",
7542 req.best_of.is_some_and(|n| n != 1),
7543 " (single choice only)",
7544 ),
7545 ]) {
7546 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7547 }
7548 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
7551 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
7552 }
7553 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7556 Ok(ms) => RequestDeadline::starting_now(ms),
7557 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7558 };
7559 let lane = match lane_for_tenant(&headers, &tenant) {
7560 Ok(l) => l,
7561 Err(resp) => return resp,
7562 };
7563 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7564 let model = req.model.clone();
7565 let stream = req.stream;
7566 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7567 let mut request = build_request_with_trace(
7568 &req,
7569 tx,
7570 lane,
7571 affinity,
7572 ttft.clone(),
7573 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
7577 );
7578 request.cache_ns = cache_ns;
7579 request.request_id = env.id.clone();
7580 if let Err((message, param)) = apply_model_request_limits(
7581 &mut request,
7582 st.openrouter_metadata.get(&model),
7583 st.caps.get(&model),
7584 ) {
7585 return with_request_id(&env.id, bad_request(&message, Some(param)));
7586 }
7587 if let Err(msg) = nonstream_deadline_gate(
7592 &request,
7593 req.stream,
7594 deadline,
7595 req.max_tokens.is_some(),
7596 st.budget_tokenizers
7597 .as_ref()
7598 .and_then(|t| t.get(&req.model))
7599 .map(Arc::as_ref),
7600 ) {
7601 return with_request_id(
7602 &env.id,
7603 error_response_coded(
7604 StatusCode::BAD_REQUEST,
7605 &msg,
7606 "invalid_request_error",
7607 Some("max_tokens"),
7608 Some("nonstream_deadline_infeasible"),
7609 ),
7610 );
7611 }
7612 if draining() {
7615 let receipt = start_request_receipt(
7616 &st,
7617 &env,
7618 &tenant,
7619 &req.model,
7620 "/v1/completions",
7621 lane,
7622 req.stream,
7623 effective_max_tokens(&request),
7624 None,
7625 None,
7626 );
7627 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7628 }
7629 let budget = match admit_tenant_budget(&st, &tenant, &mut request) {
7630 Ok(budget) => budget,
7631 Err(rejection) => {
7632 let (response, error_code) = rejection.into_response();
7633 let receipt = start_request_receipt(
7634 &st,
7635 &env,
7636 &tenant,
7637 &req.model,
7638 "/v1/completions",
7639 lane,
7640 req.stream,
7641 effective_max_tokens(&request),
7642 None,
7643 None,
7644 );
7645 return ledger_rejected(receipt, response, error_code, &env.id);
7646 }
7647 };
7648 let receipt = start_request_receipt(
7649 &st,
7650 &env,
7651 &tenant,
7652 &req.model,
7653 "/v1/completions",
7654 lane,
7655 req.stream,
7656 effective_max_tokens(&request),
7657 budget.reserved_ctx,
7658 budget.permit,
7659 );
7660 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
7661 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7664 Ok(slot) => slot,
7665 Err(resp) => {
7666 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7667 }
7668 };
7669 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7672 Ok(guard) => guard,
7673 Err((resp, outcome)) => {
7674 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7675 }
7676 };
7677 meter_admit(&env, &tenant, &model, lane);
7678 let stop_strings = request.stop_strings.clone();
7679
7680 if let Some(trace) = ttft.as_ref() {
7685 trace.mark_submitted();
7686 }
7687 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
7688 drop(pending_admit);
7689 return ledger_rejected(
7690 receipt,
7691 rl.attach(worker_unavailable_response()),
7692 "worker_unavailable",
7693 &env.id,
7694 );
7695 }
7696 pending_admit.commit();
7697 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7701 Ok(Ok(rx)) => rx,
7702 Ok(Err((resp, error_code))) => {
7703 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7704 }
7705 Err(_) => {
7706 return ledger_unbilled(
7707 receipt,
7708 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7709 "deadline_exceeded",
7710 "deadline_exceeded",
7711 &env.id,
7712 );
7713 }
7714 };
7715
7716 let resp = if stream {
7717 let rx = match peek_first_token(rx, deadline).await {
7721 Ok(rx) => rx,
7722 Err(()) => {
7723 return ledger_unbilled(
7724 receipt,
7725 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7726 "deadline_exceeded",
7727 "deadline_exceeded",
7728 &env.id,
7729 );
7730 }
7731 };
7732 sse_response_with_receipt(
7733 rx,
7734 model,
7735 false,
7736 None,
7737 env.clone(),
7738 stop_strings,
7739 Some(guard),
7740 receipt,
7741 )
7742 .into_response()
7743 } else {
7744 let mut receipt = receipt;
7750 let resp = blocking_response_with_receipt(
7751 rx,
7752 model,
7753 false,
7754 stop_strings,
7755 None,
7756 env.clone(),
7757 &mut receipt,
7758 Some(deadline),
7759 )
7760 .await;
7761 drop(guard); resp.into_response()
7763 };
7764 rl.attach(with_request_id(&env.id, resp))
7765}
7766
7767async fn chat_completions(
7768 State(st): State<AppState>,
7769 headers: axum::http::HeaderMap,
7770 trace: Option<Extension<TtftRequestTrace>>,
7771 Json(mut req): Json<ChatCompletionReq>,
7772) -> Response {
7773 let env = Envelope::new(true);
7774 match canonical_model_id(&st.models, &req.model) {
7779 Some(canonical) => req.model = canonical,
7780 None => {
7781 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7782 }
7783 }
7784 let ttft = trace.and_then(|Extension(trace)| trace.0);
7785 if let Some(trace) = ttft.as_ref() {
7786 trace.mark_parsed();
7787 trace.bind_request(&env.id, &req.model);
7788 }
7789 let tenant = match authenticate(&st.api_auth, &headers) {
7790 Ok(t) => t,
7791 Err(resp) => return with_request_id(&env.id, resp),
7792 };
7793 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7794 Ok(ns) => ns,
7795 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7796 };
7797 if req.messages.is_empty()
7798 || req.messages.iter().any(|message| {
7799 !matches!(
7800 message.role.as_str(),
7801 "system" | "developer" | "user" | "assistant" | "tool"
7802 )
7803 })
7804 {
7805 return with_request_id(
7806 &env.id,
7807 bad_request(
7808 "messages must use system/developer/user/assistant/tool roles",
7809 Some("messages"),
7810 ),
7811 );
7812 }
7813 if let Err((msg, param)) = reject_unsupported(&[
7818 (
7819 "logit_bias",
7820 req.logit_bias.is_some(),
7821 " (device-side sampling has no bias hook yet)",
7822 ),
7823 (
7824 "logprobs",
7825 req.logprobs
7826 .as_ref()
7827 .is_some_and(|v| v.as_bool() != Some(false)),
7828 "",
7829 ),
7830 ("top_logprobs", req.top_logprobs.is_some(), ""),
7831 (
7832 "n",
7833 req.n.is_some_and(|n| n != 1),
7834 " for n != 1 (single choice only)",
7835 ),
7836 ]) {
7837 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7838 }
7839 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7842 Ok(ms) => RequestDeadline::starting_now(ms),
7843 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7844 };
7845 let lane = match lane_for_tenant(&headers, &tenant) {
7846 Ok(l) => l,
7847 Err(resp) => return resp,
7848 };
7849 let model = req.model.clone();
7850 let stream = req.stream;
7851 let capture_prompt = st
7854 .metering
7855 .as_ref()
7856 .filter(|m| m.captures(&tenant.tenant))
7857 .map(|_| capture_chat_messages(&req.messages));
7858 let declared_max_tokens = req.max_tokens.is_some();
7862 let vision_preprocess_permit = if request_has_vision(&req) {
7866 match VISION_PREPROCESS_SEMAPHORE.acquire().await {
7867 Ok(permit) => Some(permit),
7868 Err(_) => {
7869 return with_request_id(
7870 &env.id,
7871 error_response(
7872 StatusCode::SERVICE_UNAVAILABLE,
7873 "vision preprocessing is unavailable",
7874 "server_error",
7875 None,
7876 ),
7877 );
7878 }
7879 }
7880 } else {
7881 None
7882 };
7883 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7884 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7885 let mut plan = match build_chat_request_with_trace(
7886 req,
7887 st.caps.get(&model),
7888 tx,
7889 lane,
7890 affinity,
7891 ttft.clone(),
7892 st.openrouter_metadata
7893 .get(&model)
7894 .and_then(|m| m.default_reasoning_effort.as_deref()),
7895 &st.sampling_defaults(&model),
7896 ) {
7897 Ok(plan) => plan,
7898 Err(err) => {
7899 return with_request_id(&env.id, bad_request(&err, None));
7900 }
7901 };
7902 plan.request.cache_ns = cache_ns;
7903 plan.request.request_id = env.id.clone();
7904 if let Err((message, param)) = apply_model_request_limits(
7905 &mut plan.request,
7906 st.openrouter_metadata.get(&model),
7907 st.caps.get(&model),
7908 ) {
7909 return with_request_id(&env.id, bad_request(&message, Some(param)));
7910 }
7911 if let Err(msg) = nonstream_deadline_gate(
7914 &plan.request,
7915 stream,
7916 deadline,
7917 declared_max_tokens,
7918 st.budget_tokenizers
7919 .as_ref()
7920 .and_then(|t| t.get(&model))
7921 .map(Arc::as_ref),
7922 ) {
7923 return with_request_id(
7924 &env.id,
7925 error_response_coded(
7926 StatusCode::BAD_REQUEST,
7927 &msg,
7928 "invalid_request_error",
7929 Some("max_tokens"),
7930 Some("nonstream_deadline_infeasible"),
7931 ),
7932 );
7933 }
7934 plan.vision_memory = match reserve_vision_memory(&plan) {
7935 Ok(permit) => permit,
7936 Err(err) => {
7937 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
7938 }
7939 };
7940 if draining() {
7943 let receipt = start_request_receipt(
7944 &st,
7945 &env,
7946 &tenant,
7947 &model,
7948 "/v1/chat/completions",
7949 lane,
7950 stream,
7951 effective_max_tokens(&plan.request),
7952 None,
7953 None,
7954 );
7955 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7956 }
7957 let budget = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
7958 Ok(budget) => budget,
7959 Err(rejection) => {
7960 let (response, error_code) = rejection.into_response();
7961 let receipt = start_request_receipt(
7962 &st,
7963 &env,
7964 &tenant,
7965 &model,
7966 "/v1/chat/completions",
7967 lane,
7968 stream,
7969 effective_max_tokens(&plan.request),
7970 None,
7971 None,
7972 );
7973 return ledger_rejected(receipt, response, error_code, &env.id);
7974 }
7975 };
7976 let receipt = start_request_receipt(
7977 &st,
7978 &env,
7979 &tenant,
7980 &model,
7981 "/v1/chat/completions",
7982 lane,
7983 stream,
7984 effective_max_tokens(&plan.request),
7985 budget.reserved_ctx,
7986 budget.permit,
7987 );
7988 let receipt = if let Some(prompt) = capture_prompt {
7989 arm_capture(receipt, move || prompt)
7990 } else {
7991 receipt
7992 };
7993 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7997 Ok(slot) => slot,
7998 Err(resp) => {
7999 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
8000 }
8001 };
8002 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
8005 Ok(guard) => guard,
8006 Err((resp, outcome)) => {
8007 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
8008 }
8009 };
8010 if let Err(err) = decode_pending_vision(&mut plan) {
8015 return ledger_rejected(
8016 receipt,
8017 rl.attach(bad_request(&err, Some("messages"))),
8018 "invalid_request_error",
8019 &env.id,
8020 );
8021 }
8022 plan.request.vision_memory = plan.vision_memory.take();
8023 drop(vision_preprocess_permit);
8024 let constraint_ready = if plan.request.grammar.is_some() {
8025 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
8026 plan.request.constraint_ready = Some(ready_tx);
8027 Some(ready_rx)
8028 } else {
8029 None
8030 };
8031 meter_admit(&env, &tenant, &model, lane);
8032 let stop_strings = plan.request.stop_strings.clone();
8033 if let Some(trace) = ttft.as_ref() {
8035 trace.mark_submitted();
8036 }
8037 if st
8038 .cmd_tx
8039 .send(Cmd::Generate(Box::new(plan.request)))
8040 .is_err()
8041 {
8042 drop(pending_admit);
8043 return ledger_rejected(
8044 receipt,
8045 rl.attach(worker_unavailable_response()),
8046 "worker_unavailable",
8047 &env.id,
8048 );
8049 }
8050 pending_admit.commit();
8051 if let Some(ready) = constraint_ready {
8057 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
8058 match tokio::time::timeout(bound, ready).await {
8059 Ok(Ok(Ok(()))) => {}
8060 Ok(Ok(Err(err))) => {
8061 return ledger_rejected(
8062 receipt,
8063 rl.attach(engine_error_response(&err)),
8064 engine_error_code(err.class),
8065 &env.id,
8066 );
8067 }
8068 Ok(Err(_)) => {
8069 return ledger_rejected(
8070 receipt,
8071 rl.attach(worker_unavailable_response()),
8072 "worker_unavailable",
8073 &env.id,
8074 );
8075 }
8076 Err(_) if deadline.remaining().is_zero() => {
8077 return ledger_unbilled(
8078 receipt,
8079 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8080 "deadline_exceeded",
8081 "deadline_exceeded",
8082 &env.id,
8083 );
8084 }
8085 Err(_) => {
8086 return ledger_rejected(
8087 receipt,
8088 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
8089 "constraint_compile_timeout",
8090 &env.id,
8091 );
8092 }
8093 }
8094 }
8095 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
8097 Ok(Ok(rx)) => rx,
8098 Ok(Err((resp, error_code))) => {
8099 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
8100 }
8101 Err(_) => {
8102 return ledger_unbilled(
8103 receipt,
8104 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
8105 "deadline_exceeded",
8106 "deadline_exceeded",
8107 &env.id,
8108 );
8109 }
8110 };
8111 let resp = if stream {
8112 let rx = match peek_first_token(rx, deadline).await {
8114 Ok(rx) => rx,
8115 Err(()) => {
8116 return ledger_unbilled(
8117 receipt,
8118 rl.attach(deadline_exceeded_response(deadline.ms, true)),
8119 "deadline_exceeded",
8120 "deadline_exceeded",
8121 &env.id,
8122 );
8123 }
8124 };
8125 sse_response_with_receipt(
8126 rx,
8127 model,
8128 true,
8129 plan.parser,
8130 env.clone(),
8131 stop_strings,
8132 Some(guard),
8133 receipt,
8134 )
8135 .into_response()
8136 } else {
8137 let mut receipt = receipt;
8140 let resp = blocking_response_with_receipt(
8141 rx,
8142 model,
8143 true,
8144 stop_strings,
8145 plan.parser,
8146 env.clone(),
8147 &mut receipt,
8148 Some(deadline),
8149 )
8150 .await;
8151 drop(guard); resp.into_response()
8153 };
8154 rl.attach(with_request_id(&env.id, resp))
8155}
8156
8157#[cfg(test)]
8166fn sse_response(
8167 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8168 model: String,
8169 chat: bool,
8170 parser: Option<ToolStreamParser>,
8171 env: Envelope,
8172 stop_strings: Vec<String>,
8173 guard: Option<InflightGuard>,
8174) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8175 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
8176}
8177
8178fn sse_response_with_receipt(
8179 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8180 model: String,
8181 chat: bool,
8182 mut parser: Option<ToolStreamParser>,
8183 env: Envelope,
8184 stop_strings: Vec<String>,
8185 guard: Option<InflightGuard>,
8186 mut receipt: Option<Box<dyn metering::Receipt>>,
8187) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8188 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
8192 .then(|| StopScrubber::new(stop_strings));
8193 let stream = async_stream::stream! {
8194 let _guard = guard;
8197 let mut call_index: usize = 0;
8198 let mut role_sent = false;
8201 macro_rules! chat_chunk {
8202 ($delta:expr, $finish:expr) => {{
8203 let mut delta = $delta;
8204 if chat && !role_sent {
8205 role_sent = true;
8206 delta["role"] = json!("assistant");
8207 }
8208 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
8209 "choices": [{ "index": 0, "delta": delta,
8210 "finish_reason": $finish }] }))
8211 .to_string()
8212 }};
8213 }
8214 macro_rules! piece_chunks {
8216 ($piece:expr) => {{
8217 let mut payloads: Vec<String> = Vec::new();
8218 match $piece {
8219 Piece::Content(text) => {
8220 let text = match scrub.as_mut() {
8221 Some(sc) => sc.push(&text),
8222 None => text,
8223 };
8224 if !text.is_empty() {
8225 payloads.push(chat_chunk!(json!({ "content": text }),
8226 serde_json::Value::Null));
8227 }
8228 }
8229 Piece::Reasoning(text) => payloads.push(
8233 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
8234 Piece::Call(call) => {
8235 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8236 "index": call_index, "id": call.id, "type": "function",
8237 "function": { "name": call.name, "arguments": "" } }] }),
8238 serde_json::Value::Null));
8239 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8240 "index": call_index,
8241 "function": { "arguments": call.arguments } }] }),
8242 serde_json::Value::Null));
8243 call_index += 1;
8244 }
8245 }
8246 payloads
8247 }};
8248 }
8249 let mut terminal = false;
8253 while let Some(ev) = rx.recv().await {
8254 match ev {
8255 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8257 if let Some(receipt) = receipt.as_mut()
8258 && let Err(err) = receipt.record_prompt_usage(
8259 n_prompt as u64,
8260 n_cached as u64,
8261 )
8262 {
8263 eprintln!(
8264 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8265 env.id
8266 );
8267 let _ = receipt.reject(500, "request_ledger_unavailable");
8270 let payload = request_ledger_error_body().to_string();
8271 if chat || openai_compat() {
8272 yield Ok(SseEvent::default().data(payload));
8273 yield Ok(SseEvent::default().data("[DONE]"));
8274 } else {
8275 yield Ok(SseEvent::default().event("error").data(payload));
8276 }
8277 terminal = true;
8278 break;
8279 }
8280 }
8281 Event::Token { id, text } => {
8282 if let Some(receipt) = receipt.as_mut()
8283 && let Err(err) = receipt.record_completion_token()
8284 {
8285 eprintln!(
8286 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8287 env.id
8288 );
8289 let _ = receipt.reject(500, "request_ledger_unavailable");
8290 let payload = request_ledger_error_body().to_string();
8291 if chat || openai_compat() {
8292 yield Ok(SseEvent::default().data(payload));
8293 yield Ok(SseEvent::default().data("[DONE]"));
8294 } else {
8295 yield Ok(SseEvent::default().event("error").data(payload));
8296 }
8297 terminal = true;
8298 break;
8299 }
8300 if let Some(receipt) = receipt.as_mut() {
8303 receipt.capture_completion_delta(&text);
8304 }
8305 if let Some(p) = parser.as_mut() {
8306 for piece in p.push(&text) {
8307 for payload in piece_chunks!(piece) {
8308 yield Ok(SseEvent::default().data(payload));
8309 }
8310 }
8311 continue;
8312 }
8313 let text = match scrub.as_mut() {
8314 Some(sc) => sc.push(&text),
8315 None => text,
8316 };
8317 if text.is_empty() && scrub.is_some() {
8318 continue; }
8320 let payload = if chat {
8321 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
8322 } else if openai_compat() {
8323 env.stamp(json!({ "object": "text_completion", "model": model,
8324 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
8325 .to_string()
8326 } else {
8327 json!({ "model": model, "id": id, "text": text }).to_string()
8328 };
8329 yield Ok(SseEvent::default().data(payload));
8330 }
8331 Event::TokenSnapshot(_) => {}
8335 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
8336 let mut finish = stop_reason_to_finish(&stop_reason);
8337 if let Some(p) = parser.as_mut() {
8338 for piece in p.finish() {
8339 for payload in piece_chunks!(piece) {
8340 yield Ok(SseEvent::default().data(payload));
8341 }
8342 }
8343 if p.n_calls() > 0 { finish = "tool_calls"; }
8344 }
8345 if let Some(sc) = scrub.as_mut() {
8347 let tail = sc.finish();
8348 if !tail.is_empty() {
8349 let payload = if chat {
8350 chat_chunk!(json!({ "content": tail }),
8351 serde_json::Value::Null)
8352 } else {
8353 env.stamp(json!({ "object": "text_completion",
8354 "model": model,
8355 "choices": [{ "index": 0, "text": tail,
8356 "finish_reason": null }] })).to_string()
8357 };
8358 yield Ok(SseEvent::default().data(payload));
8359 }
8360 }
8361 if let Some(receipt) = receipt.as_mut()
8362 && let Err(err) = receipt.complete(
8363 metering::UsageCounts {
8364 prompt_tokens: n_prompt as u64,
8365 cached_prompt_tokens: n_cached as u64,
8366 completion_tokens: n_tokens as u64,
8367 },
8368 elapsed_s,
8369 )
8370 {
8371 eprintln!(
8372 "[ledger] ERROR: request {} completion receipt failed: {err}",
8373 env.id
8374 );
8375 let _ = receipt.reject(500, "request_ledger_unavailable");
8379 let payload = request_ledger_error_body().to_string();
8380 if chat || openai_compat() {
8381 yield Ok(SseEvent::default().data(payload));
8382 yield Ok(SseEvent::default().data("[DONE]"));
8383 } else {
8384 yield Ok(SseEvent::default().event("error").data(payload));
8385 }
8386 terminal = true;
8387 break;
8388 }
8389 if chat || openai_compat() {
8390 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
8391 let fin = if chat {
8392 let mut v = env.stamp(json!({
8393 "object": "chat.completion.chunk", "model": model,
8394 "choices": [{ "index": 0, "delta": {},
8395 "finish_reason": finish }],
8396 "usage": usage }));
8397 if !role_sent {
8399 v["choices"][0]["delta"]["role"] = json!("assistant");
8400 }
8401 v
8402 } else {
8403 env.stamp(json!({ "object": "text_completion", "model": model,
8404 "choices": [{ "index": 0, "text": "",
8405 "finish_reason": finish }],
8406 "usage": usage }))
8407 }.to_string();
8408 yield Ok(SseEvent::default().data(fin));
8409 yield Ok(SseEvent::default().data("[DONE]"));
8410 } else {
8411 let payload = json!({
8412 "stop_reason": stop_reason, "n_tokens": n_tokens,
8413 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
8414 "elapsed_s": elapsed_s
8415 }).to_string();
8416 yield Ok(SseEvent::default().event("done").data(payload));
8417 }
8418 terminal = true;
8419 break;
8420 }
8421 Event::Error(err) => {
8422 let ledger_error = if let Some(receipt) = receipt.as_mut() {
8432 receipt
8433 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
8434 .err()
8435 } else {
8436 None
8437 };
8438 if let Some(ref ledger_error) = ledger_error {
8439 eprintln!(
8440 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
8441 env.id
8442 );
8443 }
8444 let payload = if ledger_error.is_some() {
8445 request_ledger_error_body().to_string()
8446 } else {
8447 engine_error_body(&err).to_string()
8448 };
8449 if chat || openai_compat() {
8450 yield Ok(SseEvent::default().data(payload));
8453 yield Ok(SseEvent::default().data("[DONE]"));
8454 } else {
8455 yield Ok(SseEvent::default().event("error").data(payload));
8458 }
8459 terminal = true;
8460 break;
8461 }
8462 }
8463 }
8464 if !terminal {
8465 let e = worker::EngineError::overloaded(
8471 "worker closed the stream without completing (worker restart in progress)",
8472 );
8473 if let Some(receipt) = receipt.as_mut()
8474 && let Err(ledger_err) = receipt.reject(
8475 class_http(e.class).0.as_u16(),
8476 engine_error_code(e.class),
8477 )
8478 {
8479 eprintln!(
8480 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8481 env.id
8482 );
8483 }
8484 let payload = engine_error_body(&e).to_string();
8485 if chat || openai_compat() {
8486 yield Ok(SseEvent::default().data(payload));
8487 yield Ok(SseEvent::default().data("[DONE]"));
8488 } else {
8489 yield Ok(SseEvent::default().event("error").data(payload));
8490 }
8491 }
8492 };
8493 Sse::new(stream).keep_alive(
8494 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
8497 )
8498}
8499
8500fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
8502 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
8503 text.truncate(offset);
8504 }
8505}
8506
8507fn partial_stop_suffix(s: &str, tag: &str) -> usize {
8510 let mut best = 0;
8511 for (k, _) in tag.char_indices().skip(1) {
8512 if k <= s.len() && s.ends_with(&tag[..k]) {
8513 best = k;
8514 }
8515 }
8516 best
8517}
8518
8519struct StopScrubber {
8525 stops: Vec<String>,
8526 buf: String,
8527 done: bool,
8528}
8529
8530impl StopScrubber {
8531 fn new(stops: Vec<String>) -> Self {
8532 Self {
8533 stops,
8534 buf: String::new(),
8535 done: false,
8536 }
8537 }
8538
8539 fn push(&mut self, text: &str) -> String {
8541 if self.done {
8542 return String::new();
8543 }
8544 self.buf.push_str(text);
8545 if let Some(i) = self
8546 .stops
8547 .iter()
8548 .filter_map(|s| self.buf.find(s.as_str()))
8549 .min()
8550 {
8551 self.done = true;
8552 let out = self.buf[..i].to_string();
8553 self.buf.clear();
8554 return out;
8555 }
8556 let keep = self
8557 .stops
8558 .iter()
8559 .map(|s| partial_stop_suffix(&self.buf, s))
8560 .max()
8561 .unwrap_or(0);
8562 let emit_to = self.buf.len() - keep;
8563 let out = self.buf[..emit_to].to_string();
8564 self.buf.drain(..emit_to);
8565 out
8566 }
8567
8568 fn finish(&mut self) -> String {
8570 if self.done {
8571 self.buf.clear();
8572 return String::new();
8573 }
8574 std::mem::take(&mut self.buf)
8575 }
8576}
8577
8578#[cfg(test)]
8579async fn blocking_response(
8580 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8581 model: String,
8582 chat: bool,
8583 stop_strings: Vec<String>,
8584 parser: Option<ToolStreamParser>,
8585 env: Envelope,
8586) -> Response {
8587 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
8588 .await
8589}
8590
8591struct BlockingPayload<'a> {
8595 env: &'a Envelope,
8596 model: String,
8597 chat: bool,
8598 finish: &'static str,
8599 text: String,
8600 reasoning: String,
8601 calls: Vec<ParsedToolCall>,
8602 tokens: Vec<u32>,
8603 stop_reason: String,
8604 n_prompt: usize,
8605 n_tokens: usize,
8606 n_cached: usize,
8607 elapsed_s: f64,
8608 spec: Option<worker::SpecUsage>,
8609 deadline_error: Option<serde_json::Value>,
8615}
8616
8617fn blocking_payload(p: BlockingPayload<'_>) -> Response {
8618 let BlockingPayload {
8619 env,
8620 model,
8621 chat,
8622 finish,
8623 text,
8624 reasoning,
8625 calls,
8626 tokens,
8627 stop_reason,
8628 n_prompt,
8629 n_tokens,
8630 n_cached,
8631 elapsed_s,
8632 spec,
8633 deadline_error,
8634 } = p;
8635 if chat {
8636 let content = if !calls.is_empty() && text.is_empty() {
8638 serde_json::Value::Null
8639 } else {
8640 serde_json::Value::String(text)
8641 };
8642 let mut message = json!({ "role": "assistant", "content": content });
8643 if !reasoning.is_empty() {
8646 message["reasoning"] = json!(reasoning);
8647 message["reasoning_details"] = json!([{
8648 "type": "reasoning.text", "text": reasoning }]);
8649 }
8650 if !calls.is_empty() {
8651 message["tool_calls"] =
8652 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
8653 }
8654 let mut body = json!({
8655 "object": "chat.completion", "model": model,
8656 "choices": [{ "index": 0,
8657 "message": message,
8658 "finish_reason": finish }],
8659 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8660 });
8661 if let Some(err) = deadline_error {
8662 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8663 body["error"] = err;
8664 }
8665 return Json(env.stamp(body)).into_response();
8666 }
8667 if openai_compat() {
8668 let mut body = json!({
8669 "object": "text_completion", "model": model,
8670 "choices": [{ "index": 0, "text": text,
8671 "finish_reason": finish }],
8672 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8673 });
8674 if let Some(err) = deadline_error {
8675 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8676 body["error"] = err;
8677 }
8678 return Json(env.stamp(body)).into_response();
8679 }
8680 Json(CompletionResp {
8681 model,
8682 text,
8683 tokens,
8684 stop_reason,
8685 error: deadline_error,
8686 n_tokens,
8687 prompt_tokens: n_prompt,
8688 cached_tokens: n_cached,
8689 elapsed_s,
8690 })
8691 .into_response()
8692}
8693
8694async fn blocking_response_with_receipt(
8719 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8720 model: String,
8721 chat: bool,
8722 stop_strings: Vec<String>,
8723 mut parser: Option<ToolStreamParser>,
8724 env: Envelope,
8725 receipt: &mut Option<Box<dyn metering::Receipt>>,
8726 deadline: Option<RequestDeadline>,
8727) -> Response {
8728 let mut text = String::new();
8729 let mut reasoning = String::new();
8730 let mut tokens: Vec<u32> = Vec::new();
8731 let mut calls: Vec<ParsedToolCall> = Vec::new();
8732 let consume = |pieces: Vec<Piece>,
8733 text: &mut String,
8734 reasoning: &mut String,
8735 calls: &mut Vec<ParsedToolCall>| {
8736 for piece in pieces {
8737 match piece {
8738 Piece::Content(t) => text.push_str(&t),
8739 Piece::Reasoning(t) => reasoning.push_str(&t),
8740 Piece::Call(c) => calls.push(c),
8741 }
8742 }
8743 };
8744 let started = std::time::Instant::now();
8746 let mut seen_prompt: usize = 0;
8747 let mut seen_cached: usize = 0;
8748 let mut seen_tokens: usize = 0;
8749 loop {
8750 let ev = match deadline {
8751 Some(d) => tokio::select! {
8752 biased;
8753 ev = rx.recv() => ev,
8754 () = tokio::time::sleep_until(d.at) => {
8755 drop(rx);
8758 if seen_tokens == 0 {
8759 if let Some(receipt) = receipt.as_mut()
8763 && let Err(err) = receipt.settle_unbilled(
8764 "deadline_exceeded",
8765 StatusCode::REQUEST_TIMEOUT.as_u16(),
8766 "deadline_exceeded",
8767 )
8768 {
8769 eprintln!(
8770 "[ledger] ERROR: request {} deadline receipt failed: {err}",
8771 env.id
8772 );
8773 return request_ledger_error_response();
8774 }
8775 return deadline_exceeded_response(d.ms, false);
8776 }
8777 if let Some(p) = parser.as_mut() {
8778 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8779 }
8780 truncate_at_stop(&mut text, &stop_strings);
8781 let elapsed_s = started.elapsed().as_secs_f64();
8782 if let Some(receipt) = receipt.as_mut()
8785 && let Err(err) = receipt.complete_deadline_partial(
8786 metering::UsageCounts {
8787 prompt_tokens: seen_prompt as u64,
8788 cached_prompt_tokens: seen_cached as u64,
8789 completion_tokens: seen_tokens as u64,
8790 },
8791 elapsed_s,
8792 )
8793 {
8794 eprintln!(
8795 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
8796 env.id
8797 );
8798 let _ = receipt.reject(500, "request_ledger_unavailable");
8799 return request_ledger_error_response();
8800 }
8801 eprintln!(
8802 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
8803 {} ms deadline (prompt {}); non-streaming caller advised to stream",
8804 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
8805 );
8806 let err_obj = json!({
8807 "message": format!(
8808 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
8809 the {} tokens produced before the cut are delivered above and are \
8810 billed. Set \"stream\": true for work this long — a stream's \
8811 deadline bounds only the time to first token — or lower max_tokens.",
8812 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
8813 ),
8814 "code": "deadline_exceeded",
8815 "metadata": { "error_type": "timeout", "provider_name": "memra" }
8816 });
8817 return blocking_payload(BlockingPayload {
8818 env: &env,
8819 model,
8820 chat,
8821 finish: "error",
8822 text,
8823 reasoning,
8824 calls,
8825 tokens,
8826 stop_reason: "Deadline".to_string(),
8827 n_prompt: seen_prompt,
8828 n_tokens: seen_tokens,
8829 n_cached: seen_cached,
8830 elapsed_s,
8831 spec: None,
8832 deadline_error: Some(err_obj),
8833 });
8834 }
8835 },
8836 None => rx.recv().await,
8837 };
8838 let Some(ev) = ev else { break };
8839 match ev {
8840 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8842 if let Some(receipt) = receipt.as_mut()
8843 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
8844 {
8845 eprintln!(
8846 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8847 env.id
8848 );
8849 let _ = receipt.reject(500, "request_ledger_unavailable");
8852 return request_ledger_error_response();
8853 }
8854 seen_prompt = n_prompt;
8855 seen_cached = n_cached;
8856 }
8857 Event::Token { id, text: delta } => {
8858 if let Some(receipt) = receipt.as_mut()
8859 && let Err(err) = receipt.record_completion_token()
8860 {
8861 eprintln!(
8862 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8863 env.id
8864 );
8865 let _ = receipt.reject(500, "request_ledger_unavailable");
8866 return request_ledger_error_response();
8867 }
8868 if let Some(receipt) = receipt.as_mut() {
8870 receipt.capture_completion_delta(&delta);
8871 }
8872 tokens.push(id);
8873 seen_tokens += 1;
8874 match parser.as_mut() {
8875 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
8876 None => text.push_str(&delta),
8877 }
8878 }
8879 Event::TokenSnapshot(ids) => tokens = ids,
8880 Event::Done {
8881 stop_reason,
8882 n_tokens,
8883 n_prompt,
8884 n_cached,
8885 elapsed_s,
8886 spec,
8887 } => {
8888 if let Some(p) = parser.as_mut() {
8889 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8890 }
8891 truncate_at_stop(&mut text, &stop_strings);
8892 let finish = if calls.is_empty() {
8893 stop_reason_to_finish(&stop_reason)
8894 } else {
8895 "tool_calls"
8896 };
8897 if let Some(receipt) = receipt.as_mut()
8898 && let Err(err) = receipt.complete(
8899 metering::UsageCounts {
8900 prompt_tokens: n_prompt as u64,
8901 cached_prompt_tokens: n_cached as u64,
8902 completion_tokens: n_tokens as u64,
8903 },
8904 elapsed_s,
8905 )
8906 {
8907 eprintln!(
8908 "[ledger] ERROR: request {} completion receipt failed: {err}",
8909 env.id
8910 );
8911 let _ = receipt.reject(500, "request_ledger_unavailable");
8914 return request_ledger_error_response();
8915 }
8916 return blocking_payload(BlockingPayload {
8917 env: &env,
8918 model,
8919 chat,
8920 finish,
8921 text,
8922 reasoning,
8923 calls,
8924 tokens,
8925 stop_reason,
8926 n_prompt,
8927 n_tokens,
8928 n_cached,
8929 elapsed_s,
8930 spec,
8931 deadline_error: None,
8932 });
8933 }
8934 Event::Error(err) => {
8935 if let Some(receipt) = receipt.as_mut()
8939 && let Err(ledger_err) = receipt.reject(
8940 class_http(err.class).0.as_u16(),
8941 engine_error_code(err.class),
8942 )
8943 {
8944 eprintln!(
8945 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
8946 env.id
8947 );
8948 return request_ledger_error_response();
8949 }
8950 return engine_error_response(&err);
8951 }
8952 }
8953 }
8954 let e = worker::EngineError::overloaded(
8959 "worker closed the stream without completing (worker restart in progress)",
8960 );
8961 if let Some(receipt) = receipt.as_mut()
8962 && let Err(ledger_err) =
8963 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
8964 {
8965 eprintln!(
8966 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8967 env.id
8968 );
8969 return request_ledger_error_response();
8970 }
8971 engine_error_response(&e)
8972}
8973
8974#[cfg(test)]
8975mod tests {
8976 use super::*;
8977
8978 #[derive(Debug, Clone, PartialEq)]
8986 enum MeterEvent {
8987 Reserve {
8988 tenant: String,
8989 principal: Option<String>,
8990 model: String,
8991 },
8992 Open {
8993 request_id: String,
8994 tenant: String,
8995 model: String,
8996 route: &'static str,
8997 stream: bool,
8998 with_permit: bool,
8999 },
9000 PromptUsage {
9001 prompt: u64,
9002 cached: u64,
9003 },
9004 Token,
9005 CapturePrompt(serde_json::Value),
9006 CaptureDelta(String),
9007 Complete {
9008 prompt: u64,
9009 cached: u64,
9010 completion: u64,
9011 },
9012 DeadlinePartial {
9013 prompt: u64,
9014 cached: u64,
9015 completion: u64,
9016 },
9017 Reject {
9018 status: u16,
9019 code: String,
9020 },
9021 Unbilled {
9022 outcome: &'static str,
9023 status: u16,
9024 code: String,
9025 },
9026 Dropped {
9029 prompt: u64,
9030 cached: u64,
9031 completion: u64,
9032 },
9033 }
9034
9035 enum ReserveScript {
9038 Admit { with_permit: bool },
9039 Insufficient,
9040 Blocked,
9041 Unenrolled,
9042 PrincipalCapped,
9043 }
9044
9045 struct MockMetering {
9046 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
9047 limits: bool,
9048 limited: bool,
9049 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
9050 captures: bool,
9051 }
9052
9053 impl MockMetering {
9054 fn admit_all() -> Arc<Self> {
9055 Arc::new(MockMetering {
9056 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9057 limits: false,
9058 limited: true,
9059 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
9060 captures: false,
9061 })
9062 }
9063
9064 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
9065 Arc::new(MockMetering {
9066 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9067 limits: true,
9068 limited: true,
9069 reserve_script: std::sync::Mutex::new(script.into()),
9070 captures: false,
9071 })
9072 }
9073
9074 fn capturing() -> Arc<Self> {
9075 Arc::new(MockMetering {
9076 events: Arc::new(std::sync::Mutex::new(Vec::new())),
9077 limits: false,
9078 limited: true,
9079 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
9080 captures: true,
9081 })
9082 }
9083
9084 fn events(&self) -> Vec<MeterEvent> {
9085 self.events.lock().unwrap().clone()
9086 }
9087 }
9088
9089 impl metering::Metering for MockMetering {
9090 fn enforces_limits(&self) -> bool {
9091 self.limits
9092 }
9093
9094 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
9095 Ok(self.limited)
9096 }
9097
9098 fn reserve(
9099 &self,
9100 tenant: &str,
9101 principal: Option<&str>,
9102 model: &str,
9103 _prompt_tokens: u64,
9104 _completion_bound: u64,
9105 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
9106 self.events.lock().unwrap().push(MeterEvent::Reserve {
9107 tenant: tenant.into(),
9108 principal: principal.map(str::to_owned),
9109 model: model.into(),
9110 });
9111 match self.reserve_script.lock().unwrap().pop_front() {
9112 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
9113 Some(ReserveScript::Admit { with_permit: true }) => {
9114 Ok(Some(Box::new(()) as metering::Permit))
9115 }
9116 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
9117 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
9118 Some(ReserveScript::Unenrolled) => Err(metering::AdmitError::Unenrolled),
9119 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
9120 }
9121 }
9122
9123 fn open(
9124 &self,
9125 meta: &metering::RequestMeta<'_>,
9126 permit: Option<metering::Permit>,
9127 ) -> Box<dyn metering::Receipt> {
9128 self.events.lock().unwrap().push(MeterEvent::Open {
9129 request_id: meta.request_id.into(),
9130 tenant: meta.tenant.into(),
9131 model: meta.model.into(),
9132 route: meta.route,
9133 stream: meta.stream,
9134 with_permit: permit.is_some(),
9135 });
9136 Box::new(MockReceipt {
9137 events: self.events.clone(),
9138 wants_capture: self.captures,
9139 prompt: 0,
9140 cached: 0,
9141 completion: 0,
9142 finalized: false,
9143 })
9144 }
9145
9146 fn captures(&self, _tenant: &str) -> bool {
9147 self.captures
9148 }
9149
9150 fn limits_health(&self) -> Option<metering::LimitsHealth> {
9151 self.limits.then_some(metering::LimitsHealth {
9152 source_reload_failed: 0,
9153 source_reload_consecutive: 0,
9154 source_available: true,
9155 })
9156 }
9157 }
9158
9159 struct MockReceipt {
9160 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
9161 wants_capture: bool,
9162 prompt: u64,
9163 cached: u64,
9164 completion: u64,
9165 finalized: bool,
9166 }
9167
9168 impl metering::Receipt for MockReceipt {
9169 fn wants_capture(&self) -> bool {
9170 self.wants_capture
9171 }
9172
9173 fn arm_capture(&mut self, prompt: serde_json::Value) {
9174 self.events
9175 .lock()
9176 .unwrap()
9177 .push(MeterEvent::CapturePrompt(prompt));
9178 }
9179
9180 fn capture_completion_delta(&mut self, text: &str) {
9181 if self.wants_capture {
9182 self.events
9183 .lock()
9184 .unwrap()
9185 .push(MeterEvent::CaptureDelta(text.into()));
9186 }
9187 }
9188
9189 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
9190 self.prompt = prompt;
9191 self.cached = cached;
9192 self.events
9193 .lock()
9194 .unwrap()
9195 .push(MeterEvent::PromptUsage { prompt, cached });
9196 Ok(())
9197 }
9198
9199 fn record_completion_token(&mut self) -> Result<(), String> {
9200 self.completion += 1;
9201 self.events.lock().unwrap().push(MeterEvent::Token);
9202 Ok(())
9203 }
9204
9205 fn complete(
9206 &mut self,
9207 usage: metering::UsageCounts,
9208 _worker_elapsed_s: f64,
9209 ) -> Result<(), String> {
9210 self.finalized = true;
9211 self.events.lock().unwrap().push(MeterEvent::Complete {
9212 prompt: usage.prompt_tokens,
9213 cached: usage.cached_prompt_tokens,
9214 completion: usage.completion_tokens,
9215 });
9216 Ok(())
9217 }
9218
9219 fn complete_deadline_partial(
9220 &mut self,
9221 usage: metering::UsageCounts,
9222 _worker_elapsed_s: f64,
9223 ) -> Result<(), String> {
9224 self.finalized = true;
9225 self.events
9226 .lock()
9227 .unwrap()
9228 .push(MeterEvent::DeadlinePartial {
9229 prompt: usage.prompt_tokens,
9230 cached: usage.cached_prompt_tokens,
9231 completion: usage.completion_tokens,
9232 });
9233 Ok(())
9234 }
9235
9236 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
9237 self.finalized = true;
9238 self.events.lock().unwrap().push(MeterEvent::Reject {
9239 status,
9240 code: error_code.into(),
9241 });
9242 Ok(())
9243 }
9244
9245 fn settle_unbilled(
9246 &mut self,
9247 outcome: &'static str,
9248 status: u16,
9249 error_code: &str,
9250 ) -> Result<(), String> {
9251 self.finalized = true;
9252 self.events.lock().unwrap().push(MeterEvent::Unbilled {
9253 outcome,
9254 status,
9255 code: error_code.into(),
9256 });
9257 Ok(())
9258 }
9259 }
9260
9261 impl Drop for MockReceipt {
9262 fn drop(&mut self) {
9263 if !self.finalized {
9264 self.events.lock().unwrap().push(MeterEvent::Dropped {
9265 prompt: self.prompt,
9266 cached: self.cached,
9267 completion: self.completion,
9268 });
9269 }
9270 }
9271 }
9272
9273 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
9278
9279 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
9282 let req: CompletionReq = serde_json::from_value(json!({
9283 "model": "qwen/qwen3.8-27b",
9284 "prompt_ids": vec![7u32; prompt_ids],
9285 }))
9286 .unwrap();
9287 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9288 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
9289 request.params.max_new = max_new;
9290 request
9291 }
9292
9293 #[test]
9300 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
9301 let prompt = 30_278u64;
9302 let deadline_ms = TIMEOUT_MS_DEFAULT;
9303 let margin = |max_new: u64| {
9304 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
9305 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
9306 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
9307 };
9308 for allowed in [64u64, 2048, 4096, 5120, 6144] {
9309 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
9310 }
9311 for refused in [8192u64, 16384, 262_144] {
9312 assert!(
9313 !margin(refused),
9314 "{refused} measured as a 408 and must be refused"
9315 );
9316 }
9317 }
9318
9319 #[test]
9320 fn the_gate_names_a_max_tokens_that_actually_fits() {
9321 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
9324 assert!(
9325 fits > 0 && fits < 7_800,
9326 "advice {fits} must fit the measured ceiling"
9327 );
9328 assert_eq!(
9330 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
9331 None
9332 );
9333 }
9334
9335 #[test]
9336 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
9337 let req = gate_request(262_144, 30_000);
9338 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
9339 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
9341 assert!(
9342 err.contains("stream"),
9343 "message must name the streaming alternative: {err}"
9344 );
9345 assert!(
9346 err.contains("max_tokens"),
9347 "message must name the knob: {err}"
9348 );
9349 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
9351 let _l = GATE_ENV_LOCK.lock().unwrap(); for off in ["0", "off", "false"] {
9358 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
9359 assert!(
9360 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
9361 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
9362 );
9363 }
9364 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
9365 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
9366 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
9367 assert!(
9368 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
9369 "unset means ON (the documented default)"
9370 );
9371 }
9372
9373 #[test]
9380 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
9381 let strip = |src: &str| -> String {
9388 src.lines()
9389 .map(|line| match line.find("//") {
9390 Some(i) => line[..i].to_string(),
9391 None => line.to_string(),
9392 })
9393 .collect::<Vec<_>>()
9394 .join("\n")
9395 };
9396 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
9398 let start = src
9399 .find(signature)
9400 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
9401 let rest = &src[start + signature.len()..];
9402 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
9403 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
9404 &rest[..end]
9405 }
9406 let main_src = strip(include_str!("lib.rs"));
9407 let surfaces_src = strip(include_str!("surfaces.rs"));
9408 for (surface, src, signature) in [
9409 ("/v1/completions", &main_src, "async fn completions("),
9410 (
9411 "/v1/chat/completions",
9412 &main_src,
9413 "async fn chat_completions(",
9414 ),
9415 (
9416 "/v1/messages + /v1/responses (shared admission)",
9417 &surfaces_src,
9418 "pub(crate) async fn admit_translated(",
9419 ),
9420 ] {
9421 let handler = body(src, signature);
9422 assert!(
9423 handler.contains("nonstream_deadline_gate("),
9424 "{surface} must CALL the feasibility gate inside {signature}"
9425 );
9426 let limits = handler
9429 .find("apply_model_request_limits(")
9430 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
9431 let gate = handler.find("nonstream_deadline_gate(").unwrap();
9432 assert!(
9433 limits < gate,
9434 "{surface}: the gate must run after apply_model_request_limits"
9435 );
9436 }
9437 }
9438
9439 #[test]
9443 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
9444 let err = json!({"code": "deadline_exceeded",
9445 "metadata": {"error_type": "timeout"}});
9446 let cut = CompletionResp {
9447 model: "m".into(),
9448 text: "partial".into(),
9449 tokens: vec![1, 2],
9450 stop_reason: "Deadline".into(),
9451 error: Some(err.clone()),
9452 n_tokens: 2,
9453 prompt_tokens: 9,
9454 cached_tokens: 0,
9455 elapsed_s: 1.0,
9456 };
9457 let v = serde_json::to_value(&cut).unwrap();
9458 assert_eq!(v["stop_reason"], "Deadline");
9459 assert_eq!(v["error"]["code"], "deadline_exceeded");
9460 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
9461 let whole = CompletionResp {
9463 error: None,
9464 stop_reason: "Eos".into(),
9465 ..cut
9466 };
9467 let v = serde_json::to_value(&whole).unwrap();
9468 assert!(
9469 v.get("error").is_none(),
9470 "a complete response must not grow an error key: {v}"
9471 );
9472 }
9473
9474 #[test]
9475 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
9476 let _l = GATE_ENV_LOCK.lock().unwrap();
9477 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
9481 assert!(
9482 nonstream_deadline_gate(
9483 &req,
9484 false,
9485 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9486 false,
9487 None,
9488 )
9489 .is_ok(),
9490 "an omitted max_tokens is never gated — context is its only limit"
9491 );
9492 let resolved = gate_request(32_768, 30_000);
9497 assert!(
9498 nonstream_deadline_gate(
9499 &resolved,
9500 false,
9501 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9502 false,
9503 None,
9504 )
9505 .is_ok(),
9506 "a resolved-but-undeclared cap is not the caller's number to be refused over"
9507 );
9508 assert!(
9510 nonstream_deadline_gate(
9511 &resolved,
9512 false,
9513 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9514 true,
9515 None,
9516 )
9517 .is_err()
9518 );
9519 }
9520
9521 #[test]
9522 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
9523 let req = gate_request(64, 1234);
9524 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
9525 let mut text = gate_request(64, 0);
9526 text.prompt_ids.clear();
9527 text.prompt_text = "x".repeat(6_000);
9528 assert_eq!(
9529 prompt_tokens_estimate(&text, None),
9530 1_000,
9531 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
9532 that would have succeeded"
9533 );
9534 }
9535
9536 #[test]
9537 fn vision_memory_reservation_is_bounded_and_released() {
9538 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
9539 let Err(capacity) = try_reserve_vision_memory(1) else {
9540 panic!("a full process vision budget admitted another request");
9541 };
9542 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
9543 let response = vision_memory_error_response(capacity, Some("messages"));
9544 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
9545 assert_eq!(response.headers()["retry-after"], "5");
9546 assert_eq!(response.headers()["retry-after-ms"], "5000");
9547 drop(permit);
9548 assert!(try_reserve_vision_memory(1).is_ok());
9549 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
9550 panic!("an over-limit vision request was admitted");
9551 };
9552 assert!(matches!(request, VisionMemoryError::Request(_)));
9553 let response = vision_memory_error_response(request, Some("messages"));
9554 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
9555 assert_eq!(response.headers()["x-should-retry"], "false");
9556 let _ = try_reserve_vision_memory(1);
9557 }
9558
9559 #[test]
9560 fn header_auth_gate_covers_only_inference_dialects() {
9561 for path in [
9562 "/v1/auth/check",
9563 "/v1/completions",
9564 "/v1/chat/completions",
9565 "/v1/messages",
9566 "/v1/responses",
9567 "/v1/embeddings",
9568 "/v1/rerank",
9569 ] {
9570 assert!(protected_inference_path(path), "{path}");
9571 }
9572 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
9573 assert!(!protected_inference_path(path), "{path}");
9574 }
9575 }
9576 #[tokio::test]
9583 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
9584 use crate::metering::Metering as _;
9585 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
9586
9587 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
9588 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
9589 tx.send(Event::PromptUsage {
9590 n_prompt: 7,
9591 n_cached: 0,
9592 })
9593 .unwrap();
9594 tx.send(Event::Token {
9595 id: 1,
9596 text: "Hel".into(),
9597 })
9598 .unwrap();
9599 tx.send(Event::Token {
9600 id: 2,
9601 text: "lo".into(),
9602 })
9603 .unwrap();
9604 tx.send(Event::Done {
9605 stop_reason: "eos".into(),
9606 n_tokens: 2,
9607 n_prompt: 7,
9608 n_cached: 0,
9609 elapsed_s: 0.05,
9610 spec: None,
9611 })
9612 .unwrap();
9613 drop(tx);
9614 let mut receipt = receipt;
9615 blocking_response_with_receipt(
9616 rx,
9617 "m".into(),
9618 true,
9619 Vec::new(),
9620 None,
9621 Envelope::new(true),
9622 &mut receipt,
9623 None,
9624 )
9625 .await
9626 };
9627
9628 let plain = MockMetering::admit_all();
9630 let receipt = plain.open(
9631 &metering::RequestMeta {
9632 request_id: "cap-unmarked",
9633 tenant: "unmarked",
9634 principal: None,
9635 model: "m",
9636 route: "/v1/chat/completions",
9637 lane: "interactive",
9638 stream: false,
9639 max_tokens: None,
9640 reserved_ctx: None,
9641 },
9642 None,
9643 );
9644 let response = drive(Some(receipt)).await;
9645 assert_eq!(response.status(), StatusCode::OK);
9646 assert!(
9647 !plain.events().iter().any(|e| matches!(
9648 e,
9649 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
9650 )),
9651 "an unarmed receipt must see no capture traffic: {:?}",
9652 plain.events()
9653 );
9654
9655 let capturing = MockMetering::capturing();
9658 let mut receipt = capturing.open(
9659 &metering::RequestMeta {
9660 request_id: "cap-marked",
9661 tenant: "marked",
9662 principal: None,
9663 model: "m",
9664 route: "/v1/chat/completions",
9665 lane: "interactive",
9666 stream: false,
9667 max_tokens: None,
9668 reserved_ctx: None,
9669 },
9670 None,
9671 );
9672 assert!(receipt.wants_capture());
9673 receipt.arm_capture(prompt.clone());
9674 let response = drive(Some(receipt)).await;
9675 assert_eq!(response.status(), StatusCode::OK);
9676 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
9677 .await
9678 .unwrap();
9679 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
9680 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
9681
9682 let events = capturing.events();
9683 assert!(
9684 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
9685 "prompt must arm byte-exact: {events:?}"
9686 );
9687 let completion: String = events
9688 .iter()
9689 .filter_map(|e| match e {
9690 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
9691 _ => None,
9692 })
9693 .collect();
9694 assert_eq!(
9695 completion, "Hello",
9696 "the deltas must reassemble the served completion byte-exact: {events:?}"
9697 );
9698 assert!(
9699 events.contains(&MeterEvent::Complete {
9700 prompt: 7,
9701 cached: 0,
9702 completion: 2,
9703 }),
9704 "worker-truth usage settles alongside the capture: {events:?}"
9705 );
9706 }
9707
9708 fn tool_caps() -> ModelCaps {
9709 ModelCaps {
9710 tools_branch: true,
9711 qwen_think: true,
9712 think_switch: true,
9713 chat_ok: true,
9714 ..Default::default()
9715 }
9716 }
9717
9718 fn ladder_caps() -> ModelCaps {
9723 ModelCaps {
9724 qwen_effort: true,
9725 ..tool_caps()
9726 }
9727 }
9728
9729 fn gemma_tool_caps() -> ModelCaps {
9730 ModelCaps {
9731 tools_branch: true,
9732 gemma_think: true,
9733 chat_ok: true,
9734 instruct_type: Some("gemma".into()),
9735 ..Default::default()
9736 }
9737 }
9738
9739 fn gemma_template(kind: &str) -> String {
9740 let file = match kind {
9741 "qat" => "qat-trunk-template.jinja",
9742 _ => "official-tooluse-template.jinja",
9743 };
9744 let path = format!(
9745 "{}/../../research/gemma4-tools-20260817/{file}",
9746 env!("CARGO_MANIFEST_DIR")
9747 );
9748 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9749 }
9750
9751 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
9756 let tools_arr = request
9757 .get("tools")
9758 .and_then(|t| t.as_array())
9759 .cloned()
9760 .unwrap_or_default();
9761 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
9762 (Vec::new(), Vec::new(), HashMap::new())
9763 } else {
9764 prepare_tools(&tools_arr).unwrap()
9765 };
9766 let effort = request
9767 .get("reasoning_effort")
9768 .and_then(|v| v.as_str())
9769 .map(String::from);
9770 let (think, _lvl, _explicit) =
9771 parse_think(&effort, &None, None, None, None, false).unwrap();
9772
9773 let mut turns: Vec<TmplTurn> = Vec::new();
9774 for msg in request["messages"].as_array().unwrap() {
9775 let role = msg["role"].as_str().unwrap();
9776 let role = if role == "developer" { "system" } else { role };
9777 let content =
9778 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9779 let tool_calls = msg
9780 .get("tool_calls")
9781 .and_then(|a| a.as_array())
9782 .map(|a| {
9783 a.iter()
9784 .map(|tc| {
9785 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9786 render_req_tool_call(&rtc).unwrap()
9787 })
9788 .collect()
9789 })
9790 .unwrap_or_default();
9791 let tool_responses = msg
9792 .get("tool_responses")
9793 .and_then(|a| a.as_array())
9794 .map(|a| {
9795 a.iter()
9796 .map(|tr| {
9797 (
9798 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
9799 json_to_val(&tr["response"]),
9800 )
9801 })
9802 .collect()
9803 })
9804 .unwrap_or_default();
9805 turns.push(TmplTurn {
9806 role: role.to_string(),
9807 content,
9808 tool_calls,
9809 reasoning: msg
9810 .get("reasoning")
9811 .and_then(|r| r.as_str())
9812 .map(String::from)
9813 .filter(|s| !s.is_empty()),
9814 tool_call_id: msg
9815 .get("tool_call_id")
9816 .and_then(|s| s.as_str())
9817 .map(String::from),
9818 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9819 tool_responses,
9820 task: None,
9821 tools: Vec::new(),
9822 });
9823 }
9824 chat::apply_chat_template_tools_ex(
9825 Some(template),
9826 &turns,
9827 true,
9828 &tools_json,
9829 &tools_struct,
9830 think,
9831 None,
9832 None,
9833 )
9834 .unwrap()
9835 }
9836
9837 #[test]
9841 fn gemma4_tools_fixtures_match_the_official_jinja() {
9842 let dir = format!(
9843 "{}/../../research/gemma4-tools-20260817/fixtures",
9844 env!("CARGO_MANIFEST_DIR")
9845 );
9846 let mut entries: Vec<_> = std::fs::read_dir(&dir)
9847 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
9848 .map(|e| e.unwrap().path())
9849 .filter(|p| p.is_dir())
9850 .collect();
9851 entries.sort();
9852 assert!(
9853 entries.len() >= 14,
9854 "expected >=14 fixtures, found {}",
9855 entries.len()
9856 );
9857 let (mut official, mut qat) = (0u32, 0u32);
9858 for d in entries {
9859 let input: serde_json::Value =
9860 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
9861 .unwrap();
9862 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
9863 let kind = input
9864 .get("template")
9865 .and_then(|t| t.as_str())
9866 .unwrap_or("official");
9867 match kind {
9868 "qat" => qat += 1,
9869 _ => official += 1,
9870 }
9871 let tmpl = gemma_template(kind);
9872 let got = render_fixture(&input["request"], &tmpl);
9873 assert_eq!(
9874 got, expected,
9875 "fixture {:?} diverged from the jinja oracle",
9876 d
9877 );
9878 }
9879 assert!(
9880 official >= 12 && qat >= 2,
9881 "coverage: {official} official, {qat} qat"
9882 );
9883 }
9884
9885 #[test]
9891 fn gemma4_tools_flow_through_build_chat_request() {
9892 let tmpl = gemma_template("official");
9893 for name in [
9894 "01-system-tools-basic",
9895 "04-single-call-cycle",
9896 "07-multi-cycle-agentic",
9897 ] {
9898 let path = format!(
9899 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
9900 env!("CARGO_MANIFEST_DIR")
9901 );
9902 let input: serde_json::Value =
9903 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
9904 let expected_path = format!(
9905 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
9906 env!("CARGO_MANIFEST_DIR")
9907 );
9908 let expected = std::fs::read_to_string(&expected_path).unwrap();
9909 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
9910 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9911 let plan = build_chat_request(
9912 req,
9913 Some(&gemma_tool_caps()),
9914 tx,
9915 lanes::Lane::Interactive,
9916 None,
9917 )
9918 .unwrap();
9919 let got = chat::apply_chat_template_tools_ex(
9920 Some(&tmpl),
9921 &plan.request.chat_turns,
9922 true,
9923 &plan.request.tools_json,
9924 &plan.request.tools_struct,
9925 plan.request.think,
9926 plan.request.reasoning_effort.as_deref(),
9927 None,
9928 )
9929 .unwrap();
9930 assert_eq!(got, expected, "pipeline render diverged for {name}");
9931 }
9932 }
9933
9934 fn dsv4_sentinel() -> String {
9942 let path = format!(
9943 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
9944 env!("CARGO_MANIFEST_DIR")
9945 );
9946 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9947 }
9948
9949 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
9954 let role = msg["role"].as_str().unwrap().to_string();
9955 let content =
9956 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9957 let reasoning = msg
9958 .get("reasoning")
9959 .or_else(|| msg.get("reasoning_content"))
9960 .and_then(|r| r.as_str())
9961 .map(String::from)
9962 .filter(|s| !s.is_empty());
9963 let tool_calls = msg
9964 .get("tool_calls")
9965 .and_then(|a| a.as_array())
9966 .map(|a| {
9967 a.iter()
9968 .map(|tc| {
9969 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9970 render_req_tool_call(&rtc).unwrap()
9971 })
9972 .collect()
9973 })
9974 .unwrap_or_default();
9975 let tools = msg
9976 .get("tools")
9977 .and_then(|a| a.as_array())
9978 .map(|a| {
9979 a.iter()
9980 .filter_map(|t| t.get("function").map(json_to_val))
9981 .collect()
9982 })
9983 .unwrap_or_default();
9984 TmplTurn {
9985 role,
9986 content,
9987 tool_calls,
9988 reasoning,
9989 tool_call_id: msg
9990 .get("tool_call_id")
9991 .and_then(|s| s.as_str())
9992 .map(String::from),
9993 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9994 tool_responses: Vec::new(),
9995 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
9996 tools,
9997 }
9998 }
9999
10000 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
10001 v.and_then(|t| t.as_array())
10002 .map(|a| {
10003 a.iter()
10004 .filter_map(|t| t.get("function").map(json_to_val))
10005 .collect()
10006 })
10007 .unwrap_or_default()
10008 }
10009
10010 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
10014 let dir = format!(
10015 "{}/../../research/dsv4-template-20260818/{subdir}",
10016 env!("CARGO_MANIFEST_DIR")
10017 );
10018 let tmpl = dsv4_sentinel();
10019 let mut entries: Vec<_> = std::fs::read_dir(&dir)
10020 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
10021 .map(|e| e.unwrap().path())
10022 .filter(|p| p.is_dir())
10023 .collect();
10024 entries.sort();
10025 assert!(
10026 entries.len() >= min_fixtures,
10027 "expected >={min_fixtures} fixtures, found {}",
10028 entries.len()
10029 );
10030 for d in &entries {
10031 let input: serde_json::Value =
10032 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
10033 .unwrap();
10034 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
10035 let turns: Vec<TmplTurn> = input["turns"]
10036 .as_array()
10037 .unwrap()
10038 .iter()
10039 .map(dsv4_turn)
10040 .collect();
10041 let think = match input["think"].as_str().unwrap() {
10042 "chat" => ThinkMode::NoThink,
10043 _ => ThinkMode::Think,
10044 };
10045 let effort = input
10046 .get("reasoning_effort")
10047 .and_then(|v| v.as_str())
10048 .map(String::from);
10049 let req_tools = dsv4_req_tools(input.get("req_tools"));
10050 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
10051 let got = chat::apply_chat_template_tools_ex(
10052 Some(&tmpl),
10053 &turns,
10054 agp,
10055 &[],
10056 &req_tools,
10057 think,
10058 effort.as_deref(),
10059 Some(encoding),
10060 )
10061 .unwrap();
10062 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
10063 }
10064 }
10065
10066 #[test]
10067 fn dsv4_template_fixtures_match_the_oracle() {
10068 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
10069 }
10070
10071 #[test]
10077 fn dsv4_0731_fixtures_match_the_oracle() {
10078 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
10079 }
10080
10081 #[test]
10082 fn dsv4_artifact_fixtures_are_byte_identical() {
10083 let base = format!(
10087 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
10088 env!("CARGO_MANIFEST_DIR")
10089 );
10090 let tmpl = dsv4_sentinel();
10091 for (n, think) in [
10092 (1u32, ThinkMode::Think),
10093 (2, ThinkMode::Think),
10094 (3, ThinkMode::Think),
10095 (4, ThinkMode::NoThink),
10096 ] {
10097 let td: serde_json::Value = serde_json::from_str(
10098 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
10099 )
10100 .unwrap();
10101 let (messages, tools) = if td.is_object() {
10102 (td["messages"].clone(), td.get("tools").cloned())
10103 } else {
10104 (td.clone(), None)
10105 };
10106 let mut turns: Vec<TmplTurn> = Vec::new();
10107 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
10108 let mut t = dsv4_turn(msg);
10109 if i == 0 {
10110 if let Some(tl) = &tools {
10111 t.tools = tl
10112 .as_array()
10113 .unwrap()
10114 .iter()
10115 .filter_map(|x| x.get("function").map(json_to_val))
10116 .collect();
10117 }
10118 }
10119 turns.push(t);
10120 }
10121 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
10122 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
10126 let got = chat::apply_chat_template_tools_ex(
10127 Some(&tmpl),
10128 &turns,
10129 true,
10130 &[],
10131 &[],
10132 think,
10133 None,
10134 Some(encoding),
10135 )
10136 .unwrap();
10137 assert_eq!(
10138 got, expected,
10139 "artifact fixture {n} diverged from the oracle under {encoding:?}"
10140 );
10141 }
10142 }
10143 }
10144
10145 #[test]
10146 fn dsv4_default_thinkmode_renders_thinking() {
10147 let tmpl = dsv4_sentinel();
10150 let turns = vec![TmplTurn {
10151 role: "user".into(),
10152 content: "Hi".into(),
10153 ..Default::default()
10154 }];
10155 let dflt = chat::apply_chat_template_tools_ex(
10156 Some(&tmpl),
10157 &turns,
10158 true,
10159 &[],
10160 &[],
10161 ThinkMode::Default,
10162 None,
10163 None,
10164 )
10165 .unwrap();
10166 let think = chat::apply_chat_template_tools_ex(
10167 Some(&tmpl),
10168 &turns,
10169 true,
10170 &[],
10171 &[],
10172 ThinkMode::Think,
10173 None,
10174 None,
10175 )
10176 .unwrap();
10177 assert_eq!(dflt, think);
10178 assert!(
10179 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
10180 "{dflt:?}"
10181 );
10182 let chat_mode = chat::apply_chat_template_tools_ex(
10183 Some(&tmpl),
10184 &turns,
10185 true,
10186 &[],
10187 &[],
10188 ThinkMode::NoThink,
10189 None,
10190 None,
10191 )
10192 .unwrap();
10193 assert!(
10194 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
10195 "{chat_mode:?}"
10196 );
10197 }
10198
10199 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
10204 let base = format!(
10205 "{}/../../research/dsv4-template-20260818",
10206 env!("CARGO_MANIFEST_DIR")
10207 );
10208 let refdir = std::path::Path::new(&base).join("ref");
10209 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
10210 .expect("load dsv4 tokenizer from ref dir");
10211 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
10212 let banked: serde_json::Value = serde_json::from_str(
10213 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
10214 .unwrap(),
10215 )
10216 .unwrap();
10217 let obj = banked.as_object().unwrap();
10218 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
10219 for (name, ids_v) in obj {
10220 let rendered =
10221 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
10222 let want: Vec<u32> = ids_v
10223 .as_array()
10224 .unwrap()
10225 .iter()
10226 .map(|v| v.as_u64().unwrap() as u32)
10227 .collect();
10228 let got = tok.encode(&rendered, true);
10229 assert_eq!(got, want, "tokenization diverged for {name}");
10230 }
10231 }
10232
10233 #[test]
10234 fn dsv4_tokenization_crosscheck_matches_official_ids() {
10235 dsv4_run_tokenization_crosscheck("fixtures");
10236 }
10237
10238 #[test]
10242 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
10243 dsv4_run_tokenization_crosscheck("fixtures-0731");
10244 }
10245
10246 #[test]
10247 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
10248 let base = format!(
10258 "{}/../../research/dsv4-template-20260818",
10259 env!("CARGO_MANIFEST_DIR")
10260 );
10261 let refdir = std::path::Path::new(&base).join("ref");
10262 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
10263 .expect("load dsv4 tokenizer from ref dir");
10264 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
10265 let tmpl = dsv4_sentinel();
10266 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
10267 {"type": "function", "function": {
10268 "name": "get_data",
10269 "description": "Fetch a blob",
10270 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
10271 "required": ["key"]}
10272 }}
10273 ])));
10274
10275 let cases: Vec<(&str, String)> = vec![
10276 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
10278 ("space-131k", " ".repeat(131_072)),
10279 ("digit-131k", "7".repeat(131_072)),
10280 (
10281 "mixed-runs",
10282 format!(
10283 "{}{}{}{}",
10284 "Z".repeat(65_536),
10285 " ".repeat(65_536),
10286 "7".repeat(65_536),
10287 "\n".repeat(65_536)
10288 ),
10289 ),
10290 ("cjk-64k", "中".repeat(65_536)),
10291 ("accented-letter-64k", "é".repeat(65_536)),
10292 ];
10293 for (name, blob) in &cases {
10294 let msgs = serde_json::json!([
10295 {"role": "system", "content": "You are a tool-using assistant."},
10296 {"role": "user", "content": "Fetch the blob."},
10297 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
10298 "tool_calls": [{"id": "call_001", "type": "function",
10299 "function": {"name": "get_data",
10300 "arguments": "{\"key\": \"blob\"}"}}]},
10301 {"role": "tool", "tool_call_id": "call_001", "content": blob}
10302 ]);
10303 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
10304 let rendered = chat::apply_chat_template_tools_ex(
10305 Some(&tmpl),
10306 &turns,
10307 true,
10308 &[],
10309 &req_tools,
10310 ThinkMode::Think,
10311 None,
10312 None,
10313 )
10314 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
10315 assert!(
10316 rendered.contains(blob.as_str()),
10317 "{name}: tool result missing from render"
10318 );
10319 let t0 = std::time::Instant::now();
10320 let ids = tok.encode(&rendered, true);
10321 let encode_dt = t0.elapsed();
10322 assert!(!ids.is_empty(), "{name}: empty encode");
10323 let back = tok.decode(&ids);
10324 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
10325 assert!(
10329 encode_dt < std::time::Duration::from_secs(60),
10330 "{name}: encode took {encode_dt:?}"
10331 );
10332 if *name == "ascii-letter-131k" {
10336 if let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR") {
10337 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
10338 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
10339 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
10340 }
10341 }
10342 }
10343 }
10344
10345 #[test]
10346 fn models_v1_entry_advertises_thinking_support() {
10347 let step_caps = ModelCaps {
10350 effort_levels: true,
10351 ..tool_caps()
10352 };
10353 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
10354 assert_eq!(entry["capabilities"]["reasoning"], true);
10355 assert_eq!(entry["capabilities"]["tools"], true);
10356
10357 let plain = ModelCaps {
10359 chat_ok: true,
10360 ..Default::default()
10361 };
10362 let entry = model_entry_v1("plain", Some(&plain), None);
10363 assert_eq!(entry["capabilities"]["reasoning"], false);
10364 assert_eq!(entry["capabilities"]["tools"], false);
10365 let entry = model_entry_v1("unknown", None, None);
10367 assert_eq!(entry["capabilities"]["reasoning"], false);
10368 assert_eq!(entry["capabilities"]["streaming"], true);
10369 }
10370
10371 #[test]
10372 fn chat_request_preserves_turns_and_openai_stop_forms() {
10373 let payload = serde_json::json!({
10374 "model": "plain_quant",
10375 "messages": [
10376 {"role": "system", "content": "rules"},
10377 {"role": "developer", "content": "dev rules"},
10378 {"role": "user", "content": "task"},
10379 {"role": "assistant", "content": "work"}
10380 ],
10381 "max_tokens": 64,
10382 "temperature": 0.0,
10383 "stop": "<stop>"
10384 });
10385 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10386 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10387 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10388 let request = plan.request;
10389 assert!(
10390 plan.parser.is_none(),
10391 "no tools -> no parser (isolation contract)"
10392 );
10393 assert!(request.tools_json.is_empty());
10394 assert_eq!(request.think, ThinkMode::Default);
10395 assert_eq!(request.model, "plain_quant");
10396 assert_eq!(request.params.max_new, 64);
10397 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10399 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
10400 }))
10401 .unwrap();
10402 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10403 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10404 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
10405 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10407 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10408 "max_completion_tokens": 7
10409 }))
10410 .unwrap();
10411 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10412 assert_eq!(
10413 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
10414 .unwrap()
10415 .request
10416 .params
10417 .max_new,
10418 7
10419 );
10420 let req: CompletionReq = serde_json::from_value(serde_json::json!({
10422 "model": "plain_quant", "prompt": "task"
10423 }))
10424 .unwrap();
10425 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10426 assert_eq!(
10427 build_request(&req, tx, lanes::Lane::Interactive, None)
10428 .params
10429 .max_new,
10430 worker::MAX_NEW_CTX_BOUNDED
10431 );
10432 let turns: Vec<(String, String)> = request
10433 .chat_turns
10434 .iter()
10435 .map(|t| (t.role.clone(), t.content.clone()))
10436 .collect();
10437 assert_eq!(
10438 turns,
10439 vec![
10440 ("system".into(), "rules".into()),
10441 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
10443 ("assistant".into(), "work".into()),
10444 ]
10445 );
10446 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
10447 assert_eq!(request.stop_strings, vec!["<stop>"]);
10448
10449 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10450 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10451 "stop": ["a", "b"]
10452 }))
10453 .unwrap();
10454 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
10455
10456 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10460 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10461 "stop": ["", "real", ""]
10462 }))
10463 .unwrap();
10464 assert_eq!(req.stop.into_vec(), vec!["real"]);
10465 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10466 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10467 "stop": ""
10468 }))
10469 .unwrap();
10470 assert!(req.stop.into_vec().is_empty());
10471
10472 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10473 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10474 "stop": null
10475 }))
10476 .unwrap();
10477 assert!(req.stop.into_vec().is_empty());
10478 }
10479
10480 #[tokio::test]
10481 async fn chat_response_has_openai_message_shape() {
10482 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10483 tx.send(Event::Token {
10484 id: 1,
10485 text: "hello".into(),
10486 })
10487 .unwrap();
10488 tx.send(Event::Done {
10489 stop_reason: "Eos".into(),
10490 n_tokens: 1,
10491 n_prompt: 42,
10492 n_cached: 30,
10493 elapsed_s: 0.5,
10494 spec: None,
10495 })
10496 .unwrap();
10497 drop(tx);
10498 let response = blocking_response(
10499 rx,
10500 "plain_quant".into(),
10501 true,
10502 Vec::new(),
10503 None,
10504 Envelope::new(true),
10505 )
10506 .await;
10507 assert_eq!(response.status(), StatusCode::OK);
10508 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10509 .await
10510 .unwrap();
10511 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10512 assert_eq!(payload["object"], "chat.completion");
10513 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
10515 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
10516 assert!(
10517 payload["system_fingerprint"]
10518 .as_str()
10519 .unwrap()
10520 .starts_with("memra-")
10521 );
10522 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
10523 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
10524 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
10525 assert_eq!(payload["usage"]["prompt_tokens"], 42);
10527 assert_eq!(payload["usage"]["completion_tokens"], 1);
10528 assert_eq!(payload["usage"]["total_tokens"], 43);
10529 assert_eq!(
10530 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
10531 30
10532 );
10533 assert!(payload["usage"].get("spec").is_none());
10536 }
10537
10538 #[tokio::test]
10539 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
10540 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10541 tx.send(Event::Token {
10543 id: 4,
10544 text: "hello".into(),
10545 })
10546 .unwrap();
10547 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
10548 tx.send(Event::Done {
10549 stop_reason: "MaxNew".into(),
10550 n_tokens: 4,
10551 n_prompt: 2,
10552 n_cached: 0,
10553 elapsed_s: 0.5,
10554 spec: None,
10555 })
10556 .unwrap();
10557 drop(tx);
10558
10559 let response = blocking_response(
10560 rx,
10561 "plain_quant".into(),
10562 false,
10563 Vec::new(),
10564 None,
10565 Envelope::new(false),
10566 )
10567 .await;
10568 assert_eq!(response.status(), StatusCode::OK);
10569 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10570 .await
10571 .unwrap();
10572 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10573 assert_eq!(payload["text"], "hello");
10574 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
10575 assert_eq!(payload["n_tokens"], 4);
10576 }
10577
10578 #[tokio::test]
10581 async fn chat_usage_carries_spec_acceptance_summary() {
10582 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10583 tx.send(Event::Token {
10584 id: 1,
10585 text: "hello".into(),
10586 })
10587 .unwrap();
10588 tx.send(Event::Done {
10589 stop_reason: "Eos".into(),
10590 n_tokens: 1,
10591 n_prompt: 42,
10592 n_cached: 0,
10593 elapsed_s: 0.5,
10594 spec: Some(worker::SpecUsage {
10595 rounds: 10,
10596 drafted: 30,
10597 accepted: 21,
10598 }),
10599 })
10600 .unwrap();
10601 drop(tx);
10602 let response = blocking_response(
10603 rx,
10604 "plain_quant".into(),
10605 true,
10606 Vec::new(),
10607 None,
10608 Envelope::new(true),
10609 )
10610 .await;
10611 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10612 .await
10613 .unwrap();
10614 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10615 let sp = &payload["usage"]["spec"];
10616 assert_eq!(sp["rounds"], 10);
10617 assert_eq!(sp["drafted"], 30);
10618 assert_eq!(sp["accepted"], 21);
10619 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
10620 assert_eq!(payload["usage"]["total_tokens"], 43);
10622 }
10623
10624 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
10625 let mut payload = serde_json::json!({
10626 "model": "m",
10627 "messages": [{"role": "user", "content": "Weather in Paris?"}],
10628 "tools": [{"type": "function", "function": {
10629 "name": "get_weather",
10630 "description": "Get current weather",
10631 "parameters": {"type": "object",
10632 "properties": {"city": {"type": "string"},
10633 "days": {"type": "integer"}},
10634 "required": ["city"]}}}],
10635 });
10636 if let Some(obj) = extra.as_object() {
10637 for (k, v) in obj {
10638 payload[k] = v.clone();
10639 }
10640 }
10641 serde_json::from_value(payload).unwrap()
10642 }
10643
10644 #[test]
10645 fn vision_decode_is_deferred_and_grid_pinned() {
10646 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10651 let req: ChatCompletionReq = serde_json::from_value(json!({
10652 "model": "m", "messages": [{"role": "user", "content": "hi"}],
10653 }))
10654 .unwrap();
10655 let mut plan = build_chat_request(
10656 req,
10657 Some(&ModelCaps {
10658 chat_ok: true,
10659 ..Default::default()
10660 }),
10661 tx,
10662 lanes::Lane::Interactive,
10663 None,
10664 )
10665 .unwrap();
10666 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
10670 let mut b = Vec::new();
10671 b.extend_from_slice(b"BM");
10672 b.extend_from_slice(&54u32.to_le_bytes());
10673 b.extend_from_slice(&0u32.to_le_bytes());
10674 b.extend_from_slice(&54u32.to_le_bytes());
10675 b.extend_from_slice(&40u32.to_le_bytes());
10676 b.extend_from_slice(&w.to_le_bytes());
10677 b.extend_from_slice(&h.to_le_bytes());
10678 b.extend_from_slice(&1u16.to_le_bytes());
10679 b.extend_from_slice(&24u16.to_le_bytes());
10680 b.extend_from_slice(&[0u8; 24]);
10681 if with_pixels {
10682 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
10683 }
10684 b
10685 };
10686 let bytes = bmp(64, 64, true);
10687 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
10688 plan.pending_images.push(PendingVisionUnit::Still {
10689 bytes: bytes.clone(),
10690 gh,
10691 gw,
10692 });
10693 decode_pending_vision(&mut plan).unwrap();
10694 assert_eq!(plan.request.images.len(), 1);
10695 assert_eq!(
10696 (
10697 plan.request.images[0].prep.gh,
10698 plan.request.images[0].prep.gw
10699 ),
10700 (gh, gw),
10701 "decoded grid must equal the header-planned grid the pad run was rendered from"
10702 );
10703 plan.request.images.clear();
10705 plan.pending_images.push(PendingVisionUnit::Still {
10706 bytes,
10707 gh: gh + 2,
10708 gw,
10709 });
10710 let err = decode_pending_vision(&mut plan).unwrap_err();
10711 assert!(err.contains("header-planned"), "got: {err}");
10712 let bomb = bmp(16_000, 16_000, false);
10715 plan.pending_images.clear();
10716 plan.pending_images.push(PendingVisionUnit::Still {
10717 bytes: bomb,
10718 gh: 2,
10719 gw: 2,
10720 });
10721 let err = decode_pending_vision(&mut plan).unwrap_err();
10722 assert!(err.contains("exceeds the decode budget"), "got: {err}");
10723 }
10724
10725 #[test]
10726 fn tools_request_renders_client_key_order_and_arms_parser() {
10727 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10728 let plan = build_chat_request(
10729 weather_request(json!({})),
10730 Some(&tool_caps()),
10731 tx,
10732 lanes::Lane::Interactive,
10733 None,
10734 )
10735 .unwrap();
10736 assert!(plan.parser.is_some());
10737 assert_eq!(plan.request.tools_json.len(), 1);
10738 assert_eq!(
10740 plan.request.tools_json[0],
10741 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
10742 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
10743 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
10744 \"integer\"}}, \"required\": [\"city\"]}}}"
10745 );
10746 }
10747
10748 #[test]
10749 fn tool_choice_none_strips_tools_and_parser() {
10750 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10751 let plan = build_chat_request(
10752 weather_request(json!({"tool_choice": "none"})),
10753 Some(&tool_caps()),
10754 tx,
10755 lanes::Lane::Interactive,
10756 None,
10757 )
10758 .unwrap();
10759 let mut p = plan
10762 .parser
10763 .expect("think-open chat arms the reasoning splitter");
10764 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
10765 assert_eq!(
10766 pieces,
10767 vec![
10768 Piece::Reasoning("x".into()),
10769 Piece::Content("<tool_call> stays prose".into()),
10770 ]
10771 );
10772 assert!(plan.request.tools_json.is_empty());
10773 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10775 assert!(
10776 build_chat_request(
10777 weather_request(json!({"tool_choice": "required"})),
10778 Some(&tool_caps()),
10779 tx,
10780 lanes::Lane::Interactive,
10781 None
10782 )
10783 .is_err()
10784 );
10785 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10786 assert!(
10787 build_chat_request(
10788 weather_request(json!({"tool_choice":
10789 {"type": "function", "function": {"name": "get_weather"}}})),
10790 Some(&tool_caps()),
10791 tx,
10792 lanes::Lane::Interactive,
10793 None
10794 )
10795 .is_err()
10796 );
10797 }
10798
10799 #[test]
10800 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
10801 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
10802 let _ = std::fs::remove_dir_all(&root);
10803
10804 let st = root.join("st_single");
10806 std::fs::create_dir_all(&st).unwrap();
10807 std::fs::write(st.join("config.json"), "{}").unwrap();
10808 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
10809 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
10810
10811 let sh = root.join("st_sharded");
10813 std::fs::create_dir_all(&sh).unwrap();
10814 std::fs::write(sh.join("config.json"), "{}").unwrap();
10815 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
10816 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
10817
10818 let rp = root.join("repack");
10820 std::fs::create_dir_all(&rp).unwrap();
10821 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
10822 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
10823
10824 let bogus = root.join("bogus");
10826 std::fs::create_dir_all(&bogus).unwrap();
10827 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
10828 assert!(
10829 err.contains("model.safetensors"),
10830 "error should say what is missing: {err}"
10831 );
10832 assert!(
10833 err.contains("manifest.json"),
10834 "error should mention the repack form: {err}"
10835 );
10836
10837 let nc = root.join("no_config");
10839 std::fs::create_dir_all(&nc).unwrap();
10840 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
10841 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
10842 assert!(
10843 err.contains("config.json"),
10844 "error should name config.json: {err}"
10845 );
10846
10847 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
10849 assert!(err.contains("does not exist"), "{err}");
10850
10851 let f = root.join("model.gguf");
10853 std::fs::write(&f, b"g").unwrap();
10854 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
10855
10856 let _ = std::fs::remove_dir_all(&root);
10857 }
10858
10859 #[test]
10860 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
10861 let caps = ModelCaps {
10864 tools_branch: false,
10865 qwen_think: false,
10866 think_switch: false,
10867 chat_ok: false,
10868 ..Default::default()
10869 };
10870 let payload = serde_json::json!({
10871 "model": "st_model",
10872 "messages": [{"role": "user", "content": "hello"}],
10873 });
10874 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10875 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10876 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
10877 Err(e) => e,
10878 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
10879 };
10880 assert!(
10881 err.contains("no chat template"),
10882 "message should name the cause: {err}"
10883 );
10884 assert!(
10885 err.contains("/v1/completions"),
10886 "message should point at the raw-prompt escape hatch: {err}"
10887 );
10888 }
10889
10890 #[test]
10891 fn tools_on_model_without_tools_branch_is_rejected() {
10892 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10893 let caps = ModelCaps {
10894 chat_ok: true,
10895 ..Default::default()
10896 };
10897 assert!(
10898 build_chat_request(
10899 weather_request(json!({})),
10900 Some(&caps),
10901 tx,
10902 lanes::Lane::Interactive,
10903 None
10904 )
10905 .is_err()
10906 );
10907 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10908 assert!(
10909 build_chat_request(
10910 weather_request(json!({})),
10911 None,
10912 tx,
10913 lanes::Lane::Interactive,
10914 None
10915 )
10916 .is_err()
10917 );
10918 }
10919
10920 #[test]
10921 fn reasoning_effort_maps_to_think_switch() {
10922 for (extra, want) in [
10928 (json!({}), ThinkMode::Default),
10929 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
10930 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
10931 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
10932 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
10933 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
10934 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
10935 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
10936 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
10937 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
10941 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
10942 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
10943 (
10947 json!({"reasoning": {"enabled": true, "effort": "none"}}),
10948 ThinkMode::Think,
10949 ),
10950 (
10951 json!({"reasoning": {"enabled": false, "effort": "high"}}),
10952 ThinkMode::NoThink,
10953 ),
10954 ] {
10955 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10956 let plan = build_chat_request(
10957 weather_request(extra.clone()),
10958 Some(&ladder_caps()),
10963 tx,
10964 lanes::Lane::Interactive,
10965 None,
10966 )
10967 .unwrap();
10968 assert_eq!(plan.request.think, want, "extra={extra}");
10969 }
10970 for extra in [
10974 json!({"reasoning_effort": "extreme"}),
10975 json!({"reasoning": {"effort": "banana"}}),
10976 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
10977 json!({"reasoning": {"enabled": true, "effort": ""}}),
10978 ] {
10979 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10980 assert!(
10981 build_chat_request(
10982 weather_request(extra.clone()),
10983 Some(&tool_caps()),
10984 tx,
10985 lanes::Lane::Interactive,
10986 None
10987 )
10988 .is_err(),
10989 "extra={extra} must be rejected by the one allowlist"
10990 );
10991 }
10992 for (raw, want) in [
10995 ("none", Some("none")),
10996 ("minimal", Some("minimal")),
10997 ("low", Some("low")),
10998 ("medium", Some("medium")),
10999 ("high", Some("high")),
11000 ("xhigh", Some("high")),
11001 ("max", Some("high")),
11002 ("ultra", Some("high")),
11003 ("banana", None),
11004 ("", None),
11005 ("HIGH", None),
11006 ] {
11007 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
11008 }
11009 for (raw, want) in [
11012 ("none", Some("none")),
11013 ("minimal", Some("minimal")),
11014 ("low", Some("low")),
11015 ("medium", Some("medium")),
11016 ("high", Some("high")),
11017 ("xhigh", Some("max")),
11018 ("max", Some("max")),
11019 ("ultra", Some("max")),
11020 ("banana", None),
11021 ("", None),
11022 ("MAX", None),
11023 ] {
11024 assert_eq!(
11025 canonical_effort_for(raw, true),
11026 want,
11027 "canonical_effort_for({raw:?}, dsv4)"
11028 );
11029 }
11030 }
11031
11032 #[test]
11033 fn dsv4_reasoning_effort_max_survives_canonicalization() {
11034 let dsv4_caps = ModelCaps {
11040 chat_ok: true,
11041 dsv4: true,
11042 ..Default::default()
11043 };
11044 let build = |caps: &ModelCaps, effort: &str| {
11045 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11046 let req: ChatCompletionReq = serde_json::from_value(json!({
11047 "model": "m",
11048 "messages": [{"role": "user", "content": "hi"}],
11049 "reasoning_effort": effort,
11050 }))
11051 .unwrap();
11052 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
11053 };
11054 for raw in ["max", "xhigh", "ultra"] {
11055 let plan = build(&dsv4_caps, raw).unwrap();
11056 assert_eq!(
11057 plan.request.reasoning_effort.as_deref(),
11058 Some("max"),
11059 "dsv4 {raw:?} must reach the renderer as the max rung"
11060 );
11061 assert_eq!(plan.request.think, chat::ThinkMode::Think);
11062 }
11063 let plan = build(&dsv4_caps, "high").unwrap();
11065 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
11066 let step_caps = ModelCaps {
11068 chat_ok: true,
11069 effort_levels: true,
11070 ..Default::default()
11071 };
11072 let plan = build(&step_caps, "max").unwrap();
11073 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
11074 }
11075
11076 #[test]
11077 fn default_reasoning_effort_flips_only_the_unset_request() {
11078 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
11085 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11086 build_chat_request_with_trace(
11087 weather_request(extra),
11088 Some(&ladder_caps()),
11089 tx,
11090 lanes::Lane::Interactive,
11091 None,
11092 None,
11093 default_effort,
11094 &ModelSamplingDefaults::default(),
11095 )
11096 .unwrap()
11097 };
11098 for (extra, want) in [
11099 (json!({}), ThinkMode::Think),
11101 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
11105 (json!({"include_reasoning": false}), ThinkMode::NoThink),
11106 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
11108 (json!({"include_reasoning": true}), ThinkMode::Think),
11109 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
11111 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
11112 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
11113 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
11115 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
11116 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
11117 ] {
11118 let plan = build(extra.clone(), Some("high"));
11119 assert_eq!(plan.request.think, want, "extra={extra}");
11120 }
11121 assert_eq!(
11123 build(json!({}), Some("none")).request.think,
11124 ThinkMode::NoThink
11125 );
11126 assert_eq!(
11127 build(json!({"reasoning_effort": "high"}), Some("none"))
11128 .request
11129 .think,
11130 ThinkMode::Think
11131 );
11132 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
11136 }
11137
11138 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
11144 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
11145 {%- else %}'<think>\\n'{%- endif %}";
11146
11147 #[test]
11148 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
11149 let build = |extra: serde_json::Value| {
11155 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11156 build_chat_request(
11157 weather_request(extra),
11158 Some(&tool_caps()),
11159 tx,
11160 lanes::Lane::Interactive,
11161 None,
11162 )
11163 };
11164 for (extra, want) in [
11165 (json!({"enable_thinking": false}), ThinkMode::NoThink),
11166 (json!({"enable_thinking": true}), ThinkMode::Think),
11167 (
11168 json!({"chat_template_kwargs": {"enable_thinking": false}}),
11169 ThinkMode::NoThink,
11170 ),
11171 (
11172 json!({"chat_template_kwargs": {"enable_thinking": true}}),
11173 ThinkMode::Think,
11174 ),
11175 (
11178 json!({"enable_thinking": false, "reasoning_effort": "high"}),
11179 ThinkMode::NoThink,
11180 ),
11181 (
11183 json!({"enable_thinking": false,
11184 "chat_template_kwargs": {"enable_thinking": false}}),
11185 ThinkMode::NoThink,
11186 ),
11187 ] {
11188 let plan = build(extra.clone()).unwrap_or_else(|e| {
11189 panic!("{extra} must be accepted and honored, got 400: {e}");
11190 });
11191 assert_eq!(
11192 plan.request.think, want,
11193 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
11194 );
11195 }
11196 let render = |extra: serde_json::Value| -> String {
11199 let plan = build(extra).unwrap();
11200 chat::apply_chat_template_tools_ex(
11201 Some(SWITCHED_QWEN_TMPL),
11202 &plan.request.chat_turns,
11203 true,
11204 &plan.request.tools_json,
11205 &plan.request.tools_struct,
11206 plan.request.think,
11207 plan.request.reasoning_effort.as_deref(),
11208 None,
11209 )
11210 .unwrap()
11211 };
11212 let off = render(json!({"enable_thinking": false}));
11213 assert!(
11214 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
11215 "enable_thinking:false must render the CLOSED think pair: {off:?}"
11216 );
11217 let on = render(json!({}));
11218 assert!(
11219 on.ends_with("<|im_start|>assistant\n<think>\n"),
11220 "an unset request must still render the template's OPEN think tail: {on:?}"
11221 );
11222 assert_eq!(
11223 off,
11224 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
11225 "both vLLM spellings must render byte-identically"
11226 );
11227 assert_eq!(
11228 off,
11229 render(json!({"reasoning_effort": "none"})),
11230 "the vLLM spelling must render byte-identically to the OpenAI spelling"
11231 );
11232 }
11233
11234 #[test]
11235 fn unknown_chat_template_kwarg_refuses_by_name() {
11236 let build = |extra: serde_json::Value| {
11239 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11240 build_chat_request(
11241 weather_request(extra),
11242 Some(&tool_caps()),
11243 tx,
11244 lanes::Lane::Interactive,
11245 None,
11246 )
11247 };
11248 let refusal = |extra: serde_json::Value, why: &str| -> String {
11249 build(extra).err().unwrap_or_else(|| panic!("{why}"))
11250 };
11251 let err = refusal(
11252 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
11253 "an unimplementable template kwarg must not be accepted",
11254 );
11255 assert!(
11256 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
11257 "the refusal must name the offending key AND the supported one: {err}"
11258 );
11259 let err = refusal(
11260 json!({"chat_template_kwargs": "enable_thinking=false"}),
11261 "a non-object chat_template_kwargs must not be accepted",
11262 );
11263 assert!(
11264 err.contains("must be an object"),
11265 "refusal must say what shape is expected: {err}"
11266 );
11267 let err = refusal(
11268 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
11269 "a stringly-typed switch must not be accepted",
11270 );
11271 assert!(
11272 err.contains("true or false"),
11273 "refusal must name the expected type: {err}"
11274 );
11275 let plan = build(json!({"chat_template_kwargs": null}))
11277 .expect("null chat_template_kwargs is the unset case");
11278 assert_eq!(plan.request.think, ThinkMode::Default);
11279 }
11280
11281 const Q38_TMPL: &str =
11296 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
11297
11298 fn render_with(
11301 tmpl: &str,
11302 caps: &ModelCaps,
11303 extra: serde_json::Value,
11304 default_effort: Option<&str>,
11305 ) -> Result<String, String> {
11306 let mut payload = serde_json::json!({
11307 "model": "m",
11308 "messages": [{"role": "user", "content": "hi"}],
11309 });
11310 if let Some(obj) = extra.as_object() {
11311 for (k, v) in obj {
11312 payload[k] = v.clone();
11313 }
11314 }
11315 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11316 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11317 let plan = build_chat_request_with_trace(
11318 req,
11319 Some(caps),
11320 tx,
11321 lanes::Lane::Interactive,
11322 None,
11323 None,
11324 default_effort,
11325 &ModelSamplingDefaults::default(),
11326 )?;
11327 Ok(chat::apply_chat_template_tools_ex(
11328 Some(tmpl),
11329 &plan.request.chat_turns,
11330 true,
11331 &plan.request.tools_json,
11332 &plan.request.tools_struct,
11333 plan.request.think,
11334 plan.request.reasoning_effort.as_deref(),
11335 None,
11336 )
11337 .unwrap())
11338 }
11339
11340 #[test]
11341 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
11342 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
11349 let xhigh = "Reasoning effort is set to xhigh.";
11350 let low = "Reasoning effort is set to low.";
11351 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
11353 assert!(
11354 r(json!({"reasoning_effort": "high"}))
11355 .unwrap()
11356 .contains(xhigh)
11357 );
11358 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
11361 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
11362 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
11364 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
11365 assert_ne!(low_p, high_p);
11366 assert_ne!(low_p, medium);
11367 assert_ne!(high_p, medium);
11368 for alias in ["xhigh", "max", "ultra"] {
11371 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
11372 }
11373 assert_eq!(r(json!({})).unwrap(), high_p);
11376 assert_eq!(
11379 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
11380 medium
11381 );
11382 let off = r(json!({"reasoning_effort": "none"})).unwrap();
11385 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
11386 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
11387 }
11388
11389 #[test]
11390 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
11391 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
11404focused, moving directly to the conclusion without unnecessary elaboration.";
11405 let expected = format!(
11406 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
11407 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
11408 );
11409 let after_fix = render_with(
11411 Q38_TMPL,
11412 &ladder_caps(),
11413 json!({"reasoning_effort": "low"}),
11414 None,
11415 )
11416 .unwrap();
11417 assert_eq!(
11418 after_fix, expected,
11419 "the shipped prompt for reasoning_effort:\"low\""
11420 );
11421 const ORNITH_TMPL: &str = include_str!(
11424 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11425 );
11426 let on_deployed_binary = render_with(
11427 ORNITH_TMPL,
11428 &tool_caps(),
11429 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
11430 {"role": "user", "content": "hi"}]}),
11431 None,
11432 )
11433 .unwrap();
11434 assert_eq!(
11435 on_deployed_binary, expected,
11436 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
11437 level, or its reasoning-volume numbers do not describe the shipped prompt"
11438 );
11439 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
11442 assert!(
11443 !ladderless_unset.contains("Reasoning effort is set to"),
11444 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
11445 );
11446 assert_eq!(
11447 ladderless_unset,
11448 render_with(
11449 Q38_TMPL,
11450 &ladder_caps(),
11451 json!({"reasoning_effort": "medium"}),
11452 None
11453 )
11454 .unwrap(),
11455 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
11456 );
11457 }
11458
11459 #[test]
11460 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
11461 let off = render_with(
11468 Q38_TMPL,
11469 &ladder_caps(),
11470 json!({"reasoning_effort": "none"}),
11471 None,
11472 )
11473 .unwrap();
11474 for extra in [
11475 json!({"include_reasoning": false}),
11476 json!({"reasoning": {"exclude": true}}),
11477 ] {
11478 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
11479 assert!(
11480 got.ends_with("<think>\n\n</think>\n\n"),
11481 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
11482 );
11483 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
11484 }
11485 for extra in [
11490 json!({"enable_thinking": true, "include_reasoning": false}),
11491 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
11492 json!({"reasoning": {"enabled": true, "exclude": true}}),
11493 ] {
11494 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
11495 .err()
11496 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
11497 assert!(e.contains("contradictory"), "{extra}: {e}");
11498 assert!(
11499 e.contains("include_reasoning") || e.contains("exclude"),
11500 "{extra}: the refusal must name the suppression field the caller sent: {e}"
11501 );
11502 }
11503 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
11506 for extra in [
11507 json!({"include_reasoning": true}),
11508 json!({"reasoning": {"exclude": false}}),
11509 ] {
11510 assert_eq!(
11511 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
11512 dflt,
11513 "{extra} must not perturb the model's default"
11514 );
11515 }
11516 let switchless = ModelCaps {
11520 think_switch: false,
11521 ..tool_caps()
11522 };
11523 let err = render_with(
11524 Q38_TMPL,
11525 &switchless,
11526 json!({"include_reasoning": false}),
11527 None,
11528 )
11529 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
11530 assert!(err.contains("cannot disable reasoning"), "{err}");
11531 }
11532
11533 #[test]
11534 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
11535 let build = |extra: serde_json::Value| {
11536 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11537 build_chat_request(
11538 weather_request(extra),
11539 Some(&ladder_caps()),
11540 tx,
11541 lanes::Lane::Interactive,
11542 None,
11543 )
11544 };
11545 let err = |extra: serde_json::Value, why: &str| -> String {
11546 build(extra).err().unwrap_or_else(|| panic!("{why}"))
11547 };
11548 let e = err(
11552 json!({"reasoning": {"max_tokens": 1024}}),
11553 "reasoning.max_tokens must not be accepted-and-ignored",
11554 );
11555 assert!(e.contains("reasoning.max_tokens"), "{e}");
11556 assert!(e.contains("ONE output budget"), "{e}");
11557 for extra in [
11562 json!({"reasoning": {"max_tokens": null}}),
11563 json!({"reasoning": {"banana": null}}),
11564 ] {
11565 let e = err(
11566 extra.clone(),
11567 "a null-valued unhonourable key must still refuse",
11568 );
11569 assert!(
11570 e.contains("max_tokens") || e.contains("banana"),
11571 "{extra}: {e}"
11572 );
11573 }
11574 let e = err(
11576 json!({"reasoning": {"budget": 5}}),
11577 "an unknown reasoning key must not be accepted",
11578 );
11579 assert!(
11580 e.contains("reasoning.budget") && e.contains("enabled"),
11581 "{e}"
11582 );
11583 for (extra, want) in [
11587 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
11588 (json!({"reasoning": {"exclude": 1}}), "true or false"),
11589 (json!({"reasoning": {"effort": 3}}), "must be a string"),
11590 ] {
11591 let e = err(
11592 extra.clone(),
11593 "a wrong-typed reasoning key must not be ignored",
11594 );
11595 assert!(e.contains(want), "{extra}: {e}");
11596 }
11597 for extra in [
11602 json!({"reasoning": {"enabled": true}}),
11603 json!({"reasoning": {"effort": "low"}}),
11604 json!({"reasoning": {"exclude": false}}),
11605 json!({"reasoning": null}),
11606 json!({"reasoning": {"effort": null}}),
11607 json!({"reasoning": {"enabled": null, "exclude": null}}),
11608 ] {
11609 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
11610 }
11611 }
11612
11613 #[test]
11614 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
11615 const ORNITH_TMPL: &str = include_str!(
11623 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11624 );
11625 let explicit_on = render_with(
11629 ORNITH_TMPL,
11630 &tool_caps(),
11631 json!({"reasoning": {"enabled": true}}),
11632 None,
11633 )
11634 .unwrap();
11635 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
11636 for extra in [
11637 json!({"reasoning_effort": "low"}),
11638 json!({"reasoning_effort": "medium"}),
11639 json!({"reasoning_effort": "high"}),
11640 json!({"reasoning_effort": "xhigh"}),
11642 json!({"reasoning": {"effort": "xhigh"}}),
11643 ] {
11644 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11645 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
11646 assert_eq!(
11647 got, explicit_on,
11648 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
11649 documented translation, not a decorative accept"
11650 );
11651 }
11652 for extra in [
11654 json!({}),
11655 json!({"reasoning_effort": "none"}),
11656 json!({"reasoning_effort": "minimal"}),
11657 json!({"enable_thinking": false}),
11658 ] {
11659 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11660 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
11661 }
11662 let minimal = render_with(
11665 ORNITH_TMPL,
11666 &tool_caps(),
11667 json!({"reasoning_effort": "minimal"}),
11668 None,
11669 )
11670 .unwrap();
11671 assert!(
11672 minimal.ends_with("<think>\n\n</think>\n\n"),
11673 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
11674 );
11675 let ladder_low = render_with(
11678 Q38_TMPL,
11679 &ladder_caps(),
11680 json!({"reasoning_effort": "low"}),
11681 None,
11682 )
11683 .unwrap();
11684 assert!(
11685 ladder_low.contains("Reasoning effort is set to low."),
11686 "{ladder_low:?}"
11687 );
11688 assert_ne!(
11689 ladder_low,
11690 render_with(
11691 Q38_TMPL,
11692 &ladder_caps(),
11693 json!({"reasoning_effort": "high"}),
11694 None
11695 )
11696 .unwrap(),
11697 "the ladder model's rungs stay distinct prompts"
11698 );
11699 }
11700
11701 #[test]
11702 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
11703 let render_chat = |body: serde_json::Value| -> Result<String, String> {
11714 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11715 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11716 let plan = build_chat_request(
11717 req,
11718 Some(&ladder_caps()),
11719 tx,
11720 lanes::Lane::Interactive,
11721 None,
11722 )?;
11723 Ok(chat::apply_chat_template_tools_ex(
11724 Some(Q38_TMPL),
11725 &plan.request.chat_turns,
11726 true,
11727 &plan.request.tools_json,
11728 &plan.request.tools_struct,
11729 plan.request.think,
11730 plan.request.reasoning_effort.as_deref(),
11731 None,
11732 )
11733 .unwrap())
11734 };
11735 for (intent, chat_body, responses_body, messages_body) in [
11740 (
11741 "reasoning OFF",
11742 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11743 "reasoning_effort": "none"}),
11744 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
11745 json!({"model": "m", "max_tokens": 16,
11746 "messages": [{"role": "user", "content": "hi"}],
11747 "thinking": {"type": "disabled"}}),
11748 ),
11749 (
11750 "reasoning ON at the top rung",
11751 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11752 "reasoning_effort": "xhigh"}),
11753 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
11754 json!({"model": "m", "max_tokens": 16,
11755 "messages": [{"role": "user", "content": "hi"}],
11756 "output_config": {"effort": "xhigh"}}),
11757 ),
11758 (
11759 "reasoning ON at the bottom rung",
11760 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11761 "reasoning_effort": "low"}),
11762 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
11763 json!({"model": "m", "max_tokens": 16,
11764 "messages": [{"role": "user", "content": "hi"}],
11765 "output_config": {"effort": "low"}}),
11766 ),
11767 (
11768 "the model's own default",
11769 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11770 json!({"model": "m", "input": "hi"}),
11771 json!({"model": "m", "max_tokens": 16,
11772 "messages": [{"role": "user", "content": "hi"}]}),
11773 ),
11774 ] {
11775 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
11776 let via_responses = responses_api::translate(&responses_body)
11777 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
11778 let via_messages = anthropic::translate(&messages_body)
11779 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
11780 for (surface, translated) in [
11781 ("/v1/responses", via_responses),
11782 ("/v1/messages", via_messages),
11783 ] {
11784 let got = render_chat(translated)
11785 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
11786 assert_eq!(
11787 got, chat,
11788 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
11789 /v1/chat/completions — the parameter is honoured on one format and not \
11790 the other"
11791 );
11792 }
11793 }
11794 let switchless = ModelCaps {
11797 think_switch: false,
11798 ..ladder_caps()
11799 };
11800 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
11801 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11802 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11803 let plan =
11804 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
11805 Ok(format!("{:?}", plan.request.think))
11806 };
11807 for (surface, body) in [
11808 (
11809 "/v1/responses",
11810 responses_api::translate(&json!({
11811 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
11812 .unwrap(),
11813 ),
11814 (
11815 "/v1/messages",
11816 anthropic::translate(&json!({
11817 "model": "m", "max_tokens": 16,
11818 "messages": [{"role": "user", "content": "hi"}],
11819 "thinking": {"type": "disabled"}}))
11820 .unwrap(),
11821 ),
11822 ] {
11823 let err = render_switchless(body)
11824 .err()
11825 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
11826 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
11827 }
11828 }
11829
11830 #[test]
11831 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
11832 let build = |extra: serde_json::Value| {
11840 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11841 build_chat_request(
11842 weather_request(extra),
11843 Some(&ladder_caps()),
11844 tx,
11845 lanes::Lane::Interactive,
11846 None,
11847 )
11848 };
11849 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
11850 .expect("preserve_thinking:true is the vendor default the renderer implements");
11851 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
11852 .err()
11853 .expect("preserve_thinking:false (the strip arm) must refuse");
11854 assert!(e.contains("preserve_thinking"), "{e}");
11855 assert!(e.contains("strip"), "{e}");
11856 assert_eq!(
11859 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
11860 .unwrap()
11861 .request
11862 .think,
11863 ThinkMode::NoThink
11864 );
11865 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
11867 .err()
11868 .expect("a stringly-typed preserve_thinking must not be accepted");
11869 assert!(e.contains("true or false"), "{e}");
11870 }
11871
11872 #[test]
11873 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
11874 let dsv4_caps = ModelCaps {
11880 qwen_think: true,
11881 think_switch: false,
11882 dsv4: true,
11883 ..tool_caps()
11884 };
11885 for extra in [
11886 json!({"reasoning_effort": "none"}),
11887 json!({"reasoning": {"enabled": false}}),
11888 json!({"enable_thinking": false}),
11889 json!({"include_reasoning": false}),
11890 ] {
11891 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11892 let plan = build_chat_request(
11893 weather_request(extra.clone()),
11894 Some(&dsv4_caps),
11895 tx,
11896 lanes::Lane::Interactive,
11897 None,
11898 )
11899 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
11900 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
11901 }
11902 }
11903
11904 #[test]
11905 fn contradictory_think_switches_refuse_instead_of_picking_one() {
11906 let build = |extra: serde_json::Value| {
11909 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11910 build_chat_request(
11911 weather_request(extra),
11912 Some(&tool_caps()),
11913 tx,
11914 lanes::Lane::Interactive,
11915 None,
11916 )
11917 };
11918 for extra in [
11919 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
11920 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
11921 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
11922 ] {
11923 match build(extra.clone()) {
11924 Err(err) => assert!(
11925 err.contains("contradictory"),
11926 "the refusal must say the switches contradict: {err}"
11927 ),
11928 Ok(plan) => panic!(
11929 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
11930 plan.request.think
11931 ),
11932 }
11933 }
11934 for extra in [
11936 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
11937 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
11938 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
11939 ] {
11940 build(extra.clone())
11941 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
11942 }
11943 }
11944
11945 #[test]
11946 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
11947 let switchless = ModelCaps {
11952 tools_branch: true,
11953 qwen_think: true,
11954 think_switch: false,
11955 chat_ok: true,
11956 ..Default::default()
11957 };
11958 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
11959 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11960 build_chat_request_with_trace(
11961 weather_request(extra),
11962 Some(caps),
11963 tx,
11964 lanes::Lane::Interactive,
11965 None,
11966 None,
11967 default_effort,
11968 &ModelSamplingDefaults::default(),
11969 )
11970 };
11971 for extra in [
11972 json!({"reasoning_effort": "none"}),
11973 json!({"reasoning_effort": "minimal"}),
11974 json!({"reasoning": {"enabled": false}}),
11975 json!({"enable_thinking": false}),
11976 json!({"chat_template_kwargs": {"enable_thinking": false}}),
11977 ] {
11978 let err = build(extra.clone(), &switchless, None)
11979 .err()
11980 .unwrap_or_else(|| {
11981 panic!(
11982 "{extra} on a switchless think template must not be accepted-and-ignored"
11983 )
11984 });
11985 assert!(
11986 err.contains("cannot disable reasoning"),
11987 "the refusal must say the model cannot disable reasoning: {err}"
11988 );
11989 }
11990 for (extra, default_effort) in [
11994 (json!({}), None),
11995 (json!({"reasoning_effort": "high"}), None),
11998 (json!({"reasoning": {"enabled": true}}), None),
11999 (json!({"enable_thinking": true}), None),
12000 (json!({}), Some("none")),
12001 (json!({}), Some("minimal")),
12002 (json!({}), Some("high")),
12003 ] {
12004 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
12005 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
12006 });
12007 }
12008 assert_eq!(
12011 build(json!({"enable_thinking": false}), &tool_caps(), None)
12012 .unwrap()
12013 .request
12014 .think,
12015 ThinkMode::NoThink
12016 );
12017 }
12018
12019 #[test]
12020 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
12021 let gemma_caps = ModelCaps {
12028 tools_branch: true,
12029 chat_ok: true,
12030 gemma_think: true,
12031 instruct_type: Some("gemma".into()),
12032 ..Default::default()
12033 };
12034 let render =
12035 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
12036 let mut payload = serde_json::json!({
12037 "model": "google/gemma-4-31b-it",
12038 "messages": [{"role": "user", "content": "Weather in Paris?"}],
12039 });
12040 if let Some(obj) = extra.as_object() {
12041 for (k, v) in obj {
12042 payload[k] = v.clone();
12043 }
12044 }
12045 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12046 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12047 let plan = build_chat_request_with_trace(
12048 req,
12049 Some(&gemma_caps),
12050 tx,
12051 lanes::Lane::Interactive,
12052 None,
12053 None,
12054 default_effort,
12055 &ModelSamplingDefaults::default(),
12056 )
12057 .unwrap();
12058 chat::apply_chat_template_tools_ex(
12059 Some(tmpl),
12060 &plan.request.chat_turns,
12061 true,
12062 &plan.request.tools_json,
12063 &plan.request.tools_struct,
12064 plan.request.think,
12065 plan.request.reasoning_effort.as_deref(),
12066 None, )
12068 .unwrap()
12069 };
12070 let official = gemma_template("official");
12071 let unset_with_knob = render(&official, json!({}), Some("high"));
12072 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
12073 assert_eq!(
12074 unset_with_knob, explicit_on,
12075 "knob render must be byte-identical to the explicit think-on render"
12076 );
12077 assert!(
12078 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
12079 "think-on injects the <|think|> system token: {unset_with_knob:?}"
12080 );
12081 assert!(
12082 unset_with_knob.ends_with("<|turn>model\n"),
12083 "think-on generation turn is OPEN: {unset_with_knob:?}"
12084 );
12085 let explicit_off_with_knob =
12089 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
12090 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
12091 assert_eq!(explicit_off_with_knob, explicit_off);
12092 assert!(
12093 !explicit_off_with_knob.contains("<|think|>")
12094 && explicit_off_with_knob.ends_with("<|turn>model\n"),
12095 "explicit off keeps the official template's thinking-off bytes: \
12096 {explicit_off_with_knob:?}"
12097 );
12098 let unset_no_knob = render(&official, json!({}), None);
12100 assert_eq!(
12101 unset_no_knob, explicit_off,
12102 "knobless unset stays the template's own thinking-off default"
12103 );
12104 assert_ne!(unset_no_knob, unset_with_knob);
12105 let qat = gemma_template("qat");
12108 assert!(
12109 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
12110 "QAT knobless unset keeps the closed-channel default"
12111 );
12112 assert_eq!(
12113 render(&qat, json!({}), Some("high")),
12114 render(&qat, json!({"reasoning_effort": "high"}), None),
12115 "QAT knob render must equal the explicit think-on render"
12116 );
12117 }
12118
12119 #[test]
12120 fn default_reasoning_effort_is_validated_at_metadata_load() {
12121 let parsed = OpenRouterMetadataFile::from_toml(
12123 r#"
12124[models.g]
12125default_reasoning_effort = "high"
12126"#,
12127 )
12128 .unwrap();
12129 assert_eq!(
12130 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
12131 Some("high")
12132 );
12133 let err = OpenRouterMetadataFile::from_toml(
12134 r#"
12135[models.g]
12136default_reasoning_effort = "always"
12137"#,
12138 )
12139 .unwrap_err();
12140 assert!(err.contains("default_reasoning_effort"), "{err}");
12141 }
12142
12143 #[test]
12144 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
12145 let effort_caps = ModelCaps {
12157 effort_levels: true,
12158 think_switch: false,
12159 ..tool_caps()
12160 };
12161 for (extra, want) in [
12162 (json!({}), None),
12163 (json!({"reasoning_effort": "low"}), Some("low")),
12164 (json!({"reasoning_effort": "medium"}), Some("medium")),
12165 (json!({"reasoning_effort": "high"}), Some("high")),
12166 (json!({"reasoning": {"effort": "high"}}), Some("high")),
12167 (json!({"reasoning_effort": "xhigh"}), Some("high")),
12169 (json!({"reasoning": {"effort": "max"}}), Some("high")),
12170 ] {
12171 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12172 let plan = build_chat_request(
12173 weather_request(extra.clone()),
12174 Some(&effort_caps),
12175 tx,
12176 lanes::Lane::Interactive,
12177 None,
12178 )
12179 .unwrap();
12180 assert_eq!(
12181 plan.request.reasoning_effort.as_deref(),
12182 want,
12183 "extra={extra}"
12184 );
12185 }
12186 for extra in [
12193 json!({"reasoning_effort": "none"}),
12194 json!({"reasoning_effort": "minimal"}),
12195 json!({"reasoning": {"enabled": false}}),
12196 json!({"enable_thinking": false}),
12197 json!({"include_reasoning": false}),
12198 ] {
12199 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12200 let err = build_chat_request(
12201 weather_request(extra.clone()),
12202 Some(&effort_caps),
12203 tx,
12204 lanes::Lane::Interactive,
12205 None,
12206 )
12207 .err()
12208 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
12209 assert!(
12210 err.contains("cannot disable reasoning"),
12211 "extra={extra}: {err}"
12212 );
12213 }
12214 for extra in [
12221 json!({"reasoning_effort": "high"}),
12222 json!({"reasoning": {"effort": "low"}}),
12223 ] {
12224 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12225 let plan = build_chat_request(
12226 weather_request(extra.clone()),
12227 Some(&tool_caps()),
12228 tx,
12229 lanes::Lane::Interactive,
12230 None,
12231 )
12232 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
12233 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
12234 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
12235 }
12236 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12238 let plan = build_chat_request(
12239 weather_request(json!({})),
12240 Some(&tool_caps()),
12241 tx,
12242 lanes::Lane::Interactive,
12243 None,
12244 )
12245 .unwrap();
12246 assert_eq!(plan.request.reasoning_effort, None);
12247 }
12248
12249 #[test]
12250 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
12251 let payload = serde_json::json!({
12252 "model": "m",
12253 "messages": [
12254 {"role": "user", "content": "Weather in Paris?"},
12255 {"role": "assistant", "content": null, "tool_calls": [
12256 {"id": "call_x", "type": "function", "function": {
12257 "name": "get_weather",
12258 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
12259 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
12260 ],
12261 });
12262 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12263 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12264 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
12265 .unwrap();
12266 let turns = &plan.request.chat_turns;
12267 assert_eq!(turns[1].tool_calls.len(), 1);
12268 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
12269 assert_eq!(
12270 turns[1].tool_calls[0].params,
12271 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
12272 );
12273 assert_eq!(turns[2].role, "tool");
12274 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
12275 let mut p = plan
12278 .parser
12279 .expect("think-open chat arms the reasoning splitter");
12280 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
12281 assert_eq!(
12282 pieces,
12283 vec![
12284 Piece::Reasoning("thought".into()),
12285 Piece::Content("answer <tool_call> is prose here".into()),
12286 ]
12287 );
12288 }
12289
12290 #[tokio::test]
12291 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
12292 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12293 tx.send(Event::Token {
12294 id: 1,
12295 text: "plan</think>\n\n".into(),
12296 })
12297 .unwrap();
12298 tx.send(Event::Token {
12299 id: 2,
12300 text: "<tool_call>\n<function=get_weather>\n\
12301<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
12302 .into(),
12303 })
12304 .unwrap();
12305 tx.send(Event::Done {
12306 stop_reason: "Eos".into(),
12307 n_tokens: 2,
12308 n_prompt: 40,
12309 n_cached: 0,
12310 elapsed_s: 0.5,
12311 spec: None,
12312 })
12313 .unwrap();
12314 drop(tx);
12315 let parser = ToolStreamParser::new(HashMap::new(), true);
12316 let response = blocking_response(
12317 rx,
12318 "m".into(),
12319 true,
12320 Vec::new(),
12321 Some(parser),
12322 Envelope::new(true),
12323 )
12324 .await;
12325 assert_eq!(response.status(), StatusCode::OK);
12326 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12327 .await
12328 .unwrap();
12329 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12330 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
12331 assert_eq!(
12334 payload["choices"][0]["message"]["content"],
12335 serde_json::Value::Null
12336 );
12337 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
12338 assert_eq!(
12339 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
12340 "plan"
12341 );
12342 let call = &payload["choices"][0]["message"]["tool_calls"][0];
12343 assert_eq!(call["type"], "function");
12344 assert_eq!(call["function"]["name"], "get_weather");
12345 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
12346 assert_eq!(payload["usage"]["prompt_tokens"], 40);
12349 assert_eq!(payload["usage"]["completion_tokens"], 2);
12350 assert_eq!(payload["usage"]["total_tokens"], 42);
12351 assert_eq!(
12352 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12353 0
12354 );
12355 }
12356
12357 #[test]
12358 fn cache_salt_plumbs_to_the_worker_namespace() {
12359 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12361 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
12362 }))
12363 .unwrap();
12364 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12365 assert_eq!(
12366 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12367 "tenant-a"
12368 );
12369
12370 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12371 "model": "m", "messages": [{"role": "user", "content": "task"}],
12372 "cache_salt": "tenant-b"
12373 }))
12374 .unwrap();
12375 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12376 assert_eq!(
12377 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12378 .unwrap()
12379 .request
12380 .cache_ns,
12381 "tenant-b"
12382 );
12383
12384 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12386 "model": "m", "prompt": "task"
12387 }))
12388 .unwrap();
12389 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12390 assert_eq!(
12391 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12392 ""
12393 );
12394 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12395 "model": "m", "messages": [{"role": "user", "content": "task"}]
12396 }))
12397 .unwrap();
12398 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12399 assert_eq!(
12400 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12401 .unwrap()
12402 .request
12403 .cache_ns,
12404 ""
12405 );
12406 }
12407
12408 #[test]
12409 fn cache_salt_validation_rejects_oversized_value() {
12410 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
12411 assert_eq!(
12412 validate_cache_namespace(&salt, false),
12413 Err("cache_salt must be at most 64 bytes")
12414 );
12415 }
12416
12417 #[test]
12418 fn cache_salt_validation_rejects_reserved_open_namespace() {
12419 let salt = Some("t:acme\u{1f}private".to_string());
12420 assert_eq!(
12421 validate_cache_namespace(&salt, false),
12422 Err("cache_salt must not use the reserved t: prefix without a keyring")
12423 );
12424 }
12425
12426 #[test]
12427 fn cache_salt_validation_accepts_normal_value() {
12428 let salt = Some("tenant-A_7.c2VjcmV0LXNjb3Bl+/=".to_string());
12429 assert_eq!(
12430 validate_cache_namespace(&salt, false).unwrap(),
12431 salt.unwrap()
12432 );
12433 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
12434 let max = Some("a".repeat(CACHE_SALT_MAX_BYTES));
12435 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max.unwrap());
12436 }
12437
12438 #[test]
12439 fn cache_salt_validation_rejects_unsupported_characters() {
12440 let salt = Some("tenant salt".to_string());
12441 assert_eq!(
12442 validate_cache_namespace(&salt, false),
12443 Err("cache_salt contains unsupported characters")
12444 );
12445 }
12446
12447 #[test]
12448 fn affinity_key_honors_both_client_conventions_in_priority_order() {
12449 use axum::http::HeaderMap;
12450 let hdr = |v: &str| {
12451 let mut h = HeaderMap::new();
12452 h.insert("x-session-id", v.parse().unwrap());
12453 h
12454 };
12455 let empty = HeaderMap::new();
12456 let s = |v: &str| Some(v.to_string());
12457 assert_eq!(affinity_key(&s("explicit"), &None, &empty), s("explicit"));
12459 assert_eq!(
12460 affinity_key(&None, &s("openai-user"), &empty),
12461 s("openai-user")
12462 );
12463 assert_eq!(affinity_key(&None, &None, &hdr("hdr-id")), s("hdr-id"));
12464 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")), s("a"));
12467 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")), s("b"));
12468 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")), None);
12471 assert_eq!(affinity_key(&s(""), &s("real"), &empty), s("real"));
12472 assert_eq!(affinity_key(&s(" padded "), &None, &empty), s("padded"));
12474 assert_eq!(affinity_key(&None, &None, &empty), None);
12476 }
12477
12478 #[test]
12479 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
12480 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12481 "model": "m", "prompt": "task", "session_id": "conv-1"
12482 }))
12483 .unwrap();
12484 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12485 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12486 assert_eq!(
12487 build_request(&req, tx, lanes::Lane::Interactive, key)
12488 .affinity
12489 .as_deref(),
12490 Some("conv-1")
12491 );
12492 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12494 "model": "m", "messages": [{"role": "user", "content": "task"}],
12495 "user": "conv-2"
12496 }))
12497 .unwrap();
12498 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12499 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12500 assert_eq!(
12501 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
12502 .unwrap()
12503 .request
12504 .affinity
12505 .as_deref(),
12506 Some("conv-2")
12507 );
12508 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12510 "model": "m", "prompt": "task"
12511 }))
12512 .unwrap();
12513 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12514 assert!(
12515 build_request(&req, tx, lanes::Lane::Interactive, None)
12516 .affinity
12517 .is_none()
12518 );
12519 }
12520
12521 async fn sse_data_lines(resp: Response) -> Vec<String> {
12523 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12524 .await
12525 .unwrap();
12526 String::from_utf8(bytes.to_vec())
12527 .unwrap()
12528 .lines()
12529 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
12530 .collect()
12531 }
12532
12533 #[tokio::test]
12534 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
12535 let feed = |think: bool| {
12541 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12542 let body = if think {
12543 "a plan</think>\n\nanswer"
12544 } else {
12545 "answer"
12546 };
12547 tx.send(Event::Token {
12548 id: 1,
12549 text: body.into(),
12550 })
12551 .unwrap();
12552 tx.send(Event::Done {
12553 stop_reason: "Eos".into(),
12554 n_tokens: 3,
12555 n_prompt: 10,
12556 n_cached: 0,
12557 elapsed_s: 0.1,
12558 spec: None,
12559 })
12560 .unwrap();
12561 drop(tx);
12562 rx
12563 };
12564 let resp = blocking_response(
12566 feed(true),
12567 "m".into(),
12568 true,
12569 Vec::new(),
12570 Some(ToolStreamParser::reasoning_only()),
12571 Envelope::new(true),
12572 )
12573 .await;
12574 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12575 .await
12576 .unwrap();
12577 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12578 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
12579 assert_eq!(
12580 v["choices"][0]["message"]["reasoning_details"][0]["text"],
12581 "a plan"
12582 );
12583 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12584 let resp = blocking_response(
12587 feed(false),
12588 "m".into(),
12589 true,
12590 Vec::new(),
12591 None,
12592 Envelope::new(true),
12593 )
12594 .await;
12595 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12596 .await
12597 .unwrap();
12598 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12599 assert!(
12600 v["choices"][0]["message"].get("reasoning").is_none(),
12601 "a reasoning-off response must carry no reasoning field: {v}"
12602 );
12603 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12604 let resp = sse_response(
12606 feed(true),
12607 "m".into(),
12608 true,
12609 Some(ToolStreamParser::reasoning_only()),
12610 Envelope::new(true),
12611 Vec::new(),
12612 None,
12613 )
12614 .into_response();
12615 let lines = sse_data_lines(resp).await;
12616 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12617 .iter()
12618 .map(|l| serde_json::from_str(l).unwrap())
12619 .collect();
12620 let reasoning: String = chunks
12621 .iter()
12622 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
12623 .collect();
12624 assert_eq!(
12625 reasoning, "a plan",
12626 "think text must stream as delta.reasoning"
12627 );
12628 let content: String = chunks
12629 .iter()
12630 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
12631 .collect();
12632 assert_eq!(content, "answer", "content must exclude the think segment");
12633 let resp = sse_response(
12635 feed(false),
12636 "m".into(),
12637 true,
12638 None,
12639 Envelope::new(true),
12640 Vec::new(),
12641 None,
12642 )
12643 .into_response();
12644 let lines = sse_data_lines(resp).await;
12645 for l in &lines[..lines.len() - 1] {
12646 let c: serde_json::Value = serde_json::from_str(l).unwrap();
12647 assert!(
12648 c["choices"][0]["delta"].get("reasoning").is_none(),
12649 "a reasoning-off stream must carry no reasoning deltas: {c}"
12650 );
12651 }
12652 }
12653
12654 #[tokio::test]
12655 async fn stream_chunks_carry_envelope_and_first_delta_role() {
12656 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12657 tx.send(Event::Token {
12658 id: 1,
12659 text: "he".into(),
12660 })
12661 .unwrap();
12662 tx.send(Event::Token {
12663 id: 2,
12664 text: "llo".into(),
12665 })
12666 .unwrap();
12667 tx.send(Event::Done {
12668 stop_reason: "Eos".into(),
12669 n_tokens: 2,
12670 n_prompt: 10,
12671 n_cached: 0,
12672 elapsed_s: 0.1,
12673 spec: None,
12674 })
12675 .unwrap();
12676 drop(tx);
12677 let resp = sse_response(
12678 rx,
12679 "m".into(),
12680 true,
12681 None,
12682 Envelope::new(true),
12683 Vec::new(),
12684 None,
12685 )
12686 .into_response();
12687 let lines = sse_data_lines(resp).await;
12688 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12689 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12690 .iter()
12691 .map(|l| serde_json::from_str(l).unwrap())
12692 .collect();
12693 let id = chunks[0]["id"].as_str().unwrap().to_string();
12695 assert!(id.starts_with("chatcmpl-"));
12696 for c in &chunks {
12697 assert_eq!(c["id"], id.as_str());
12698 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
12699 assert!(
12700 c["system_fingerprint"]
12701 .as_str()
12702 .unwrap()
12703 .starts_with("memra-")
12704 );
12705 assert_eq!(c["object"], "chat.completion.chunk");
12706 }
12707 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
12709 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
12710 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
12711 let fin = chunks.last().unwrap();
12713 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
12714 assert_eq!(fin["usage"]["prompt_tokens"], 10);
12715 }
12716
12717 #[tokio::test]
12718 async fn stream_token_events_equal_usage_on_every_finish_path() {
12719 for (stop_reason, expected_finish) in [
12720 ("Eos", "stop"),
12721 ("Callback", "stop"),
12722 ("MaxNew", "length"),
12723 ("ContextFull", "length"),
12724 ] {
12725 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12726 tx.send(Event::Token {
12729 id: 248_046,
12730 text: String::new(),
12731 })
12732 .unwrap();
12733 tx.send(Event::Done {
12734 stop_reason: stop_reason.into(),
12735 n_tokens: 1,
12736 n_prompt: 8,
12737 n_cached: 8,
12738 elapsed_s: 0.1,
12739 spec: None,
12740 })
12741 .unwrap();
12742 drop(tx);
12743
12744 let resp = sse_response(
12745 rx,
12746 "m".into(),
12747 true,
12748 None,
12749 Envelope::new(true),
12750 Vec::new(),
12751 None,
12752 )
12753 .into_response();
12754 let lines = sse_data_lines(resp).await;
12755 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12756 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12757 .iter()
12758 .map(|line| serde_json::from_str(line).unwrap())
12759 .collect();
12760 let token_events = chunks
12761 .iter()
12762 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
12763 .count();
12764 let terminal = chunks.last().unwrap();
12765 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
12766 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
12767 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
12768 }
12769 }
12770
12771 #[tokio::test]
12772 async fn stream_excludes_stop_text_like_non_stream_does() {
12773 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12777 tx.send(Event::Token {
12778 id: 1,
12779 text: "answer\nPro".into(),
12780 })
12781 .unwrap();
12782 tx.send(Event::Token {
12783 id: 2,
12784 text: "blem: leaked prompt".into(),
12785 })
12786 .unwrap();
12787 tx.send(Event::Done {
12788 stop_reason: "Callback".into(),
12789 n_tokens: 2,
12790 n_prompt: 8,
12791 n_cached: 0,
12792 elapsed_s: 0.1,
12793 spec: None,
12794 })
12795 .unwrap();
12796 drop(tx);
12797 let resp = sse_response(
12798 rx,
12799 "m".into(),
12800 true,
12801 None,
12802 Envelope::new(true),
12803 vec!["Problem:".into()],
12804 None,
12805 )
12806 .into_response();
12807 let lines = sse_data_lines(resp).await;
12808 let content: String = lines
12809 .iter()
12810 .filter(|l| *l != "[DONE]")
12811 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12812 .filter_map(|c| {
12813 c["choices"][0]["delta"]["content"]
12814 .as_str()
12815 .map(str::to_string)
12816 })
12817 .collect();
12818 assert_eq!(content, "answer\n");
12819
12820 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12822 tx.send(Event::Token {
12823 id: 1,
12824 text: "ends in Pro".into(),
12825 })
12826 .unwrap();
12827 tx.send(Event::Done {
12828 stop_reason: "Eos".into(),
12829 n_tokens: 1,
12830 n_prompt: 8,
12831 n_cached: 0,
12832 elapsed_s: 0.1,
12833 spec: None,
12834 })
12835 .unwrap();
12836 drop(tx);
12837 let resp = sse_response(
12838 rx,
12839 "m".into(),
12840 true,
12841 None,
12842 Envelope::new(true),
12843 vec!["Problem:".into()],
12844 None,
12845 )
12846 .into_response();
12847 let lines = sse_data_lines(resp).await;
12848 let content: String = lines
12849 .iter()
12850 .filter(|l| *l != "[DONE]")
12851 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12852 .filter_map(|c| {
12853 c["choices"][0]["delta"]["content"]
12854 .as_str()
12855 .map(str::to_string)
12856 })
12857 .collect();
12858 assert_eq!(content, "ends in Pro");
12859 }
12860
12861 #[tokio::test]
12862 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
12863 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12864 tx.send(Event::Error(worker::EngineError::engine("boom")))
12865 .unwrap();
12866 drop(tx);
12867 let resp = sse_response(
12868 rx,
12869 "m".into(),
12870 true,
12871 None,
12872 Envelope::new(true),
12873 Vec::new(),
12874 None,
12875 )
12876 .into_response();
12877 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12878 .await
12879 .unwrap();
12880 let body = String::from_utf8(bytes.to_vec()).unwrap();
12881 assert!(
12883 !body.contains("event: error"),
12884 "named SSE event leaked: {body}"
12885 );
12886 let lines: Vec<&str> = body
12887 .lines()
12888 .filter_map(|l| l.strip_prefix("data: "))
12889 .collect();
12890 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
12891 assert_eq!(err["error"]["message"], "boom");
12892 assert_eq!(err["error"]["type"], "server_error");
12893 assert_eq!(err["error"]["code"], "engine_error");
12894 assert_eq!(lines.last(), Some(&"[DONE]"));
12895 }
12896
12897 #[test]
12898 fn ttft_sse_marker_ignores_keepalive_comments() {
12899 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
12900 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
12901 assert!(is_sse_data_frame(
12902 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
12903 ));
12904 }
12905
12906 #[tokio::test]
12907 async fn error_bodies_use_the_openai_object_shape() {
12908 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12909 tx.send(Event::Error(worker::EngineError::model_not_found(
12910 "unknown model \"x\"",
12911 )))
12912 .unwrap();
12913 drop(tx);
12914 let response =
12915 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
12916 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
12917 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12918 .await
12919 .unwrap();
12920 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12921 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
12923 assert_eq!(payload["error"]["type"], "invalid_request_error");
12924 assert_eq!(payload["error"]["param"], "model");
12925 assert_eq!(payload["error"]["code"], "model_not_found");
12926 }
12927
12928 fn retry_after(resp: &Response) -> Option<String> {
12935 resp.headers()
12936 .get(axum::http::header::RETRY_AFTER)
12937 .and_then(|v| v.to_str().ok())
12938 .map(str::to_string)
12939 }
12940
12941 async fn body_value(resp: Response) -> serde_json::Value {
12944 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12945 .await
12946 .expect("body");
12947 serde_json::from_slice(&bytes).expect("json body")
12948 }
12949
12950 #[test]
12951 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
12952 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
12954 assert_eq!(
12955 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
12956 TIMEOUT_MS_DEFAULT
12957 );
12958 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
12962 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
12963 }
12964 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
12966 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
12967 assert!(err.contains("timeout_ms"), "{err}");
12968 assert!(
12969 err.contains(&TIMEOUT_MS_MIN.to_string())
12970 && err.contains(&TIMEOUT_MS_MAX.to_string()),
12971 "the message must state the range: {err}"
12972 );
12973 assert!(
12974 err.contains("stream"),
12975 "the message must point at streaming for longer work: {err}"
12976 );
12977 }
12978 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
12980 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
12981 assert!(
12982 err.contains("timeout_ms") && err.contains("stream"),
12983 "{err}"
12984 );
12985 }
12986 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
12988 }
12989
12990 #[tokio::test]
12993 async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
12994 let _l = DRAIN_LOCK.lock().unwrap();
12995 let st = fake_worker_state();
12996
12997 let comp = completions(
12998 State(st.clone()),
12999 HeaderMap::new(),
13000 None,
13001 Json(
13002 serde_json::from_value(json!({
13003 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
13004 .unwrap(),
13005 ),
13006 )
13007 .await;
13008 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
13009 let chat = chat_completions(
13010 State(st.clone()),
13011 HeaderMap::new(),
13012 None,
13013 Json(
13014 serde_json::from_value(json!({
13015 "model": "m", "messages": [{"role": "user", "content": "t"}],
13016 "timeout_ms": 90_001}))
13017 .unwrap(),
13018 ),
13019 )
13020 .await;
13021 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
13022 let resp_api = responses_api::responses(
13023 State(st.clone()),
13024 HeaderMap::new(),
13025 None,
13026 axum::body::Bytes::from(
13027 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
13028 ),
13029 )
13030 .await;
13031 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
13032 let msgs = anthropic::messages(
13033 State(st.clone()),
13034 HeaderMap::new(),
13035 None,
13036 axum::body::Bytes::from(
13037 json!({"model": "m", "max_tokens": 16,
13038 "messages": [{"role": "user", "content": "t"}],
13039 "timeout_ms": 90_001})
13040 .to_string(),
13041 ),
13042 )
13043 .await;
13044 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
13045
13046 for (surface, resp) in [
13048 ("/v1/completions", comp),
13049 ("/v1/chat/completions", chat),
13050 ("/v1/responses", resp_api),
13051 ] {
13052 let body = body_value(resp).await;
13053 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
13054 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
13055 let m = body["error"]["message"].as_str().unwrap();
13056 assert!(
13057 m.contains("90000") && m.contains("stream"),
13058 "{surface}: {m}"
13059 );
13060 }
13061 let body = body_value(msgs).await;
13063 assert_eq!(body["error"]["type"], "invalid_request_error");
13064 let m = body["error"]["message"].as_str().unwrap();
13065 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
13066 }
13067
13068 #[tokio::test]
13071 async fn a_non_integer_timeout_ms_is_a_named_400() {
13072 let _l = DRAIN_LOCK.lock().unwrap();
13073 let st = fake_worker_state();
13074 let resp = chat_completions(
13075 State(st),
13076 HeaderMap::new(),
13077 None,
13078 Json(
13079 serde_json::from_value(json!({
13080 "model": "m", "messages": [{"role": "user", "content": "t"}],
13081 "timeout_ms": "30s"}))
13082 .unwrap(),
13083 ),
13084 )
13085 .await;
13086 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13087 let body = body_value(resp).await;
13088 assert_eq!(body["error"]["param"], "timeout_ms");
13089 }
13090
13091 #[tokio::test]
13097 async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
13098 let _l = DRAIN_LOCK.lock().unwrap();
13099 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13104 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
13105 let worker_cancel = cancel_seen.clone();
13106 let health = health::WorkerHealth::new();
13107 let h = health.clone();
13108 std::thread::spawn(move || {
13109 h.mark_ready();
13110 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
13111 worker::release_pending_admit();
13112 worker::release_admission_reservation(req.lane);
13113 let _ = req.tx.send(Event::PromptUsage {
13114 n_prompt: 1,
13115 n_cached: 0,
13116 });
13117 let _ = req.tx.send(Event::Token {
13118 id: 1,
13119 text: "partial".into(),
13120 });
13121 for _ in 0..5_000 {
13124 if req.tx.is_closed() {
13125 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
13126 break;
13127 }
13128 std::thread::sleep(std::time::Duration::from_millis(1));
13129 }
13130 }
13131 });
13132 for _ in 0..2_000 {
13133 if health.live().is_ok() {
13134 break;
13135 }
13136 std::thread::sleep(std::time::Duration::from_millis(1));
13137 }
13138 let mut st = fake_worker_state();
13139 st.cmd_tx = cmd_tx;
13140 st.health = health;
13141 let mock = MockMetering::admit_all();
13142 st.metering = Some(mock.clone());
13143
13144 let resp = chat_completions(
13145 State(st),
13146 HeaderMap::new(),
13147 None,
13148 Json(
13149 serde_json::from_value(json!({
13150 "model": "m", "messages": [{"role": "user", "content": "t"}],
13151 "timeout_ms": 1_000}))
13152 .unwrap(),
13153 ),
13154 )
13155 .await;
13156
13157 assert_eq!(resp.status(), StatusCode::OK);
13162 let body = body_value(resp).await;
13163 assert!(
13164 body["choices"][0]["message"]["content"]
13165 .as_str()
13166 .unwrap()
13167 .contains("partial"),
13168 "the tokens generated before the cut must be delivered: {body}"
13169 );
13170 assert_eq!(body["choices"][0]["finish_reason"], "error");
13174 assert_eq!(
13175 body["choices"][0]["native_finish_reason"],
13176 "deadline_exceeded"
13177 );
13178 assert_eq!(body["error"]["code"], "deadline_exceeded");
13179 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
13180 let message = body["error"]["message"].as_str().unwrap();
13181 assert!(
13182 message.contains("1000") && message.contains("stream"),
13183 "the partial must name the deadline and the streaming alternative: {message}"
13184 );
13185 assert_eq!(body["usage"]["completion_tokens"], 1);
13186
13187 let mut cancelled = false;
13192 for _ in 0..500 {
13193 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
13194 cancelled = true;
13195 break;
13196 }
13197 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
13198 }
13199 assert!(
13200 cancelled,
13201 "the deadline must CANCEL generation (worker's event channel closed)"
13202 );
13203
13204 let events = mock.events();
13209 assert!(
13210 events.contains(&MeterEvent::DeadlinePartial {
13211 prompt: 1,
13212 cached: 0,
13213 completion: 1,
13214 }),
13215 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
13216 );
13217 assert!(
13218 !events
13219 .iter()
13220 .any(|e| matches!(e, MeterEvent::Complete { .. })),
13221 "a deadline cut must stay distinguishable from a full answer: {events:?}"
13222 );
13223 }
13224
13225 #[tokio::test]
13229 async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
13230 let _l = DRAIN_LOCK.lock().unwrap();
13231 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13232 let health = health::WorkerHealth::new();
13233 let h = health.clone();
13234 std::thread::spawn(move || {
13235 h.mark_ready();
13236 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
13239 worker::release_pending_admit();
13240 worker::release_admission_reservation(req.lane);
13241 let _ = req.tx.send(Event::PromptUsage {
13242 n_prompt: 1,
13243 n_cached: 0,
13244 });
13245 for _ in 0..5_000 {
13246 if req.tx.is_closed() {
13247 break;
13248 }
13249 std::thread::sleep(std::time::Duration::from_millis(1));
13250 }
13251 }
13252 });
13253 for _ in 0..2_000 {
13254 if health.live().is_ok() {
13255 break;
13256 }
13257 std::thread::sleep(std::time::Duration::from_millis(1));
13258 }
13259 let mut st = fake_worker_state();
13260 st.cmd_tx = cmd_tx;
13261 st.health = health;
13262 let mock = MockMetering::admit_all();
13263 st.metering = Some(mock.clone());
13264 let resp = chat_completions(
13265 State(st),
13266 HeaderMap::new(),
13267 None,
13268 Json(
13269 serde_json::from_value(json!({
13270 "model": "m", "messages": [{"role": "user", "content": "t"}],
13271 "timeout_ms": 1_000}))
13272 .unwrap(),
13273 ),
13274 )
13275 .await;
13276 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
13277 assert!(resp.headers().get("x-should-retry").is_none());
13279 assert_eq!(retry_after(&resp), None);
13280 let body = body_value(resp).await;
13281 assert_eq!(body["error"]["code"], "deadline_exceeded");
13282 assert!(
13283 body["error"]["message"]
13284 .as_str()
13285 .unwrap()
13286 .contains("not billed"),
13287 "the zero-token 408 keeps the billing promise: {body}"
13288 );
13289 let events = mock.events();
13290 assert!(
13291 events.contains(&MeterEvent::Unbilled {
13292 outcome: "deadline_exceeded",
13293 status: 408,
13294 code: "deadline_exceeded".into(),
13295 }),
13296 "the named zero-debit census outcome, not the generic reject — every sibling \
13297 deadline path settles this one: {events:?}"
13298 );
13299 }
13300
13301 #[tokio::test]
13304 async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
13305 let _l = DRAIN_LOCK.lock().unwrap();
13306 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13308 let health = health::WorkerHealth::new();
13309 let h = health.clone();
13310 std::thread::spawn(move || {
13311 h.mark_ready();
13312 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
13313 worker::release_pending_admit();
13314 worker::release_admission_reservation(req.lane);
13315 let _ = req.tx.send(Event::PromptUsage {
13316 n_prompt: 1,
13317 n_cached: 0,
13318 });
13319 while !req.tx.is_closed() {
13320 std::thread::sleep(std::time::Duration::from_millis(1));
13321 }
13322 }
13323 });
13324 for _ in 0..2_000 {
13325 if health.live().is_ok() {
13326 break;
13327 }
13328 std::thread::sleep(std::time::Duration::from_millis(1));
13329 }
13330 let mut st = fake_worker_state();
13331 st.cmd_tx = cmd_tx;
13332 st.health = health;
13333 let mock = MockMetering::admit_all();
13334 st.metering = Some(mock.clone());
13335
13336 let resp = chat_completions(
13337 State(st),
13338 HeaderMap::new(),
13339 None,
13340 Json(
13341 serde_json::from_value(json!({
13342 "model": "m", "messages": [{"role": "user", "content": "t"}],
13343 "stream": true, "timeout_ms": 1_000}))
13344 .unwrap(),
13345 ),
13346 )
13347 .await;
13348 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
13351 let body = body_value(resp).await;
13352 assert_eq!(body["error"]["code"], "deadline_exceeded");
13353 assert!(
13354 body["error"]["message"]
13355 .as_str()
13356 .unwrap()
13357 .contains("first token"),
13358 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
13359 );
13360 let events = mock.events();
13361 assert!(
13362 events.contains(&MeterEvent::Unbilled {
13363 outcome: "deadline_exceeded",
13364 status: 408,
13365 code: "deadline_exceeded".into(),
13366 }),
13367 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
13368 );
13369 }
13370
13371 #[tokio::test]
13375 async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
13376 let _l = DRAIN_LOCK.lock().unwrap();
13377 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
13380 let mock = MockMetering::admit_all();
13381 st.metering = Some(mock.clone());
13382 let resp = chat_completions(
13383 State(st),
13384 HeaderMap::new(),
13385 None,
13386 Json(
13387 serde_json::from_value(json!({
13388 "model": "m", "messages": [{"role": "user", "content": "t"}],
13389 "stream": true, "timeout_ms": 1_000}))
13390 .unwrap(),
13391 ),
13392 )
13393 .await;
13394 assert_eq!(
13395 resp.status(),
13396 StatusCode::OK,
13397 "TTFT was met — 200 is correct"
13398 );
13399 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13400 .await
13401 .expect("the stream must run to completion past the deadline");
13402 let text = String::from_utf8(bytes.to_vec()).unwrap();
13403 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
13404 let events = mock.events();
13405 assert!(
13406 events
13407 .iter()
13408 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
13409 "a stream past its deadline after first token still settles as COMPLETE with \
13410 all four tokens: {events:?}"
13411 );
13412 }
13413
13414 #[test]
13417 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
13418 let st = fake_worker_state();
13419 let cap = lane_cap(lanes::Lane::Interactive);
13421 {
13422 let mut m = st.metrics.lock().unwrap();
13423 m.completed = 10;
13424 m.tokens_out = 1_000; m.step_p50_ms = 10.0; m.queued_requests = (cap * 4 + 1) as u64;
13427 }
13428 let rl = RateLimit {
13429 limit: cap,
13430 remaining: 0,
13431 reset_s: 1,
13432 };
13433 let (resp, outcome) = admission_backpressure(
13434 &st,
13435 lanes::Lane::Interactive,
13436 &rl,
13437 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13438 )
13439 .expect_err("a backlog past the bound must shed");
13440 assert_eq!(outcome, "shed_queue");
13441 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13442 assert!(
13443 retry_after(&resp).is_some(),
13444 "a shed must carry Retry-After so the router's spill can act on it"
13445 );
13446 let stamped = rl.attach(resp);
13448 for h in [
13449 "x-ratelimit-limit",
13450 "x-ratelimit-remaining",
13451 "x-ratelimit-reset",
13452 ] {
13453 assert!(stamped.headers().get(h).is_some(), "missing {h}");
13454 }
13455 }
13456
13457 #[test]
13461 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
13462 let st = fake_worker_state();
13463 let cap = lane_cap(lanes::Lane::Interactive);
13464 {
13465 let mut m = st.metrics.lock().unwrap();
13466 m.completed = 10;
13467 m.tokens_out = 1_000;
13468 m.step_p50_ms = 10.0; m.queued_requests = cap as u64; }
13471 let rl = RateLimit {
13472 limit: cap,
13473 remaining: 0,
13474 reset_s: 1,
13475 };
13476 assert!(
13478 admission_backpressure(
13479 &st,
13480 lanes::Lane::Interactive,
13481 &rl,
13482 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13483 )
13484 .is_ok(),
13485 "a request whose deadline covers the estimate must be admitted"
13486 );
13487 let (resp, outcome) = admission_backpressure(
13489 &st,
13490 lanes::Lane::Interactive,
13491 &rl,
13492 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13493 )
13494 .expect_err("a deadline shorter than the estimated wait must shed");
13495 assert_eq!(outcome, "shed_deadline");
13496 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13497 assert!(retry_after(&resp).is_some());
13498 }
13499
13500 #[test]
13503 fn admission_backpressure_is_interactive_only_and_silent_with_free_slots() {
13504 let st = fake_worker_state();
13505 let cap = lane_cap(lanes::Lane::Interactive);
13506 {
13507 let mut m = st.metrics.lock().unwrap();
13508 m.completed = 10;
13509 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
13511 m.queued_requests = (cap * 100) as u64;
13512 }
13513 let free = RateLimit {
13515 limit: cap,
13516 remaining: 1,
13517 reset_s: 0,
13518 };
13519 assert!(
13520 admission_backpressure(
13521 &st,
13522 lanes::Lane::Interactive,
13523 &free,
13524 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13525 )
13526 .is_ok()
13527 );
13528 let full = RateLimit {
13530 limit: cap,
13531 remaining: 0,
13532 reset_s: 5,
13533 };
13534 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
13535 assert!(
13536 admission_backpressure(
13537 &st,
13538 lane,
13539 &full,
13540 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13541 )
13542 .is_ok(),
13543 "{lane:?} must not be shed by the interactive gate"
13544 );
13545 }
13546 }
13547
13548 #[test]
13549 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
13550 let st = fake_worker_state();
13551 let cap = lane_cap(lanes::Lane::Interactive);
13552 let bound = max_queue_depth(cap);
13553 assert!(bound > 0, "the queue bound must admit at least one request");
13554 let rl = RateLimit {
13555 limit: cap,
13556 remaining: 0,
13557 reset_s: 1,
13558 };
13559 let _ = worker::PENDING_ADMITS.fetch_update(
13560 std::sync::atomic::Ordering::AcqRel,
13561 std::sync::atomic::Ordering::Acquire,
13562 |_| Some(0),
13563 );
13564 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
13565 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
13566 let guard = reserve_pending_admit(
13567 &st,
13568 lanes::Lane::Interactive,
13569 &rl,
13570 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13571 )
13572 .expect("the final queue slot should be reservable");
13573 assert_eq!(
13574 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13575 1
13576 );
13577 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
13578 drop(guard);
13579 assert_eq!(
13580 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13581 0
13582 );
13583 assert_eq!(
13584 counter.load(std::sync::atomic::Ordering::Acquire),
13585 bound - 1
13586 );
13587
13588 counter.store(bound, std::sync::atomic::Ordering::Release);
13589 let rejected = reserve_pending_admit(
13590 &st,
13591 lanes::Lane::Interactive,
13592 &rl,
13593 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13594 );
13595 assert!(matches!(rejected, Err((_, "shed_queue"))));
13596 counter.store(0, std::sync::atomic::Ordering::Release);
13597 }
13598
13599 #[test]
13600 fn admission_reservations_are_lane_scoped() {
13601 let st = fake_worker_state();
13602 let harvest = lanes::Lane::Harvest;
13603 let interactive = lanes::Lane::Interactive;
13604 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
13605 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
13606 harvest_counter.store(
13607 max_queue_depth(lane_cap(harvest)),
13608 std::sync::atomic::Ordering::Release,
13609 );
13610 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
13611 let free = RateLimit {
13612 limit: lane_cap(interactive),
13613 remaining: 1,
13614 reset_s: 0,
13615 };
13616 let guard = reserve_pending_admit(
13617 &st,
13618 interactive,
13619 &free,
13620 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13621 )
13622 .expect("a full harvest queue must not consume interactive capacity");
13623 drop(guard);
13624 let harvest_rl = RateLimit {
13625 limit: lane_cap(harvest),
13626 remaining: 0,
13627 reset_s: 1,
13628 };
13629 assert!(matches!(
13630 reserve_pending_admit(
13631 &st,
13632 harvest,
13633 &harvest_rl,
13634 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13635 ),
13636 Err((_, "shed_queue"))
13637 ));
13638 harvest_counter.store(0, std::sync::atomic::Ordering::Release);
13639 }
13640
13641 #[test]
13642 fn taxonomy_maps_every_class_to_its_status_and_code() {
13643 use worker::{EngineError as E, ErrClass as C};
13644 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
13645 (
13646 E::invalid_param("bad json", "response_format"),
13647 StatusCode::BAD_REQUEST,
13648 "invalid_request_error",
13649 "",
13650 ),
13651 (
13652 E::context_length("prompt (9000 tok) >= context cap (8192)"),
13653 StatusCode::BAD_REQUEST,
13654 "invalid_request_error",
13655 "context_length_exceeded",
13656 ),
13657 (
13658 E::model_not_found("unknown model \"nope\""),
13659 StatusCode::BAD_REQUEST,
13660 "invalid_request_error",
13661 "model_not_found",
13662 ),
13663 (
13664 E::rate_limit("lane judge is at capacity, retry"),
13665 StatusCode::TOO_MANY_REQUESTS,
13666 "rate_limit_error",
13667 "rate_limit_exceeded",
13668 ),
13669 (
13670 E::overloaded("no VRAM for a new session"),
13671 StatusCode::SERVICE_UNAVAILABLE,
13672 "server_error",
13673 "overloaded",
13674 ),
13675 (
13676 E::engine("graph step failed: launch error"),
13677 StatusCode::INTERNAL_SERVER_ERROR,
13678 "server_error",
13679 "engine_error",
13680 ),
13681 ];
13682 for (err, want_status, want_type, want_code) in cases {
13683 let (status, etype, code) = class_http(err.class);
13684 assert_eq!(status, want_status, "{:?}", err);
13685 assert_eq!(etype, want_type, "{:?}", err);
13686 if !want_code.is_empty() {
13687 assert_eq!(code, Some(want_code), "{:?}", err);
13688 }
13689 let body = engine_error_body(&err);
13691 assert_eq!(body["error"]["message"], err.message);
13692 assert_eq!(body["error"]["type"], want_type);
13693 }
13694 for c in [
13696 C::InvalidRequest,
13697 C::ContextLength,
13698 C::ModelNotFound,
13699 C::RateLimit,
13700 C::Overloaded,
13701 C::Engine,
13702 ] {
13703 let (s, t, _) = class_http(c);
13704 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
13705 assert!(!t.is_empty());
13706 }
13707 }
13708
13709 #[test]
13710 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
13711 let e = worker::EngineError::engine(
13716 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
13717 );
13718 let resp = engine_error_response(&e);
13719 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13720 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13721 }
13722
13723 #[test]
13724 fn retry_headers_follow_the_sdk_contract() {
13725 for e in [
13729 worker::EngineError::rate_limit("shed"),
13730 worker::EngineError::overloaded("no VRAM"),
13731 ] {
13732 let resp = engine_error_response(&e);
13733 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
13734 let secs: u64 = ra
13735 .parse()
13736 .expect("Retry-After must be integer delay-seconds");
13737 assert!(
13738 secs > 0 && secs <= 60,
13739 "Retry-After {secs}s outside the honored window"
13740 );
13741 let ms = resp
13742 .headers()
13743 .get("retry-after-ms")
13744 .unwrap()
13745 .to_str()
13746 .unwrap();
13747 assert_eq!(
13748 ms.parse::<u64>().unwrap(),
13749 secs * 1000,
13750 "the two headers disagree"
13751 );
13752 assert!(
13753 resp.headers().get("x-should-retry").is_none(),
13754 "a retryable class must not say x-should-retry: false"
13755 );
13756 }
13757 }
13758
13759 #[tokio::test]
13766 async fn admit_predict_reject_matches_shed_contract() {
13767 let shed = retry_contract_response(
13769 (
13770 StatusCode::TOO_MANY_REQUESTS,
13771 Json(error_body(
13772 "interactive queue is at its bound",
13773 "rate_limit_error",
13774 None,
13775 Some("shed_queue"),
13776 )),
13777 )
13778 .into_response(),
13779 Some(7),
13780 );
13781 let predict = engine_error_response(&worker::EngineError::rate_limit_after(
13784 "predicted KV-to-completion exceeds the box budget; retry",
13785 7,
13786 ));
13787 assert_eq!(shed.status(), predict.status());
13788 for header in ["retry-after", "retry-after-ms"] {
13789 assert_eq!(
13790 shed.headers().get(header),
13791 predict.headers().get(header),
13792 "header {header} must be byte-identical to the shed contract"
13793 );
13794 }
13795 let shed_body: serde_json::Value = serde_json::from_slice(
13796 &axum::body::to_bytes(shed.into_body(), usize::MAX)
13797 .await
13798 .unwrap(),
13799 )
13800 .unwrap();
13801 let predict_body: serde_json::Value = serde_json::from_slice(
13802 &axum::body::to_bytes(predict.into_body(), usize::MAX)
13803 .await
13804 .unwrap(),
13805 )
13806 .unwrap();
13807 assert_eq!(shed_body["error"]["type"], predict_body["error"]["type"]);
13808 assert_eq!(predict_body["error"]["type"], "rate_limit_error");
13809 let shed_keys: Vec<&String> = shed_body["error"].as_object().unwrap().keys().collect();
13810 let predict_keys: Vec<&String> =
13811 predict_body["error"].as_object().unwrap().keys().collect();
13812 assert_eq!(shed_keys, predict_keys, "same body schema, key for key");
13813 assert_eq!(predict_body["error"]["code"], "rate_limit_exceeded");
13814
13815 let clamped = engine_error_response(&worker::EngineError::rate_limit_after("m", 400));
13817 assert_eq!(retry_after(&clamped).as_deref(), Some("60"));
13818 let plain = engine_error_response(&worker::EngineError::rate_limit("m"));
13820 assert_eq!(retry_after(&plain).as_deref(), Some("2"));
13821 }
13822
13823 #[tokio::test]
13824 async fn command_send_failure_obeys_the_retry_contract() {
13825 let _l = DRAIN_LOCK.lock().unwrap();
13826 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
13827 let mut st = fake_worker_state();
13828 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13829 drop(cmd_rx);
13830 st.cmd_tx = cmd_tx;
13831
13832 let completion = completions(
13833 State(st.clone()),
13834 axum::http::HeaderMap::new(),
13835 None,
13836 Json(
13837 serde_json::from_value(serde_json::json!({
13838 "model": "m", "prompt": "test"
13839 }))
13840 .unwrap(),
13841 ),
13842 )
13843 .await;
13844 let chat = chat_completions(
13845 State(st),
13846 axum::http::HeaderMap::new(),
13847 None,
13848 Json(
13849 serde_json::from_value(serde_json::json!({
13850 "model": "m", "messages": [{"role": "user", "content": "test"}]
13851 }))
13852 .unwrap(),
13853 ),
13854 )
13855 .await;
13856
13857 for resp in [completion, chat] {
13858 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13859 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13860 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
13861 assert_ne!(
13862 resp.headers()
13863 .get("x-should-retry")
13864 .and_then(|v| v.to_str().ok()),
13865 Some("false")
13866 );
13867 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13868 .await
13869 .unwrap();
13870 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13871 assert_eq!(payload["error"]["type"], "server_error");
13872 assert_eq!(payload["error"]["code"], "overloaded");
13873 }
13874 }
13875
13876 #[test]
13877 fn unfixable_client_errors_say_x_should_retry_false() {
13878 for e in [
13881 worker::EngineError::model_not_found("unknown model \"x\""),
13882 worker::EngineError::context_length("prompt too long"),
13883 worker::EngineError::invalid_param("bad", "messages"),
13884 ] {
13885 let resp = engine_error_response(&e);
13886 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13887 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
13888 assert!(
13889 retry_after(&resp).is_none(),
13890 "a 400 must not promise a retry window"
13891 );
13892 }
13893 }
13894
13895 #[tokio::test]
13896 async fn a_closed_worker_channel_is_503_not_500() {
13897 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
13901 drop(tx);
13902 let resp =
13903 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
13904 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13905 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13906 }
13907
13908 #[tokio::test]
13909 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
13910 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13913 tx.send(Event::Error(worker::EngineError::rate_limit(
13914 "lane judge shed: interactive p99 over budget, retry",
13915 )))
13916 .unwrap();
13917 let (resp, error_code) = peek_admission(rx)
13918 .await
13919 .expect_err("a shed must not be forwarded into the stream");
13920 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13921 assert_eq!(error_code, "rate_limit_exceeded");
13922 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13923 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13924 .await
13925 .unwrap();
13926 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13927 assert!(
13928 payload["error"].is_object(),
13929 "bare-string error body: {payload}"
13930 );
13931 assert_eq!(payload["error"]["type"], "rate_limit_error");
13932 assert!(
13933 payload["error"]["message"]
13934 .as_str()
13935 .unwrap()
13936 .contains("shed")
13937 );
13938 }
13939
13940 #[tokio::test]
13941 async fn interactive_admission_error_is_a_preheader_429() {
13942 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13945 tx.send(Event::Error(worker::EngineError::rate_limit(
13946 "KV capacity unavailable",
13947 )))
13948 .unwrap();
13949 let (resp, error_code) = peek_admission(rx)
13950 .await
13951 .expect_err("admission error must stay pre-header");
13952 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13953 assert_eq!(error_code, "rate_limit_exceeded");
13954 }
13955
13956 #[tokio::test]
13957 async fn admission_peek_preserves_context_error_for_the_ledger() {
13958 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13959 tx.send(Event::Error(worker::EngineError::context_length(
13960 "prompt exceeds configured model maximum",
13961 )))
13962 .unwrap();
13963 let (resp, error_code) = peek_admission(rx)
13964 .await
13965 .expect_err("context rejection must stay pre-header");
13966 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13967 assert_eq!(error_code, "context_length_exceeded");
13968 }
13969
13970 #[tokio::test]
13971 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
13972 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13973 tx.send(Event::PromptUsage {
13974 n_prompt: 262_143,
13975 n_cached: 0,
13976 })
13977 .unwrap();
13978 let mut replay = peek_admission(rx).await.expect("successful admission");
13979 assert!(matches!(
13980 replay.recv().await,
13981 Some(Event::PromptUsage {
13982 n_prompt: 262_143,
13983 n_cached: 0
13984 }),
13985 ));
13986 }
13987
13988 #[test]
13989 fn penalties_plumb_from_http_to_sampler_config() {
13990 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13993 "model": "m", "messages": [{"role": "user", "content": "task"}],
13994 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
13995 }))
13996 .unwrap();
13997 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13998 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13999 .unwrap()
14000 .request
14001 .sampler_cfg;
14002 assert_eq!(cfg.penalty_freq, 0.5);
14003 assert_eq!(cfg.penalty_present, 0.25);
14004 assert_eq!(cfg.penalty_repeat, 1.1);
14005 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14006
14007 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14008 "model": "m", "prompt": "task", "frequency_penalty": 1.5
14009 }))
14010 .unwrap();
14011 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14012 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
14013 assert_eq!(cfg.penalty_freq, 1.5);
14014 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14015
14016 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14018 "model": "m", "prompt": "task"
14019 }))
14020 .unwrap();
14021 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14022 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
14023 assert_eq!(cfg.penalty_last_n, 0);
14024 assert_eq!(cfg.penalty_repeat, 1.0);
14025 }
14026
14027 #[test]
14028 fn omitted_temperature_is_openai_default_not_greedy() {
14029 let chat_temp = |body: serde_json::Value| {
14045 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14046 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14047 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
14048 .unwrap()
14049 .request
14050 .sampler_cfg
14051 .temperature
14052 };
14053 let comp_temp = |body: serde_json::Value| {
14054 let req: CompletionReq = serde_json::from_value(body).unwrap();
14055 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14056 build_request(&req, tx, lanes::Lane::Interactive, None)
14057 .sampler_cfg
14058 .temperature
14059 };
14060
14061 assert_eq!(
14063 chat_temp(serde_json::json!({
14064 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
14065 1.0,
14066 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
14067 );
14068 assert_eq!(
14069 comp_temp(serde_json::json!({
14070 "model": "m", "prompt": "t"})),
14071 1.0,
14072 "omitted completions temperature must be the OpenAI 1.0 default"
14073 );
14074
14075 assert_eq!(
14077 chat_temp(serde_json::json!({
14078 "model": "m", "messages": [{"role": "user", "content": "t"}],
14079 "temperature": 0.0})),
14080 0.0,
14081 "explicit temperature 0 must stay greedy"
14082 );
14083 assert_eq!(
14084 comp_temp(serde_json::json!({
14085 "model": "m", "prompt": "t", "temperature": 0})),
14086 0.0,
14087 "explicit temperature 0 must stay greedy"
14088 );
14089 assert!(
14091 memra_engine::sampler::Sampler::new(sampler_config(
14092 0.0,
14093 0,
14094 1.0,
14095 0.0,
14096 0.0,
14097 0.0,
14098 1.0,
14099 Some(0)
14100 ))
14101 .is_greedy()
14102 );
14103 assert!(
14104 !memra_engine::sampler::Sampler::new(sampler_config(
14105 1.0,
14106 0,
14107 1.0,
14108 0.0,
14109 0.0,
14110 0.0,
14111 1.0,
14112 Some(0)
14113 ))
14114 .is_greedy()
14115 );
14116
14117 assert_eq!(
14119 chat_temp(serde_json::json!({
14120 "model": "m", "messages": [{"role": "user", "content": "t"}],
14121 "temperature": 0.7})),
14122 0.7
14123 );
14124
14125 let req: CompletionReq = serde_json::from_value(serde_json::json!({
14129 "model": "m", "prompt": "t"}))
14130 .unwrap();
14131 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14132 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
14133 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
14134 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
14135 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
14136 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
14137 assert!(
14142 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
14143 "the omitted-temperature default must ride sampled spec's pure-temp regime"
14144 );
14145 }
14146
14147 #[test]
14148 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
14149 let caps = ModelCaps {
14150 chat_temperature_default: Some(0.5),
14151 chat_top_p_default: Some(0.9),
14152 chat_ok: true,
14153 ..Default::default()
14154 };
14155 let cfg = |extra: serde_json::Value| {
14156 let mut body = serde_json::json!({
14157 "model": "step35",
14158 "messages": [{"role": "user", "content": "task"}]
14159 });
14160 body.as_object_mut()
14161 .unwrap()
14162 .extend(extra.as_object().unwrap().clone());
14163 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14164 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14165 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
14166 .unwrap()
14167 .request
14168 .sampler_cfg
14169 };
14170
14171 let omitted = cfg(serde_json::json!({}));
14172 assert_eq!(omitted.temperature, 0.5);
14173 assert_eq!(omitted.top_p, 0.9);
14174
14175 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
14176 assert_eq!(explicit_temp.temperature, 0.7);
14177 assert_eq!(
14178 explicit_temp.top_p, 0.9,
14179 "omitting top_p must retain StepFun's nucleus default"
14180 );
14181
14182 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
14183 assert_eq!(
14184 explicit.temperature, 0.0,
14185 "explicit greedy must remain authoritative"
14186 );
14187 assert_eq!(
14188 explicit.top_p, 1.0,
14189 "explicit untruncated sampling must remain authoritative"
14190 );
14191 }
14192
14193 fn qwen38_vendor_defaults() -> SamplingDefaults {
14197 SamplingDefaults {
14198 temperature: Some(1.0),
14199 top_p: Some(0.95),
14200 top_k: Some(20),
14201 min_p: Some(0.0),
14202 presence_penalty: Some(0.0),
14203 repetition_penalty: Some(1.0),
14204 frequency_penalty: None,
14205 }
14206 }
14207
14208 fn gemma4_vendor_defaults() -> SamplingDefaults {
14213 SamplingDefaults {
14214 temperature: Some(1.0),
14215 top_p: Some(0.95),
14216 top_k: Some(64),
14217 ..Default::default()
14218 }
14219 }
14220
14221 #[test]
14222 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
14223 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
14228 let chat = |extra: serde_json::Value| {
14229 let mut body = serde_json::json!({
14230 "model": "google/gemma-4-31b-it",
14231 "messages": [{"role": "user", "content": "task"}],
14232 "seed": 7
14234 });
14235 body.as_object_mut()
14236 .unwrap()
14237 .extend(extra.as_object().unwrap().clone());
14238 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14239 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14240 build_chat_request_with_trace(
14241 req,
14242 Some(&ModelCaps {
14243 chat_ok: true,
14244 ..Default::default()
14245 }),
14246 tx,
14247 lanes::Lane::Interactive,
14248 None,
14249 None,
14250 None,
14251 &d,
14252 )
14253 .unwrap()
14254 .request
14255 .sampler_cfg
14256 };
14257
14258 let omitted = chat(serde_json::json!({}));
14260 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
14261 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
14262 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
14263 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
14265 assert_eq!(omitted.penalty_repeat, 1.0);
14266 assert_eq!(omitted.penalty_freq, 0.0);
14267 assert_eq!(omitted.penalty_present, 0.0);
14268 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
14269 assert!(
14270 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
14271 "the vendor default must NOT be greedy — that is the whole point of the lane"
14272 );
14273
14274 let greedy = chat(serde_json::json!({"temperature": 0}));
14277 assert_eq!(
14278 greedy.temperature, 0.0,
14279 "explicit temperature 0 stays greedy"
14280 );
14281 assert!(
14282 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
14283 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
14284 spec/graph exactness arms"
14285 );
14286
14287 let one_field = chat(serde_json::json!({"top_k": 3}));
14289 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
14290 assert_eq!(
14291 one_field.temperature, 1.0,
14292 "omitting temperature still takes the vendor value"
14293 );
14294 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
14295
14296 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
14299 assert_eq!(
14300 disabled.top_k, 0,
14301 "an explicit top_k 0 means KEEP ALL, not 'unset'"
14302 );
14303 assert_eq!(
14304 disabled.top_p, 1.0,
14305 "an explicit top_p 1.0 means untruncated"
14306 );
14307
14308 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
14310 assert_eq!(penal.penalty_present, 1.5);
14311 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14312 }
14313
14314 #[test]
14315 fn vendor_sampling_defaults_are_identical_on_every_surface() {
14316 let d = qwen38_vendor_defaults();
14328 let md = ModelSamplingDefaults::single(d);
14329 let comp = |extra: serde_json::Value| {
14330 let mut body = serde_json::json!({
14331 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
14332 body.as_object_mut()
14333 .unwrap()
14334 .extend(extra.as_object().unwrap().clone());
14335 let req: CompletionReq = serde_json::from_value(body).unwrap();
14336 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14337 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
14338 };
14339 let chat = |extra: serde_json::Value| {
14340 let mut body = serde_json::json!({
14341 "model": "qwen/qwen3.8-27b",
14342 "messages": [{"role": "user", "content": "task"}],
14343 "seed": 11 });
14344 body.as_object_mut()
14345 .unwrap()
14346 .extend(extra.as_object().unwrap().clone());
14347 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14348 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14349 build_chat_request_with_trace(
14350 req,
14351 Some(&ModelCaps {
14352 chat_ok: true,
14353 ..Default::default()
14354 }),
14355 tx,
14356 lanes::Lane::Interactive,
14357 None,
14358 None,
14359 None,
14360 &md,
14361 )
14362 .unwrap()
14363 .request
14364 .sampler_cfg
14365 };
14366
14367 for extra in [
14368 serde_json::json!({}),
14369 serde_json::json!({"temperature": 0}),
14370 serde_json::json!({"temperature": 0.0}),
14371 serde_json::json!({"temperature": 0.7}),
14372 serde_json::json!({"top_p": 1.0}),
14373 serde_json::json!({"top_k": 0}),
14374 serde_json::json!({"min_p": 0.05}),
14375 serde_json::json!({"repetition_penalty": 1.1}),
14376 serde_json::json!({"frequency_penalty": 0.5}),
14377 serde_json::json!({"presence_penalty": 1.5}),
14378 serde_json::json!({
14379 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
14380 "frequency_penalty": 0.1, "presence_penalty": 0.2,
14381 "repetition_penalty": 1.05 }),
14382 ] {
14383 let c = comp(extra.clone());
14384 let h = chat(extra.clone());
14385 assert_eq!(
14386 (
14387 c.temperature,
14388 c.top_p,
14389 c.top_k,
14390 c.min_p,
14391 c.penalty_repeat,
14392 c.penalty_freq,
14393 c.penalty_present,
14394 c.penalty_last_n,
14395 c.seed
14396 ),
14397 (
14398 h.temperature,
14399 h.top_p,
14400 h.top_k,
14401 h.min_p,
14402 h.penalty_repeat,
14403 h.penalty_freq,
14404 h.penalty_present,
14405 h.penalty_last_n,
14406 h.seed
14407 ),
14408 "/v1/completions and /v1/chat/completions disagree on {extra} — \
14409 standard-surface-law violation"
14410 );
14411 }
14412
14413 let omitted = comp(serde_json::json!({}));
14415 assert_eq!(
14416 omitted.temperature, 1.0,
14417 "qwen3.8 card thinking temperature"
14418 );
14419 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
14420 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
14421 assert!(
14423 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
14424 .is_greedy()
14425 );
14426 }
14427
14428 #[tokio::test]
14441 async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
14442 let _l = DRAIN_LOCK.lock().unwrap();
14443 let step_caps = ModelCaps {
14444 chat_ok: true,
14445 chat_temperature_default: Some(0.5),
14446 chat_top_p_default: Some(0.9),
14447 ..Default::default()
14448 };
14449 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14450 let st = fake_worker_state_full(
14451 1,
14452 std::time::Duration::ZERO,
14453 HashMap::from([("m".to_string(), step_caps)]),
14454 Some(cfg_tx),
14455 );
14456 let fields = |saw: &WorkerSaw| {
14461 let c = &saw.sampler_cfg;
14462 (
14463 c.temperature,
14464 c.top_p,
14465 c.top_k,
14466 c.min_p,
14467 c.penalty_repeat,
14468 c.penalty_freq,
14469 c.penalty_present,
14470 c.penalty_last_n,
14471 )
14472 };
14473 let worker_saw = |surface: &str| {
14474 cfg_rx
14475 .recv_timeout(std::time::Duration::from_secs(10))
14476 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
14477 };
14478
14479 let resp = completions(
14480 State(st.clone()),
14481 axum::http::HeaderMap::new(),
14482 None,
14483 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
14484 )
14485 .await;
14486 assert_eq!(
14487 resp.status(),
14488 StatusCode::OK,
14489 "/v1/completions rejected the omitted-sampling request"
14490 );
14491 let comp = worker_saw("/v1/completions");
14492
14493 let resp = chat_completions(
14494 State(st.clone()),
14495 axum::http::HeaderMap::new(),
14496 None,
14497 Json(
14498 serde_json::from_value(serde_json::json!({
14499 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
14500 .unwrap(),
14501 ),
14502 )
14503 .await;
14504 assert_eq!(
14505 resp.status(),
14506 StatusCode::OK,
14507 "/v1/chat/completions rejected the omitted-sampling request"
14508 );
14509 let chat = worker_saw("/v1/chat/completions");
14510
14511 let resp = anthropic::messages(
14512 State(st.clone()),
14513 axum::http::HeaderMap::new(),
14514 None,
14515 axum::body::Bytes::from(
14516 serde_json::json!({
14517 "model": "m", "max_tokens": 16,
14518 "messages": [{"role": "user", "content": "t"}]})
14519 .to_string(),
14520 ),
14521 )
14522 .await;
14523 assert_eq!(
14524 resp.status(),
14525 StatusCode::OK,
14526 "/v1/messages rejected the omitted-sampling request"
14527 );
14528 let msg = worker_saw("/v1/messages");
14529
14530 let resp = responses_api::responses(
14531 State(st.clone()),
14532 axum::http::HeaderMap::new(),
14533 None,
14534 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
14535 )
14536 .await;
14537 assert_eq!(
14538 resp.status(),
14539 StatusCode::OK,
14540 "/v1/responses rejected the omitted-sampling request"
14541 );
14542 let rsp = worker_saw("/v1/responses");
14543
14544 for (surface, cfg) in [
14545 ("/v1/completions", &comp),
14546 ("/v1/messages", &msg),
14547 ("/v1/responses", &rsp),
14548 ] {
14549 assert_eq!(
14550 fields(cfg),
14551 fields(&chat),
14552 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
14553 for the same omitted-sampling request — standard-surface-law violation \
14554 (hermes d991b51699218285)"
14555 );
14556 }
14557 assert_eq!(
14560 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
14561 (0.5, 0.9),
14562 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
14563 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
14564 );
14565 }
14566
14567 #[tokio::test]
14577 async fn same_effort_value_resolves_identically_on_every_surface() {
14578 let _l = DRAIN_LOCK.lock().unwrap();
14579 let caps = ModelCaps {
14582 chat_ok: true,
14583 effort_levels: true,
14584 ..Default::default()
14585 };
14586 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14587 let st = fake_worker_state_full(
14588 1,
14589 std::time::Duration::ZERO,
14590 HashMap::from([("m".to_string(), caps)]),
14591 Some(saw_tx),
14592 );
14593 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
14594 match surface {
14595 "/v1/chat/completions" => {
14596 chat_completions(
14597 State(st),
14598 axum::http::HeaderMap::new(),
14599 None,
14600 Json(
14601 serde_json::from_value(serde_json::json!({
14602 "model": "m", "max_tokens": 8,
14603 "reasoning_effort": effort,
14604 "messages": [{"role": "user", "content": "t"}]}))
14605 .unwrap(),
14606 ),
14607 )
14608 .await
14609 }
14610 "/v1/responses" => {
14611 responses_api::responses(
14612 State(st),
14613 axum::http::HeaderMap::new(),
14614 None,
14615 axum::body::Bytes::from(
14616 serde_json::json!({
14617 "model": "m", "max_output_tokens": 8, "input": "t",
14618 "reasoning": {"effort": effort}})
14619 .to_string(),
14620 ),
14621 )
14622 .await
14623 }
14624 "/v1/messages" => {
14625 anthropic::messages(
14626 State(st),
14627 axum::http::HeaderMap::new(),
14628 None,
14629 axum::body::Bytes::from(
14630 serde_json::json!({
14631 "model": "m", "max_tokens": 8,
14632 "messages": [{"role": "user", "content": "t"}],
14633 "output_config": {"effort": effort}})
14634 .to_string(),
14635 ),
14636 )
14637 .await
14638 }
14639 other => panic!("unknown surface {other}"),
14640 }
14641 };
14642 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
14643
14644 for (effort, want_think, want_level) in [
14647 ("none", ThinkMode::NoThink, Some("low")),
14648 ("minimal", ThinkMode::NoThink, Some("low")),
14649 ("low", ThinkMode::Think, Some("low")),
14650 ("medium", ThinkMode::Think, Some("medium")),
14651 ("high", ThinkMode::Think, Some("high")),
14652 ("xhigh", ThinkMode::Think, Some("high")),
14654 ] {
14655 for surface in SURFACES {
14656 let resp = send(st.clone(), surface, effort).await;
14657 assert_eq!(
14658 resp.status(),
14659 StatusCode::OK,
14660 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
14661 diverged again (issue #31)"
14662 );
14663 let saw = saw_rx
14664 .recv_timeout(std::time::Duration::from_secs(10))
14665 .unwrap_or_else(|_| {
14666 panic!("{surface}: effort {effort:?} request never reached the worker")
14667 });
14668 assert_eq!(
14669 (saw.think, saw.reasoning_effort.as_deref()),
14670 (want_think, want_level),
14671 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
14672 reasoning surface — the parameter was dropped or remapped before \
14673 parse_think (issue #31 regression)"
14674 );
14675 }
14676 }
14677
14678 for effort in ["bogus", "banana", ""] {
14681 for surface in SURFACES {
14682 let resp = send(st.clone(), surface, effort).await;
14683 assert_eq!(
14684 resp.status(),
14685 StatusCode::BAD_REQUEST,
14686 "{surface} accepted effort {effort:?} — silent-accept regression \
14687 (issue #31: the value never reached parse_think's allowlist)"
14688 );
14689 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
14691 .await
14692 .unwrap();
14693 let v: serde_json::Value = serde_json::from_slice(&body)
14694 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
14695 match surface {
14696 "/v1/messages" => {
14697 assert_eq!(v["type"], "error", "{surface} error envelope");
14698 assert_eq!(
14699 v["error"]["type"], "invalid_request_error",
14700 "{surface} error type"
14701 );
14702 }
14703 _ => {
14704 assert!(
14705 v["error"]["message"].is_string(),
14706 "{surface} OpenAI-shaped error body: {v}"
14707 );
14708 }
14709 }
14710 }
14711 }
14712
14713 let resp = anthropic::messages(
14717 State(st.clone()),
14718 axum::http::HeaderMap::new(),
14719 None,
14720 axum::body::Bytes::from(
14721 serde_json::json!({
14722 "model": "m", "max_tokens": 8,
14723 "messages": [{"role": "user", "content": "t"}],
14724 "thinking": {"type": "enabled"},
14725 "output_config": {"effort": "none"}})
14726 .to_string(),
14727 ),
14728 )
14729 .await;
14730 assert_eq!(resp.status(), StatusCode::OK);
14731 let saw = saw_rx
14732 .recv_timeout(std::time::Duration::from_secs(10))
14733 .expect("thinking+effort request never reached the worker");
14734 assert_eq!(
14735 saw.think,
14736 ThinkMode::Think,
14737 "thinking.type (the documented Anthropic lever) must win the switch over \
14738 output_config.effort"
14739 );
14740 let resp = anthropic::messages(
14741 State(st.clone()),
14742 axum::http::HeaderMap::new(),
14743 None,
14744 axum::body::Bytes::from(
14745 serde_json::json!({
14746 "model": "m", "max_tokens": 8,
14747 "messages": [{"role": "user", "content": "t"}],
14748 "thinking": {"type": "enabled"},
14749 "output_config": {"effort": "banana"}})
14750 .to_string(),
14751 ),
14752 )
14753 .await;
14754 assert_eq!(
14755 resp.status(),
14756 StatusCode::BAD_REQUEST,
14757 "an invalid effort must 400 even next to an explicit thinking.type — \
14758 precedence must not re-open the silent-accept hole"
14759 );
14760 }
14761
14762 #[test]
14763 fn vendor_sampling_defaults_are_boot_validated() {
14764 let parsed = OpenRouterMetadataFile::from_toml(
14767 r#"
14768[models.g]
14769default_temperature = 1.0
14770default_top_p = 0.95
14771default_top_k = 64
14772default_min_p = 0.0
14773default_presence_penalty = 0.0
14774default_frequency_penalty = 0.0
14775default_repetition_penalty = 1.0
14776"#,
14777 )
14778 .unwrap();
14779 let g = parsed.get("g").unwrap();
14780 assert_eq!(g.default_temperature, Some(1.0));
14781 assert_eq!(g.default_top_p, Some(0.95));
14782 assert_eq!(g.default_top_k, Some(64));
14783
14784 let err = OpenRouterMetadataFile::from_toml(
14788 r#"
14789[models.g]
14790default_temperature = 0.0
14791"#,
14792 )
14793 .unwrap_err();
14794 assert!(err.contains("default_temperature"), "{err}");
14795 assert!(
14796 err.contains("greedy"),
14797 "the refusal must say WHY a zero default is refused: {err}"
14798 );
14799
14800 for bad in [
14801 "default_temperature = 2.5",
14802 "default_temperature = -1.0",
14803 "default_top_p = 0.0",
14804 "default_top_p = 1.5",
14805 "default_min_p = 1.0",
14806 "default_min_p = -0.1",
14807 "default_presence_penalty = 3.0",
14808 "default_frequency_penalty = -2.5",
14809 "default_repetition_penalty = 0.0",
14810 ] {
14811 let err =
14812 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
14813 let key = bad.split(' ').next().unwrap();
14814 assert!(err.contains(key), "{bad} must be refused by name: {err}");
14815 }
14816
14817 let err = OpenRouterMetadataFile::from_toml(
14821 r#"
14822[models.g]
14823default_temperture = 1.0
14824"#,
14825 )
14826 .unwrap_err();
14827 assert!(
14828 err.contains("unknown field"),
14829 "an unknown key must be fatal, which is what makes binary-first ordering \
14830 mandatory: {err}"
14831 );
14832 }
14833
14834 #[test]
14835 fn non_thinking_sampling_arm_is_boot_validated() {
14836 let parsed = OpenRouterMetadataFile::from_toml(
14840 r#"
14841[models.q]
14842default_temperature = 1.0
14843default_top_p = 0.95
14844default_top_k = 20
14845
14846[models.q.non_thinking_sampling]
14847temperature = 0.7
14848top_p = 0.8
14849top_k = 20
14850presence_penalty = 1.5
14851"#,
14852 )
14853 .unwrap();
14854 let arm = parsed
14855 .get("q")
14856 .unwrap()
14857 .non_thinking_sampling
14858 .as_ref()
14859 .unwrap();
14860 assert_eq!(arm.temperature, Some(0.7));
14861 assert_eq!(arm.top_p, Some(0.8));
14862 assert_eq!(arm.top_k, Some(20));
14863 assert_eq!(arm.presence_penalty, Some(1.5));
14864 assert_eq!(
14865 arm.min_p, None,
14866 "undeclared arm fields stay undeclared, never invented"
14867 );
14868
14869 let err = OpenRouterMetadataFile::from_toml(
14873 r#"
14874[models.q]
14875[models.q.non_thinking_sampling]
14876temperature = 0.0
14877"#,
14878 )
14879 .unwrap_err();
14880 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
14881 assert!(err.contains("greedy"), "{err}");
14882
14883 let err = OpenRouterMetadataFile::from_toml(
14886 r#"
14887[models.q]
14888[models.q.non_thinking_sampling]
14889"#,
14890 )
14891 .unwrap_err();
14892 assert!(err.contains("non_thinking_sampling"), "{err}");
14893 assert!(err.contains("declare"), "{err}");
14894
14895 for bad in [
14897 "temperature = 2.5",
14898 "top_p = 0.0",
14899 "top_p = 1.5",
14900 "min_p = 1.0",
14901 "presence_penalty = 3.0",
14902 "frequency_penalty = -2.5",
14903 "repetition_penalty = 0.0",
14904 ] {
14905 let err = OpenRouterMetadataFile::from_toml(&format!(
14906 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
14907 ))
14908 .unwrap_err();
14909 let key = bad.split(' ').next().unwrap();
14910 assert!(
14911 err.contains(&format!("non_thinking_sampling.{key}")),
14912 "the refusal for {bad:?} must name the nested key: {err}"
14913 );
14914 }
14915
14916 let err = OpenRouterMetadataFile::from_toml(
14920 r#"
14921[models.q]
14922[models.q.non_thinking_sampling]
14923temperture = 0.7
14924"#,
14925 )
14926 .unwrap_err();
14927 assert!(err.contains("unknown field"), "{err}");
14928 }
14929
14930 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
14935 SamplingDefaults {
14936 temperature: Some(0.7),
14937 top_p: Some(0.8),
14938 top_k: Some(20),
14939 presence_penalty: Some(1.5),
14940 ..Default::default()
14941 }
14942 }
14943
14944 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
14945 ModelSamplingDefaults {
14946 thinking: qwen38_vendor_defaults(),
14947 non_thinking: Some(qwen38_non_thinking_defaults()),
14948 }
14949 }
14950
14951 fn qwen38_caps() -> ModelCaps {
14955 ModelCaps {
14956 chat_ok: true,
14957 qwen_think: true,
14958 think_switch: true,
14959 ..Default::default()
14960 }
14961 }
14962
14963 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
14966 (
14967 c.temperature,
14968 c.top_p,
14969 c.top_k,
14970 c.min_p,
14971 c.penalty_present,
14972 c.penalty_freq,
14973 c.penalty_repeat,
14974 c.penalty_last_n,
14975 c.seed,
14976 )
14977 }
14978
14979 fn build_with_arms(
14980 defaults: &ModelSamplingDefaults,
14981 caps: &ModelCaps,
14982 default_effort: Option<&str>,
14983 extra: serde_json::Value,
14984 ) -> Request {
14985 let mut body = serde_json::json!({
14986 "model": "m",
14987 "messages": [{"role": "user", "content": "task"}],
14988 "seed": 3
14990 });
14991 body.as_object_mut()
14992 .unwrap()
14993 .extend(extra.as_object().unwrap().clone());
14994 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14995 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14996 build_chat_request_with_trace(
14997 req,
14998 Some(caps),
14999 tx,
15000 lanes::Lane::Interactive,
15001 None,
15002 None,
15003 default_effort,
15004 defaults,
15005 )
15006 .unwrap()
15007 .request
15008 }
15009
15010 #[test]
15011 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
15012 let two_arm = qwen38_two_arm_defaults();
15017 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
15018 let caps = qwen38_caps();
15019
15020 let off_spellings = [
15022 serde_json::json!({"reasoning_effort": "none"}),
15023 serde_json::json!({"enable_thinking": false}),
15024 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
15025 serde_json::json!({"reasoning": {"enabled": false}}),
15026 ];
15027 for extra in &off_spellings {
15028 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
15029 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
15030 let c = &r.sampler_cfg;
15031 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
15032 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
15033 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
15034 assert_eq!(
15035 c.penalty_present, 1.5,
15036 "{extra}: non-thinking presence_penalty"
15037 );
15038 assert_eq!(
15039 c.penalty_last_n,
15040 memra_engine::spec::PEN_WINDOW_MAX,
15041 "{extra}: the arm's presence penalty uses the cross-path history window"
15042 );
15043 assert_eq!(
15044 c.min_p, 0.0,
15045 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
15046 );
15047
15048 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
15051 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
15052 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
15053 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
15054 assert_eq!(
15055 s.sampler_cfg.penalty_present, 0.0,
15056 "{extra}: single-arm model"
15057 );
15058 }
15059
15060 for extra in [
15063 serde_json::json!({}),
15064 serde_json::json!({"enable_thinking": true}),
15065 serde_json::json!({"reasoning_effort": "high"}),
15066 serde_json::json!({"reasoning": {"enabled": true}}),
15067 ] {
15068 for defaults in [&two_arm, &single_arm] {
15069 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
15070 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
15071 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
15072 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
15073 assert_eq!(
15074 c.penalty_present, 0.0,
15075 "{extra}: thinking arm has no presence"
15076 );
15077 }
15078 }
15079
15080 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
15083 assert_eq!(
15084 c.temperature, 0.7,
15085 "deployment-default off = non-thinking arm"
15086 );
15087 let c = build_with_arms(
15089 &two_arm,
15090 &caps,
15091 Some("none"),
15092 serde_json::json!({"enable_thinking": true}),
15093 )
15094 .sampler_cfg;
15095 assert_eq!(
15096 c.temperature, 1.0,
15097 "explicit ON beats the deployment default"
15098 );
15099
15100 let c = build_with_arms(
15102 &two_arm,
15103 &caps,
15104 None,
15105 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
15106 )
15107 .sampler_cfg;
15108 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
15109 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
15110 let c = build_with_arms(
15111 &two_arm,
15112 &caps,
15113 None,
15114 serde_json::json!({
15115 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
15116 )
15117 .sampler_cfg;
15118 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
15119 assert_eq!(
15120 c.penalty_present, 0.0,
15121 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
15122 is a value, not an absence"
15123 );
15124 assert_eq!(
15125 c.penalty_last_n, 0,
15126 "all penalties off => no history window"
15127 );
15128 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
15129
15130 let c = build_with_arms(
15133 &two_arm,
15134 &caps,
15135 None,
15136 serde_json::json!({"enable_thinking": false, "temperature": 0}),
15137 )
15138 .sampler_cfg;
15139 assert!(
15140 memra_engine::sampler::Sampler::new(c).is_greedy(),
15141 "explicit temperature 0 must stay greedy on the non-thinking arm"
15142 );
15143
15144 let c = build_with_arms(
15147 &single_arm,
15148 &caps,
15149 None,
15150 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
15151 )
15152 .sampler_cfg;
15153 assert_eq!(c.temperature, 0.55);
15154 assert_eq!(
15155 c.top_p, 0.95,
15156 "single-arm model: unset top_p takes its one arm"
15157 );
15158 }
15159
15160 #[test]
15161 fn sampling_arms_never_blend_field_by_field() {
15162 let parsed = OpenRouterMetadataFile::from_toml(
15167 r#"
15168[models.m]
15169default_temperature = 1.0
15170default_min_p = 0.05
15171
15172[models.m.non_thinking_sampling]
15173temperature = 0.6
15174"#,
15175 )
15176 .unwrap();
15177 let caps = ModelCaps {
15178 chat_temperature_default: Some(0.5),
15179 chat_top_p_default: Some(0.9),
15180 ..Default::default()
15181 };
15182 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
15183 let client = ClientSampling {
15184 seed: Some(1),
15185 ..Default::default()
15186 };
15187
15188 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
15189 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
15190 assert_eq!(
15191 off.min_p, 0.0,
15192 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
15193 );
15194 assert_eq!(
15195 off.top_p, 1.0,
15196 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
15197 );
15198
15199 for mode in [ThinkMode::Default, ThinkMode::Think] {
15201 let on = resolve_sampler_config(client, d.for_mode(mode));
15202 assert_eq!(on.temperature, 1.0);
15203 assert_eq!(on.min_p, 0.05);
15204 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
15205 }
15206 }
15207
15208 #[test]
15209 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
15210 let caps = qwen38_caps();
15219 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
15220 let two_arm = qwen38_two_arm_defaults();
15221
15222 let bodies = [
15223 serde_json::json!({}),
15224 serde_json::json!({"enable_thinking": true}),
15225 serde_json::json!({"reasoning_effort": "high"}),
15226 serde_json::json!({"reasoning_effort": "none"}),
15227 serde_json::json!({"enable_thinking": false}),
15228 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
15229 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
15230 serde_json::json!({"enable_thinking": false, "temperature": 0}),
15231 ];
15232 for extra in &bodies {
15233 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
15235 let mut client = ClientSampling {
15236 seed: Some(3),
15237 ..Default::default()
15238 };
15239 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
15240 client.temperature = Some(t as f32);
15241 }
15242 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
15243 client.top_p = Some(p as f32);
15244 }
15245 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
15246 assert_eq!(
15247 sampler_key(&r.sampler_cfg),
15248 sampler_key(&pre_arm),
15249 "{extra}: single-arm model diverged from the pre-arm resolution law"
15250 );
15251
15252 if r.think != ThinkMode::NoThink {
15255 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
15256 assert_eq!(t.think, r.think, "{extra}");
15257 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
15258 assert_eq!(
15259 sampler_key(&t.sampler_cfg),
15260 sampler_key(&r.sampler_cfg),
15261 "{extra}: a thinking-on request must not feel the non-thinking arm"
15262 );
15263 }
15264 }
15265 }
15266
15267 #[test]
15268 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
15269 let r = build_with_arms(
15275 &qwen38_two_arm_defaults(),
15276 &qwen38_caps(),
15277 None,
15278 serde_json::json!({"response_format": {"type": "json_object"}}),
15279 );
15280 assert_eq!(
15281 r.think,
15282 ThinkMode::NoThink,
15283 "constraint forces the switch off"
15284 );
15285 assert_eq!(
15286 r.sampler_cfg.temperature, 0.7,
15287 "and the arm follows the real mode"
15288 );
15289 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
15290 }
15291
15292 #[test]
15293 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
15294 let caps = ModelCaps {
15299 chat_temperature_default: Some(0.5),
15300 chat_top_p_default: Some(0.9),
15301 chat_ok: true,
15302 ..Default::default()
15303 };
15304 let metadata = OpenRouterModelMetadata {
15305 default_temperature: Some(1.0),
15306 default_top_p: Some(0.95),
15307 default_top_k: Some(64),
15308 ..Default::default()
15309 };
15310
15311 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
15312 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
15313 assert_eq!(caps_only.top_p, Some(0.9));
15314 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
15315
15316 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
15317 assert_eq!(
15318 both.temperature,
15319 Some(1.0),
15320 "metadata outranks the arch cap"
15321 );
15322 assert_eq!(both.top_p, Some(0.95));
15323 assert_eq!(both.top_k, Some(64));
15324
15325 let partial = SamplingDefaults::resolve(
15327 Some(&OpenRouterModelMetadata {
15328 default_temperature: Some(0.7),
15329 ..Default::default()
15330 }),
15331 Some(&caps),
15332 );
15333 assert_eq!(partial.temperature, Some(0.7));
15334 assert_eq!(
15335 partial.top_p,
15336 Some(0.9),
15337 "an undeclared metadata field must fall through to the cap, not to 1.0"
15338 );
15339
15340 assert_eq!(
15342 SamplingDefaults::resolve(None, None),
15343 SamplingDefaults::default()
15344 );
15345 }
15346
15347 #[test]
15348 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
15349 let resolved = |d: &SamplingDefaults| {
15362 resolve_sampler_config(
15363 ClientSampling {
15364 seed: Some(1),
15365 ..Default::default()
15366 },
15367 d,
15368 )
15369 };
15370
15371 assert!(
15373 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
15374 .is_spec_sampling(),
15375 "the API-standard default must stay in the fast pure-temp regime"
15376 );
15377
15378 for (name, d) in [
15379 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
15380 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
15381 ] {
15382 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
15383 assert!(
15384 !sampler.is_greedy(),
15385 "{name}: vendor default must not be greedy"
15386 );
15387 assert!(
15388 !sampler.is_spec_sampling(),
15389 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
15390 starts passing, either the vendor numbers changed or the in-graph draft \
15391 learned filters, and the perf note in docs/SERVING.md needs revisiting"
15392 );
15393 }
15394
15395 let opted_out = resolve_sampler_config(
15397 ClientSampling {
15398 top_p: Some(1.0),
15399 top_k: Some(0),
15400 seed: Some(1),
15401 ..Default::default()
15402 },
15403 &qwen38_vendor_defaults(),
15404 );
15405 assert!(
15406 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
15407 "explicitly disabling the filters must restore the pure-temp regime"
15408 );
15409 }
15410
15411 #[test]
15412 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
15413 let comp_seed = |body: serde_json::Value| {
15420 let req: CompletionReq = serde_json::from_value(body).unwrap();
15421 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15422 build_request(&req, tx, lanes::Lane::Interactive, None)
15423 .sampler_cfg
15424 .seed
15425 };
15426 let chat_seed = |body: serde_json::Value| {
15427 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15428 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15429 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15430 .unwrap()
15431 .request
15432 .sampler_cfg
15433 .seed
15434 };
15435
15436 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15439 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15440 let c = chat_seed(serde_json::json!({
15441 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
15442 assert_ne!(
15443 a, 0,
15444 "omitted seed must not be the pinned 0 that caused the loop"
15445 );
15446 assert_ne!(b, 0);
15447 assert_ne!(c, 0);
15448 assert_ne!(
15449 a, b,
15450 "two seed-omitting requests must get DIFFERENT streams"
15451 );
15452 assert_ne!(a, c);
15453
15454 assert_eq!(
15457 comp_seed(serde_json::json!({
15458 "model": "m", "prompt": "t", "seed": 0})),
15459 0,
15460 "explicit seed 0 must stay 0 — the determinism gates depend on it"
15461 );
15462 assert_eq!(
15463 comp_seed(serde_json::json!({
15464 "model": "m", "prompt": "t", "seed": 12345})),
15465 12345
15466 );
15467 assert_eq!(
15468 chat_seed(serde_json::json!({
15469 "model": "m", "messages": [{"role": "user", "content": "t"}],
15470 "seed": 777})),
15471 777
15472 );
15473 assert_eq!(
15475 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
15476 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
15477 );
15478
15479 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
15482 assert_eq!(
15483 seeds.len(),
15484 256,
15485 "fresh_seed must not collide across rapid calls"
15486 );
15487 assert!(!seeds.contains(&0));
15488 }
15489
15490 #[test]
15491 fn response_format_builds_grammar_only_when_present() {
15492 let mk = |rf: Option<serde_json::Value>| {
15496 let mut body = serde_json::json!({
15497 "model": "m", "messages": [{"role": "user", "content": "t"}]});
15498 if let Some(rf) = rf {
15499 body["response_format"] = rf;
15500 }
15501 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15502 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15503 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15504 };
15505 assert!(mk(None).unwrap().request.grammar.is_none());
15506 assert!(
15507 mk(Some(serde_json::json!({"type": "text"})))
15508 .unwrap()
15509 .request
15510 .grammar
15511 .is_none()
15512 );
15513 assert!(matches!(
15514 mk(Some(serde_json::json!({"type": "json_object"})))
15515 .unwrap()
15516 .request
15517 .grammar,
15518 Some(constrained::GrammarSpec::JsonObject)
15519 ));
15520 assert!(matches!(
15521 mk(Some(serde_json::json!({"type": "json_schema",
15522 "json_schema": {"schema": {"type": "object"}}})))
15523 .unwrap()
15524 .request
15525 .grammar,
15526 Some(constrained::GrammarSpec::JsonSchema(_))
15527 ));
15528 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
15530 }
15531
15532 #[test]
15542 fn response_format_think_table_switch_postthink_refusal() {
15543 let mk = |caps: &ModelCaps| {
15544 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15545 "model": "m", "messages": [{"role": "user", "content": "t"}],
15546 "response_format": {"type": "json_object"}}))
15547 .unwrap();
15548 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15549 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
15550 };
15551 let switch = ModelCaps {
15553 chat_ok: true,
15554 qwen_think: true,
15555 think_switch: true,
15556 ..Default::default()
15557 };
15558 let plan = mk(&switch).unwrap();
15559 assert_eq!(
15560 plan.request.think,
15561 memra_tokenizer::chat::ThinkMode::NoThink,
15562 "switch-carrying template must keep the grammar-from-token-1 path"
15563 );
15564 assert!(plan.request.grammar.is_some());
15565
15566 let postthink = ModelCaps {
15568 chat_ok: true,
15569 qwen_think: true,
15570 think_switch: false,
15571 think_close: vec![128799],
15572 ..Default::default()
15573 };
15574 let plan = mk(&postthink).unwrap();
15575 assert_ne!(
15576 plan.request.think,
15577 memra_tokenizer::chat::ThinkMode::NoThink,
15578 "post-think constrained request must keep the think channel ON"
15579 );
15580 assert!(plan.request.grammar.is_some());
15581
15582 let no_contract = ModelCaps {
15584 chat_ok: true,
15585 qwen_think: true,
15586 think_switch: false,
15587 think_close: Vec::new(),
15588 ..Default::default()
15589 };
15590 let err = match mk(&no_contract) {
15591 Err(err) => err,
15592 Ok(_) => panic!("think-forced template with no close contract must refuse"),
15593 };
15594 assert!(
15595 err.contains("think-close"),
15596 "refusal must name the missing close contract: {err}"
15597 );
15598 }
15599
15600 #[test]
15601 fn unsupported_semantic_params_are_named_rejections() {
15602 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15604 "model": "m", "messages": [{"role": "user", "content": "t"}],
15605 "response_format": {"type": "json_object"}
15606 }))
15607 .unwrap();
15608 assert!(req.response_format.is_some());
15609 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15610 "model": "m", "messages": [{"role": "user", "content": "t"}],
15611 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
15612 "user": "u-1", "stream_options": {"include_usage": true}
15613 }))
15614 .unwrap();
15615 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
15617 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
15618 assert_eq!(req.n, Some(1));
15619 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
15621 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
15622 assert_eq!(param, "logit_bias");
15623 assert_eq!(msg, "logit_bias is not supported (why)");
15624 }
15625
15626 #[test]
15627 fn completions_accept_openai_stop_forms() {
15628 for (value, expected) in [
15629 (serde_json::json!("Problem:"), vec!["Problem:"]),
15630 (
15631 serde_json::json!(["Question:", "Problem:"]),
15632 vec!["Question:", "Problem:"],
15633 ),
15634 (serde_json::Value::Null, Vec::<&str>::new()),
15635 ] {
15636 let req: CompletionReq = serde_json::from_value(serde_json::json!({
15637 "model": "plain_quant", "prompt": "task", "stop": value
15638 }))
15639 .unwrap();
15640 assert_eq!(req.stop.into_vec(), expected);
15641 }
15642 }
15643
15644 fn fake_worker_state() -> AppState {
15651 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
15652 }
15653
15654 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
15655 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
15656 }
15657
15658 struct WorkerSaw {
15663 sampler_cfg: SamplerConfig,
15664 think: ThinkMode,
15665 reasoning_effort: Option<String>,
15666 }
15667
15668 fn fake_worker_state_full(
15675 steps: usize,
15676 step_delay: std::time::Duration,
15677 caps: HashMap<String, ModelCaps>,
15678 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
15679 ) -> AppState {
15680 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15681 let health = health::WorkerHealth::new();
15682 let h = health.clone();
15683 std::thread::spawn(move || {
15684 h.mark_ready();
15685 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
15686 if let Some(tx) = &saw_tx {
15687 let _ = tx.send(WorkerSaw {
15688 sampler_cfg: req.sampler_cfg.clone(),
15689 think: req.think,
15690 reasoning_effort: req.reasoning_effort.clone(),
15691 });
15692 }
15693 worker::release_pending_admit();
15697 worker::release_admission_reservation(req.lane);
15698 h.beat_busy();
15699 if let Some(ready) = req.constraint_ready.take() {
15700 let _ = ready.send(Ok(()));
15701 }
15702 let _ = req.tx.send(Event::PromptUsage {
15703 n_prompt: 1,
15704 n_cached: 0,
15705 });
15706 for step in 0..steps {
15707 h.beat_busy();
15708 let text = if steps == 1 { "ok" } else { "x" };
15709 let _ = req.tx.send(Event::Token {
15710 id: step as u32 + 1,
15711 text: text.into(),
15712 });
15713 if !step_delay.is_zero() {
15714 std::thread::sleep(step_delay);
15715 }
15716 }
15717 let _ = req.tx.send(Event::Done {
15718 stop_reason: "Eos".into(),
15719 n_tokens: steps,
15720 n_prompt: 1,
15721 n_cached: 0,
15722 elapsed_s: 0.01,
15723 spec: None,
15724 });
15725 h.set_phase(health::PHASE_IDLE);
15726 }
15727 });
15728 for _ in 0..2000 {
15731 if health.live().is_ok() {
15732 break;
15733 }
15734 std::thread::sleep(std::time::Duration::from_millis(1));
15735 }
15736 AppState {
15737 cmd_tx,
15738 models: Arc::new(vec!["m".into()]),
15739 caps: Arc::new(caps),
15740 openrouter_metadata: Arc::new(HashMap::new()),
15741 provider_metadata: Arc::new(None),
15742 metering: None,
15743
15744 budget_tokenizers: None,
15745 api_auth: ApiAuth::default(),
15746 metrics_auth: MetricsAuth::default(),
15747 metrics: SharedMetrics::default(),
15748 started: 1,
15749 inflight: Arc::new(Default::default()),
15750 tenant_inflight: Arc::new(Default::default()),
15751 health,
15752 bg: None,
15753 }
15754 }
15755
15756 #[tokio::test]
15757 async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
15758 let _l = DRAIN_LOCK.lock().unwrap();
15759 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
15760 let normal_state = st.clone();
15761 let normal = tokio::spawn(async move {
15762 chat_completions(
15763 State(normal_state),
15764 axum::http::HeaderMap::new(),
15765 None,
15766 Json(
15767 serde_json::from_value(serde_json::json!({
15768 "model": "m",
15769 "messages": [{"role": "user", "content": "keep decoding"}],
15770 }))
15771 .unwrap(),
15772 ),
15773 )
15774 .await
15775 });
15776 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
15777
15778 let mut deep = serde_json::json!({"type": "string"});
15779 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
15780 deep = serde_json::json!({"allOf": [deep]});
15781 }
15782 let bad = chat_completions(
15783 State(st.clone()),
15784 axum::http::HeaderMap::new(),
15785 None,
15786 Json(
15787 serde_json::from_value(serde_json::json!({
15788 "model": "m",
15789 "messages": [{"role": "user", "content": "bad schema"}],
15790 "response_format": {
15791 "type": "json_schema",
15792 "json_schema": {"schema": deep},
15793 },
15794 }))
15795 .unwrap(),
15796 ),
15797 )
15798 .await;
15799 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
15800 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
15801 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
15802 .await
15803 .unwrap();
15804 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15805 assert!(
15806 payload["error"]["message"]
15807 .as_str()
15808 .unwrap()
15809 .contains("maximum nesting depth")
15810 );
15811 assert!(
15812 !normal.is_finished(),
15813 "bad schema stalled or replaced the normal decode"
15814 );
15815
15816 let normal_response = normal.await.unwrap();
15817 assert_eq!(normal_response.status(), StatusCode::OK);
15818 let snapshot = st.health.snapshot();
15819 assert!(
15820 st.health.live().is_ok(),
15821 "normal decode left health stalled"
15822 );
15823 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
15824 }
15825
15826 #[tokio::test]
15827 async fn valid_response_format_preflight_preserves_generation() {
15828 let _l = DRAIN_LOCK.lock().unwrap();
15829 let response = chat_completions(
15830 State(fake_worker_state()),
15831 axum::http::HeaderMap::new(),
15832 None,
15833 Json(
15834 serde_json::from_value(serde_json::json!({
15835 "model": "m",
15836 "messages": [{"role": "user", "content": "valid schema"}],
15837 "response_format": {"type": "json_object"},
15838 }))
15839 .unwrap(),
15840 ),
15841 )
15842 .await;
15843 assert_eq!(response.status(), StatusCode::OK);
15844 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15845 .await
15846 .unwrap();
15847 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15848 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
15849 }
15850
15851 #[tokio::test]
15852 async fn unknown_model_refuses_model_not_found_before_admission() {
15853 let _l = DRAIN_LOCK.lock().unwrap();
15854 let response = chat_completions(
15859 State(fake_worker_state()),
15860 axum::http::HeaderMap::new(),
15861 None,
15862 Json(
15863 serde_json::from_value(serde_json::json!({
15864 "model": "qwen/qwen3.8-27b-typo",
15865 "messages": [{"role": "user", "content": "hi"}],
15866 }))
15867 .unwrap(),
15868 ),
15869 )
15870 .await;
15871 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15872 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15873 .await
15874 .unwrap();
15875 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15876 assert_eq!(payload["error"]["code"], "model_not_found");
15877 assert_eq!(payload["error"]["type"], "invalid_request_error");
15878
15879 let response = completions(
15881 State(fake_worker_state()),
15882 axum::http::HeaderMap::new(),
15883 None,
15884 Json(
15885 serde_json::from_value(serde_json::json!({
15886 "model": "nope",
15887 "prompt": "hi",
15888 }))
15889 .unwrap(),
15890 ),
15891 )
15892 .await;
15893 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15894 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15895 .await
15896 .unwrap();
15897 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15898 assert_eq!(payload["error"]["code"], "model_not_found");
15899 }
15900
15901 const METRICS_KEY_ACME: &str = "completion-acme-secret";
15902 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
15903
15904 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
15905 let spec = format!(
15906 "acme:{},blue:{}",
15907 auth::sha256_hex(METRICS_KEY_ACME),
15908 auth::sha256_hex(METRICS_KEY_BLUE),
15909 );
15910 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
15911 let mut st = fake_worker_state();
15912 st.api_auth.keyring = Some(keyring);
15913 st.metrics_auth = MetricsAuth::new(
15914 true,
15915 st.api_auth.configured(),
15916 metrics_token.map(str::to_string),
15917 );
15918 {
15919 let mut metrics = st.metrics.lock().unwrap();
15920 metrics.admitted = 17;
15921 metrics.prompt_tokens_in = 400;
15922 metrics.cached_tokens_in = 60;
15923 metrics.prefix_hits = 2;
15924 metrics.prefix_misses = 3;
15925 metrics.prefix_inserts = 5;
15926 metrics.prefix_evictions = 7;
15927 metrics.prefix_skips_budget = 9;
15928 metrics.prefix_skips_pinned = 10;
15929 metrics.prefix_hit_tokens = 11;
15930 metrics.lcp_hist[4] = 13;
15931 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
15932 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
15933 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
15934 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
15935 metrics.prefix_entries = 29;
15936 metrics.prefix_bytes = 31;
15937 metrics.active_sessions = 3;
15938 metrics.queued_requests = 5;
15939 metrics.admission_inflight.insert("m".into(), 4);
15940 metrics
15941 .admission_booked_bytes
15942 .insert("m".into(), 41_000_000);
15943 metrics.continuation_pool_entries = 7;
15944 metrics.spec_pool_entries = 11;
15945 metrics.cuda_driver_free_bytes = 13;
15946 metrics.cuda_pool_reserved_bytes = 17;
15947 metrics.cuda_pool_used_bytes = 19;
15948 metrics.cuda_pool_cached_bytes = 23;
15949 metrics.batch_size_last = 37;
15950 metrics.spec.insert(
15951 "m".into(),
15952 memra_engine::spec::SpecTelemetry {
15953 rounds: 2,
15954 drafted: 6,
15955 accepted: 4,
15956 ..Default::default()
15957 },
15958 );
15959 let mut spec_window = memra_engine::spec::SpecTelemetry {
15960 rounds: 4,
15961 drafted: 12,
15962 accepted: 6,
15963 ..Default::default()
15964 };
15965 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
15966 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
15967 metrics.spec_window.insert("m".into(), spec_window);
15968 metrics.constraint_compiler_fail_closed.insert(
15969 "m".into(),
15970 Arc::new(std::sync::atomic::AtomicBool::new(true)),
15971 );
15972 }
15973 st
15974 }
15975
15976 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15977 let mut headers = HeaderMap::new();
15978 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15979 let response = get_metrics(State(st), headers).await;
15980 assert_eq!(response.status(), StatusCode::OK);
15981 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15982 .await
15983 .unwrap();
15984 serde_json::from_slice(&bytes).unwrap()
15985 }
15986
15987 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15988 let mut headers = HeaderMap::new();
15989 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15990 let response = yield_metrics(State(st), headers).await;
15991 assert_eq!(response.status(), StatusCode::OK);
15992 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15993 .await
15994 .unwrap();
15995 serde_json::from_slice(&bytes).unwrap()
15996 }
15997
15998 #[test]
15999 fn exposed_open_bind_is_refused_before_server_start() {
16000 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
16001 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
16002
16003 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
16004 assert!(err.contains("refusing unauthenticated non-loopback bind"));
16005 assert!(err.contains("MEMRA_API_KEY"));
16006 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
16007 assert!(validate_bind_security("[::]:8000", false, false).is_err());
16008
16009 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
16010 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
16011 }
16012
16013 #[tokio::test]
16014 async fn keyed_metrics_require_and_accept_api_bearer() {
16015 let mut st = fake_worker_state();
16016 st.api_auth.single_key = Some(Arc::from("completion-secret"));
16017 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
16018
16019 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
16020 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
16021 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
16022 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
16023
16024 let mut headers = HeaderMap::new();
16025 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
16026 assert_eq!(
16027 get_metrics(State(st.clone()), headers.clone())
16028 .await
16029 .status(),
16030 StatusCode::OK,
16031 );
16032 let body = metrics_json(st.clone(), "completion-secret").await;
16033 assert!(
16034 body.get("admitted").is_some(),
16035 "the legacy single-key domain keeps cumulative counters",
16036 );
16037 assert!(
16038 body.get("active_sessions").is_none(),
16039 "a static completion key is not an operator metrics principal",
16040 );
16041 assert_eq!(
16042 yield_metrics(State(st), headers).await.status(),
16043 StatusCode::OK
16044 );
16045 }
16046
16047 #[tokio::test]
16048 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
16049 let st = multi_key_metrics_state(None);
16050 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
16051 assert_eq!(
16052 body.as_object().unwrap().len(),
16053 2,
16054 "completion metrics must contain only tenant-scoped rows",
16055 );
16056 let tenants = body["tenants"].as_object().unwrap();
16057 assert_eq!(tenants.len(), 1);
16058 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
16059 assert!(!tenants.contains_key("t:blue"));
16060 let adsd = body["adsd_suspect_total"].as_object().unwrap();
16061 assert_eq!(adsd.len(), 1);
16062 assert_eq!(adsd["t:acme"], 1);
16063 assert!(!adsd.contains_key("t:blue"));
16064
16065 let mut headers = HeaderMap::new();
16066 headers.insert(
16067 "authorization",
16068 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
16069 );
16070 assert_eq!(
16071 yield_metrics(State(st), headers).await.status(),
16072 StatusCode::FORBIDDEN,
16073 "the process-wide yield view requires an operator metrics token",
16074 );
16075 }
16076
16077 #[tokio::test]
16078 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
16079 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
16080 for operator_only in [
16081 "prefix_cache_entries",
16082 "prefix_cache_bytes",
16083 "prefix_cache_skips_budget",
16084 "prefix_cache_skips_pinned",
16085 "active_sessions",
16086 "queued_requests",
16087 "admission_inflight",
16088 "admission_booked_bytes",
16089 "continuation_pool_entries",
16090 "spec_pool_entries",
16091 "cuda_driver_free_bytes",
16092 "cuda_pool_reserved_bytes",
16093 "cuda_pool_used_bytes",
16094 "cuda_pool_cached_bytes",
16095 "constraint_compiler_fail_closed",
16096 "serve_idle_seconds",
16097 "spec",
16098 "spec_tau",
16099 "spec_accept_by_position",
16100 "dual_pp",
16101 "peer_probe_bypassed",
16102 "peer_probe_boundary_copies",
16103 "peer_probe_runtime_reprobes",
16104 "peer_probe_runtime_failures",
16105 "peer_probe_deferred_total",
16106 "peer_probe_integrity_degraded",
16107 "peer_probe_degraded_to_host_bounce",
16108 ] {
16109 assert!(
16110 body.get(operator_only).is_none(),
16111 "tenant metrics must not expose operator field {operator_only}",
16112 );
16113 }
16114 }
16115
16116 #[test]
16117 fn populated_spec_acceptance_metrics_are_operator_only() {
16118 for scope in [
16119 MetricsScope::CompletionDomain,
16120 MetricsScope::Tenant("t:acme".into()),
16121 ] {
16122 let mut body = json!({});
16123 insert_spec_acceptance_metrics(&mut body, &scope, || {
16124 panic!("tenant scope evaluated the process-wide spec snapshot")
16125 });
16126 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
16127 assert!(
16128 body.get("spec_accept_by_position").is_none(),
16129 "{scope:?} leaked the accept histogram"
16130 );
16131 }
16132
16133 let mut telemetry = memra_engine::spec::SpecTelemetry {
16134 rounds: 4,
16135 drafted: 12,
16136 accepted: 6,
16137 ..Default::default()
16138 };
16139 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
16140 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
16141 let mut body = json!({});
16142 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
16143 HashMap::from([("model-a".to_string(), telemetry)])
16144 });
16145 assert_eq!(body["spec_tau"]["model-a"], 1.5);
16146 let histogram = &body["spec_accept_by_position"]["model-a"];
16147 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
16148 assert_eq!(histogram["rounds"], 4);
16149 assert_eq!(histogram["offered"], json!([4, 4, 4]));
16150 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
16151 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
16152 }
16153
16154 #[test]
16155 fn populated_dual_pp_metrics_are_operator_only() {
16156 let populated = DualPpMetricsSnapshot {
16157 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
16158 stage_samples: [1, 1, 1, 1],
16159 dropped_timing_samples: 0,
16160 overlaps: 17,
16161 slot_pairs: 19,
16162 slot_uses: [19, 19],
16163 slot_collisions: 0,
16164 };
16165 for scope in [
16166 MetricsScope::CompletionDomain,
16167 MetricsScope::Tenant("t:acme".into()),
16168 ] {
16169 let mut body = json!({});
16170 insert_dual_pp_metrics(&mut body, &scope, || populated);
16171 assert!(
16172 body.get("dual_pp").is_none(),
16173 "{scope:?} leaked dual PP topology"
16174 );
16175 }
16176
16177 let mut body = json!({});
16178 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
16179 assert_eq!(body["dual_pp"]["overlaps"], 17);
16180 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
16181 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
16182 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
16183 assert_eq!(
16184 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
16185 1.0
16186 );
16187 }
16188
16189 #[test]
16190 fn peer_probe_metrics_are_operator_only() {
16191 let populated = memra_engine::pp::PeerProbeMetrics {
16192 bypassed: 1,
16193 boundary_copies: 8_192,
16194 runtime_probes: 1,
16195 runtime_failures: 0,
16196 deferred_total: 4,
16197 integrity_degraded: true,
16198 degraded_to_host_bounce: true,
16199 };
16200 for scope in [
16201 MetricsScope::CompletionDomain,
16202 MetricsScope::Tenant("t:acme".into()),
16203 ] {
16204 let mut body = json!({});
16205 insert_peer_probe_metrics(&mut body, &scope, || populated);
16206 assert!(body.get("peer_probe_bypassed").is_none());
16207 }
16208
16209 let mut body = json!({});
16210 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
16211 assert_eq!(body["peer_probe_bypassed"], 1);
16212 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
16213 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
16214 assert_eq!(body["peer_probe_runtime_failures"], 0);
16215 assert_eq!(body["peer_probe_deferred_total"], 4);
16216 assert_eq!(body["peer_probe_integrity_degraded"], true);
16217 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
16218 }
16219
16220 #[tokio::test]
16221 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
16222 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
16223 for operator_only in [
16224 "lcp_histogram",
16225 "cache_hit_token_ratio",
16226 "prefix_cache_hits",
16227 "prefix_cache_misses",
16228 "prefix_cache_inserts",
16229 "prefix_cache_evictions",
16230 "prefix_cache_skips_budget",
16231 "prefix_cache_skips_pinned",
16232 "prefix_cache_hit_tokens",
16233 ] {
16234 assert!(
16235 tenant_body.get(operator_only).is_none(),
16236 "tenant metrics must not expose global prefix field {operator_only}",
16237 );
16238 }
16239 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
16240 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
16241 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
16242 assert_eq!(
16243 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
16244 0.4
16245 );
16246
16247 let operator_body = metrics_json(
16248 multi_key_metrics_state(Some("scrape-secret")),
16249 "scrape-secret",
16250 )
16251 .await;
16252 assert_eq!(operator_body["prefix_cache_hits"], 2);
16253 assert_eq!(operator_body["prefix_cache_misses"], 3);
16254 assert_eq!(operator_body["prefix_cache_inserts"], 5);
16255 assert_eq!(operator_body["prefix_cache_evictions"], 7);
16256 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
16257 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
16258 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
16259 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
16260 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
16261 }
16262
16263 #[tokio::test]
16264 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
16265 let st = multi_key_metrics_state(Some("scrape-secret"));
16266 let mut completion_headers = HeaderMap::new();
16267 completion_headers.insert(
16268 "authorization",
16269 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
16270 );
16271 assert_eq!(
16272 get_metrics(State(st.clone()), completion_headers.clone())
16273 .await
16274 .status(),
16275 StatusCode::FORBIDDEN,
16276 );
16277 assert_eq!(
16278 yield_metrics(State(st.clone()), completion_headers)
16279 .await
16280 .status(),
16281 StatusCode::FORBIDDEN,
16282 );
16283
16284 let body = metrics_json(st.clone(), "scrape-secret").await;
16285 let tenants = body["tenants"].as_object().unwrap();
16286 assert_eq!(tenants.len(), 2);
16287 assert!(tenants.contains_key("t:acme"));
16288 assert!(tenants.contains_key("t:blue"));
16289 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
16290 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
16291 assert_eq!(body["active_sessions"], 3);
16292 assert_eq!(body["queued_requests"], 5);
16293 assert_eq!(body["admission_inflight"]["m"], 4);
16295 assert_eq!(body["admission_booked_bytes"]["m"], 41_000_000);
16296 assert_eq!(body["prefix_cache_bytes"], 31);
16297 assert_eq!(body["cuda_driver_free_bytes"], 13);
16298 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
16299 assert_eq!(body["spec"]["m"]["drafted"], 6);
16300 assert_eq!(body["spec_tau"]["m"], 1.5);
16301 assert_eq!(
16302 body["spec_accept_by_position"]["m"]["accepted"],
16303 json!([3, 2, 1])
16304 );
16305 let yield_body = yield_metrics_json(st, "scrape-secret").await;
16306 assert_eq!(yield_body["batch_size_last"], 37);
16307 }
16308
16309 #[tokio::test]
16310 async fn metrics_token_protects_public_override_without_api_keys() {
16311 let mut st = fake_worker_state();
16312 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
16313
16314 assert_eq!(
16315 get_metrics(State(st.clone()), HeaderMap::new())
16316 .await
16317 .status(),
16318 StatusCode::UNAUTHORIZED,
16319 );
16320 let mut headers = HeaderMap::new();
16321 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
16322 assert_eq!(
16323 get_metrics(State(st.clone()), headers.clone())
16324 .await
16325 .status(),
16326 StatusCode::OK,
16327 );
16328 assert_eq!(
16329 yield_metrics(State(st), headers).await.status(),
16330 StatusCode::OK
16331 );
16332 }
16333
16334 #[tokio::test]
16335 async fn no_key_loopback_metrics_remain_open_for_development() {
16336 let mut st = fake_worker_state();
16337 st.metrics_auth = MetricsAuth::new(true, false, None);
16338 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
16339 assert_eq!(response.status(), StatusCode::OK);
16340 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16341 .await
16342 .unwrap();
16343 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16344 assert!(
16345 body.get("active_sessions").is_some(),
16346 "no-key loopback development keeps full operator visibility",
16347 );
16348 assert_eq!(
16349 yield_metrics(State(st), HeaderMap::new()).await.status(),
16350 StatusCode::OK,
16351 );
16352 }
16353
16354 #[test]
16355 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
16356 let metrics = SharedMetrics::default();
16357 let rl = RateLimit::compute(4, 1, &metrics);
16359 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
16360 let rl = RateLimit::compute(4, 3, &metrics);
16361 assert_eq!(rl.remaining, 1);
16362 let rl = RateLimit::compute(4, 4, &metrics);
16364 assert_eq!(rl.remaining, 0);
16365 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
16366 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
16368 let m = worker::Metrics {
16370 completed: 2,
16371 tokens_out: 200,
16372 step_p50_ms: 20.0,
16373 ..Default::default()
16374 };
16375 assert_eq!(reset_estimate_s(&m), 2); }
16377
16378 #[test]
16379 fn inflight_guard_counts_up_and_frees_on_drop() {
16380 let counts: InflightCounts = Arc::new(Default::default());
16381 let tenants: TenantGauge = Arc::new(Default::default());
16382 let (g1, n1, t1) = InflightGuard::try_acquire(
16383 counts.clone(),
16384 lanes::Lane::Interactive,
16385 tenants.clone(),
16386 "acme",
16387 None,
16388 )
16389 .unwrap();
16390 let (g2, n2, t2) = InflightGuard::try_acquire(
16391 counts.clone(),
16392 lanes::Lane::Interactive,
16393 tenants.clone(),
16394 "acme",
16395 None,
16396 )
16397 .unwrap();
16398 assert_eq!((n1, n2), (1, 2));
16399 assert_eq!((t1, t2), (1, 2));
16401 let (gj, nj, tj) = InflightGuard::try_acquire(
16403 counts.clone(),
16404 lanes::Lane::Judge,
16405 tenants.clone(),
16406 "blue",
16407 None,
16408 )
16409 .unwrap();
16410 assert_eq!((nj, tj), (1, 1));
16411 drop(g1);
16412 drop(gj);
16413 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
16414 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
16415 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
16416 assert!(tenants.lock().unwrap().get("blue").is_none());
16418 drop(g2);
16419 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
16420 assert!(tenants.lock().unwrap().is_empty());
16421 }
16422
16423 #[test]
16424 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
16425 let counts: InflightCounts = Arc::new(Default::default());
16426 let tenants: TenantGauge = Arc::new(Default::default());
16427 let start = Arc::new(std::sync::Barrier::new(3));
16428 let attempted = Arc::new(std::sync::Barrier::new(3));
16429 let mut joins = Vec::new();
16430 for _ in 0..2 {
16431 let counts = counts.clone();
16432 let tenants = tenants.clone();
16433 let start = start.clone();
16434 let attempted = attempted.clone();
16435 joins.push(std::thread::spawn(move || {
16436 start.wait();
16437 let result = InflightGuard::try_acquire(
16438 counts,
16439 lanes::Lane::Interactive,
16440 tenants,
16441 "preview_001",
16442 Some(1),
16443 );
16444 let won = result.is_ok();
16445 attempted.wait(); drop(result);
16447 won
16448 }));
16449 }
16450 start.wait();
16451 attempted.wait();
16452 let wins = joins
16453 .into_iter()
16454 .map(|join| join.join().unwrap())
16455 .filter(|won| *won)
16456 .count();
16457 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
16458 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
16459 assert!(tenants.lock().unwrap().is_empty());
16460 }
16461
16462 #[tokio::test]
16463 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
16464 let st = fake_worker_state();
16465 let tenant = auth::TenantCtx {
16466 tenant: "preview_001".into(),
16467 lane_class: auth::LaneClass::Interactive,
16468 rate_limit: Some(1),
16469 key_prefix: None,
16470 };
16471 let first_env = Envelope::new(true);
16472 let (guard, first_rl) =
16473 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
16474 Ok(slot) => slot,
16475 Err(_) => panic!("the first request must acquire the tenant slot"),
16476 };
16477 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
16478
16479 let second_env = Envelope::new(true);
16480 let response =
16481 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
16482 Err(response) => response,
16483 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
16484 };
16485 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
16486 assert_eq!(response.headers()["retry-after"], "2");
16487 assert_eq!(response.headers()["retry-after-ms"], "2000");
16488 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
16489 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
16490 assert_eq!(response.headers()["x-request-id"], second_env.id);
16491 assert_eq!(
16492 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16493 1,
16494 "rejected request must not consume a lane slot"
16495 );
16496 assert_eq!(
16497 st.tenant_inflight
16498 .lock()
16499 .unwrap()
16500 .get("preview_001")
16501 .copied(),
16502 Some(1),
16503 "rejected request must not increment the tenant gauge"
16504 );
16505 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16506 .await
16507 .unwrap();
16508 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16509 assert_eq!(payload["error"]["type"], "rate_limit_error");
16510 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
16511 assert!(
16512 payload["error"]["message"]
16513 .as_str()
16514 .unwrap()
16515 .contains("concurrent request limit")
16516 );
16517
16518 drop(guard);
16519 let _ = InflightGuard::try_acquire(
16520 st.inflight.clone(),
16521 lanes::Lane::Interactive,
16522 st.tenant_inflight.clone(),
16523 "preview_001",
16524 Some(1),
16525 )
16526 .expect("slot must reopen after the in-flight request completes");
16527 }
16528
16529 #[test]
16530 fn tenant_rate_limit_override_is_min_with_global_cap() {
16531 let metrics = SharedMetrics::default();
16532 let unlimited = auth::TenantCtx::default_tenant();
16533 let capped = auth::TenantCtx {
16534 tenant: "acme".into(),
16535 lane_class: auth::LaneClass::Interactive,
16536 rate_limit: Some(2),
16537 key_prefix: None,
16538 };
16539 let global = lane_cap(lanes::Lane::Interactive);
16540 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
16542 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16543 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
16545 assert_eq!((rl.limit, rl.remaining), (2, 1));
16546 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
16547 assert_eq!(rl.remaining, 0);
16548 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
16549 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
16553 assert_eq!(rl.remaining, 0);
16554 let wide = auth::TenantCtx {
16555 rate_limit: Some(global + 100),
16556 ..capped.clone()
16557 };
16558 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
16559 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16560 }
16561
16562 #[test]
16563 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
16564 let batch = auth::TenantCtx {
16565 tenant: "bulk".into(),
16566 lane_class: auth::LaneClass::Batch,
16567 rate_limit: None,
16568 key_prefix: None,
16569 };
16570 let interactive = auth::TenantCtx::default_tenant();
16571 let hdr = |v: Option<&str>| {
16572 let mut h = axum::http::HeaderMap::new();
16573 if let Some(v) = v {
16574 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16575 }
16576 h
16577 };
16578 assert_eq!(
16580 lane_for_tenant(&hdr(None), &interactive).unwrap(),
16581 lanes::Lane::Interactive
16582 );
16583 assert_eq!(
16584 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
16585 lanes::Lane::Judge
16586 );
16587 assert_eq!(
16589 lane_for_tenant(&hdr(None), &batch).unwrap(),
16590 lanes::Lane::Harvest
16591 );
16592 assert_eq!(
16593 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
16594 lanes::Lane::Judge
16595 );
16596 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
16597 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16598 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
16600 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16601 }
16602
16603 #[tokio::test]
16604 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
16605 let hdr = |v: &str| {
16610 let mut h = axum::http::HeaderMap::new();
16611 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16612 h
16613 };
16614 let body = |resp: Response| async move {
16615 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16616 .await
16617 .unwrap();
16618 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
16619 };
16620
16621 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
16622 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16623 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16624 let payload = body(resp).await;
16625 assert!(
16626 payload["error"].is_object(),
16627 "bare-string error body: {payload}"
16628 );
16629 assert_eq!(payload["error"]["type"], "invalid_request_error");
16630 assert_eq!(payload["error"]["param"], "x-lane");
16631 assert_eq!(payload["error"]["code"], "invalid_lane");
16632
16633 let batch = auth::TenantCtx {
16634 tenant: "bulk".into(),
16635 lane_class: auth::LaneClass::Batch,
16636 rate_limit: None,
16637 key_prefix: None,
16638 };
16639 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
16640 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16641 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16642 let payload = body(resp).await;
16643 assert_eq!(payload["error"]["type"], "authentication_error");
16644 assert_eq!(payload["error"]["param"], "x-lane");
16645 }
16646
16647 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
16650
16651 #[tokio::test]
16652 async fn responses_carry_rate_limit_headers_and_slot_frees() {
16653 let _l = DRAIN_LOCK.lock().unwrap();
16654 let st = fake_worker_state();
16655 let resp = chat_completions(
16658 State(st.clone()),
16659 axum::http::HeaderMap::new(),
16660 None,
16661 Json(
16662 serde_json::from_value(serde_json::json!({
16663 "model": "m", "messages": [{"role": "user", "content": "t"}]
16664 }))
16665 .unwrap(),
16666 ),
16667 )
16668 .await;
16669 assert_eq!(resp.status(), StatusCode::OK);
16670 let h = resp.headers();
16671 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
16672 let remaining: usize = h["x-ratelimit-remaining"]
16673 .to_str()
16674 .unwrap()
16675 .parse()
16676 .unwrap();
16677 assert_eq!(remaining, limit - 1);
16678 assert_eq!(h["x-ratelimit-reset"], "0");
16679 assert_eq!(
16680 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16681 0,
16682 "slot must free at completion"
16683 );
16684 let resp = completions(
16687 State(st.clone()),
16688 axum::http::HeaderMap::new(),
16689 None,
16690 Json(
16691 serde_json::from_value(serde_json::json!({
16692 "model": "m", "prompt": "t", "stream": true
16693 }))
16694 .unwrap(),
16695 ),
16696 )
16697 .await;
16698 assert_eq!(resp.status(), StatusCode::OK);
16699 assert!(resp.headers().contains_key("x-ratelimit-limit"));
16700 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
16701 assert!(resp.headers().contains_key("x-ratelimit-reset"));
16702 assert_eq!(
16703 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16704 1,
16705 "stream in flight holds the slot"
16706 );
16707 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
16708 .await
16709 .unwrap();
16710 assert_eq!(
16711 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16712 0,
16713 "slot must free when the stream completes"
16714 );
16715 }
16716
16717 #[tokio::test]
16718 async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
16719 let _l = DRAIN_LOCK.lock().unwrap();
16720 let mut st = fake_worker_state();
16721 let mock = MockMetering::admit_all();
16722 st.metering = Some(mock.clone());
16723
16724 let nonstream = chat_completions(
16725 State(st.clone()),
16726 HeaderMap::new(),
16727 None,
16728 Json(
16729 serde_json::from_value(json!({
16730 "model": "m",
16731 "messages": [{"role": "user", "content": "t"}],
16732 }))
16733 .unwrap(),
16734 ),
16735 )
16736 .await;
16737 assert_eq!(nonstream.status(), StatusCode::OK);
16738 let nonstream_id = nonstream.headers()["x-request-id"]
16739 .to_str()
16740 .unwrap()
16741 .to_string();
16742
16743 let stream = completions(
16744 State(st),
16745 HeaderMap::new(),
16746 None,
16747 Json(
16748 serde_json::from_value(json!({
16749 "model": "m",
16750 "prompt": "t",
16751 "stream": true,
16752 }))
16753 .unwrap(),
16754 ),
16755 )
16756 .await;
16757 assert_eq!(stream.status(), StatusCode::OK);
16758 let stream_id = stream.headers()["x-request-id"]
16759 .to_str()
16760 .unwrap()
16761 .to_string();
16762 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
16763 .await
16764 .unwrap();
16765
16766 let events = mock.events();
16770 let opened: Vec<&str> = events
16771 .iter()
16772 .filter_map(|e| match e {
16773 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
16774 _ => None,
16775 })
16776 .collect();
16777 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
16778 let completes = events
16779 .iter()
16780 .filter(|e| {
16781 matches!(
16782 e,
16783 MeterEvent::Complete {
16784 prompt: 1,
16785 cached: 0,
16786 completion: 1,
16787 }
16788 )
16789 })
16790 .count();
16791 assert_eq!(
16792 completes, 2,
16793 "both surfaces settle complete with worker-truth usage: {events:?}"
16794 );
16795 }
16796
16797 #[tokio::test]
16798 async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
16799 let _l = DRAIN_LOCK.lock().unwrap();
16800 let mock = MockMetering::with_limits(vec![
16806 ReserveScript::Insufficient,
16807 ReserveScript::Admit { with_permit: false },
16808 ReserveScript::Blocked,
16809 ReserveScript::Admit { with_permit: true },
16810 ]);
16811 let mut st = fake_worker_state();
16812 st.metering = Some(mock.clone());
16813
16814 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
16816 assert_eq!(metrics.status(), StatusCode::OK);
16817 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
16818 .await
16819 .unwrap();
16820 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
16821 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
16822 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
16823 assert_eq!(metrics_body["budget_source_available"], true);
16824
16825 let request = || {
16826 Json(
16827 serde_json::from_value::<CompletionReq>(json!({
16828 "model": "m",
16829 "prompt_ids": [1],
16830 "max_tokens": 1,
16831 }))
16832 .unwrap(),
16833 )
16834 };
16835
16836 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16837 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
16838 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
16839 .await
16840 .unwrap();
16841 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
16842 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
16843 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
16844
16845 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16846 assert_eq!(included.status(), StatusCode::OK);
16847
16848 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16851 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
16852
16853 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16854 assert_eq!(admitted.status(), StatusCode::OK);
16855
16856 let events = mock.events();
16857 let terminal: Vec<&MeterEvent> = events
16858 .iter()
16859 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
16860 .collect();
16861 assert_eq!(
16862 terminal.len(),
16863 4,
16864 "four requests, four terminal settles: {events:?}"
16865 );
16866 assert!(matches!(
16867 terminal[0],
16868 MeterEvent::Reject { status: 402, .. }
16869 ));
16870 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
16871 assert!(matches!(
16872 terminal[2],
16873 MeterEvent::Reject { status: 402, .. }
16874 ));
16875 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
16876 let permits: Vec<bool> = events
16878 .iter()
16879 .filter_map(|e| match e {
16880 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
16881 _ => None,
16882 })
16883 .collect();
16884 assert_eq!(
16885 permits,
16886 vec![false, false, false, true],
16887 "the permit rides the receipt exactly when reserve minted one: {events:?}"
16888 );
16889 }
16890
16891 #[tokio::test]
16895 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
16896 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
16897 let mut st = fake_worker_state();
16898 st.metering = Some(mock.clone());
16899 let tenant = auth::TenantCtx {
16900 tenant: "acme".into(),
16901 lane_class: auth::LaneClass::Interactive,
16902 rate_limit: None,
16903 key_prefix: Some("mk-acme-testprefix00".into()),
16904 };
16905 let mut request = gate_request(1, 1);
16906 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
16907 .expect_err("a capped key must be refused at admission");
16908 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
16909 let (response, outcome) = rejection.into_response();
16910 assert_eq!(outcome, "key_spend_cap_reached");
16911 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
16912 let body = body_value(response).await;
16913 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
16914 assert!(
16915 body["error"]["message"].as_str().unwrap().contains("cap"),
16916 "the 402 must point at the KEY's cap, not tenant credit: {body}"
16917 );
16918 let events = mock.events();
16919 assert!(
16920 events.contains(&MeterEvent::Reserve {
16921 tenant: "acme".into(),
16922 principal: Some("mk-acme-testprefix00".into()),
16923 model: "qwen/qwen3.8-27b".into(),
16924 }),
16925 "the key prefix must reach reserve: {events:?}"
16926 );
16927 }
16928
16929 #[tokio::test]
16930 async fn streaming_client_disconnect_records_partial_usage_and_cost() {
16931 let _l = DRAIN_LOCK.lock().unwrap();
16932 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
16933 let mock = MockMetering::admit_all();
16934 st.metering = Some(mock.clone());
16935
16936 let response = completions(
16937 State(st),
16938 HeaderMap::new(),
16939 None,
16940 Json(
16941 serde_json::from_value(json!({
16942 "model": "m",
16943 "prompt": "disconnect after one delta",
16944 "stream": true,
16945 }))
16946 .unwrap(),
16947 ),
16948 )
16949 .await;
16950 assert_eq!(response.status(), StatusCode::OK);
16951 let request_id = response.headers()["x-request-id"]
16952 .to_str()
16953 .unwrap()
16954 .to_string();
16955 let mut body = Box::pin(response.into_body().into_data_stream());
16956 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
16957 .await
16958 .expect("stream ended before first delta")
16959 .expect("stream body failed");
16960 assert!(
16961 is_sse_data_frame(&first),
16962 "first frame was not SSE data: {first:?}"
16963 );
16964 drop(body);
16965
16966 let mut dropped = None;
16969 for _ in 0..500 {
16970 if let Some(event) = mock
16971 .events()
16972 .into_iter()
16973 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
16974 {
16975 dropped = Some(event);
16976 break;
16977 }
16978 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
16979 }
16980 let events = mock.events();
16981 assert!(
16982 events
16983 .iter()
16984 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
16985 "the receipt was opened under the caller-visible request id: {events:?}"
16986 );
16987 assert_eq!(
16988 dropped,
16989 Some(MeterEvent::Dropped {
16990 prompt: 1,
16991 cached: 0,
16992 completion: 1,
16993 }),
16994 "a client disconnect must leave the partial counts on the dropped receipt \
16995 (the implementation prices that drop): {events:?}"
16996 );
16997 }
16998
16999 #[tokio::test]
17000 async fn draining_rejects_new_requests_with_503_and_retry_after() {
17001 let _l = DRAIN_LOCK.lock().unwrap();
17002 let st = fake_worker_state();
17003 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
17004 let resp = chat_completions(
17006 State(st.clone()),
17007 axum::http::HeaderMap::new(),
17008 None,
17009 Json(
17010 serde_json::from_value(serde_json::json!({
17011 "model": "m", "messages": [{"role": "user", "content": "t"}]
17012 }))
17013 .unwrap(),
17014 ),
17015 )
17016 .await;
17017 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17018 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
17023 let ra_s: u64 = ra
17024 .parse()
17025 .expect("Retry-After must be integer delay-seconds");
17026 assert!(
17027 ra_s > 0 && ra_s <= 60,
17028 "Retry-After {ra_s}s is outside the honored window"
17029 );
17030 let ra_ms: u64 = resp.headers()["retry-after-ms"]
17031 .to_str()
17032 .unwrap()
17033 .parse()
17034 .unwrap();
17035 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
17036 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17037 .await
17038 .unwrap();
17039 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17040 assert!(
17041 payload["error"]["message"]
17042 .as_str()
17043 .unwrap()
17044 .contains("draining")
17045 );
17046 assert_eq!(payload["error"]["type"], "server_error");
17047 assert_eq!(payload["error"]["code"], "draining");
17048 let resp = completions(
17049 State(st.clone()),
17050 axum::http::HeaderMap::new(),
17051 None,
17052 Json(
17053 serde_json::from_value(serde_json::json!({
17054 "model": "m", "prompt": "t"
17055 }))
17056 .unwrap(),
17057 ),
17058 )
17059 .await;
17060 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17061 assert!(resp.headers().contains_key("retry-after"));
17062 assert_eq!(
17063 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
17064 0,
17065 "rejected requests must not hold slots"
17066 );
17067 let resp = health_live(State(st.clone())).await.into_response();
17070 assert_eq!(
17071 resp.status(),
17072 StatusCode::OK,
17073 "a drain must not look like a liveness fault"
17074 );
17075 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17076 .await
17077 .unwrap();
17078 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17079 assert_eq!(payload["status"], "draining");
17080 let resp = health_ready(State(st.clone())).await.into_response();
17082 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17083 let retry_s = drain_deadline_s().clamp(1, 60);
17084 let retry_s_text = retry_s.to_string();
17085 let retry_ms_text = (retry_s * 1000).to_string();
17086 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
17087 assert_eq!(
17088 resp.headers().get("retry-after-ms").unwrap(),
17089 retry_ms_text.as_str()
17090 );
17091 assert_ne!(
17092 resp.headers()
17093 .get("x-should-retry")
17094 .and_then(|v| v.to_str().ok()),
17095 Some("false")
17096 );
17097 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17098 .await
17099 .unwrap();
17100 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17101 assert_eq!(payload["status"], "not_ready");
17102 assert!(payload["detail"].as_str().unwrap().contains("draining"));
17103 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
17104 let resp = chat_completions(
17106 State(st.clone()),
17107 axum::http::HeaderMap::new(),
17108 None,
17109 Json(
17110 serde_json::from_value(serde_json::json!({
17111 "model": "m", "messages": [{"role": "user", "content": "t"}]
17112 }))
17113 .unwrap(),
17114 ),
17115 )
17116 .await;
17117 assert_eq!(resp.status(), StatusCode::OK);
17118 }
17119
17120 #[tokio::test]
17123 async fn health_is_green_only_while_the_worker_is_alive() {
17124 let _l = DRAIN_LOCK.lock().unwrap();
17127 let st = fake_worker_state();
17128 let resp = health_live(State(st.clone())).await.into_response();
17131 assert_eq!(resp.status(), StatusCode::OK);
17132 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17133 .await
17134 .unwrap();
17135 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17136 assert_eq!(payload["status"], "ok");
17137 assert_eq!(payload["worker"]["phase"], "idle");
17138 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
17139 let ready = health_ready(State(st.clone())).await.into_response();
17140 assert_eq!(ready.status(), StatusCode::OK);
17141
17142 st.health.mark_dead("worker thread panicked: test-injected");
17146 let resp = health_live(State(st.clone())).await.into_response();
17147 assert_eq!(
17148 resp.status(),
17149 StatusCode::SERVICE_UNAVAILABLE,
17150 "a dead worker MUST NOT report a healthy liveness"
17151 );
17152 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17153 .await
17154 .unwrap();
17155 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17156 assert_eq!(payload["status"], "unhealthy");
17157 assert!(
17159 payload["detail"]
17160 .as_str()
17161 .unwrap()
17162 .contains("test-injected"),
17163 "cause not surfaced: {payload}"
17164 );
17165 let ready = health_ready(State(st.clone())).await.into_response();
17166 assert_eq!(
17167 ready.status(),
17168 StatusCode::SERVICE_UNAVAILABLE,
17169 "dead is also not ready"
17170 );
17171
17172 st.health.mark_ready();
17175 assert_eq!(
17176 health_live(State(st.clone()))
17177 .await
17178 .into_response()
17179 .status(),
17180 StatusCode::OK,
17181 "mark_ready must clear the latch (a successful respawn)"
17182 );
17183 }
17184
17185 #[tokio::test]
17186 async fn readyz_peer_probe_integrity_is_present_and_advisory() {
17187 let _l = DRAIN_LOCK.lock().unwrap();
17188 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
17189 let st = fake_worker_state();
17190
17191 let ready = health_ready(State(st.clone())).await.into_response();
17192 assert_eq!(ready.status(), StatusCode::OK);
17193 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
17194 .await
17195 .unwrap();
17196 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17197 assert_eq!(payload["peer_probe_integrity"], "ok");
17198
17199 st.health.note_peer_probe_deferral(2, false);
17200 let deferred = health_ready(State(st.clone())).await.into_response();
17201 assert_eq!(deferred.status(), StatusCode::OK);
17202 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
17203 .await
17204 .unwrap();
17205 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17206 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
17207
17208 st.health.note_peer_probe_deferral(4, true);
17209 let degraded = health_ready(State(st.clone())).await.into_response();
17210 assert_eq!(
17211 degraded.status(),
17212 StatusCode::OK,
17213 "peer degradation is advisory while plain serving remains healthy"
17214 );
17215 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
17216 .await
17217 .unwrap();
17218 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17219 assert_eq!(payload["peer_probe_integrity"], "degraded");
17220
17221 st.health.mark_dead("test-injected worker failure");
17222 let unready = health_ready(State(st)).await.into_response();
17223 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
17224 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
17225 .await
17226 .unwrap();
17227 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17228 assert_eq!(
17229 payload["peer_probe_integrity"], "degraded",
17230 "the advisory field must also survive an unrelated readiness failure"
17231 );
17232 }
17233
17234 #[tokio::test]
17235 async fn liveness_failure_obeys_the_retry_contract() {
17236 let _l = DRAIN_LOCK.lock().unwrap();
17240 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
17241 let st = fake_worker_state();
17242 st.health
17243 .mark_dead("worker thread panicked: retry-contract-test");
17244
17245 let resp = health_live(State(st)).await.into_response();
17246 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17247 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
17248 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
17249 assert_ne!(
17250 resp.headers()
17251 .get("x-should-retry")
17252 .and_then(|v| v.to_str().ok()),
17253 Some("false")
17254 );
17255 }
17256
17257 #[tokio::test]
17258 async fn readiness_failure_obeys_the_retry_contract() {
17259 let _l = DRAIN_LOCK.lock().unwrap();
17260 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
17261 let st = fake_worker_state();
17262 st.health
17263 .mark_dead("worker thread panicked: retry-contract-test");
17264
17265 let resp = health_ready(State(st)).await.into_response();
17266 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17267 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
17268 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
17269 assert_ne!(
17270 resp.headers()
17271 .get("x-should-retry")
17272 .and_then(|v| v.to_str().ok()),
17273 Some("false")
17274 );
17275 }
17276
17277 #[tokio::test]
17278 async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
17279 let _l = DRAIN_LOCK.lock().unwrap();
17289 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
17290 let st = fake_worker_state();
17291 assert_eq!(
17292 health_live(State(st.clone()))
17293 .await
17294 .into_response()
17295 .status(),
17296 StatusCode::OK
17297 );
17298 st.health
17299 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
17300 let resp = health_live(State(st.clone())).await.into_response();
17301 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17302 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17303 .await
17304 .unwrap();
17305 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17306 assert!(
17307 payload["detail"]
17308 .as_str()
17309 .unwrap()
17310 .contains("probe exceeded")
17311 );
17312 st.health.mark_ready();
17315 assert_eq!(
17316 health_live(State(st.clone()))
17317 .await
17318 .into_response()
17319 .status(),
17320 StatusCode::SERVICE_UNAVAILABLE,
17321 "a GPU fault must not be cleared by an in-process respawn"
17322 );
17323 }
17324
17325 #[test]
17326 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
17327 let caps = ModelCaps {
17329 tools_branch: true,
17330 qwen_think: true,
17331 think_switch: true,
17332 chat_ok: true,
17333 context_length: 262144,
17334 tokenizer: "qwen2".into(),
17335 instruct_type: Some("chatml".into()),
17336 effort_levels: false,
17337 qwen_effort: false,
17338 gemma_think: false,
17339 dsv4: false,
17340 chat_temperature_default: None,
17341 chat_top_p_default: None,
17342 n_vocab: 151_936,
17343 think_close: Vec::new(),
17344 };
17345 let e = model_entry_v1("main", Some(&caps), None);
17346 assert_eq!(e["id"], "main");
17347 assert_eq!(e["name"], "main");
17348 assert_eq!(e["object"], "model");
17349 assert_eq!(e["context_length"], 262144);
17350 assert!(e["pricing"]["input"].is_null());
17352 assert!(e["pricing"]["output"].is_null());
17353
17354 let meta = OpenRouterModelMetadata {
17357 pricing: OpenRouterPricing {
17358 prompt: Some("0.00000038".into()),
17359 cached_prompt: Some("0.0000002".into()),
17360 completion: Some("0.0000026".into()),
17361 ..Default::default()
17362 },
17363 input_modalities: vec!["image".into(), "video".into()],
17364 max_output_length: Some(32768),
17365 ..Default::default()
17366 };
17367 let e = model_entry_v1("main", Some(&caps), Some(&meta));
17368 assert_eq!(e["pricing"]["currency"], "USD");
17371 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
17372 assert_eq!(e["pricing"]["input"], "0.38");
17373 assert_eq!(e["pricing"]["output"], "2.60");
17374 assert_eq!(e["pricing"]["cached_input"], "0.20");
17375 assert!(e["pricing"]["cache_write"].is_null());
17376 assert_eq!(e["pricing"]["minimum_request"], "0");
17377 assert_eq!(e["owned_by"], "main");
17378 assert_eq!(e["type"], "chat");
17379 assert_eq!(e["max_output_tokens"], 32768);
17380 assert_eq!(e["endpoints"], json!(["chat/completions"]));
17381 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
17382 assert_eq!(e["output_modalities"], json!(["text"]));
17383 assert_eq!(e["capabilities"]["streaming"], true);
17384 assert_eq!(e["capabilities"]["tools"], true);
17385 assert_eq!(e["lifecycle"]["status"], "active");
17386 assert!(e["lifecycle"]["deprecation_at"].is_null());
17387 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
17388 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
17389 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
17393 keys.sort_unstable();
17394 assert_eq!(
17395 keys,
17396 [
17397 "capabilities",
17398 "context_length",
17399 "endpoints",
17400 "id",
17401 "input_modalities",
17402 "lifecycle",
17403 "max_output_tokens",
17404 "name",
17405 "object",
17406 "output_modalities",
17407 "owned_by",
17408 "pricing",
17409 "reliability",
17410 "type",
17411 ],
17412 "unexpected /v1/models entry keys"
17413 );
17414 let mut price_keys: Vec<&str> = e["pricing"]
17415 .as_object()
17416 .unwrap()
17417 .keys()
17418 .map(String::as_str)
17419 .collect();
17420 price_keys.sort_unstable();
17421 assert_eq!(
17422 price_keys,
17423 [
17424 "cache_write",
17425 "cached_input",
17426 "currency",
17427 "input",
17428 "minimum_request",
17429 "output",
17430 "unit",
17431 ],
17432 "unexpected /v1/models pricing keys"
17433 );
17434
17435 let e = model_entry_v1("m", None, None);
17437 assert!(e["context_length"].is_null());
17438 assert!(e["max_output_tokens"].is_null());
17439 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
17441 assert!(e["context_length"].is_null());
17442 }
17443
17444 #[test]
17450 fn catalog_row_follows_the_declared_surface() {
17451 let caps = ModelCaps {
17452 tools_branch: true,
17453 ..Default::default()
17454 };
17455
17456 let embed = OpenRouterModelMetadata {
17457 surface: Some("embedding".into()),
17458 max_output_length: Some(1),
17459 ..Default::default()
17460 };
17461 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
17462 assert_eq!(e["type"], "embedding");
17463 assert_eq!(e["endpoints"], json!(["embeddings"]));
17464 assert_eq!(e["output_modalities"], json!(["embeddings"]));
17465 assert_eq!(e["capabilities"]["streaming"], false);
17466 assert_eq!(
17467 e["capabilities"]["tools"], false,
17468 "an embedder has no tools"
17469 );
17470 assert_eq!(e["capabilities"]["reasoning"], false);
17471 assert_eq!(e["capabilities"]["structured_output"], false);
17472 assert_eq!(e["capabilities"]["prompt_caching"], false);
17473 assert!(
17474 e["max_output_tokens"].is_null(),
17475 "a surface that emits no completion tokens must not advertise a ceiling"
17476 );
17477
17478 let rerank = OpenRouterModelMetadata {
17479 surface: Some("rerank".into()),
17480 ..Default::default()
17481 };
17482 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
17483 assert_eq!(r["type"], "rerank");
17484 assert_eq!(r["endpoints"], json!(["rerank"]));
17485 assert_eq!(r["output_modalities"], json!(["rerank"]));
17486 assert_eq!(r["capabilities"]["tools"], false);
17487 assert_eq!(r["capabilities"]["reasoning"], false);
17488
17489 let chat = OpenRouterModelMetadata {
17492 max_output_length: Some(32768),
17493 ..Default::default()
17494 };
17495 let c = model_entry_v1("main", Some(&caps), Some(&chat));
17496 assert_eq!(c["type"], "chat");
17497 assert_eq!(c["endpoints"], json!(["chat/completions"]));
17498 assert_eq!(c["output_modalities"], json!(["text"]));
17499 assert_eq!(c["capabilities"]["tools"], true);
17500 assert_eq!(c["max_output_tokens"], 32768);
17501 }
17502
17503 #[test]
17506 fn unknown_surface_is_rejected_at_config_load() {
17507 let bad = OpenRouterModelMetadata {
17508 surface: Some("embeddings".into()), ..Default::default()
17510 };
17511 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
17512 .expect_err("an unknown surface must not load");
17513 assert!(err.contains("surface"), "{err}");
17514
17515 for good in ["chat", "embedding", "rerank"] {
17516 let ok = OpenRouterModelMetadata {
17517 surface: Some(good.into()),
17518 ..Default::default()
17519 };
17520 assert!(
17521 validate_openrouter_metadata("m", &ok).is_ok(),
17522 "{good} must load"
17523 );
17524 }
17525 }
17526
17527 #[test]
17528 fn per_million_price_is_exact_decimal_shift() {
17529 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
17531 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
17532 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
17533 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
17534 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
17535 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
17536 assert_eq!(per_million_price("not-a-price"), None);
17537 assert_eq!(per_million_price(""), None);
17538 }
17539
17540 #[test]
17541 fn metadata_provider_block_parses_and_validates() {
17542 let (_, provider) = OpenRouterMetadataFile::parse(
17543 r#"
17544 [provider]
17545 id = "tiyuvta"
17546 status_url = "https://status.tiyuvta.ai"
17547 support_contact = "mailto:support@tiyuvta.ai"
17548 incident_contact = "mailto:incidents@tiyuvta.ai"
17549 regions = ["eu-central"]
17550 "#,
17551 )
17552 .unwrap();
17553 let provider = provider.unwrap();
17554 assert_eq!(provider.id, "tiyuvta");
17555 assert_eq!(provider.regions, vec!["eu-central"]);
17556 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
17558 assert!(err.contains("provider.id"), "{err}");
17559 let err = OpenRouterMetadataFile::parse(
17561 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
17562 )
17563 .unwrap_err();
17564 assert!(err.contains("must be a URI"), "{err}");
17565 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
17567 assert!(provider.is_none());
17568 }
17569
17570 #[test]
17571 fn models_openai_default_body_stays_byte_identical() {
17572 let body = models_openai_body(&["main".into(), "judge".into()]);
17573 let bytes = serde_json::to_vec(&body).unwrap();
17574 assert_eq!(
17575 bytes,
17576 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
17577 );
17578 }
17579
17580 #[test]
17581 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
17582 let loaded = vec![
17584 "qwen/qwen3.6-27b".to_string(),
17585 "qwen/qwen3.6-35b-a3b".to_string(),
17586 ];
17587 assert_eq!(
17588 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
17589 Some("qwen/qwen3.6-35b-a3b"),
17590 );
17591 assert_eq!(
17592 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
17593 Some("qwen/qwen3.6-27b"),
17594 );
17595 assert_eq!(
17597 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
17598 Some("qwen/qwen3.6-35b-a3b"),
17599 );
17600 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
17602 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
17603 assert_eq!(canonical_model_id(&loaded, ""), None);
17604 }
17605
17606 #[test]
17607 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
17608 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
17611 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
17612 assert_eq!(
17614 canonical_model_id(&loaded, "a/shared-name").as_deref(),
17615 Some("a/shared-name")
17616 );
17617 assert_eq!(
17618 canonical_model_id(&loaded, "b/shared-name").as_deref(),
17619 Some("b/shared-name")
17620 );
17621 let bare = vec!["solo".to_string()];
17623 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
17624 }
17625
17626 #[test]
17627 fn openrouter_models_entry_serializes_complete_metadata() {
17628 let metadata = OpenRouterMetadataFile::from_toml(
17629 r#"
17630[models.main]
17631hugging_face_id = "Qwen/Qwen3.6-27B"
17632created = 1786032000
17633quantization = "nvfp4"
17634description = "Qwen3.6 27B served by memra."
17635max_prompt_length = 245760
17636max_output_length = 16384
17637default_output_length = 4096
17638is_ready = true
17639is_free = false
17640discount_to_user = 0.1
17641openrouter_slug = "qwen/qwen3.6-27b"
17642datacenters = [{ country_code = "US", region = "us-east-1" }]
17643zdr = true
17644hipaa = false
17645
17646[models.main.pricing]
17647prompt = "0.000000234"
17648cached_prompt = "0.0000000585"
17649cache_write = "0.000000234"
17650completion = "0.000001872"
17651internal_reasoning = "0.000001872"
17652request = "0.01"
17653
17654[models.main.capacity]
17655prompt_tpm = 1000000
17656cached_prompt_tpm = 2000000
17657completion_tpm = 500000
17658request_rpm = 1000
17659concurrency = 64
17660"#,
17661 )
17662 .unwrap();
17663 let caps = ModelCaps {
17664 tools_branch: true,
17665 qwen_think: true,
17666 think_switch: true,
17667 chat_ok: true,
17668 context_length: 262144,
17669 tokenizer: "qwen2".into(),
17670 instruct_type: Some("chatml".into()),
17671 ..Default::default()
17672 };
17673 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
17674
17675 assert_eq!(entry["schema_version"], "2.4");
17676 assert_eq!(entry["id"], "main");
17677 assert_eq!(entry["name"], "main");
17678 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
17679 assert_eq!(entry["created"], 1786032000u64);
17680 assert_eq!(entry["quantization"], "nvfp4");
17681 assert_eq!(entry["tokenizer"], "qwen2");
17682 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
17683 assert!(
17684 entry.get("object").is_none(),
17685 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
17686 );
17687
17688 let input = &entry["input_modalities"][0];
17689 assert_eq!(input["type"], "text");
17690 assert_eq!(
17691 input["supported_inputs"]["max_context_length"]["value"],
17692 262144
17693 );
17694 assert_eq!(
17695 input["supported_inputs"]["max_prompt_length"]["value"],
17696 245760
17697 );
17698 let input_prices = input["pricing"].as_array().unwrap();
17699 let input_price = |kind: &str| {
17700 input_prices
17701 .iter()
17702 .find(|price| price["type"] == kind)
17703 .unwrap()
17704 };
17705 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
17706 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
17707 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
17708 assert_eq!(input["capacity"][0]["value"], 1000000);
17709 assert_eq!(input["capacity"][1]["value"], 2000000);
17710
17711 let output = &entry["output_modalities"][0];
17712 assert_eq!(output["type"], "text");
17713 assert_eq!(output["max_length"]["value"], 16384);
17714 assert_eq!(output["streaming"], true);
17715 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
17716 assert_eq!(
17717 output["supported_parameters"]["structured_outputs"]["type"],
17718 "boolean"
17719 );
17720 assert_eq!(
17721 output["supported_parameters"]["reasoning"]["type"],
17722 "boolean"
17723 );
17724 assert_eq!(output["pricing"][0]["type"], "completion");
17725 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
17726 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
17727 assert_eq!(output["capacity"][0]["value"], 500000);
17728 assert_eq!(output["capacity"][1]["type"], "concurrency");
17729 assert_eq!(output["capacity"][1]["value"], 64);
17730
17731 assert_eq!(entry["pricing"][0]["type"], "request");
17732 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
17733 assert_eq!(entry["capacity"][0]["value"], 1000);
17734 assert_eq!(entry["is_ready"], true);
17735 assert_eq!(entry["is_free"], false);
17736 assert_eq!(entry["discount_to_user"], 0.1);
17737 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
17738 assert_eq!(entry["datacenters"][0]["country_code"], "US");
17739 assert_eq!(entry["compliance"]["zdr"], true);
17740 assert_eq!(entry["compliance"]["hipaa"], false);
17741 }
17742
17743 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
17748[models."qwen/qwen3.6-35b-a3b"]
17749hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
17750created = 1777260255
17751quantization = "int4"
17752description = "Qwen3.6 35B-A3B fixture entry."
17753max_prompt_length = 262144
17754max_output_length = 262144
17755default_output_length = 8192
17756is_ready = true
17757is_free = false
17758discount_to_user = 0.0
17759openrouter_slug = "qwen/qwen3.6-35b-a3b"
17760zdr = false
17761hipaa = false
17762
17763[[models."qwen/qwen3.6-35b-a3b".datacenters]]
17764country_code = "CA"
17765region = "Ontario"
17766
17767[models."qwen/qwen3.6-35b-a3b".pricing]
17768prompt = "0.0000000931"
17769cached_prompt = "0.0000000652"
17770completion = "0.0000009025"
17771
17772[models."qwen/qwen3.6-35b-a3b".capacity]
17773prompt_tpm = 780000
17774cached_prompt_tpm = 310000
17775completion_tpm = 9600
17776request_rpm = 160
17777concurrency = 16
17778
17779[planned_models."qwen/qwen3.8-27b"]
17780description = "Planned fixture entry; must never be emitted."
17781max_prompt_length = 262144
17782max_output_length = 262144
17783default_output_length = 8192
17784is_ready = false
17785is_free = false
17786discount_to_user = 0.0
17787openrouter_slug = "qwen/qwen3.8-27b"
17788zdr = false
17789hipaa = false
17790
17791[planned_models."qwen/qwen3.8-27b".pricing]
17792prompt = "0.0000002745"
17793cached_prompt = "0.0000001922"
17794completion = "0.0000022800"
17795
17796[planned_models."google/gemma-4-26b-a4b-it"]
17797hugging_face_id = "google/gemma-4-26B-A4B-it"
17798created = 1775227989
17799quantization = "int4"
17800description = "Planned fixture entry; must never be emitted."
17801max_prompt_length = 262144
17802max_output_length = 262144
17803default_output_length = 8192
17804is_ready = false
17805is_free = false
17806discount_to_user = 0.0
17807openrouter_slug = "google/gemma-4-26b-a4b-it"
17808zdr = false
17809hipaa = false
17810
17811[planned_models."google/gemma-4-26b-a4b-it".pricing]
17812prompt = "0.0000000665"
17813cached_prompt = "0.0000000466"
17814completion = "0.0000003230"
17815"#;
17816
17817 #[test]
17818 fn gateway_registry_generates_the_staged_active_shape() {
17819 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17820 let caps = ModelCaps {
17821 tools_branch: true,
17822 qwen_think: true,
17823 think_switch: true,
17824 chat_ok: true,
17825 context_length: 262144,
17826 tokenizer: "qwen2".into(),
17827 instruct_type: Some("chatml".into()),
17828 ..Default::default()
17829 };
17830 let q35_entry = model_entry_openrouter(
17831 "qwen/qwen3.6-35b-a3b",
17832 Some(&caps),
17833 metadata.get("qwen/qwen3.6-35b-a3b"),
17834 );
17835 assert_eq!(q35_entry["created"], 1777260255u64);
17836 assert_eq!(q35_entry["quantization"], "int4");
17837 assert_eq!(q35_entry["is_ready"], true);
17838 assert_eq!(
17839 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
17840 262144
17841 );
17842 assert_eq!(
17843 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
17844 262144
17845 );
17846 assert_eq!(
17847 q35_entry["output_modalities"][0]["max_length"]["value"],
17848 262144
17849 );
17850 let prices = q35_entry["input_modalities"][0]["pricing"]
17851 .as_array()
17852 .unwrap();
17853 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
17854 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
17855 assert_eq!(
17859 q35_entry["input_modalities"][0]["capacity"][0]["value"],
17860 780000
17861 );
17862 assert_eq!(
17863 q35_entry["input_modalities"][0]["capacity"][1]["value"],
17864 310000
17865 );
17866 assert_eq!(
17867 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
17868 262144
17869 );
17870 assert_eq!(
17871 q35_entry["output_modalities"][0]["capacity"][0]["value"],
17872 9600
17873 );
17874 assert_eq!(
17875 q35_entry["output_modalities"][0]["capacity"][1]["value"],
17876 16
17877 );
17878 assert_eq!(
17879 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
17880 "0.0000009025"
17881 );
17882 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
17884
17885 assert_eq!(
17886 metadata.len(),
17887 1,
17888 "planned models must never enter the active map"
17889 );
17890 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
17891 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
17892 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
17893
17894 let openmodels = model_entry_openmodels(
17895 "qwen/qwen3.6-35b-a3b",
17896 Some(&caps),
17897 metadata.get("qwen/qwen3.6-35b-a3b"),
17898 )
17899 .unwrap();
17900 assert_eq!(openmodels["currency"], "USD");
17901 assert_eq!(openmodels["max_output_length"], 262144);
17902 assert_eq!(openmodels["is_ready"], true);
17903 assert_eq!(openmodels["is_free"], false);
17904 assert_eq!(openmodels["discount_to_user"], 0.0);
17905 }
17906
17907 #[test]
17908 fn gateway_registry_limits_are_live_request_limits() {
17909 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17910 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
17911 let caps = ModelCaps {
17912 context_length: 262_144,
17913 ..Default::default()
17914 };
17915 let build = |value: serde_json::Value| {
17916 let req: CompletionReq = serde_json::from_value(value).unwrap();
17917 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
17918 build_request(&req, tx, lanes::Lane::Interactive, None)
17919 };
17920
17921 let mut omitted = build(json!({
17922 "model": "qwen/qwen3.6-35b-a3b",
17923 "prompt_ids": [1, 2, 3]
17924 }));
17925 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
17926 assert_eq!(omitted.params.max_new, 8_192);
17927 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
17928
17929 let mut field_top = build(json!({
17930 "model": "qwen/qwen3.6-35b-a3b",
17931 "prompt_ids": [1],
17932 "max_tokens": 262144
17933 }));
17934 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
17935 assert_eq!(field_top.params.max_new, 262_144);
17936 assert_eq!(
17937 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
17938 262_044,
17939 "the field-top output request is accepted but bounded by remaining trained context",
17940 );
17941
17942 let mut too_much_output = build(json!({
17943 "model": "qwen/qwen3.6-35b-a3b",
17944 "prompt_ids": [1],
17945 "max_tokens": 262145
17946 }));
17947 let (message, param) =
17948 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
17949 .unwrap_err();
17950 assert_eq!(param, "max_tokens");
17951 assert!(message.contains("262145"));
17952
17953 let mut oversized_allocation = build(json!({
17954 "model": "qwen/qwen3.6-35b-a3b",
17955 "prompt_ids": [1],
17956 "max_tokens": 1,
17957 "max_ctx": 262145
17958 }));
17959 let (_, param) =
17960 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
17961 .unwrap_err();
17962 assert_eq!(param, "max_ctx");
17963 }
17964
17965 #[test]
17966 fn planned_registry_entries_are_validated_but_never_activated() {
17967 let parsed = OpenRouterMetadataFile::from_toml(
17968 r#"
17969[planned_models.future]
17970max_output_length = 262144
17971default_output_length = 8192
17972
17973[planned_models.future.pricing]
17974prompt = "0.0000001"
17975"#,
17976 )
17977 .unwrap();
17978 assert!(parsed.is_empty());
17979
17980 let error = OpenRouterMetadataFile::from_toml(
17981 r#"
17982[planned_models.future]
17983default_output_length = 8192
17984"#,
17985 )
17986 .unwrap_err();
17987 assert!(error.contains("requires max_output_length"));
17988 }
17989
17990 #[test]
17995 fn every_catalog_feed_honours_the_declared_surface() {
17996 let metadata = OpenRouterMetadataFile::from_toml(
17997 r#"
17998[models."qwen/qwen3-embedding-8b"]
17999surface = "embedding"
18000created = 1787961600
18001max_output_length = 1
18002is_ready = true
18003is_free = false
18004discount_to_user = 0.0
18005
18006[models."qwen/qwen3-embedding-8b".pricing]
18007prompt = "0.00000001"
18008cached_prompt = "0.0"
18009completion = "0.0"
18010
18011[models."main"]
18012created = 1787443200
18013max_output_length = 32768
18014is_ready = true
18015is_free = false
18016discount_to_user = 0.0
18017
18018[models."main".pricing]
18019prompt = "0.00000025"
18020cached_prompt = "0.00000009"
18021completion = "0.0000012"
18022"#,
18023 )
18024 .unwrap();
18025 let caps = ModelCaps {
18026 tools_branch: true,
18027 qwen_think: true,
18028 chat_ok: true,
18029 context_length: 32768,
18030 ..Default::default()
18031 };
18032 let embed = metadata.get("qwen/qwen3-embedding-8b");
18033 let chat = metadata.get("main");
18034
18035 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
18037 let out = &or["output_modalities"][0];
18038 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
18039 assert!(
18040 out.get("streaming").is_none(),
18041 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
18042 );
18043 let params = &out["supported_parameters"];
18048 assert_eq!(
18049 params.as_object().map(|o| o.len()),
18050 Some(0),
18051 "no completion parameter belongs on an embedder row: {params}"
18052 );
18053 for field in [
18054 "tools",
18055 "tool_choice",
18056 "reasoning",
18057 "max_tokens",
18058 "json_mode",
18059 "structured_outputs",
18060 "stop",
18061 "temperature",
18062 "seed",
18063 ] {
18064 assert!(params[field].is_null(), "{field} leaked onto an embedder");
18065 }
18066 assert!(
18067 out["max_length"].is_null(),
18068 "a surface emitting no completion tokens advertises no ceiling: {out}"
18069 );
18070
18071 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
18073 .expect("openmodels entry builds");
18074 assert_eq!(om["output_modalities"], json!(["embeddings"]));
18075 let features = om["supported_features"].as_array().unwrap();
18076 assert!(
18077 !features
18078 .iter()
18079 .any(|f| f == "tool_calling" || f == "reasoning"),
18080 "chat-only features leaked onto an embedder: {features:?}"
18081 );
18082
18083 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
18085 assert_eq!(v1["type"], "embedding");
18086 assert_eq!(v1["capabilities"]["tools"], false);
18087
18088 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
18090 let out_chat = &or_chat["output_modalities"][0];
18091 assert_eq!(out_chat["type"], "text");
18092 assert_eq!(out_chat["streaming"], true);
18093 assert!(!out_chat["supported_parameters"]["tools"].is_null());
18094 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
18095 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
18096 assert_eq!(out_chat["max_length"]["value"], 32768u64);
18097 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
18098 assert_eq!(om_chat["output_modalities"], json!(["text"]));
18099 assert!(
18100 om_chat["supported_features"]
18101 .as_array()
18102 .unwrap()
18103 .iter()
18104 .any(|f| f == "tool_calling")
18105 );
18106 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
18107 }
18108
18109 #[test]
18116 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
18117 let raw = std::fs::read_to_string(concat!(
18118 env!("CARGO_MANIFEST_DIR"),
18119 "/../../research/gateway-20260812/raw/sources/",
18120 "openrouter-provider-schema-v2.4-20260812.json"
18121 ))
18122 .expect("vendored Provider Monitor 2.4 schema is in-tree");
18123 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
18124 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
18125 .as_array()
18126 .expect("OutputModality is a oneOf");
18127
18128 let metadata = OpenRouterMetadataFile::from_toml(
18129 r#"
18130[models."embed"]
18131surface = "embedding"
18132created = 1787961600
18133max_output_length = 1
18134is_ready = true
18135is_free = false
18136discount_to_user = 0.0
18137
18138[models."embed".pricing]
18139prompt = "0.00000001"
18140cached_prompt = "0.0"
18141completion = "0.0"
18142
18143[models."rr"]
18144surface = "rerank"
18145created = 1787961600
18146max_output_length = 1
18147is_ready = true
18148is_free = false
18149discount_to_user = 0.0
18150
18151[models."rr".pricing]
18152prompt = "0.00000003"
18153cached_prompt = "0.0"
18154completion = "0.0"
18155
18156[models."chatty"]
18157created = 1787443200
18158max_output_length = 32768
18159is_ready = true
18160is_free = false
18161discount_to_user = 0.0
18162
18163[models."chatty".pricing]
18164prompt = "0.00000025"
18165cached_prompt = "0.00000009"
18166completion = "0.0000012"
18167"#,
18168 )
18169 .unwrap();
18170 let caps = ModelCaps {
18171 tools_branch: true,
18172 qwen_think: true,
18173 chat_ok: true,
18174 context_length: 32768,
18175 ..Default::default()
18176 };
18177
18178 for (alias, want_type) in [
18179 ("embed", "embeddings"),
18180 ("rr", "rerank"),
18181 ("chatty", "text"),
18182 ] {
18183 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
18184 let modality = &row["output_modalities"][0];
18185 assert_eq!(modality["type"], want_type, "{alias}: {row}");
18186
18187 let branch = branches
18189 .iter()
18190 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
18191 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
18192 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
18193 .as_object()
18194 .expect("branch properties")
18195 .keys()
18196 .map(String::as_str)
18197 .collect();
18198 for key in modality.as_object().expect("modality object").keys() {
18199 assert!(
18200 allowed.contains(key.as_str()),
18201 "{alias}: {key:?} is not a property of the {want_type:?} branch \
18202 (additionalProperties:false); allowed = {allowed:?}"
18203 );
18204 }
18205 for req in branch["required"].as_array().into_iter().flatten() {
18206 let req = req.as_str().expect("required entry is a string");
18207 assert!(
18208 modality.get(req).is_some(),
18209 "{alias}: required property {req:?} missing from the {want_type:?} branch"
18210 );
18211 }
18212 }
18213 }
18214
18215 #[test]
18216 fn openrouter_models_entry_omits_undeclared_optional_fields() {
18217 let entry = model_entry_openrouter("minimal", None, None);
18218 let object = entry.as_object().unwrap();
18219 for field in [
18220 "hugging_face_id",
18221 "created",
18222 "quantization",
18223 "tokenizer",
18224 "description",
18225 "pricing",
18226 "capacity",
18227 "is_ready",
18228 "is_free",
18229 "discount_to_user",
18230 "openrouter",
18231 "datacenters",
18232 "compliance",
18233 ] {
18234 assert!(
18235 !object.contains_key(field),
18236 "optional field {field} must be absent, not null"
18237 );
18238 }
18239 assert_eq!(entry["schema_version"], "2.4");
18240 assert_eq!(entry["input_modalities"][0]["type"], "text");
18241 assert!(
18242 entry["input_modalities"][0]
18243 .get("supported_inputs")
18244 .is_none()
18245 );
18246 assert!(entry["input_modalities"][0].get("pricing").is_none());
18247 assert!(entry["input_modalities"][0].get("capacity").is_none());
18248 assert_eq!(entry["output_modalities"][0]["type"], "text");
18249 assert_eq!(entry["output_modalities"][0]["streaming"], true);
18250 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
18251 assert!(entry["output_modalities"][0].get("max_length").is_none());
18252 assert!(entry["output_modalities"][0].get("pricing").is_none());
18253 assert!(entry["output_modalities"][0].get("capacity").is_none());
18254 }
18255
18256 #[test]
18257 fn openmodels_entry_serializes_standard_provider_shape() {
18258 let metadata = OpenRouterMetadataFile::from_toml(
18259 r#"
18260[models."qwen/qwen3.6-27b"]
18261created = 1786032000
18262max_output_length = 16384
18263is_ready = true
18264is_free = false
18265discount_to_user = 0.05
18266
18267[models."qwen/qwen3.6-27b".pricing]
18268prompt = "0.000000291"
18269cached_prompt = "0.000000291"
18270completion = "0.000002763"
18271request = "0"
18272"#,
18273 )
18274 .unwrap();
18275 let caps = ModelCaps {
18276 tools_branch: true,
18277 qwen_think: true,
18278 chat_ok: true,
18279 context_length: 262144,
18280 ..Default::default()
18281 };
18282 let entry = model_entry_openmodels(
18283 "qwen/qwen3.6-27b",
18284 Some(&caps),
18285 metadata.get("qwen/qwen3.6-27b"),
18286 )
18287 .unwrap();
18288
18289 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
18290 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
18291 assert_eq!(entry["created"], 1786032000u64);
18292 assert_eq!(entry["input_modalities"], json!(["text"]));
18293 assert_eq!(entry["output_modalities"], json!(["text"]));
18294 assert_eq!(entry["context_length"], 262144u64);
18295 assert_eq!(entry["max_output_length"], 16384u64);
18296 assert_eq!(entry["currency"], "USD");
18297 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
18298 assert_eq!(entry["pricing"]["completion"], "0.000002763");
18299 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
18300 assert_eq!(entry["pricing"]["request"], "0");
18301 assert_eq!(
18302 entry["supported_features"],
18303 json!(["tool_calling", "reasoning"])
18304 );
18305 assert_eq!(entry["is_ready"], true);
18306 assert_eq!(entry["is_free"], false);
18307 assert_eq!(entry["discount_to_user"], 0.05);
18308 assert!(entry.get("schema_version").is_none());
18309 assert!(entry.get("quantization").is_none());
18310 }
18311
18312 #[test]
18313 fn openmodels_entry_rejects_missing_operator_metadata() {
18314 let caps = ModelCaps {
18315 context_length: 262144,
18316 ..Default::default()
18317 };
18318 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
18319 assert_eq!(
18320 error,
18321 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
18322 );
18323 }
18324
18325 #[tokio::test]
18326 async fn blocking_response_excludes_stop_text_across_token_events() {
18327 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
18328 tx.send(Event::Token {
18329 id: 1,
18330 text: "answer\nPro".into(),
18331 })
18332 .unwrap();
18333 tx.send(Event::Token {
18334 id: 2,
18335 text: "blem: leaked prompt".into(),
18336 })
18337 .unwrap();
18338 tx.send(Event::Done {
18339 stop_reason: "Callback".into(),
18340 n_tokens: 2,
18341 n_prompt: 8,
18342 n_cached: 0,
18343 elapsed_s: 0.5,
18344 spec: None,
18345 })
18346 .unwrap();
18347 drop(tx);
18348 let response = blocking_response(
18349 rx,
18350 "plain_quant".into(),
18351 false,
18352 vec!["Problem:".into()],
18353 None,
18354 Envelope::new(false),
18355 )
18356 .await;
18357 assert_eq!(response.status(), StatusCode::OK);
18358 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18359 .await
18360 .unwrap();
18361 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18362 assert_eq!(payload["text"], "answer\n");
18363 assert_eq!(payload["stop_reason"], "Callback");
18364 }
18365
18366 #[test]
18370 fn step_walker_expansion_and_separator_law() {
18371 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
18373 let uri = format!("data:image/png;base64,{PNG64}");
18374 let content = serde_json::json!([
18375 {"type": "text", "text": "look at"},
18376 {"type": "text", "text": "this:"},
18377 {"type": "image_url", "image_url": {"url": uri}},
18378 {"type": "text", "text": "what is it?"},
18379 ]);
18380 let mut pending: Vec<PendingStepImage> = Vec::new();
18381 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
18382 let mut expansion = String::from("<im_start>");
18383 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
18384 expansion.push_str("<im_patch>");
18385 }
18386 expansion.push_str("<im_end>");
18387 assert_eq!(out, format!("look at this:{expansion}what is it?"));
18390 assert_eq!(pending.len(), 1);
18391 assert_eq!(pending[0].plan.n_tiles, 0);
18392 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
18393
18394 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
18396 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
18397 let http = serde_json::json!([
18398 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
18399 ]);
18400 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
18401 }
18402}