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 anthropic;
75mod dsv4_serve;
76mod embed_api;
77pub mod metering;
84mod responses_api;
85mod surfaces;
86mod toolcall;
87mod ttft;
88mod worker;
89
90use std::collections::HashMap;
91use std::net::{SocketAddr, ToSocketAddrs};
92use std::sync::Arc;
93use std::sync::mpsc::Sender;
94
95use axum::{
96 Extension, Json, Router,
97 body::Body,
98 extract::{DefaultBodyLimit, Query, Request as AxumRequest, State},
99 http::{
100 HeaderMap, StatusCode,
101 header::{CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING},
102 },
103 middleware::{self, Next},
104 response::{
105 IntoResponse, Response,
106 sse::{Event as SseEvent, Sse},
107 },
108 routing::{get, post},
109};
110use futures_core::Stream as _;
111use serde::{Deserialize, Serialize};
112use serde_json::json;
113
114use memra_engine::decode::GenParams;
115use memra_engine::sampler::SamplerConfig;
116use memra_tokenizer::{
117 Tokenizer,
118 chat::{self, ThinkMode, ToolCall as TmplToolCall, Turn as TmplTurn},
119};
120use toolcall::{ParsedToolCall, Piece, ToolStreamParser};
121use worker::{Cmd, Event, ModelCaps, Request, SharedMetrics};
122
123const MAX_BODY_BYTES: usize = 192 * 1024 * 1024;
140const MAX_BODY_ADMISSIONS: usize = 4;
141const MAX_SMALL_BODY_ADMISSIONS: usize = 32;
142const BODY_ADMISSION_BYPASS_BYTES: usize = 1 * 1024 * 1024;
146const BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
147const BODY_READ_RATE_BYTES_PER_SEC: u64 = 2 * 1024 * 1024;
148const BODY_READ_TIMEOUT_MAX: std::time::Duration = std::time::Duration::from_secs(180);
149const BODY_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
150const BODY_ADMISSION_RETRY_AFTER_S: u64 = 1;
151
152fn body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
153 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
154 SEMAPHORE
155 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_BODY_ADMISSIONS)))
156 .clone()
157}
158
159fn small_body_admission_semaphore() -> Arc<tokio::sync::Semaphore> {
160 static SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();
161 SEMAPHORE
162 .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_SMALL_BODY_ADMISSIONS)))
163 .clone()
164}
165
166fn declared_body_length(req: &AxumRequest) -> Option<usize> {
167 req.headers()
168 .get(CONTENT_LENGTH)
169 .and_then(|value| value.to_str().ok())
170 .and_then(|value| value.parse().ok())
171}
172
173fn body_requires_admission(req: &AxumRequest) -> bool {
174 if req.headers().contains_key(TRANSFER_ENCODING) {
178 return true;
179 }
180 declared_body_length(req).map_or(true, |length| length > BODY_ADMISSION_BYPASS_BYTES)
181}
182
183fn body_read_timeout(req: &AxumRequest) -> std::time::Duration {
187 let Some(length) = declared_body_length(req) else {
188 return BODY_READ_TIMEOUT;
189 };
190 let bytes = length as u64;
191 let extra_seconds =
192 bytes.saturating_add(BODY_READ_RATE_BYTES_PER_SEC - 1) / BODY_READ_RATE_BYTES_PER_SEC;
193 let seconds = BODY_READ_TIMEOUT
194 .as_secs()
195 .saturating_add(extra_seconds)
196 .min(BODY_READ_TIMEOUT_MAX.as_secs());
197 std::time::Duration::from_secs(seconds)
198}
199
200async fn shape_payload_too_large(req: AxumRequest, next: Next) -> Response {
205 let resp = next.run(req).await;
206 if resp.status() != StatusCode::PAYLOAD_TOO_LARGE {
207 return resp;
208 }
209 error_response_coded(
210 StatusCode::PAYLOAD_TOO_LARGE,
211 &format!(
212 "request body exceeds the {} MiB limit",
213 MAX_BODY_BYTES / (1024 * 1024)
214 ),
215 "invalid_request_error",
216 None,
217 Some("request_too_large"),
218 )
219}
220
221fn apply_body_limit(app: Router) -> Router {
224 app.layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
225 .layer(middleware::from_fn(shape_payload_too_large))
226}
227
228fn protected_inference_path(path: &str) -> bool {
229 matches!(
230 path,
231 "/v1/auth/check"
232 | "/v1/completions"
233 | "/v1/chat/completions"
234 | "/v1/messages"
235 | "/v1/responses"
236 | "/v1/embeddings"
237 | "/v1/rerank"
238 )
239}
240
241async fn shape_inference_early_response(path: &str, response: Response) -> Response {
245 let request_id = Envelope::new(path != "/v1/completions");
246 if path == "/v1/messages" {
247 anthropic::with_anthropic_request_id(
248 &request_id.id,
249 anthropic::reshape_error(response, &request_id.id).await,
250 )
251 } else {
252 with_request_id(&request_id.id, response)
253 }
254}
255
256async fn authenticate_inference_before_body(
261 State(st): State<AppState>,
262 mut req: AxumRequest,
263 next: Next,
264) -> Response {
265 if !protected_inference_path(req.uri().path()) {
266 return next.run(req).await;
267 }
268 let path = req.uri().path().to_string();
269 if declared_body_length(&req).is_some_and(|length| length > MAX_BODY_BYTES) {
273 return shape_inference_early_response(
274 &path,
275 error_response_coded(
276 StatusCode::PAYLOAD_TOO_LARGE,
277 &format!(
278 "request body exceeds the {} MiB limit",
279 MAX_BODY_BYTES / (1024 * 1024)
280 ),
281 "invalid_request_error",
282 None,
283 Some("request_too_large"),
284 ),
285 )
286 .await;
287 }
288 let headers = req.headers();
289 let bearer = bearer_token(headers);
290 let auth = if matches!(path.as_str(), "/v1/messages" | "/v1/auth/check") {
291 let api_key = headers
292 .get("x-api-key")
293 .and_then(|value| value.to_str().ok());
294 surfaces::authenticate_candidates(&st.api_auth, &[bearer, api_key])
295 } else {
296 surfaces::authenticate_candidates(&st.api_auth, &[bearer])
297 };
298 if let Err(why) = auth {
299 return shape_inference_early_response(&path, authentication_error(why)).await;
300 }
301 let body_deadline = tokio::time::Instant::now() + body_read_timeout(&req);
309 let body_admission = if body_requires_admission(&req) {
310 body_admission_semaphore()
311 } else {
312 small_body_admission_semaphore()
313 };
314 let body_permit = match body_admission.try_acquire_owned() {
315 Ok(permit) => Some(permit),
316 Err(tokio::sync::TryAcquireError::Closed) => {
317 let response = retry_contract_response(
318 error_response_coded(
319 StatusCode::SERVICE_UNAVAILABLE,
320 "request body admission is unavailable",
321 "server_error",
322 None,
323 Some("body_admission_unavailable"),
324 ),
325 Some(BODY_ADMISSION_RETRY_AFTER_S),
326 );
327 return shape_inference_early_response(&path, response).await;
328 }
329 Err(tokio::sync::TryAcquireError::NoPermits) => {
330 let response = retry_contract_response(
331 error_response_coded(
332 StatusCode::TOO_MANY_REQUESTS,
333 "request body admission is busy",
334 "rate_limit_error",
335 None,
336 Some("body_admission_busy"),
337 ),
338 Some(BODY_ADMISSION_RETRY_AFTER_S),
339 );
340 return shape_inference_early_response(&path, response).await;
341 }
342 };
343 let body = std::mem::replace(req.body_mut(), Body::empty());
347 let mut body = Box::pin(body.into_data_stream());
348 let body_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
349 let body_timed_out_flag = body_timed_out.clone();
350 let guarded_body = async_stream::stream! {
351 loop {
352 let remaining = body_deadline.saturating_duration_since(tokio::time::Instant::now());
353 if remaining.is_zero() {
354 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
355 yield Err(std::io::Error::new(
356 std::io::ErrorKind::TimedOut,
357 "request body read deadline exceeded",
358 ));
359 break;
360 }
361 let poll = std::future::poll_fn(|cx| body.as_mut().poll_next(cx));
362 let frame = match tokio::time::timeout(BODY_IDLE_TIMEOUT.min(remaining), poll).await {
363 Ok(frame) => frame,
364 Err(_) => {
365 body_timed_out_flag.store(true, std::sync::atomic::Ordering::Release);
366 yield Err(std::io::Error::new(
367 std::io::ErrorKind::TimedOut,
368 "request body idle timeout exceeded",
369 ));
370 break;
371 }
372 };
373 match frame {
374 Some(Ok(bytes)) => yield Ok(bytes),
375 Some(Err(error)) => {
376 yield Err(std::io::Error::other(error.to_string()));
377 break;
378 }
379 None => break,
380 }
381 }
382 drop(body_permit);
383 };
384 *req.body_mut() = Body::from_stream(guarded_body);
385 let response = next.run(req).await;
386 if body_timed_out.load(std::sync::atomic::Ordering::Acquire) {
387 let request_id = Envelope::new(path != "/v1/completions");
388 let timeout = error_response_coded(
389 StatusCode::REQUEST_TIMEOUT,
390 "request body read timed out",
391 "invalid_request_error",
392 None,
393 Some("request_body_timeout"),
394 );
395 return if path == "/v1/messages" {
396 anthropic::with_anthropic_request_id(
397 &request_id.id,
398 anthropic::reshape_error(timeout, &request_id.id).await,
399 )
400 } else {
401 with_request_id(&request_id.id, timeout)
402 };
403 }
404 if path == "/v1/messages" && response.status() == StatusCode::PAYLOAD_TOO_LARGE {
405 let request_id = Envelope::new(true);
406 return anthropic::with_anthropic_request_id(
407 &request_id.id,
408 anthropic::reshape_error(response, &request_id.id).await,
409 );
410 }
411 response
412}
413
414#[cfg(test)]
415mod body_limit_tests {
416 use super::*;
417 use tower::ServiceExt as _;
418
419 static BODY_ADMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
420
421 fn test_app() -> Router {
425 let app = Router::new()
426 .route(
427 "/bytes",
428 post(|b: axum::body::Bytes| async move { b.len().to_string() }),
429 )
430 .route(
431 "/json",
432 post(|Json(v): Json<serde_json::Value>| async move {
433 v["pad"].as_str().unwrap_or("").len().to_string()
434 }),
435 );
436 apply_body_limit(app)
437 }
438
439 fn streamed_body(chunks: usize) -> Body {
440 let chunk = axum::body::Bytes::from(vec![b'x'; 1024 * 1024]);
443 Body::from_stream(async_stream::stream! {
444 for _ in 0..chunks {
445 yield Ok::<_, std::io::Error>(chunk.clone());
446 }
447 })
448 }
449
450 #[tokio::test]
451 async fn bodies_past_the_old_2mib_default_are_accepted() {
452 for (path, body) in [
455 ("/bytes", Body::from(vec![b'x'; 3 * 1024 * 1024])),
456 (
457 "/json",
458 Body::from(
459 serde_json::to_vec(&json!({ "pad": "x".repeat(3 * 1024 * 1024) })).unwrap(),
460 ),
461 ),
462 ] {
463 let resp = test_app()
464 .oneshot(
465 axum::http::Request::post(path)
466 .header(CONTENT_TYPE, "application/json")
467 .body(body)
468 .unwrap(),
469 )
470 .await
471 .unwrap();
472 assert_eq!(resp.status(), StatusCode::OK, "{path}");
473 }
474 }
475
476 #[tokio::test]
477 async fn body_at_exactly_the_limit_is_accepted() {
478 let resp = test_app()
479 .oneshot(
480 axum::http::Request::post("/bytes")
481 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024)))
482 .unwrap(),
483 )
484 .await
485 .unwrap();
486 assert_eq!(resp.status(), StatusCode::OK);
487 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
488 .await
489 .unwrap();
490 assert_eq!(body.as_ref(), MAX_BODY_BYTES.to_string().as_bytes());
491 }
492
493 #[tokio::test]
494 async fn oversize_body_is_a_clean_413_in_our_error_shape() {
495 for path in ["/bytes", "/json"] {
499 let resp = test_app()
500 .oneshot(
501 axum::http::Request::post(path)
502 .header(CONTENT_TYPE, "application/json")
503 .body(streamed_body(MAX_BODY_BYTES / (1024 * 1024) + 1))
504 .unwrap(),
505 )
506 .await
507 .unwrap();
508 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE, "{path}");
509 assert_eq!(
510 resp.headers().get("x-should-retry").map(|v| v.as_bytes()),
511 Some(b"false".as_ref()),
512 "{path}: retrying identical bytes cannot fix a 413"
513 );
514 let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
515 .await
516 .unwrap();
517 let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON error shape");
518 assert_eq!(v["error"]["type"], "invalid_request_error", "{path}");
519 assert_eq!(v["error"]["code"], "request_too_large", "{path}");
520 assert!(
521 v["error"]["message"].as_str().unwrap().contains("192 MiB"),
522 "{path}: message names the limit"
523 );
524 }
525 }
526
527 #[tokio::test]
528 async fn authenticated_body_admission_is_finite() {
529 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
530 let semaphore = body_admission_semaphore();
531 let mut permits = Vec::new();
532 for _ in 0..MAX_BODY_ADMISSIONS {
533 permits.push(semaphore.clone().acquire_owned().await.unwrap());
534 }
535 assert!(
536 tokio::time::timeout(std::time::Duration::from_millis(20), semaphore.acquire())
537 .await
538 .is_err(),
539 "body parser admission must not be unbounded"
540 );
541 drop(permits);
542 assert!(semaphore.acquire().await.is_ok());
543 }
544
545 #[tokio::test]
546 async fn small_body_admission_is_finite_and_separate() {
547 let _test_lock = BODY_ADMISSION_TEST_LOCK.lock().await;
548 let large = body_admission_semaphore();
549 let small = small_body_admission_semaphore();
550 let mut small_permits = Vec::new();
551 for _ in 0..MAX_SMALL_BODY_ADMISSIONS {
552 small_permits.push(small.clone().acquire_owned().await.unwrap());
553 }
554 assert!(
555 tokio::time::timeout(std::time::Duration::from_millis(20), small.acquire())
556 .await
557 .is_err(),
558 "small body parser admission must be bounded"
559 );
560 assert!(
561 large.clone().try_acquire().is_ok(),
562 "small uploads must not consume large-upload permits"
563 );
564 drop(small_permits);
565 assert!(small.acquire().await.is_ok());
566 }
567
568 #[test]
569 fn small_declared_bodies_bypass_large_upload_admission() {
570 let request = axum::http::Request::post("/v1/chat/completions")
571 .header(CONTENT_LENGTH, "2048")
572 .body(Body::empty())
573 .unwrap();
574 assert!(!body_requires_admission(&request));
575
576 let request = axum::http::Request::post("/v1/chat/completions")
577 .header(
578 CONTENT_LENGTH,
579 (BODY_ADMISSION_BYPASS_BYTES + 1).to_string(),
580 )
581 .body(Body::empty())
582 .unwrap();
583 assert!(body_requires_admission(&request));
584
585 let request = axum::http::Request::post("/v1/chat/completions")
586 .header(CONTENT_LENGTH, "2048")
587 .header(TRANSFER_ENCODING, "chunked")
588 .body(Body::empty())
589 .unwrap();
590 assert!(body_requires_admission(&request));
591 }
592
593 #[test]
594 fn declared_body_timeout_scales_with_upload_size_and_has_a_cap() {
595 let unknown = axum::http::Request::post("/v1/chat/completions")
596 .body(Body::empty())
597 .unwrap();
598 assert_eq!(body_read_timeout(&unknown), BODY_READ_TIMEOUT);
599
600 let large = axum::http::Request::post("/v1/chat/completions")
601 .header(CONTENT_LENGTH, MAX_BODY_BYTES.to_string())
602 .body(Body::empty())
603 .unwrap();
604 assert!(body_read_timeout(&large) > BODY_READ_TIMEOUT);
605 assert_eq!(body_read_timeout(&large), BODY_READ_TIMEOUT_MAX);
606
607 let absurd = axum::http::Request::post("/v1/chat/completions")
608 .header(CONTENT_LENGTH, u64::MAX.to_string())
609 .body(Body::empty())
610 .unwrap();
611 assert_eq!(body_read_timeout(&absurd), BODY_READ_TIMEOUT_MAX);
612 }
613
614 #[tokio::test]
615 async fn early_body_refusals_keep_dialect_ids_and_retry_contracts() {
616 let too_large = shape_inference_early_response(
617 "/v1/messages",
618 error_response_coded(
619 StatusCode::PAYLOAD_TOO_LARGE,
620 "request body exceeds the 192 MiB limit",
621 "invalid_request_error",
622 None,
623 Some("request_too_large"),
624 ),
625 )
626 .await;
627 assert_eq!(too_large.status(), StatusCode::PAYLOAD_TOO_LARGE);
628 let house_id = too_large.headers()["x-request-id"].clone();
629 assert_eq!(too_large.headers()["request-id"], house_id);
630 assert_eq!(too_large.headers()["x-should-retry"], "false");
631 let body = axum::body::to_bytes(too_large.into_body(), usize::MAX)
632 .await
633 .unwrap();
634 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
635 assert_eq!(payload["type"], "error");
636 assert_eq!(payload["request_id"], house_id.to_str().unwrap());
637
638 let busy = shape_inference_early_response(
639 "/v1/chat/completions",
640 retry_contract_response(
641 error_response_coded(
642 StatusCode::TOO_MANY_REQUESTS,
643 "request body admission is busy",
644 "rate_limit_error",
645 None,
646 Some("body_admission_busy"),
647 ),
648 Some(BODY_ADMISSION_RETRY_AFTER_S),
649 ),
650 )
651 .await;
652 assert_eq!(busy.status(), StatusCode::TOO_MANY_REQUESTS);
653 assert!(!busy.headers()["x-request-id"].is_empty());
654 assert_eq!(busy.headers()["retry-after"], "1");
655 assert_eq!(busy.headers()["retry-after-ms"], "1000");
656 assert!(busy.headers().get("x-should-retry").is_none());
657 let body = axum::body::to_bytes(busy.into_body(), usize::MAX)
658 .await
659 .unwrap();
660 let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
661 assert_eq!(payload["error"]["code"], "body_admission_busy");
662 }
663}
664
665#[derive(Clone, Default)]
666struct TtftRequestTrace(Option<Arc<ttft::Trace>>);
667
668fn is_sse_data_frame(bytes: &[u8]) -> bool {
669 bytes
670 .windows(b"data:".len())
671 .any(|window| window == b"data:")
672}
673
674async fn ttft_request_start(mut req: AxumRequest, next: Next) -> Response {
675 let trace = ttft::start(req.uri().path());
676 req.extensions_mut().insert(TtftRequestTrace(trace.clone()));
677 let response = next.run(req).await;
678 let Some(trace) = trace else {
679 return response;
680 };
681 let is_sse = response
682 .headers()
683 .get(CONTENT_TYPE)
684 .and_then(|value| value.to_str().ok())
685 .is_some_and(|value| value.starts_with("text/event-stream"));
686 if !is_sse {
687 return response;
688 }
689
690 let (parts, body) = response.into_parts();
693 let mut body = Box::pin(body.into_data_stream());
694 let stream = async_stream::stream! {
695 while let Some(frame) =
696 std::future::poll_fn(|cx| body.as_mut().poll_next(cx)).await
697 {
698 if frame
699 .as_ref()
700 .is_ok_and(|bytes| is_sse_data_frame(bytes))
701 {
702 trace.mark_first_sse_byte();
703 }
704 yield frame;
705 }
706 };
707 Response::from_parts(parts, Body::from_stream(stream))
708}
709
710const OPENROUTER_SCHEMA_VERSION: &str = "2.4";
711const JSON_SAFE_INTEGER_MAX: u64 = 9_007_199_254_740_991;
712
713#[derive(Debug, Clone, Default, Deserialize)]
714#[serde(deny_unknown_fields)]
715struct OpenRouterMetadataFile {
716 #[serde(default)]
717 models: HashMap<String, OpenRouterModelMetadata>,
718 #[serde(default)]
721 planned_models: HashMap<String, OpenRouterModelMetadata>,
722 #[serde(default)]
725 provider: Option<ProviderMetadata>,
726}
727
728#[derive(Debug, Clone, Deserialize)]
732#[serde(deny_unknown_fields)]
733struct ProviderMetadata {
734 id: String,
735 #[serde(default)]
736 status_url: Option<String>,
737 #[serde(default)]
738 support_contact: Option<String>,
739 #[serde(default)]
740 incident_contact: Option<String>,
741 #[serde(default)]
742 regions: Vec<String>,
743}
744
745#[derive(Debug, Clone, Default, Deserialize)]
747#[serde(deny_unknown_fields)]
748struct LifecycleMetadata {
749 #[serde(default)]
750 status: Option<String>,
751 #[serde(default)]
752 deprecation_at: Option<String>,
753 #[serde(default)]
754 retirement_at: Option<String>,
755 #[serde(default)]
756 replacement_model_id: Option<String>,
757}
758
759#[derive(Debug, Clone, Default, Deserialize)]
761#[serde(deny_unknown_fields)]
762struct ReliabilityMetadata {
763 #[serde(default)]
764 first_token_timeout_seconds: Option<u64>,
765 #[serde(default)]
766 completion_timeout_seconds: Option<u64>,
767 #[serde(default)]
768 stream_idle_timeout_seconds: Option<u64>,
769 #[serde(default)]
770 capacity_scope: Option<String>,
771}
772
773#[derive(Debug, Clone, Default, Deserialize)]
774#[serde(deny_unknown_fields)]
775struct OpenRouterModelMetadata {
776 #[serde(default)]
778 owned_by: Option<String>,
779 #[serde(default)]
780 lifecycle: Option<LifecycleMetadata>,
781 #[serde(default)]
782 reliability: Option<ReliabilityMetadata>,
783 #[serde(default)]
784 hugging_face_id: Option<String>,
785 #[serde(default)]
786 created: Option<u64>,
787 #[serde(default)]
788 quantization: Option<String>,
789 #[serde(default)]
790 description: Option<String>,
791 #[serde(default)]
792 max_prompt_length: Option<u64>,
793 #[serde(default)]
794 max_output_length: Option<u64>,
795 #[serde(default)]
798 default_output_length: Option<u64>,
799 #[serde(default)]
800 pricing: OpenRouterPricing,
801 #[serde(default)]
802 capacity: OpenRouterCapacity,
803 #[serde(default)]
804 is_ready: Option<bool>,
805 #[serde(default)]
806 is_free: Option<bool>,
807 #[serde(default)]
808 discount_to_user: Option<f64>,
809 #[serde(default)]
810 openrouter_slug: Option<String>,
811 #[serde(default)]
812 datacenters: Vec<OpenRouterDatacenter>,
813 #[serde(default)]
817 input_modalities: Vec<String>,
818 #[serde(default)]
833 surface: Option<String>,
834 #[serde(default)]
835 zdr: Option<bool>,
836 #[serde(default)]
837 hipaa: Option<bool>,
838 #[serde(default)]
848 default_reasoning_effort: Option<String>,
849 #[serde(default)]
867 default_temperature: Option<f32>,
868 #[serde(default)]
869 default_top_p: Option<f32>,
870 #[serde(default)]
872 default_top_k: Option<usize>,
873 #[serde(default)]
874 default_min_p: Option<f32>,
875 #[serde(default)]
876 default_presence_penalty: Option<f32>,
877 #[serde(default)]
878 default_frequency_penalty: Option<f32>,
879 #[serde(default)]
881 default_repetition_penalty: Option<f32>,
882 #[serde(default)]
903 non_thinking_sampling: Option<SamplingArmMetadata>,
904}
905
906#[derive(Debug, Clone, Default, Deserialize)]
912#[serde(deny_unknown_fields)]
913struct SamplingArmMetadata {
914 #[serde(default)]
915 temperature: Option<f32>,
916 #[serde(default)]
917 top_p: Option<f32>,
918 #[serde(default)]
919 top_k: Option<usize>,
920 #[serde(default)]
921 min_p: Option<f32>,
922 #[serde(default)]
923 presence_penalty: Option<f32>,
924 #[serde(default)]
925 frequency_penalty: Option<f32>,
926 #[serde(default)]
927 repetition_penalty: Option<f32>,
928}
929
930impl SamplingArmMetadata {
931 fn is_empty(&self) -> bool {
932 self.temperature.is_none()
933 && self.top_p.is_none()
934 && self.top_k.is_none()
935 && self.min_p.is_none()
936 && self.presence_penalty.is_none()
937 && self.frequency_penalty.is_none()
938 && self.repetition_penalty.is_none()
939 }
940}
941
942#[derive(Debug, Clone, Default, Deserialize)]
943#[serde(deny_unknown_fields)]
944struct OpenRouterPricing {
945 #[serde(default)]
946 prompt: Option<String>,
947 #[serde(default)]
948 cached_prompt: Option<String>,
949 #[serde(default)]
950 cache_write: Option<String>,
951 #[serde(default)]
952 completion: Option<String>,
953 #[serde(default)]
954 internal_reasoning: Option<String>,
955 #[serde(default)]
956 request: Option<String>,
957}
958
959#[derive(Debug, Clone, Default, Deserialize)]
960#[serde(deny_unknown_fields)]
961struct OpenRouterCapacity {
962 #[serde(default)]
963 prompt_tpm: Option<u64>,
964 #[serde(default)]
965 cached_prompt_tpm: Option<u64>,
966 #[serde(default)]
967 completion_tpm: Option<u64>,
968 #[serde(default)]
969 request_rpm: Option<u64>,
970 #[serde(default)]
971 concurrency: Option<u64>,
972}
973
974#[derive(Debug, Clone, Deserialize, Serialize)]
975#[serde(deny_unknown_fields)]
976struct OpenRouterDatacenter {
977 country_code: String,
978 #[serde(default, skip_serializing_if = "Option::is_none")]
979 region: Option<String>,
980}
981
982impl OpenRouterMetadataFile {
983 fn parse(
984 text: &str,
985 ) -> Result<
986 (
987 HashMap<String, OpenRouterModelMetadata>,
988 Option<ProviderMetadata>,
989 ),
990 String,
991 > {
992 let file: Self =
993 toml::from_str(text).map_err(|e| format!("models metadata TOML parse: {e}"))?;
994 for (alias, metadata) in &file.models {
995 validate_openrouter_metadata(alias, metadata)?;
996 }
997 for (alias, metadata) in &file.planned_models {
998 validate_openrouter_metadata(alias, metadata)?;
999 if file.models.contains_key(alias) {
1000 return Err(format!(
1001 "model alias {alias:?} appears in both models and planned_models"
1002 ));
1003 }
1004 }
1005 if let Some(provider) = &file.provider {
1006 if provider.id.is_empty() {
1007 return Err("provider.id must be a non-empty slug".into());
1008 }
1009 for (field, value) in [
1011 ("provider.support_contact", &provider.support_contact),
1012 ("provider.incident_contact", &provider.incident_contact),
1013 ] {
1014 if let Some(value) = value {
1015 if !value.contains(':') {
1016 return Err(format!(
1017 "{field} must be a URI (mailto:… or https://…), got {value:?}"
1018 ));
1019 }
1020 }
1021 }
1022 }
1023 Ok((file.models, file.provider))
1024 }
1025
1026 #[cfg(test)]
1027 fn from_toml(text: &str) -> Result<HashMap<String, OpenRouterModelMetadata>, String> {
1028 Self::parse(text).map(|(models, _)| models)
1029 }
1030}
1031
1032fn per_million_price(per_token: &str) -> Option<String> {
1036 if !valid_price_string(per_token) {
1037 return None;
1038 }
1039 let (whole, frac) = match per_token.split_once('.') {
1040 Some((whole, frac)) => (whole, frac),
1041 None => (per_token, ""),
1042 };
1043 let mut digits = format!("{whole}{frac}");
1044 let point = whole.len() + 6;
1045 while digits.len() < point {
1046 digits.push('0');
1047 }
1048 let (int_part, frac_part) = digits.split_at(point);
1049 let int_part = int_part.trim_start_matches('0');
1050 let int_part = if int_part.is_empty() { "0" } else { int_part };
1051 let mut frac_out = frac_part.trim_end_matches('0').to_string();
1052 while frac_out.len() < 2 {
1053 frac_out.push('0');
1054 }
1055 Some(format!("{int_part}.{frac_out}"))
1056}
1057
1058fn valid_price_string(value: &str) -> bool {
1059 let mut parts = value.split('.');
1060 let whole = parts.next().unwrap_or_default();
1061 let fraction = parts.next();
1062 !whole.is_empty()
1063 && whole.bytes().all(|b| b.is_ascii_digit())
1064 && fraction.is_none_or(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
1065 && parts.next().is_none()
1066}
1067
1068fn validate_openrouter_metadata(
1069 alias: &str,
1070 metadata: &OpenRouterModelMetadata,
1071) -> Result<(), String> {
1072 if alias.is_empty() {
1073 return Err("models metadata contains an empty model alias".into());
1074 }
1075 if let Some(effort) = metadata.default_reasoning_effort.as_deref()
1078 && !matches!(effort, "none" | "minimal" | "low" | "medium" | "high")
1079 {
1080 return Err(format!(
1081 "model {alias:?}: default_reasoning_effort {effort:?} is not a \
1082 reasoning_effort level (none|minimal|low|medium|high)"
1083 ));
1084 }
1085 validate_sampling_defaults(alias, metadata)?;
1086 for m in &metadata.input_modalities {
1087 if m != "image" && m != "video" {
1088 return Err(format!(
1089 "model {alias:?}: input_modalities entry {m:?} not served (image/video)"
1090 ));
1091 }
1092 }
1093 if let Some(sfc) = metadata.surface.as_deref()
1094 && !matches!(sfc, "chat" | "embedding" | "rerank")
1095 {
1096 return Err(format!(
1097 "model {alias:?}: surface {sfc:?} is not a served surface (chat|embedding|rerank)"
1098 ));
1099 }
1100 if let Some(q) = metadata.quantization.as_deref()
1101 && !matches!(
1102 q,
1103 "int4"
1104 | "int8"
1105 | "fp4"
1106 | "mxfp4"
1107 | "nvfp4"
1108 | "fp6"
1109 | "fp8"
1110 | "mxfp8"
1111 | "fp16"
1112 | "bf16"
1113 | "fp32"
1114 )
1115 {
1116 return Err(format!(
1117 "model {alias:?}: quantization {q:?} is not in the OpenRouter schema 2.4 enum"
1118 ));
1119 }
1120 for (field, value) in [
1121 ("pricing.prompt", metadata.pricing.prompt.as_deref()),
1122 (
1123 "pricing.cached_prompt",
1124 metadata.pricing.cached_prompt.as_deref(),
1125 ),
1126 (
1127 "pricing.cache_write",
1128 metadata.pricing.cache_write.as_deref(),
1129 ),
1130 ("pricing.completion", metadata.pricing.completion.as_deref()),
1131 (
1132 "pricing.internal_reasoning",
1133 metadata.pricing.internal_reasoning.as_deref(),
1134 ),
1135 ("pricing.request", metadata.pricing.request.as_deref()),
1136 ] {
1137 if let Some(value) = value
1138 && !valid_price_string(value)
1139 {
1140 return Err(format!(
1141 "model {alias:?}: {field} must be a non-negative per-unit USD decimal string"
1142 ));
1143 }
1144 }
1145 for (field, value) in [
1146 ("created", metadata.created),
1147 ("max_prompt_length", metadata.max_prompt_length),
1148 ("max_output_length", metadata.max_output_length),
1149 ("default_output_length", metadata.default_output_length),
1150 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1151 (
1152 "capacity.cached_prompt_tpm",
1153 metadata.capacity.cached_prompt_tpm,
1154 ),
1155 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1156 ("capacity.request_rpm", metadata.capacity.request_rpm),
1157 ("capacity.concurrency", metadata.capacity.concurrency),
1158 ] {
1159 if let Some(value) = value
1160 && value > JSON_SAFE_INTEGER_MAX
1161 {
1162 return Err(format!(
1163 "model {alias:?}: {field} exceeds OpenRouter's JSON safe-integer maximum"
1164 ));
1165 }
1166 }
1167 for (field, value) in [
1168 ("max_prompt_length", metadata.max_prompt_length),
1169 ("max_output_length", metadata.max_output_length),
1170 ("default_output_length", metadata.default_output_length),
1171 ("capacity.prompt_tpm", metadata.capacity.prompt_tpm),
1172 (
1173 "capacity.cached_prompt_tpm",
1174 metadata.capacity.cached_prompt_tpm,
1175 ),
1176 ("capacity.completion_tpm", metadata.capacity.completion_tpm),
1177 ("capacity.request_rpm", metadata.capacity.request_rpm),
1178 ("capacity.concurrency", metadata.capacity.concurrency),
1179 ] {
1180 if value == Some(0) {
1181 return Err(format!(
1182 "model {alias:?}: {field} must be greater than zero when declared"
1183 ));
1184 }
1185 }
1186 if let (Some(default), Some(maximum)) =
1187 (metadata.default_output_length, metadata.max_output_length)
1188 && default > maximum
1189 {
1190 return Err(format!(
1191 "model {alias:?}: default_output_length {default} exceeds max_output_length {maximum}"
1192 ));
1193 }
1194 if metadata.default_output_length.is_some() && metadata.max_output_length.is_none() {
1195 return Err(format!(
1196 "model {alias:?}: default_output_length requires max_output_length"
1197 ));
1198 }
1199 if let Some(discount) = metadata.discount_to_user
1200 && (!discount.is_finite() || discount >= 1.0)
1201 {
1202 return Err(format!(
1203 "model {alias:?}: discount_to_user must be finite and less than 1"
1204 ));
1205 }
1206 if metadata
1207 .openrouter_slug
1208 .as_deref()
1209 .is_some_and(str::is_empty)
1210 {
1211 return Err(format!(
1212 "model {alias:?}: openrouter_slug must not be empty when declared"
1213 ));
1214 }
1215 for dc in &metadata.datacenters {
1216 if dc.country_code.len() != 2 || !dc.country_code.bytes().all(|b| b.is_ascii_uppercase()) {
1217 return Err(format!(
1218 "model {alias:?}: datacenter country_code {:?} must be two uppercase ASCII letters",
1219 dc.country_code
1220 ));
1221 }
1222 }
1223 Ok(())
1224}
1225
1226fn validate_sampling_defaults(
1241 alias: &str,
1242 metadata: &OpenRouterModelMetadata,
1243) -> Result<(), String> {
1244 validate_sampling_arm(
1245 alias,
1246 &[
1247 "default_temperature",
1248 "default_top_p",
1249 "default_min_p",
1250 "default_presence_penalty",
1251 "default_frequency_penalty",
1252 "default_repetition_penalty",
1253 ],
1254 metadata.default_temperature,
1255 metadata.default_top_p,
1256 metadata.default_min_p,
1257 metadata.default_presence_penalty,
1258 metadata.default_frequency_penalty,
1259 metadata.default_repetition_penalty,
1260 )?;
1261 if let Some(arm) = &metadata.non_thinking_sampling {
1262 if arm.is_empty() {
1266 return Err(format!(
1267 "model {alias:?}: non_thinking_sampling declares no fields — declare at \
1268 least one vendor recommendation or delete the table"
1269 ));
1270 }
1271 validate_sampling_arm(
1272 alias,
1273 &[
1274 "non_thinking_sampling.temperature",
1275 "non_thinking_sampling.top_p",
1276 "non_thinking_sampling.min_p",
1277 "non_thinking_sampling.presence_penalty",
1278 "non_thinking_sampling.frequency_penalty",
1279 "non_thinking_sampling.repetition_penalty",
1280 ],
1281 arm.temperature,
1282 arm.top_p,
1283 arm.min_p,
1284 arm.presence_penalty,
1285 arm.frequency_penalty,
1286 arm.repetition_penalty,
1287 )?;
1288 }
1289 Ok(())
1290}
1291
1292#[allow(clippy::too_many_arguments)]
1298fn validate_sampling_arm(
1299 alias: &str,
1300 keys: &[&str; 6],
1301 temperature: Option<f32>,
1302 top_p: Option<f32>,
1303 min_p: Option<f32>,
1304 presence_penalty: Option<f32>,
1305 frequency_penalty: Option<f32>,
1306 repetition_penalty: Option<f32>,
1307) -> Result<(), String> {
1308 if let Some(t) = temperature {
1309 if !t.is_finite() || t <= 0.0 || t > 2.0 {
1310 return Err(format!(
1311 "model {alias:?}: {} {t} must be finite and in (0, 2]. \
1312 A zero DEFAULT would make greedy decoding the deployment-wide behavior for \
1313 every request that omits temperature (owner ruling 2026-08-19: we serve the \
1314 vendor recommendation, not greedy); clients reach greedy by sending an \
1315 explicit temperature 0.",
1316 keys[0]
1317 ));
1318 }
1319 }
1320 if let Some(p) = top_p
1321 && (!p.is_finite() || p <= 0.0 || p > 1.0)
1322 {
1323 return Err(format!(
1324 "model {alias:?}: {} {p} must be finite and in (0, 1] (1.0 = disabled)",
1325 keys[1]
1326 ));
1327 }
1328 if let Some(m) = min_p
1329 && (!m.is_finite() || !(0.0..1.0).contains(&m))
1330 {
1331 return Err(format!(
1332 "model {alias:?}: {} {m} must be finite and in [0, 1) (0.0 = disabled)",
1333 keys[2]
1334 ));
1335 }
1336 for (field, value) in [(keys[3], presence_penalty), (keys[4], frequency_penalty)] {
1337 if let Some(v) = value
1338 && (!v.is_finite() || !(-2.0..=2.0).contains(&v))
1339 {
1340 return Err(format!(
1341 "model {alias:?}: {field} {v} must be finite and in [-2, 2]"
1342 ));
1343 }
1344 }
1345 if let Some(r) = repetition_penalty
1346 && (!r.is_finite() || r <= 0.0)
1347 {
1348 return Err(format!(
1349 "model {alias:?}: {} {r} must be finite and \
1350 greater than zero (1.0 = off)",
1351 keys[5]
1352 ));
1353 }
1354 Ok(())
1355}
1356
1357fn load_openrouter_metadata(
1358 models: &[(String, String, Option<String>)],
1359) -> Result<
1360 (
1361 HashMap<String, OpenRouterModelMetadata>,
1362 Option<ProviderMetadata>,
1363 ),
1364 String,
1365> {
1366 let path = match std::env::var("MEMRA_MODEL_METADATA") {
1367 Ok(path) => path,
1368 Err(_) => return Ok((HashMap::new(), None)),
1369 };
1370 let p = std::path::Path::new(&path);
1371 if !p.is_file() {
1372 return Err(format!(
1373 "MEMRA_MODEL_METADATA={path:?} is not an existing TOML file"
1374 ));
1375 }
1376 let text =
1377 std::fs::read_to_string(p).map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1378 let (metadata, provider) = OpenRouterMetadataFile::parse(&text)
1379 .map_err(|e| format!("MEMRA_MODEL_METADATA {path:?}: {e}"))?;
1380 for alias in metadata.keys() {
1381 if !models.iter().any(|(name, _, _)| name == alias) {
1382 return Err(format!(
1383 "MEMRA_MODEL_METADATA {path:?}: model alias {alias:?} is not present in MEMRA_MODELS"
1384 ));
1385 }
1386 }
1387 eprintln!(
1388 "[server] OpenRouter metadata loaded: {} model(s) from {path}",
1389 metadata.len()
1390 );
1391 Ok((metadata, provider))
1392}
1393
1394#[derive(Clone)]
1395struct AppState {
1396 cmd_tx: Sender<Cmd>,
1397 models: Arc<Vec<String>>,
1398 caps: Arc<HashMap<String, ModelCaps>>,
1399 openrouter_metadata: Arc<HashMap<String, OpenRouterModelMetadata>>,
1400 provider_metadata: Arc<Option<ProviderMetadata>>,
1402 metering: Option<Arc<dyn metering::Metering>>,
1408 budget_tokenizers: Option<Arc<HashMap<String, Arc<Tokenizer>>>>,
1411 api_auth: ApiAuth,
1414 metrics_auth: MetricsAuth,
1416 metrics: SharedMetrics,
1417 started: u64,
1420 inflight: InflightCounts,
1424 tenant_inflight: TenantGauge,
1427 health: health::SharedHealth,
1431 bg: Option<(Arc<darklane::BgJobState>, &'static str)>,
1435}
1436
1437impl AppState {
1438 fn sampling_defaults(&self, model: &str) -> ModelSamplingDefaults {
1452 ModelSamplingDefaults::resolve(self.openrouter_metadata.get(model), self.caps.get(model))
1453 }
1454}
1455
1456#[derive(Clone, Default)]
1457struct ApiAuth {
1458 keyring: Option<&'static auth::KeyStore>,
1459 single_key: Option<Arc<str>>,
1460}
1461
1462impl ApiAuth {
1463 fn from_env() -> Result<ApiAuth, String> {
1464 let single_key = match std::env::var("MEMRA_API_KEY") {
1465 Ok(key) if key.is_empty() => return Err("MEMRA_API_KEY must not be empty".into()),
1466 Ok(key) => Some(Arc::from(key)),
1467 Err(std::env::VarError::NotPresent) => None,
1468 Err(std::env::VarError::NotUnicode(_)) => {
1469 return Err("MEMRA_API_KEY must be valid UTF-8".into());
1470 }
1471 };
1472 Ok(ApiAuth {
1473 keyring: auth::global(),
1474 single_key,
1475 })
1476 }
1477
1478 fn configured(&self) -> bool {
1479 self.keyring.is_some() || self.single_key.is_some()
1480 }
1481}
1482
1483#[derive(Clone, Default)]
1484struct MetricsAuth {
1485 required: bool,
1486 token: Option<Arc<str>>,
1487}
1488
1489impl MetricsAuth {
1490 fn new(bind_loopback: bool, api_auth_configured: bool, token: Option<String>) -> MetricsAuth {
1491 let token = token.map(Arc::from);
1492 MetricsAuth {
1493 required: !bind_loopback || api_auth_configured || token.is_some(),
1494 token,
1495 }
1496 }
1497}
1498
1499fn resolve_bind_addr(addr: &str) -> Result<(SocketAddr, bool), String> {
1500 let mut resolved = addr
1501 .to_socket_addrs()
1502 .map_err(|e| format!("MEMRA_ADDR={addr:?} cannot be resolved: {e}"))?;
1503 let first = resolved
1504 .next()
1505 .ok_or_else(|| format!("MEMRA_ADDR={addr:?} resolved to no socket addresses"))?;
1506 let mut loopback = first.ip().to_canonical().is_loopback();
1507 for socket in resolved {
1508 loopback &= socket.ip().to_canonical().is_loopback();
1509 }
1510 Ok((first, loopback))
1511}
1512
1513fn bind_is_loopback(addr: &str) -> Result<bool, String> {
1514 resolve_bind_addr(addr).map(|(_, loopback)| loopback)
1515}
1516
1517fn validate_bind_security(
1518 addr: &str,
1519 api_auth_configured: bool,
1520 allow_open_bind: bool,
1521) -> Result<bool, String> {
1522 let loopback = bind_is_loopback(addr)?;
1523 if !loopback && !api_auth_configured && !allow_open_bind {
1524 return Err(format!(
1525 "refusing unauthenticated non-loopback bind {addr:?}; configure MEMRA_API_KEY or \
1526 MEMRA_API_KEYS, or set MEMRA_ALLOW_OPEN_BIND=1 for an explicit development override"
1527 ));
1528 }
1529 Ok(loopback)
1530}
1531
1532type InflightCounts = Arc<[std::sync::atomic::AtomicUsize; 3]>;
1550
1551type TenantGauge = Arc<std::sync::Mutex<HashMap<String, usize>>>;
1554
1555struct InflightGuard {
1559 counts: InflightCounts,
1560 idx: usize,
1561 tenants: TenantGauge,
1562 tenant: String,
1563}
1564
1565impl InflightGuard {
1566 fn try_acquire(
1570 counts: InflightCounts,
1571 lane: lanes::Lane,
1572 tenants: TenantGauge,
1573 tenant: &str,
1574 tenant_cap: Option<usize>,
1575 ) -> Result<(Self, usize, usize), usize> {
1576 let idx = lane.idx();
1577 let nt = {
1578 let mut m = tenants.lock().unwrap();
1579 let e = m.entry(tenant.to_string()).or_insert(0);
1580 if tenant_cap.is_some_and(|cap| *e >= cap) {
1581 return Err(*e);
1582 }
1583 *e += 1;
1584 *e
1585 };
1586 let n = counts[idx].fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1587 Ok((
1588 InflightGuard {
1589 counts,
1590 idx,
1591 tenants,
1592 tenant: tenant.to_string(),
1593 },
1594 n,
1595 nt,
1596 ))
1597 }
1598}
1599
1600impl Drop for InflightGuard {
1601 fn drop(&mut self) {
1602 self.counts[self.idx].fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1603 let mut m = self.tenants.lock().unwrap();
1604 if let Some(e) = m.get_mut(&self.tenant) {
1605 *e -= 1;
1606 if *e == 0 {
1607 m.remove(&self.tenant);
1608 }
1609 }
1610 }
1611}
1612
1613fn lane_cap(lane: lanes::Lane) -> usize {
1617 static CAPS: std::sync::OnceLock<[usize; 3]> = std::sync::OnceLock::new();
1618 CAPS.get_or_init(|| {
1619 let batching = std::env::var("MEMRA_SERVE_BATCH")
1620 .map(|v| v != "0")
1621 .unwrap_or(true);
1622 let interactive = if batching {
1623 std::env::var("MEMRA_MAX_SESSIONS")
1624 .ok()
1625 .and_then(|v| v.parse().ok())
1626 .unwrap_or(64)
1627 } else {
1628 worker::MAX_ACTIVE
1629 };
1630 let p = lanes::LanePolicy::from_env();
1631 [interactive, p.max_sessions[1], p.max_sessions[2]]
1632 })[lane.idx()]
1633}
1634
1635fn reset_estimate_s(m: &worker::Metrics) -> u64 {
1638 if m.completed > 0 && m.step_p50_ms > 0.0 {
1639 let mean_toks = m.tokens_out as f64 / m.completed as f64;
1640 return ((mean_toks * m.step_p50_ms as f64 / 1000.0).ceil() as u64).clamp(1, 600);
1641 }
1642 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
1643 *D.get_or_init(|| {
1644 std::env::var("MEMRA_RL_RESET_S")
1645 .ok()
1646 .and_then(|v| v.parse().ok())
1647 .unwrap_or(2)
1648 })
1649}
1650
1651pub(crate) const TIMEOUT_MS_MIN: u64 = 1_000;
1668pub(crate) const TIMEOUT_MS_MAX: u64 = 90_000;
1669pub(crate) const TIMEOUT_MS_DEFAULT: u64 = 90_000;
1670
1671pub(crate) fn parse_timeout_ms(v: Option<&serde_json::Value>) -> Result<u64, String> {
1675 let Some(v) = v.filter(|v| !v.is_null()) else {
1676 return Ok(TIMEOUT_MS_DEFAULT);
1677 };
1678 let Some(ms) = v.as_u64() else {
1679 return Err(format!(
1680 "timeout_ms must be an integer number of milliseconds in \
1681 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX}, got {v}; for work longer than \
1682 {TIMEOUT_MS_MAX} ms use \"stream\": true — the deadline then bounds only the \
1683 time to first token and the stream may run as long as it needs"
1684 ));
1685 };
1686 if !(TIMEOUT_MS_MIN..=TIMEOUT_MS_MAX).contains(&ms) {
1687 return Err(format!(
1688 "timeout_ms {ms} is outside the accepted range \
1689 {TIMEOUT_MS_MIN}..={TIMEOUT_MS_MAX} (milliseconds). {TIMEOUT_MS_MAX} is a \
1690 platform ceiling, not a preference: the fronting proxy fails a non-streaming \
1691 response whose headers take ~100 s (HTTP 524), so promising more would be a \
1692 lie. For work longer than {TIMEOUT_MS_MAX} ms use \"stream\": true — the \
1693 deadline then bounds only the time to first token and the stream may run as \
1694 long as it needs"
1695 ));
1696 }
1697 Ok(ms)
1698}
1699
1700#[derive(Clone, Copy)]
1703pub(crate) struct RequestDeadline {
1704 pub(crate) at: tokio::time::Instant,
1705 pub(crate) ms: u64,
1706}
1707
1708impl RequestDeadline {
1709 pub(crate) fn starting_now(ms: u64) -> Self {
1710 Self {
1711 at: tokio::time::Instant::now() + std::time::Duration::from_millis(ms),
1712 ms,
1713 }
1714 }
1715
1716 pub(crate) fn remaining(&self) -> std::time::Duration {
1717 self.at
1718 .saturating_duration_since(tokio::time::Instant::now())
1719 }
1720}
1721
1722pub(crate) fn deadline_exceeded_response(ms: u64, stream: bool) -> Response {
1728 let what = if stream {
1729 "the first token was produced"
1730 } else {
1731 "the response completed"
1732 };
1733 let msg = format!(
1734 "deadline of {ms} ms (timeout_ms; default {TIMEOUT_MS_DEFAULT}) elapsed before \
1735 {what}; generation was cancelled and this request is not billed"
1736 );
1737 error_response_coded(
1738 StatusCode::REQUEST_TIMEOUT,
1739 &msg,
1740 "timeout",
1741 Some("timeout_ms"),
1742 Some("deadline_exceeded"),
1743 )
1744}
1745
1746pub(crate) const PREFILL_FLOOR_TOK_S: u64 = 2_000;
1786
1787pub(crate) const DECODE_FLOOR_TOK_S: u64 = 60;
1792
1793pub(crate) const DEADLINE_INFEASIBLE_MARGIN_PCT: u64 = 150;
1797
1798fn env_flag_on(name: &'static str, default_on: bool) -> bool {
1804 match std::env::var(name) {
1805 Ok(v) => !matches!(
1806 v.trim().to_ascii_lowercase().as_str(),
1807 "0" | "off" | "false"
1808 ),
1809 Err(_) => default_on,
1810 }
1811}
1812
1813fn env_u64(name: &'static str, default: u64) -> u64 {
1816 std::env::var(name)
1817 .ok()
1818 .and_then(|v| v.parse::<u64>().ok())
1819 .filter(|v| *v > 0)
1820 .unwrap_or(default)
1821}
1822
1823const CHARS_PER_TOKEN_FLOOR: usize = 6;
1834
1835pub(crate) fn prompt_tokens_estimate(
1836 request: &worker::Request,
1837 tokenizer: Option<&Tokenizer>,
1838) -> u64 {
1839 if !request.prompt_ids.is_empty() {
1840 return request.prompt_ids.len() as u64;
1841 }
1842 let mut text = String::new();
1843 text.push_str(&request.prompt_text);
1844 for turn in &request.chat_turns {
1845 text.push_str(&turn.content);
1846 }
1847 for tool in &request.tools_json {
1848 text.push_str(tool);
1849 }
1850 if let Some(tokenizer) = tokenizer {
1851 return tokenizer.encode(text.as_str(), false).len() as u64;
1852 }
1853 (text.len() / CHARS_PER_TOKEN_FLOOR) as u64
1854}
1855
1856pub(crate) fn deadline_fitting_max_tokens(prompt_tokens: u64, remaining_ms: u64) -> Option<u64> {
1860 let prefill_ms = prompt_tokens
1861 .saturating_mul(1_000)
1862 .checked_div(env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S))
1863 .unwrap_or(u64::MAX);
1864 let decode_ms = remaining_ms.checked_sub(prefill_ms)?;
1865 if decode_ms == 0 {
1866 return None;
1867 }
1868 Some(decode_ms.saturating_mul(env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S)) / 1_000)
1869}
1870
1871pub(crate) fn nonstream_deadline_gate(
1881 request: &worker::Request,
1882 stream: bool,
1883 deadline: RequestDeadline,
1884 caller_declared_max_tokens: bool,
1885 tokenizer: Option<&Tokenizer>,
1886) -> Result<(), String> {
1887 if stream || !env_flag_on("MEMRA_NONSTREAM_DEADLINE_GATE", true) {
1888 return Ok(());
1889 }
1890 let max_new = request.params.max_new as u64;
1891 if !caller_declared_max_tokens || max_new == worker::MAX_NEW_CTX_BOUNDED as u64 || max_new == 0
1898 {
1899 return Ok(());
1900 }
1901 let prompt_tokens = prompt_tokens_estimate(request, tokenizer);
1902 let remaining_ms = deadline.remaining().as_millis() as u64;
1903 let prefill_ms = prompt_tokens.saturating_mul(1_000)
1904 / env_u64("MEMRA_PREFILL_FLOOR_TOK_S", PREFILL_FLOOR_TOK_S).max(1);
1905 let decode_ms = max_new.saturating_mul(1_000)
1906 / env_u64("MEMRA_DECODE_FLOOR_TOK_S", DECODE_FLOOR_TOK_S).max(1);
1907 let est_ms = prefill_ms.saturating_add(decode_ms);
1908 let bound_ms = remaining_ms.saturating_mul(DEADLINE_INFEASIBLE_MARGIN_PCT) / 100;
1909 if est_ms <= bound_ms {
1910 return Ok(());
1911 }
1912 let fits = deadline_fitting_max_tokens(prompt_tokens, remaining_ms);
1913 let advice = match fits {
1914 Some(fits) if fits > 0 => format!(
1915 "lower max_tokens to about {fits} for this prompt, or set \"stream\": true — a \
1916 stream's deadline bounds only the time to first token, so it may run as long \
1917 as it needs"
1918 ),
1919 _ => format!(
1920 "this prompt ({prompt_tokens} tok) needs most of the deadline before the first \
1921 token, so no max_tokens fits: set \"stream\": true"
1922 ),
1923 };
1924 Err(format!(
1925 "a non-streaming request for {max_new} tokens on a ~{prompt_tokens}-token prompt \
1926 needs an estimated ~{}s, which does not fit the {remaining_ms} ms timeout_ms \
1927 deadline (max {TIMEOUT_MS_MAX} ms — a platform ceiling: the fronting proxy fails \
1928 a non-streaming response whose headers take ~100 s). Refused before any GPU work \
1929 rather than after the deadline: {advice}",
1930 est_ms / 1_000,
1931 ))
1932}
1933
1934fn max_queue_depth(cap: usize) -> usize {
1938 static D: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1939 D.get_or_init(|| {
1940 std::env::var("MEMRA_MAX_QUEUE_DEPTH")
1941 .ok()
1942 .and_then(|v| v.parse().ok())
1943 })
1944 .unwrap_or(cap.saturating_mul(4))
1945}
1946
1947pub(crate) fn admission_backpressure(
1960 st: &AppState,
1961 lane: lanes::Lane,
1962 rl: &RateLimit,
1963 deadline: RequestDeadline,
1964) -> Result<(), (Response, &'static str)> {
1965 if lane != lanes::Lane::Interactive || rl.remaining > 0 {
1966 return Ok(());
1967 }
1968 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
1969 let backlog = m.queued_requests as usize
1970 + worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire);
1971 let cap = lane_cap(lane).max(1);
1972 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
1973 let bound = max_queue_depth(cap);
1974 if backlog >= bound {
1975 let msg = format!(
1976 "interactive queue is at its bound ({backlog} queued, bound {bound}); this \
1977 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
1978 coarse estimate, not a promise)"
1979 );
1980 let resp = retry_contract_response(
1981 (
1982 StatusCode::TOO_MANY_REQUESTS,
1983 Json(error_body(
1984 &msg,
1985 "rate_limit_error",
1986 None,
1987 Some("shed_queue"),
1988 )),
1989 )
1990 .into_response(),
1991 Some(est_wait_s),
1992 );
1993 return Err((resp, "shed_queue"));
1994 }
1995 let remaining_ms = deadline.remaining().as_millis() as u64;
1996 if est_wait_s.saturating_mul(1_000) > remaining_ms {
1997 let msg = format!(
1998 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
1999 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2000 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2001 estimate, not a promise)"
2002 );
2003 let resp = retry_contract_response(
2004 (
2005 StatusCode::TOO_MANY_REQUESTS,
2006 Json(error_body(
2007 &msg,
2008 "rate_limit_error",
2009 None,
2010 Some("shed_deadline"),
2011 )),
2012 )
2013 .into_response(),
2014 Some(est_wait_s),
2015 );
2016 return Err((resp, "shed_deadline"));
2017 }
2018 Ok(())
2019}
2020
2021pub(crate) struct PendingAdmissionGuard {
2026 reserved: bool,
2027 lane: lanes::Lane,
2028}
2029
2030impl PendingAdmissionGuard {
2031 pub(crate) fn commit(mut self) {
2035 self.reserved = false;
2036 std::mem::forget(self);
2037 }
2038}
2039
2040impl Drop for PendingAdmissionGuard {
2041 fn drop(&mut self) {
2042 if self.reserved {
2043 worker::release_pending_admit();
2044 worker::release_admission_reservation(self.lane);
2045 }
2046 }
2047}
2048
2049pub(crate) fn reserve_pending_admit(
2050 st: &AppState,
2051 lane: lanes::Lane,
2052 rl: &RateLimit,
2053 deadline: RequestDeadline,
2054) -> Result<PendingAdmissionGuard, (Response, &'static str)> {
2055 let cap = lane_cap(lane).max(1);
2060 let bound = max_queue_depth(cap);
2061 let reservations_for_lane = &worker::ADMISSION_RESERVATIONS[lane.idx()];
2062 loop {
2063 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
2064 let reservations = reservations_for_lane.load(std::sync::atomic::Ordering::Acquire);
2065 let backlog = reservations;
2069 let est_wait_s = reset_estimate_s(&m).saturating_mul((backlog / cap + 1) as u64);
2070 if backlog >= bound {
2071 let msg = format!(
2072 "{} queue is at its bound ({backlog} queued, bound {bound}); this \
2073 request was not admitted and is not billed; retry after ~{est_wait_s}s (a \
2074 coarse estimate, not a promise)",
2075 lane.as_str()
2076 );
2077 let resp = retry_contract_response(
2078 (
2079 StatusCode::TOO_MANY_REQUESTS,
2080 Json(error_body(
2081 &msg,
2082 "rate_limit_error",
2083 None,
2084 Some("shed_queue"),
2085 )),
2086 )
2087 .into_response(),
2088 Some(est_wait_s),
2089 );
2090 return Err((resp, "shed_queue"));
2091 }
2092 let remaining_ms = deadline.remaining().as_millis() as u64;
2093 let waits_for_capacity = rl.remaining == 0 || backlog > 0;
2097 if lane == lanes::Lane::Interactive
2098 && waits_for_capacity
2099 && est_wait_s.saturating_mul(1_000) > remaining_ms
2100 {
2101 let msg = format!(
2102 "estimated queue wait ~{est_wait_s}s exceeds this request's remaining \
2103 timeout_ms deadline ({remaining_ms} ms); this request was not admitted and \
2104 is not billed; retry after ~{est_wait_s}s or raise timeout_ms (a coarse \
2105 estimate, not a promise)"
2106 );
2107 let resp = retry_contract_response(
2108 (
2109 StatusCode::TOO_MANY_REQUESTS,
2110 Json(error_body(
2111 &msg,
2112 "rate_limit_error",
2113 None,
2114 Some("shed_deadline"),
2115 )),
2116 )
2117 .into_response(),
2118 Some(est_wait_s),
2119 );
2120 return Err((resp, "shed_deadline"));
2121 }
2122 if reservations_for_lane
2123 .compare_exchange(
2124 reservations,
2125 reservations.saturating_add(1),
2126 std::sync::atomic::Ordering::AcqRel,
2127 std::sync::atomic::Ordering::Acquire,
2128 )
2129 .is_ok()
2130 {
2131 worker::PENDING_ADMITS.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2135 return Ok(PendingAdmissionGuard {
2136 reserved: true,
2137 lane,
2138 });
2139 }
2140 }
2141}
2142
2143static DRAINING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
2154
2155fn draining() -> bool {
2156 DRAINING.load(std::sync::atomic::Ordering::SeqCst)
2157}
2158
2159fn drain_deadline_s() -> u64 {
2161 static D: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
2162 *D.get_or_init(|| {
2163 std::env::var("MEMRA_DRAIN_S")
2164 .ok()
2165 .and_then(|v| v.parse().ok())
2166 .unwrap_or(30)
2167 })
2168}
2169
2170fn drain_response() -> Response {
2180 let resp = (
2181 StatusCode::SERVICE_UNAVAILABLE,
2182 Json(error_body(
2183 "server is draining (shutdown in progress); retry",
2184 "server_error",
2185 None,
2186 Some("draining"),
2187 )),
2188 )
2189 .into_response();
2190 retry_contract_response(resp, Some(drain_deadline_s()))
2191}
2192
2193struct RateLimit {
2195 limit: usize,
2196 remaining: usize,
2197 reset_s: u64,
2198}
2199
2200impl RateLimit {
2201 fn at_admit(
2206 lane: lanes::Lane,
2207 n_inflight: usize,
2208 metrics: &SharedMetrics,
2209 tenant: &auth::TenantCtx,
2210 n_tenant: usize,
2211 ) -> Self {
2212 let global = lane_cap(lane);
2213 let Some(t) = tenant.rate_limit.filter(|&t| t < global) else {
2214 return Self::compute(global, n_inflight, metrics);
2215 };
2216 let headroom = t
2217 .saturating_sub(n_tenant)
2218 .min(global.saturating_sub(n_inflight));
2219 Self::compute(t, t - headroom, metrics)
2221 }
2222
2223 fn compute(limit: usize, n_inflight: usize, metrics: &SharedMetrics) -> Self {
2224 let remaining = limit.saturating_sub(n_inflight);
2225 let reset_s = if remaining > 0 {
2226 0
2227 } else {
2228 let m = metrics.lock().map(|m| m.clone()).unwrap_or_default();
2229 reset_estimate_s(&m)
2230 };
2231 RateLimit {
2232 limit,
2233 remaining,
2234 reset_s,
2235 }
2236 }
2237
2238 fn attach(&self, mut resp: Response) -> Response {
2240 let h = resp.headers_mut();
2241 for (k, v) in [
2242 ("x-ratelimit-limit", self.limit as u64),
2243 ("x-ratelimit-remaining", self.remaining as u64),
2244 ("x-ratelimit-reset", self.reset_s),
2245 ] {
2246 if let Ok(v) = axum::http::HeaderValue::from_str(&v.to_string()) {
2247 h.insert(axum::http::HeaderName::from_static(k), v);
2248 }
2249 }
2250 resp
2251 }
2252}
2253
2254fn acquire_request_slot(
2258 st: &AppState,
2259 lane: lanes::Lane,
2260 tenant: &auth::TenantCtx,
2261 env: &Envelope,
2262) -> Result<(InflightGuard, RateLimit), Response> {
2263 let global = lane_cap(lane);
2264 let tenant_cap = tenant.rate_limit.filter(|&cap| cap < global);
2265 match InflightGuard::try_acquire(
2266 st.inflight.clone(),
2267 lane,
2268 st.tenant_inflight.clone(),
2269 &tenant.tenant,
2270 tenant_cap,
2271 ) {
2272 Ok((guard, n_inflight, n_tenant)) => {
2273 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2274 Ok((guard, rl))
2275 }
2276 Err(n_tenant) => {
2277 let n_inflight = st.inflight[lane.idx()].load(std::sync::atomic::Ordering::SeqCst);
2278 let rl = RateLimit::at_admit(lane, n_inflight, &st.metrics, tenant, n_tenant);
2279 let error =
2280 worker::EngineError::rate_limit("api key concurrent request limit reached; retry");
2281 Err(rl.attach(with_request_id(&env.id, engine_error_response(&error))))
2282 }
2283 }
2284}
2285
2286#[derive(Deserialize)]
2288struct CompletionReq {
2289 model: String,
2290 #[serde(default)]
2291 prompt: String,
2292 #[serde(default)]
2294 prompt_ids: Vec<u32>,
2295 #[serde(default)]
2298 max_tokens: Option<usize>,
2299 #[serde(default)]
2312 temperature: Option<f32>,
2313 #[serde(default)]
2314 top_p: Option<f32>,
2315 #[serde(default)]
2317 top_k: Option<usize>,
2318 #[serde(default)]
2320 min_p: Option<f32>,
2321 #[serde(default)]
2323 frequency_penalty: Option<f32>,
2324 #[serde(default)]
2325 presence_penalty: Option<f32>,
2326 #[serde(default)]
2328 repetition_penalty: Option<f32>,
2329 #[serde(default)]
2335 seed: Option<u64>,
2336 #[serde(default)]
2337 stop: StopSequences,
2338 #[serde(default)]
2341 logit_bias: Option<serde_json::Value>,
2342 #[serde(default)]
2343 logprobs: Option<serde_json::Value>,
2344 #[serde(default)]
2345 n: Option<usize>,
2346 #[serde(default)]
2347 best_of: Option<usize>,
2348 #[serde(default)]
2350 chat: bool,
2351 #[serde(default)]
2353 stream: bool,
2354 #[serde(default)]
2356 max_ctx: Option<usize>,
2357 #[serde(default)]
2359 trace_id: Option<String>,
2360 #[serde(default)]
2364 cache_salt: Option<String>,
2365 #[serde(default)]
2369 session_id: Option<String>,
2370 #[serde(default)]
2371 user: Option<String>,
2372 #[serde(default)]
2376 timeout_ms: Option<serde_json::Value>,
2377}
2378
2379#[derive(Deserialize)]
2380struct ChatMessage {
2381 role: String,
2382 #[serde(default)]
2384 content: serde_json::Value,
2385 #[serde(default)]
2387 tool_calls: Vec<ReqToolCall>,
2388 #[serde(default)]
2391 tool_call_id: Option<String>,
2392 #[serde(default)]
2395 name: Option<String>,
2396 #[serde(default, alias = "reasoning_content")]
2406 reasoning: Option<String>,
2407}
2408
2409#[derive(Deserialize)]
2410struct ReqToolCall {
2411 #[serde(default)]
2412 #[allow(dead_code)]
2413 id: Option<String>,
2414 function: ReqToolFunction,
2415}
2416
2417#[derive(Deserialize)]
2418struct ReqToolFunction {
2419 name: String,
2420 #[serde(default)]
2422 arguments: serde_json::Value,
2423}
2424
2425#[derive(Clone, Default, Deserialize)]
2426#[serde(untagged)]
2427enum StopSequences {
2428 One(String),
2429 Many(Vec<String>),
2430 #[default]
2431 None,
2432}
2433
2434impl StopSequences {
2435 fn into_vec(self) -> Vec<String> {
2441 let stops = match self {
2442 Self::One(stop) => vec![stop],
2443 Self::Many(stops) => stops,
2444 Self::None => Vec::new(),
2445 };
2446 stops.into_iter().filter(|s| !s.is_empty()).collect()
2447 }
2448}
2449
2450#[derive(Deserialize)]
2456struct ChatCompletionReq {
2457 model: String,
2458 messages: Vec<ChatMessage>,
2459 #[serde(default, alias = "max_completion_tokens")]
2462 max_tokens: Option<usize>,
2463 #[serde(default)]
2466 temperature: Option<f32>,
2467 #[serde(default)]
2468 top_p: Option<f32>,
2469 #[serde(default)]
2473 top_k: Option<usize>,
2474 #[serde(default)]
2476 min_p: Option<f32>,
2477 #[serde(default)]
2479 frequency_penalty: Option<f32>,
2480 #[serde(default)]
2481 presence_penalty: Option<f32>,
2482 #[serde(default)]
2484 repetition_penalty: Option<f32>,
2485 #[serde(default)]
2487 seed: Option<u64>,
2488 #[serde(default)]
2489 stop: StopSequences,
2490 #[serde(default)]
2491 stream: bool,
2492 #[serde(default)]
2493 max_ctx: Option<usize>,
2494 #[serde(default)]
2499 response_format: Option<serde_json::Value>,
2500 #[serde(default)]
2501 logit_bias: Option<serde_json::Value>,
2502 #[serde(default)]
2503 logprobs: Option<serde_json::Value>,
2504 #[serde(default)]
2505 top_logprobs: Option<usize>,
2506 #[serde(default)]
2507 n: Option<usize>,
2508 #[serde(default)]
2510 tools: Vec<serde_json::Value>,
2511 #[serde(default)]
2513 tool_choice: Option<serde_json::Value>,
2514 #[serde(default)]
2520 reasoning_effort: Option<String>,
2521 #[serde(default)]
2530 reasoning: Option<serde_json::Value>,
2531 #[serde(default)]
2546 include_reasoning: Option<bool>,
2547 #[serde(default)]
2556 enable_thinking: Option<bool>,
2557 #[serde(default)]
2563 chat_template_kwargs: Option<serde_json::Value>,
2564 #[serde(default)]
2568 cache_salt: Option<String>,
2569 #[serde(default)]
2571 session_id: Option<String>,
2572 #[serde(default)]
2573 user: Option<String>,
2574 #[serde(default)]
2579 timeout_ms: Option<serde_json::Value>,
2580}
2581fn one() -> f32 {
2582 1.0
2583}
2584fn default_temperature() -> f32 {
2590 1.0
2591}
2592
2593#[derive(Debug, Clone, Copy, Default, PartialEq)]
2613struct SamplingDefaults {
2614 temperature: Option<f32>,
2615 top_p: Option<f32>,
2616 top_k: Option<usize>,
2617 min_p: Option<f32>,
2618 frequency_penalty: Option<f32>,
2619 presence_penalty: Option<f32>,
2620 repetition_penalty: Option<f32>,
2621}
2622
2623impl SamplingDefaults {
2624 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2627 SamplingDefaults {
2628 temperature: metadata
2629 .and_then(|m| m.default_temperature)
2630 .or_else(|| caps.and_then(|c| c.chat_temperature_default)),
2631 top_p: metadata
2632 .and_then(|m| m.default_top_p)
2633 .or_else(|| caps.and_then(|c| c.chat_top_p_default)),
2634 top_k: metadata.and_then(|m| m.default_top_k),
2635 min_p: metadata.and_then(|m| m.default_min_p),
2636 frequency_penalty: metadata.and_then(|m| m.default_frequency_penalty),
2637 presence_penalty: metadata.and_then(|m| m.default_presence_penalty),
2638 repetition_penalty: metadata.and_then(|m| m.default_repetition_penalty),
2639 }
2640 }
2641}
2642
2643#[derive(Debug, Clone, Copy, Default, PartialEq)]
2656struct ModelSamplingDefaults {
2657 thinking: SamplingDefaults,
2658 non_thinking: Option<SamplingDefaults>,
2659}
2660
2661impl ModelSamplingDefaults {
2662 fn resolve(metadata: Option<&OpenRouterModelMetadata>, caps: Option<&ModelCaps>) -> Self {
2663 ModelSamplingDefaults {
2664 thinking: SamplingDefaults::resolve(metadata, caps),
2665 non_thinking: metadata
2670 .and_then(|m| m.non_thinking_sampling.as_ref())
2671 .map(|arm| SamplingDefaults {
2672 temperature: arm.temperature,
2673 top_p: arm.top_p,
2674 top_k: arm.top_k,
2675 min_p: arm.min_p,
2676 frequency_penalty: arm.frequency_penalty,
2677 presence_penalty: arm.presence_penalty,
2678 repetition_penalty: arm.repetition_penalty,
2679 }),
2680 }
2681 }
2682
2683 fn for_mode(&self, think: ThinkMode) -> &SamplingDefaults {
2696 match (think, &self.non_thinking) {
2697 (ThinkMode::NoThink, Some(non_thinking)) => non_thinking,
2698 _ => &self.thinking,
2699 }
2700 }
2701
2702 fn single(thinking: SamplingDefaults) -> Self {
2705 ModelSamplingDefaults {
2706 thinking,
2707 non_thinking: None,
2708 }
2709 }
2710}
2711
2712#[derive(Debug, Clone, Copy, Default)]
2718struct ClientSampling {
2719 temperature: Option<f32>,
2720 top_p: Option<f32>,
2721 top_k: Option<usize>,
2722 min_p: Option<f32>,
2723 frequency_penalty: Option<f32>,
2724 presence_penalty: Option<f32>,
2725 repetition_penalty: Option<f32>,
2726 seed: Option<u64>,
2727}
2728
2729impl From<&CompletionReq> for ClientSampling {
2730 fn from(r: &CompletionReq) -> Self {
2731 ClientSampling {
2732 temperature: r.temperature,
2733 top_p: r.top_p,
2734 top_k: r.top_k,
2735 min_p: r.min_p,
2736 frequency_penalty: r.frequency_penalty,
2737 presence_penalty: r.presence_penalty,
2738 repetition_penalty: r.repetition_penalty,
2739 seed: r.seed,
2740 }
2741 }
2742}
2743
2744impl From<&ChatCompletionReq> for ClientSampling {
2745 fn from(r: &ChatCompletionReq) -> Self {
2746 ClientSampling {
2747 temperature: r.temperature,
2748 top_p: r.top_p,
2749 top_k: r.top_k,
2750 min_p: r.min_p,
2751 frequency_penalty: r.frequency_penalty,
2752 presence_penalty: r.presence_penalty,
2753 repetition_penalty: r.repetition_penalty,
2754 seed: r.seed,
2755 }
2756 }
2757}
2758
2759fn resolve_sampler_config(client: ClientSampling, defaults: &SamplingDefaults) -> SamplerConfig {
2765 sampler_config(
2766 client
2767 .temperature
2768 .or(defaults.temperature)
2769 .unwrap_or_else(default_temperature),
2770 client.top_k.or(defaults.top_k).unwrap_or(0),
2771 client.top_p.or(defaults.top_p).unwrap_or_else(one),
2772 client.min_p.or(defaults.min_p).unwrap_or(0.0),
2773 client
2774 .frequency_penalty
2775 .or(defaults.frequency_penalty)
2776 .unwrap_or(0.0),
2777 client
2778 .presence_penalty
2779 .or(defaults.presence_penalty)
2780 .unwrap_or(0.0),
2781 client
2782 .repetition_penalty
2783 .or(defaults.repetition_penalty)
2784 .unwrap_or_else(one),
2785 client.seed,
2786 )
2787}
2788
2789#[derive(Serialize)]
2790struct CompletionResp {
2791 model: String,
2792 text: String,
2793 tokens: Vec<u32>,
2794 stop_reason: String,
2798 #[serde(default, skip_serializing_if = "Option::is_none")]
2803 error: Option<serde_json::Value>,
2804 n_tokens: usize,
2805 prompt_tokens: usize,
2808 cached_tokens: usize,
2809 elapsed_s: f64,
2810}
2811
2812fn usage_json(
2820 n_prompt: usize,
2821 n_tokens: usize,
2822 n_cached: usize,
2823 elapsed_s: f64,
2824 spec: Option<worker::SpecUsage>,
2825) -> serde_json::Value {
2826 let mut u = json!({
2827 "prompt_tokens": n_prompt,
2828 "completion_tokens": n_tokens,
2829 "total_tokens": n_prompt + n_tokens,
2830 "prompt_tokens_details": { "cached_tokens": n_cached },
2831 "elapsed_s": elapsed_s,
2832 });
2833 if let Some(sp) = spec {
2834 u["spec"] = json!({
2835 "rounds": sp.rounds,
2836 "drafted": sp.drafted,
2837 "accepted": sp.accepted,
2838 "acceptance_rate": if sp.drafted > 0 {
2839 sp.accepted as f64 / sp.drafted as f64 } else { 0.0 },
2840 });
2841 }
2842 u
2843}
2844
2845const SYSTEM_FINGERPRINT: &str = concat!("memra-", env!("MEMRA_BUILD_SHA"));
2859
2860fn gen_hex128() -> String {
2863 use std::hash::{BuildHasher, Hasher};
2864 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2865 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2866 let t = std::time::SystemTime::now()
2867 .duration_since(std::time::UNIX_EPOCH)
2868 .map(|d| d.as_nanos() as u64)
2869 .unwrap_or(0);
2870 let mut h1 = std::collections::hash_map::RandomState::new().build_hasher();
2871 h1.write_u64(n);
2872 h1.write_u64(t);
2873 let mut h2 = std::collections::hash_map::RandomState::new().build_hasher();
2874 h2.write_u64(t.rotate_left(17));
2875 h2.write_u64(n);
2876 format!("{:016x}{:016x}", h1.finish(), h2.finish())
2877}
2878
2879#[derive(Clone)]
2882struct Envelope {
2883 id: String,
2884 created: u64,
2885}
2886
2887impl Envelope {
2888 fn new(chat: bool) -> Self {
2889 Envelope {
2890 id: format!(
2891 "{}-{}",
2892 if chat { "chatcmpl" } else { "cmpl" },
2893 gen_hex128()
2894 ),
2895 created: std::time::SystemTime::now()
2896 .duration_since(std::time::UNIX_EPOCH)
2897 .map(|d| d.as_secs())
2898 .unwrap_or(0),
2899 }
2900 }
2901
2902 fn stamp(&self, mut v: serde_json::Value) -> serde_json::Value {
2904 v["id"] = json!(self.id);
2905 v["created"] = json!(self.created);
2906 v["system_fingerprint"] = json!(SYSTEM_FINGERPRINT);
2907 v
2908 }
2909}
2910
2911fn with_request_id(id: &str, mut resp: Response) -> Response {
2913 if let Ok(v) = axum::http::HeaderValue::from_str(id) {
2914 resp.headers_mut()
2915 .insert(axum::http::HeaderName::from_static("x-request-id"), v);
2916 }
2917 resp
2918}
2919
2920fn openai_compat() -> bool {
2928 static C: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2929 *C.get_or_init(|| match std::env::var("MEMRA_COMPAT").as_deref() {
2930 Ok("openai") => true,
2931 Ok(_) => false,
2932 Err(_) => std::env::var("MEMRA_API_KEY").is_ok(),
2933 })
2934}
2935
2936fn cache_namespace(cache_salt: &Option<String>) -> String {
2948 cache_salt.clone().unwrap_or_default()
2949}
2950
2951const CACHE_SALT_MAX_BYTES: usize = 64;
2952
2953fn validate_cache_namespace(
2954 cache_salt: &Option<String>,
2955 keyring_configured: bool,
2956) -> Result<String, &'static str> {
2957 let raw = cache_namespace(cache_salt);
2958 if raw.len() > CACHE_SALT_MAX_BYTES {
2959 return Err("cache_salt must be at most 64 bytes");
2960 }
2961 if !keyring_configured && raw.starts_with("t:") {
2962 return Err("cache_salt must not use the reserved t: prefix without a keyring");
2963 }
2964 if !raw
2965 .bytes()
2966 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'+' | b'/' | b'='))
2967 {
2968 return Err("cache_salt contains unsupported characters");
2969 }
2970 Ok(raw)
2971}
2972
2973fn affinity_key(
2989 session_id: &Option<String>,
2990 user: &Option<String>,
2991 headers: &axum::http::HeaderMap,
2992) -> Option<String> {
2993 let clean = |s: &str| -> Option<String> {
2994 let t = s.trim();
2995 if t.is_empty() {
2996 None
2997 } else {
2998 Some(t.to_string())
2999 }
3000 };
3001 session_id
3002 .as_deref()
3003 .and_then(clean)
3004 .or_else(|| user.as_deref().and_then(clean))
3005 .or_else(|| {
3006 headers
3007 .get("x-session-id")
3008 .and_then(|v| v.to_str().ok())
3009 .and_then(clean)
3010 })
3011}
3012
3013fn error_body(
3018 message: &str,
3019 etype: &str,
3020 param: Option<&str>,
3021 code: Option<&str>,
3022) -> serde_json::Value {
3023 json!({ "error": {
3024 "message": message,
3025 "type": etype,
3026 "param": param,
3027 "code": code,
3028 } })
3029}
3030
3031fn error_response(status: StatusCode, message: &str, etype: &str, param: Option<&str>) -> Response {
3032 error_response_coded(status, message, etype, param, None)
3033}
3034
3035fn error_response_coded(
3040 status: StatusCode,
3041 message: &str,
3042 etype: &str,
3043 param: Option<&str>,
3044 code: Option<&str>,
3045) -> Response {
3046 let mut resp = (status, Json(error_body(message, etype, param, code))).into_response();
3047 if status.is_client_error()
3048 && status != StatusCode::TOO_MANY_REQUESTS
3049 && status != StatusCode::REQUEST_TIMEOUT
3050 && status != StatusCode::CONFLICT
3051 {
3052 resp.headers_mut().insert(
3053 "x-should-retry",
3054 axum::http::HeaderValue::from_static("false"),
3055 );
3056 }
3057 resp
3058}
3059
3060fn bad_request(message: &str, param: Option<&str>) -> Response {
3061 error_response(
3062 StatusCode::BAD_REQUEST,
3063 message,
3064 "invalid_request_error",
3065 param,
3066 )
3067}
3068
3069const 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>) {
3098 use worker::ErrClass as C;
3099 match class {
3100 C::InvalidRequest => (StatusCode::BAD_REQUEST, "invalid_request_error", None),
3101 C::ContextLength => (
3102 StatusCode::BAD_REQUEST,
3103 "invalid_request_error",
3104 Some("context_length_exceeded"),
3105 ),
3106 C::ModelNotFound => (
3107 StatusCode::BAD_REQUEST,
3108 "invalid_request_error",
3109 Some("model_not_found"),
3110 ),
3111 C::RateLimit => (
3112 StatusCode::TOO_MANY_REQUESTS,
3113 "rate_limit_error",
3114 Some("rate_limit_exceeded"),
3115 ),
3116 C::Overloaded => (
3117 StatusCode::SERVICE_UNAVAILABLE,
3118 "server_error",
3119 Some("overloaded"),
3120 ),
3121 C::Engine => (
3122 StatusCode::INTERNAL_SERVER_ERROR,
3123 "server_error",
3124 Some("engine_error"),
3125 ),
3126 }
3127}
3128
3129fn class_retry_after_s(class: worker::ErrClass) -> Option<u64> {
3131 use worker::ErrClass as C;
3132 match class {
3133 C::RateLimit => Some(RETRY_AFTER_S_RATE_LIMIT),
3134 C::Overloaded => Some(RETRY_AFTER_S_OVERLOADED),
3135 C::Engine | C::InvalidRequest | C::ContextLength | C::ModelNotFound => None,
3139 }
3140}
3141
3142fn engine_error_body(e: &worker::EngineError) -> serde_json::Value {
3145 let (_, etype, code) = class_http(e.class);
3146 error_body(&e.message, etype, e.param, code)
3147}
3148
3149fn engine_error_response(e: &worker::EngineError) -> Response {
3151 engine_error_response_with_retry_after(e, class_retry_after_s(e.class))
3152}
3153
3154fn engine_error_response_with_retry_after(
3155 e: &worker::EngineError,
3156 retry_after_s: Option<u64>,
3157) -> Response {
3158 let (status, _, _) = class_http(e.class);
3159 let resp = (status, Json(engine_error_body(e))).into_response();
3160 retry_contract_response(resp, retry_after_s)
3161}
3162
3163fn retry_contract_response(mut resp: Response, retry_after_s: Option<u64>) -> Response {
3165 let status = resp.status();
3166 let h = resp.headers_mut();
3167 match retry_after_s {
3168 Some(secs) => {
3169 let secs = secs.clamp(1, 60);
3171 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
3172 h.insert(axum::http::header::RETRY_AFTER, v);
3173 }
3174 if let Ok(v) = axum::http::HeaderValue::from_str(&(secs * 1000).to_string()) {
3175 h.insert("retry-after-ms", v);
3176 }
3177 }
3178 None if status.is_client_error() => {
3179 h.insert(
3182 "x-should-retry",
3183 axum::http::HeaderValue::from_static("false"),
3184 );
3185 }
3186 None => {}
3187 }
3188 resp
3189}
3190
3191fn worker_unavailable_response() -> Response {
3192 engine_error_response_with_retry_after(
3193 &worker::EngineError::overloaded("worker unavailable"),
3194 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
3195 )
3196}
3197
3198fn stop_reason_to_finish(r: &str) -> &'static str {
3199 match r {
3200 "Eos" | "Callback" => "stop",
3201 "MaxNew" | "ContextFull" => "length",
3202 _ => "stop",
3203 }
3204}
3205
3206fn content_to_text(v: &serde_json::Value) -> Result<String, String> {
3210 match v {
3211 serde_json::Value::Null => Ok(String::new()),
3212 serde_json::Value::String(s) => Ok(s.clone()),
3213 serde_json::Value::Array(parts) => {
3214 let mut out = String::new();
3215 for p in parts {
3216 match p.get("type").and_then(|t| t.as_str()) {
3217 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3218 Some(t) => out.push_str(t),
3219 None => return Err("content part has no text field".into()),
3220 },
3221 Some(other) => {
3222 return Err(format!(
3223 "unsupported content part type {other:?} (text only)"
3224 ));
3225 }
3226 }
3227 }
3228 Ok(out)
3229 }
3230 _ => Err("content must be a string, null, or an array of text parts".into()),
3231 }
3232}
3233
3234fn vision_enabled() -> bool {
3237 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3238 *ON.get_or_init(|| {
3239 std::env::var("MEMRA_VISION_DIR").is_ok()
3240 && std::env::var("MEMRA_VISION").as_deref() != Ok("0")
3241 })
3242}
3243
3244fn gemma_vision_enabled() -> bool {
3249 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3250 *ON.get_or_init(|| {
3251 std::env::var("MEMRA_GEMMA_VISION").as_deref() == Ok("1")
3252 && std::env::var("MEMRA_GEMMA_MMPROJ").is_ok()
3253 })
3254}
3255
3256fn step_vision_enabled() -> bool {
3262 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3263 *ON.get_or_init(|| {
3264 std::env::var("MEMRA_STEP_VISION_DIR").is_ok()
3265 && std::env::var("MEMRA_STEP_VISION").as_deref() != Ok("0")
3266 })
3267}
3268
3269const VISION_MAX_IMAGES: usize = 8;
3271
3272pub(crate) const MAX_VISION_PATCH_BYTES: usize = 1 << 30; static VISION_PATCH_BYTES_IN_USE: std::sync::atomic::AtomicUsize =
3278 std::sync::atomic::AtomicUsize::new(0);
3279pub(crate) static VISION_PREPROCESS_SEMAPHORE: tokio::sync::Semaphore =
3283 tokio::sync::Semaphore::const_new(1);
3284
3285pub(crate) struct VisionMemoryPermit {
3286 bytes: usize,
3287}
3288
3289#[derive(Debug)]
3290pub(crate) enum VisionMemoryError {
3291 Request(String),
3292 Capacity(String),
3293}
3294
3295impl Drop for VisionMemoryPermit {
3296 fn drop(&mut self) {
3297 if self.bytes != 0 {
3298 VISION_PATCH_BYTES_IN_USE.fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
3299 }
3300 }
3301}
3302
3303fn try_reserve_vision_memory(
3304 bytes: usize,
3305) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
3306 if bytes == 0 {
3307 return Ok(None);
3308 }
3309 if bytes > MAX_VISION_PATCH_BYTES {
3310 return Err(VisionMemoryError::Request(format!(
3311 "vision preprocessing requires {bytes} bytes of patch memory, exceeding the {} MiB request limit",
3312 MAX_VISION_PATCH_BYTES / (1024 * 1024)
3313 )));
3314 }
3315 let mut in_use = VISION_PATCH_BYTES_IN_USE.load(std::sync::atomic::Ordering::Acquire);
3316 loop {
3317 let Some(next) = in_use.checked_add(bytes) else {
3318 return Err(VisionMemoryError::Capacity(
3319 "vision patch memory reservation overflowed".into(),
3320 ));
3321 };
3322 if next > MAX_VISION_PATCH_BYTES {
3323 return Err(VisionMemoryError::Capacity(format!(
3324 "vision preprocessing is at capacity ({} MiB reserved; request needs {} MiB)",
3325 in_use / (1024 * 1024),
3326 bytes / (1024 * 1024)
3327 )));
3328 }
3329 match VISION_PATCH_BYTES_IN_USE.compare_exchange_weak(
3330 in_use,
3331 next,
3332 std::sync::atomic::Ordering::AcqRel,
3333 std::sync::atomic::Ordering::Acquire,
3334 ) {
3335 Ok(_) => return Ok(Some(VisionMemoryPermit { bytes })),
3336 Err(actual) => in_use = actual,
3337 }
3338 }
3339}
3340
3341pub(crate) fn vision_memory_error_response(
3342 error: VisionMemoryError,
3343 param: Option<&str>,
3344) -> Response {
3345 match error {
3346 VisionMemoryError::Request(message) => bad_request(&message, param),
3347 VisionMemoryError::Capacity(message) => retry_contract_response(
3348 error_response_coded(
3349 StatusCode::SERVICE_UNAVAILABLE,
3350 &message,
3351 "server_error",
3352 None,
3353 Some("vision_memory_busy"),
3354 ),
3355 Some(RETRY_AFTER_S_OVERLOADED),
3356 ),
3357 }
3358}
3359
3360enum PendingVisionUnit {
3366 Still {
3367 bytes: Vec<u8>,
3368 gh: usize,
3369 gw: usize,
3370 },
3371 Video {
3372 bytes: Vec<u8>,
3373 groups: Vec<memra_engine::vision_pre::PlannedVideoGroup>,
3374 video: usize,
3375 },
3376}
3377
3378struct PendingGemmaImage {
3380 bytes: Vec<u8>,
3381 gw: usize,
3382 gh: usize,
3383}
3384
3385struct PendingStepImage {
3388 bytes: Vec<u8>,
3389 plan: memra_engine::vision_step::StepImagePlan,
3390}
3391
3392fn content_to_text_vision_step(
3402 v: &serde_json::Value,
3403 step_images: &mut Vec<PendingStepImage>,
3404) -> Result<String, String> {
3405 use memra_engine::vision_step::{SV_MAIN_ROWS, SV_TILE_ROWS};
3406 let parts = match v {
3407 serde_json::Value::Array(parts) => parts,
3408 _ => return content_to_text(v),
3409 };
3410 let mut out = String::new();
3411 let mut needs_sep = false;
3412 for p in parts {
3413 match p.get("type").and_then(|t| t.as_str()) {
3414 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3415 Some(t) => {
3416 if needs_sep {
3417 out.push(' ');
3418 }
3419 out.push_str(t);
3420 needs_sep = true;
3421 }
3422 None => return Err("content part has no text field".into()),
3423 },
3424 Some("image_url") => {
3425 let url = p
3426 .get("image_url")
3427 .and_then(|u| {
3428 if u.is_string() {
3429 u.as_str()
3430 } else {
3431 u.get("url").and_then(|x| x.as_str())
3432 }
3433 })
3434 .ok_or("image_url part has no url")?;
3435 if !url.starts_with("data:") {
3436 return Err(
3437 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3438 );
3439 }
3440 if step_images.len() >= VISION_MAX_IMAGES {
3441 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3442 }
3443 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3447 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3448 let plan = memra_engine::vision_step::step_plan_image(&bytes)
3449 .map_err(|e| format!("image {}: {e}", step_images.len() + 1))?;
3450 for i in 0..plan.n_tiles {
3451 out.push_str("<patch_start>");
3452 for _ in 0..SV_TILE_ROWS {
3453 out.push_str("<im_patch>");
3454 }
3455 out.push_str("<patch_end>");
3456 if plan.newline_mask[i] {
3457 out.push_str("<patch_newline>");
3458 }
3459 }
3460 out.push_str("<im_start>");
3461 for _ in 0..SV_MAIN_ROWS {
3462 out.push_str("<im_patch>");
3463 }
3464 out.push_str("<im_end>");
3465 step_images.push(PendingStepImage { bytes, plan });
3466 needs_sep = false;
3467 }
3468 Some("video_url") => {
3469 return Err("step37 has no video input (image-only processor)".into());
3470 }
3471 Some(other) => {
3472 return Err(format!("unsupported content part type {other:?}"));
3473 }
3474 }
3475 }
3476 Ok(out)
3477}
3478
3479fn content_to_text_vision(
3488 v: &serde_json::Value,
3489 images: &mut Vec<PendingVisionUnit>,
3490 gemma_images: &mut Vec<PendingGemmaImage>,
3491 step_images: &mut Vec<PendingStepImage>,
3492 next_video: &mut usize,
3493) -> Result<String, String> {
3494 if step_vision_enabled() {
3498 return content_to_text_vision_step(v, step_images);
3499 }
3500 let parts = match v {
3501 serde_json::Value::Array(parts) => parts,
3502 _ => return content_to_text(v),
3503 };
3504 let mut out = String::new();
3505 for p in parts {
3506 match p.get("type").and_then(|t| t.as_str()) {
3507 Some("text") | None => match p.get("text").and_then(|t| t.as_str()) {
3508 Some(t) => out.push_str(t),
3509 None => return Err("content part has no text field".into()),
3510 },
3511 Some("image_url") if gemma_vision_enabled() => {
3512 let url = p
3513 .get("image_url")
3514 .and_then(|u| {
3515 if u.is_string() {
3516 u.as_str()
3517 } else {
3518 u.get("url").and_then(|x| x.as_str())
3519 }
3520 })
3521 .ok_or("image_url part has no url")?;
3522 if !url.starts_with("data:") {
3523 return Err(
3524 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3525 );
3526 }
3527 if gemma_images.len() >= VISION_MAX_IMAGES {
3528 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3529 }
3530 let bytes = memra_engine::vision_gemma::gemma_decode_data_uri(url)
3534 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3535 let (gw, gh) = memra_engine::vision_gemma::gemma_plan_image(&bytes)
3536 .map_err(|e| format!("image {}: {e}", gemma_images.len() + 1))?;
3537 out.push_str("<|image>");
3539 for _ in 0..memra_engine::vision_gemma::n_soft_for_grid(gw, gh) {
3540 out.push_str("<|image|>");
3541 }
3542 out.push_str("<image|>");
3543 gemma_images.push(PendingGemmaImage { bytes, gw, gh });
3544 }
3545 Some("image_url") => {
3546 if !vision_enabled() {
3547 return Err("image input is not enabled on this deployment".into());
3548 }
3549 let url = p
3550 .get("image_url")
3551 .and_then(|u| {
3552 if u.is_string() {
3553 u.as_str()
3554 } else {
3555 u.get("url").and_then(|x| x.as_str())
3556 }
3557 })
3558 .ok_or("image_url part has no url")?;
3559 if !url.starts_with("data:") {
3560 return Err(
3561 "image_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3562 );
3563 }
3564 if images
3565 .iter()
3566 .filter(|u| matches!(u, PendingVisionUnit::Still { .. }))
3567 .count()
3568 >= VISION_MAX_IMAGES
3569 {
3570 return Err(format!("too many images (max {VISION_MAX_IMAGES})"));
3571 }
3572 let bytes = memra_engine::vision_pre::decode_data_uri(url)
3577 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3578 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes)
3579 .map_err(|e| format!("image {}: {e}", images.len() + 1))?;
3580 out.push_str("<|vision_start|>");
3581 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(gh, gw) {
3582 out.push_str("<|image_pad|>");
3583 }
3584 out.push_str("<|vision_end|>");
3585 images.push(PendingVisionUnit::Still { bytes, gh, gw });
3586 }
3587 Some("video_url") if gemma_vision_enabled() => {
3588 return Err("gemma-4 has no video input (image-only projector)".into());
3589 }
3590 Some("video_url") => {
3591 if !vision_enabled() {
3592 return Err("video input is not enabled on this deployment".into());
3593 }
3594 let url = p
3595 .get("video_url")
3596 .and_then(|u| {
3597 if u.is_string() {
3598 u.as_str()
3599 } else {
3600 u.get("url").and_then(|x| x.as_str())
3601 }
3602 })
3603 .ok_or("video_url part has no url")?;
3604 if !url.starts_with("data:") {
3605 return Err(
3606 "video_url must be a base64 data URI (http(s) fetch is disabled)".into(),
3607 );
3608 }
3609 if *next_video >= 2 {
3610 return Err("too many videos (max 2)".into());
3611 }
3612 let bytes = memra_engine::vision_pre::decode_data_uri(url)?;
3615 let vid = memra_engine::vision_pre::plan_video_gif(&bytes)
3616 .map_err(|e| format!("video: {e}"))?;
3617 let vidx = *next_video;
3618 *next_video += 1;
3619 for group in &vid.groups {
3621 out.push_str(&format!("<{:.1} seconds>", group.timestamp));
3622 out.push_str("<|vision_start|>");
3623 for _ in 0..memra_engine::vision_pre::n_tokens_for_grid(group.gh, group.gw) {
3624 out.push_str("<|video_pad|>");
3625 }
3626 out.push_str("<|vision_end|>");
3627 }
3628 images.push(PendingVisionUnit::Video {
3631 bytes,
3632 groups: vid.groups,
3633 video: vidx,
3634 });
3635 }
3636 Some(other) => {
3637 return Err(format!("unsupported content part type {other:?}"));
3638 }
3639 }
3640 }
3641 Ok(out)
3642}
3643
3644fn pyjson(v: &serde_json::Value, out: &mut String) {
3648 match v {
3649 serde_json::Value::Object(m) => {
3650 out.push('{');
3651 for (i, (k, val)) in m.iter().enumerate() {
3652 if i > 0 {
3653 out.push_str(", ");
3654 }
3655 out.push_str(&serde_json::Value::String(k.clone()).to_string());
3656 out.push_str(": ");
3657 pyjson(val, out);
3658 }
3659 out.push('}');
3660 }
3661 serde_json::Value::Array(a) => {
3662 out.push('[');
3663 for (i, val) in a.iter().enumerate() {
3664 if i > 0 {
3665 out.push_str(", ");
3666 }
3667 pyjson(val, out);
3668 }
3669 out.push(']');
3670 }
3671 scalar => out.push_str(&scalar.to_string()),
3672 }
3673}
3674
3675fn pyjson_str(v: &serde_json::Value) -> String {
3676 let mut s = String::new();
3677 pyjson(v, &mut s);
3678 s
3679}
3680
3681fn sampler_config(
3688 temperature: f32,
3689 top_k: usize,
3690 top_p: f32,
3691 min_p: f32,
3692 frequency_penalty: f32,
3693 presence_penalty: f32,
3694 repetition_penalty: f32,
3695 seed: Option<u64>,
3696) -> SamplerConfig {
3697 let penalties_on =
3698 frequency_penalty != 0.0 || presence_penalty != 0.0 || repetition_penalty != 1.0;
3699 SamplerConfig {
3700 temperature,
3701 top_k,
3702 top_p,
3703 min_p,
3704 penalty_last_n: if penalties_on {
3705 memra_engine::spec::PEN_WINDOW_MAX
3706 } else {
3707 0
3708 },
3709 penalty_repeat: repetition_penalty,
3710 penalty_freq: frequency_penalty,
3711 penalty_present: presence_penalty,
3712 seed: seed.unwrap_or_else(fresh_seed),
3715 }
3716}
3717
3718fn fresh_seed() -> u64 {
3723 use std::sync::atomic::{AtomicU64, Ordering};
3724 static COUNTER: AtomicU64 = AtomicU64::new(0);
3725 let n = COUNTER.fetch_add(1, Ordering::Relaxed);
3726 let nanos = std::time::SystemTime::now()
3727 .duration_since(std::time::UNIX_EPOCH)
3728 .map(|d| d.as_nanos() as u64)
3729 .unwrap_or(0);
3730 let mut z = nanos
3731 .wrapping_add(n.wrapping_mul(0x9E3779B97F4A7C15))
3732 .wrapping_add(0x9E3779B97F4A7C15);
3733 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
3734 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
3735 z ^= z >> 31;
3736 if z == 0 { 0x9E3779B97F4A7C15 } else { z }
3739}
3740
3741fn reject_unsupported(fields: &[(&str, bool, &str)]) -> Result<(), (String, String)> {
3746 for (param, present, why) in fields {
3747 if *present {
3748 return Err((format!("{param} is not supported{why}"), param.to_string()));
3749 }
3750 }
3751 Ok(())
3752}
3753
3754#[derive(PartialEq)]
3755enum ToolChoice {
3756 Auto,
3757 None,
3758}
3759
3760fn parse_tool_choice(v: &Option<serde_json::Value>) -> Result<ToolChoice, String> {
3761 match v {
3762 None | Some(serde_json::Value::Null) => Ok(ToolChoice::Auto),
3763 Some(serde_json::Value::String(s)) => match s.as_str() {
3764 "auto" => Ok(ToolChoice::Auto),
3765 "none" => Ok(ToolChoice::None),
3766 "required" => Err("tool_choice \"required\" is not supported (no constrained \
3767 decoding); use \"auto\""
3768 .into()),
3769 other => Err(format!("bad tool_choice {other:?} (auto|none)")),
3770 },
3771 Some(serde_json::Value::Object(_)) => {
3772 Err("named-function tool_choice is not supported; use \"auto\"".into())
3773 }
3774 Some(other) => Err(format!("bad tool_choice: {other}")),
3775 }
3776}
3777
3778fn parse_think(
3828 reasoning_effort: &Option<String>,
3829 reasoning: &Option<serde_json::Value>,
3830 vllm_switch: Option<bool>,
3831 suppress_switch: Option<bool>,
3832 default_effort: Option<&str>,
3833 dsv4: bool,
3834) -> Result<(ThinkMode, Option<String>, bool), String> {
3835 let mut effort = reasoning_effort.clone();
3836 let ReasoningObject {
3837 mut enabled,
3838 effort: object_effort,
3839 exclude,
3840 } = parse_reasoning_object(reasoning)?;
3841 if let Some(e) = object_effort {
3842 effort = Some(e);
3843 }
3844 match (enabled, vllm_switch) {
3849 (Some(a), Some(b)) if a != b => {
3850 return Err(format!(
3851 "contradictory reasoning switches: reasoning.enabled={a} and \
3852 enable_thinking={b} — send one"
3853 ));
3854 }
3855 (None, Some(b)) => enabled = Some(b),
3856 _ => {}
3857 }
3858 let suppress = match (exclude, suppress_switch) {
3870 (Some(true), _) | (_, Some(false)) => Some(false),
3871 _ => None,
3872 };
3873 match (enabled, suppress) {
3874 (Some(true), Some(false)) => {
3875 return Err(
3876 "contradictory reasoning switches: reasoning is enabled but \
3877 include_reasoning:false / reasoning.exclude:true asks for no reasoning — \
3878 on this server not delivering reasoning means not generating it, so send one"
3879 .into(),
3880 );
3881 }
3882 (None, Some(b)) => enabled = Some(b),
3883 _ => {}
3884 }
3885 let client_explicit = effort.is_some() || enabled.is_some();
3889 if effort.is_none() && enabled.is_none() {
3894 effort = default_effort.map(str::to_string);
3895 }
3896 let effort_arm = match effort.as_deref() {
3901 None => None,
3902 Some(raw) => {
3903 let level = canonical_effort_for(raw, dsv4).ok_or_else(|| {
3904 format!(
3905 "bad reasoning_effort {raw:?} \
3906 (none|minimal|low|medium|high; xhigh/max/ultra clamp to the \
3907 highest level this model's template distinguishes)"
3908 )
3909 })?;
3910 Some(match level {
3911 "none" | "minimal" => (ThinkMode::NoThink, "low"),
3912 "low" => (ThinkMode::Think, "low"),
3913 "medium" => (ThinkMode::Think, "medium"),
3914 "max" => (ThinkMode::Think, "max"),
3915 _ => (ThinkMode::Think, "high"),
3916 })
3917 }
3918 };
3919 let (think, level) = match (enabled, effort_arm) {
3920 (Some(false), _) => (ThinkMode::NoThink, Some("low".to_string())),
3923 (Some(true), arm) => (ThinkMode::Think, arm.map(|(_, level)| level.to_string())),
3924 (None, Some((think, level))) => (think, Some(level.to_string())),
3925 (None, None) => (ThinkMode::Default, None),
3926 };
3927 Ok((think, level, client_explicit))
3928}
3929
3930struct ReasoningObject {
3932 enabled: Option<bool>,
3933 effort: Option<String>,
3934 exclude: Option<bool>,
3935}
3936
3937fn parse_reasoning_object(
3956 reasoning: &Option<serde_json::Value>,
3957) -> Result<ReasoningObject, String> {
3958 let mut out = ReasoningObject {
3959 enabled: None,
3960 effort: None,
3961 exclude: None,
3962 };
3963 let Some(v) = reasoning else { return Ok(out) };
3964 let obj = match v {
3965 serde_json::Value::Null => return Ok(out),
3966 serde_json::Value::Object(obj) => obj,
3967 _ => return Err("reasoning must be an object".into()),
3968 };
3969 for (key, value) in obj {
3970 match key.as_str() {
3978 "enabled" => {
3979 if !value.is_null() {
3980 out.enabled = Some(
3981 value
3982 .as_bool()
3983 .ok_or("reasoning.enabled must be true or false")?,
3984 );
3985 }
3986 }
3987 "exclude" => {
3988 if !value.is_null() {
3989 out.exclude = Some(
3990 value
3991 .as_bool()
3992 .ok_or("reasoning.exclude must be true or false")?,
3993 );
3994 }
3995 }
3996 "effort" => {
3997 if !value.is_null() {
3998 out.effort = Some(
3999 value
4000 .as_str()
4001 .ok_or("reasoning.effort must be a string")?
4002 .to_string(),
4003 );
4004 }
4005 }
4006 "max_tokens" => {
4007 return Err(
4008 "reasoning.max_tokens is not supported by this server: reasoning tokens \
4009 are output tokens here, and max_tokens is the ONE output budget covering \
4010 reasoning and content together — there is no separate reasoning budget to \
4011 spend against, so honouring this field is impossible rather than merely \
4012 unimplemented. Use max_tokens for the budget, and reasoning.effort (or \
4013 reasoning.enabled:false) to spend less of it on reasoning"
4014 .into(),
4015 );
4016 }
4017 other => {
4018 return Err(format!(
4019 "reasoning.{other} is not a field this server implements (it would change \
4020 nothing about the request); the supported keys are enabled, effort and \
4021 exclude"
4022 ));
4023 }
4024 }
4025 }
4026 Ok(out)
4027}
4028
4029fn parse_template_kwargs(kwargs: &Option<serde_json::Value>) -> Result<Option<bool>, String> {
4050 let Some(v) = kwargs else { return Ok(None) };
4051 let obj = match v {
4052 serde_json::Value::Null => return Ok(None),
4053 serde_json::Value::Object(obj) => obj,
4054 _ => return Err("chat_template_kwargs must be an object".into()),
4055 };
4056 let mut switch = None;
4057 for (key, value) in obj {
4058 match key.as_str() {
4059 "enable_thinking" => {
4060 switch = Some(
4061 value
4062 .as_bool()
4063 .ok_or("chat_template_kwargs.enable_thinking must be true or false")?,
4064 );
4065 }
4066 "preserve_thinking" => {
4067 let preserve = value
4068 .as_bool()
4069 .ok_or("chat_template_kwargs.preserve_thinking must be true or false")?;
4070 if !preserve {
4071 return Err(
4072 "chat_template_kwargs.preserve_thinking:false is not supported by this \
4073 server: the renderer implements the vendor DEFAULT (replay every prior \
4074 assistant turn's <think> block, empty when no reasoning was sent) but \
4075 not the strip arm — serving replay bytes under a strip request would \
4076 misdescribe the prompt. Omit the flag or send true"
4077 .into(),
4078 );
4079 }
4080 }
4082 other => {
4083 return Err(format!(
4084 "chat_template_kwargs.{other} is not supported by this server's \
4085 template renderer (it would change nothing about the prompt); the only \
4086 supported key is enable_thinking (preserve_thinking is RECOGNISED but \
4087 refuses in both directions — see its own message)"
4088 ));
4089 }
4090 }
4091 }
4092 Ok(switch)
4093}
4094
4095fn resolve_vllm_think_switch(
4099 enable_thinking: Option<bool>,
4100 kwargs: &Option<serde_json::Value>,
4101) -> Result<Option<bool>, String> {
4102 let from_kwargs = parse_template_kwargs(kwargs)?;
4103 match (enable_thinking, from_kwargs) {
4104 (Some(a), Some(b)) if a != b => Err(format!(
4105 "contradictory reasoning switches: enable_thinking={a} and \
4106 chat_template_kwargs.enable_thinking={b} — send one"
4107 )),
4108 (Some(a), _) => Ok(Some(a)),
4109 (None, b) => Ok(b),
4110 }
4111}
4112
4113pub(crate) fn canonical_effort_for(value: &str, dsv4_max: bool) -> Option<&'static str> {
4133 match value {
4134 "none" => Some("none"),
4135 "minimal" => Some("minimal"),
4136 "low" => Some("low"),
4137 "medium" => Some("medium"),
4138 "high" => Some("high"),
4139 "xhigh" | "max" | "ultra" => Some(if dsv4_max { "max" } else { "high" }),
4140 _ => None,
4141 }
4142}
4143
4144pub(crate) fn canonical_effort(value: &str) -> Option<&'static str> {
4147 canonical_effort_for(value, false)
4148}
4149
4150fn json_to_val(v: &serde_json::Value) -> chat::Val {
4153 match v {
4154 serde_json::Value::Null => chat::Val::Null,
4155 serde_json::Value::Bool(b) => chat::Val::Bool(*b),
4156 serde_json::Value::Number(n) => chat::Val::Num(n.to_string()),
4157 serde_json::Value::String(s) => chat::Val::Str(s.clone()),
4158 serde_json::Value::Array(a) => chat::Val::Arr(a.iter().map(json_to_val).collect()),
4159 serde_json::Value::Object(o) => chat::Val::Obj(
4162 o.iter()
4163 .map(|(k, val)| (k.clone(), json_to_val(val)))
4164 .collect(),
4165 ),
4166 }
4167}
4168
4169#[allow(clippy::type_complexity)]
4173fn prepare_tools(
4174 tools: &[serde_json::Value],
4175) -> Result<
4176 (
4177 Vec<String>,
4178 Vec<chat::Val>,
4179 HashMap<String, HashMap<String, String>>,
4180 ),
4181 String,
4182> {
4183 let mut tools_json = Vec::with_capacity(tools.len());
4184 let mut tools_struct = Vec::with_capacity(tools.len());
4185 let mut schemas: HashMap<String, HashMap<String, String>> = HashMap::new();
4186 for t in tools {
4187 let f = t
4188 .get("function")
4189 .ok_or("each tool needs a function object")?;
4190 let name = f
4191 .get("name")
4192 .and_then(|n| n.as_str())
4193 .ok_or("each tool needs function.name")?;
4194 let mut params: HashMap<String, String> = HashMap::new();
4195 if let Some(props) = f
4196 .get("parameters")
4197 .and_then(|p| p.get("properties"))
4198 .and_then(|p| p.as_object())
4199 {
4200 for (p, def) in props {
4201 if let Some(ty) = def.get("type").and_then(|t| t.as_str()) {
4202 params.insert(p.clone(), ty.to_string());
4203 }
4204 }
4205 }
4206 schemas.insert(name.to_string(), params);
4207 tools_json.push(pyjson_str(t));
4208 tools_struct.push(json_to_val(f));
4210 }
4211 Ok((tools_json, tools_struct, schemas))
4212}
4213
4214fn render_req_tool_call(tc: &ReqToolCall) -> Result<TmplToolCall, String> {
4219 let parsed: serde_json::Value = match &tc.function.arguments {
4220 serde_json::Value::Null => json!({}),
4221 serde_json::Value::String(s) if s.trim().is_empty() => json!({}),
4222 serde_json::Value::String(s) => serde_json::from_str(s)
4223 .map_err(|e| format!("tool_calls arguments is not valid JSON: {e}"))?,
4224 v @ serde_json::Value::Object(_) => v.clone(),
4225 _ => return Err("tool_calls arguments must be a JSON object".into()),
4226 };
4227 let obj = parsed
4228 .as_object()
4229 .ok_or("tool_calls arguments must decode to a JSON object")?;
4230 let params = obj
4231 .iter()
4232 .map(|(k, v)| {
4233 let rendered = match v {
4234 serde_json::Value::String(s) => s.clone(),
4235 v @ (serde_json::Value::Object(_) | serde_json::Value::Array(_)) => pyjson_str(v),
4236 scalar => scalar.to_string(),
4237 };
4238 (k.clone(), rendered)
4239 })
4240 .collect();
4241 let args = obj
4244 .iter()
4245 .map(|(k, v)| (k.clone(), json_to_val(v)))
4246 .collect();
4247 Ok(TmplToolCall {
4248 name: tc.function.name.clone(),
4249 params,
4250 args,
4251 id: tc.id.clone(),
4252 })
4253}
4254
4255fn tool_call_json(c: &ParsedToolCall) -> serde_json::Value {
4257 json!({ "id": c.id, "type": "function",
4258 "function": { "name": c.name, "arguments": c.arguments } })
4259}
4260
4261#[tokio::main]
4265pub async fn serve_main() -> Result<(), Box<dyn std::error::Error>> {
4266 serve_with(ServerWiring::stock()).await
4267}
4268
4269enum MeteringWiring {
4271 Stock,
4275 Custom(metering::MeteringFactory),
4280}
4281
4282pub struct ServerWiring {
4286 metering: MeteringWiring,
4287 on_ready: Option<Box<dyn FnOnce(RuntimeHandles) + Send>>,
4290 claimed_env: Vec<&'static str>,
4294}
4295
4296impl ServerWiring {
4297 pub fn stock() -> Self {
4299 ServerWiring {
4300 metering: MeteringWiring::Stock,
4301 on_ready: None,
4302 claimed_env: Vec::new(),
4303 }
4304 }
4305
4306 pub fn with_metering(factory: metering::MeteringFactory) -> Self {
4309 ServerWiring {
4310 metering: MeteringWiring::Custom(factory),
4311 on_ready: None,
4312 claimed_env: Vec::new(),
4313 }
4314 }
4315
4316 pub fn claiming(mut self, var: &'static str) -> Self {
4320 self.claimed_env.push(var);
4321 self
4322 }
4323
4324 pub fn on_ready(mut self, hook: impl FnOnce(RuntimeHandles) + Send + 'static) -> Self {
4325 self.on_ready = Some(Box::new(hook));
4326 self
4327 }
4328}
4329
4330pub struct RuntimeHandles {
4333 pub trim: TrimHandle,
4334 pub shutdown: tokio::sync::watch::Receiver<bool>,
4339}
4340
4341#[derive(Clone)]
4344pub struct TrimHandle {
4345 cmd_tx: Sender<Cmd>,
4346}
4347
4348impl TrimHandle {
4349 pub async fn trim(&self) -> Result<serde_json::Value, String> {
4351 let (tx, rx) = tokio::sync::oneshot::channel();
4352 if self.cmd_tx.send(Cmd::TrimPools(tx)).is_err() {
4353 return Err("worker is down".into());
4354 }
4355 match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
4356 Ok(Ok(report)) => Ok(json!(report)),
4357 _ => Err("worker did not answer the trim within 30s".into()),
4358 }
4359 }
4360}
4361
4362pub async fn serve_with(wiring: ServerWiring) -> Result<(), Box<dyn std::error::Error>> {
4363 let args: Vec<String> = std::env::args().skip(1).collect();
4366 if let Some(code) = auth::run_cli(&args) {
4367 std::process::exit(code);
4368 }
4369 auth::init_from_env();
4372 let api_auth = match ApiAuth::from_env() {
4373 Ok(auth) => auth,
4374 Err(err) => {
4375 eprintln!("[server] FATAL: {err}");
4376 std::process::exit(1);
4377 }
4378 };
4379 let addr = std::env::var("MEMRA_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into());
4380 let allow_open_bind = std::env::var("MEMRA_ALLOW_OPEN_BIND").as_deref() == Ok("1");
4381 let (bind_addr, bind_loopback) = match resolve_bind_addr(&addr) {
4382 Ok(resolved) => resolved,
4383 Err(err) => {
4384 eprintln!("[server] FATAL: {err}");
4385 std::process::exit(1);
4386 }
4387 };
4388 if !bind_loopback && !api_auth.configured() && !allow_open_bind {
4389 let message = format!(
4390 "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"
4391 );
4392 eprintln!("[server] FATAL: {message}");
4393 std::process::exit(1);
4394 }
4395 if !bind_loopback && !api_auth.configured() {
4396 eprintln!(
4397 "[server] WARNING: MEMRA_ALLOW_OPEN_BIND=1 permits open completion routes on {addr}; \
4398 metrics remain bearer-protected"
4399 );
4400 }
4401 let metrics_token = match std::env::var("MEMRA_METRICS_TOKEN") {
4402 Ok(token) if token.is_empty() => {
4403 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must not be empty");
4404 std::process::exit(1);
4405 }
4406 Ok(token) => Some(token),
4407 Err(std::env::VarError::NotPresent) => None,
4408 Err(std::env::VarError::NotUnicode(_)) => {
4409 eprintln!("[server] FATAL: MEMRA_METRICS_TOKEN must be valid UTF-8");
4410 std::process::exit(1);
4411 }
4412 };
4413 let metrics_auth = MetricsAuth::new(bind_loopback, api_auth.configured(), metrics_token);
4414
4415 let models = parse_models_config();
4416 let (openrouter_metadata, provider_metadata) = match load_openrouter_metadata(&models) {
4417 Ok(loaded) => loaded,
4418 Err(err) => {
4419 eprintln!("[server] FATAL: {err}");
4420 std::process::exit(1);
4421 }
4422 };
4423 let metering_obj: Option<Arc<dyn metering::Metering>> = {
4429 let factory = match wiring.metering {
4430 MeteringWiring::Stock => None,
4431 MeteringWiring::Custom(factory) => Some(factory),
4432 };
4433 for deployment_only in [
4434 "MEMRA_REQUEST_LEDGER",
4435 "MEMRA_TENANT_BUDGETS",
4436 "MEMRA_ADMIN_ADDR",
4437 "MEMRA_ADMIN_TOKEN_FILE",
4438 "MEMRA_CAPTURE_DIR",
4439 ] {
4440 if std::env::var_os(deployment_only).is_some()
4441 && !wiring.claimed_env.contains(&deployment_only)
4442 {
4443 eprintln!(
4444 "[server] FATAL: {deployment_only} is a deployment-binary surface; this \
4445 build ships no accounting/admin/capture. Wire a Metering implementation \
4446 through ServerWiring and claim the vars it consumes."
4447 );
4448 std::process::exit(1);
4449 }
4450 }
4451 match factory {
4452 None => None,
4453 Some(factory) => {
4454 let model_ids: Vec<String> =
4455 models.iter().map(|(name, _, _)| name.clone()).collect();
4456 match factory(&metering::MeteringInit { models: &model_ids }) {
4457 Ok(metering_obj) => metering_obj,
4458 Err(err) => {
4459 eprintln!("[server] FATAL: metering wiring: {err}");
4460 std::process::exit(1);
4461 }
4462 }
4463 }
4464 }
4465 };
4466 let budget_tokenizers = if metering_obj
4467 .as_ref()
4468 .is_some_and(|manager| manager.enforces_limits())
4469 {
4470 match load_budget_tokenizers(&models) {
4471 Ok(tokenizers) => Some(tokenizers),
4472 Err(err) => {
4473 eprintln!("[server] FATAL: prepaid reservation tokenizers: {err}");
4474 std::process::exit(1);
4475 }
4476 }
4477 } else {
4478 None
4479 };
4480 eprintln!("[server] starting; models config = {models:?}");
4481
4482 let health_state = health::WorkerHealth::new();
4487 health::spawn_gpu_watch(health_state.clone());
4491 health::spawn_sd_watchdog(health_state.clone());
4492
4493 let (cmd_tx, model_names, caps, metrics, worker_thread) =
4495 match worker::spawn(models, health_state.clone()) {
4496 Ok(v) => v,
4497 Err(err) => {
4498 eprintln!("[server] FATAL: worker init failed: {err}");
4499 health_state.mark_dead(format!("worker init failed: {err}"));
4500 health::sd_notify(&format!("STATUS=worker init failed: {err}"));
4501 std::process::exit(1);
4502 }
4503 };
4504 eprintln!("[server] worker ready; serving models: {model_names:?}");
4505
4506 let (drain_shutdown_tx, drain_shutdown_rx) = tokio::sync::watch::channel(false);
4513 if let Some(on_ready) = wiring.on_ready {
4514 on_ready(RuntimeHandles {
4515 trim: TrimHandle {
4516 cmd_tx: cmd_tx.clone(),
4517 },
4518 shutdown: drain_shutdown_rx.clone(),
4519 });
4520 }
4521
4522 let bg_handle = darklane::spawn_from_env(health_state.clone());
4525 let bg_state = bg_handle.as_ref().map(|h| {
4526 let mode = darklane::BgConfig::from_env()
4527 .map(|c| c.yield_mode.as_str())
4528 .unwrap_or("stop");
4529 (h.state.clone(), mode)
4530 });
4531
4532 let state = AppState {
4533 cmd_tx,
4534 models: model_names,
4535 caps,
4536 openrouter_metadata: Arc::new(openrouter_metadata),
4537 provider_metadata: Arc::new(provider_metadata),
4538 metering: metering_obj,
4539 budget_tokenizers,
4540 api_auth,
4541 metrics_auth,
4542 metrics,
4543 started: std::time::SystemTime::now()
4544 .duration_since(std::time::UNIX_EPOCH)
4545 .map(|d| d.as_secs())
4546 .unwrap_or(0),
4547 inflight: Arc::new(Default::default()),
4548 tenant_inflight: Arc::new(Default::default()),
4549 health: health_state.clone(),
4550 bg: bg_state,
4551 };
4552 let inflight_handle = state.inflight.clone();
4553 let drain_metering = state.metering.clone();
4556 worker::register_http_inflight(state.inflight.clone());
4561 let app = Router::new()
4562 .route("/health", get(health_live))
4567 .route("/livez", get(health_live))
4568 .route("/readyz", get(health_ready))
4569 .route("/models", get(list_models))
4570 .route("/v1/models", get(list_models_v1))
4571 .route("/v1/auth/check", get(auth_check))
4572 .route("/v1/completions", post(completions))
4573 .route("/v1/embeddings", post(embed_api::embeddings))
4574 .route("/v1/rerank", post(embed_api::rerank))
4575 .route("/v1/chat/completions", post(chat_completions))
4576 .route("/v1/messages", post(anthropic::messages))
4580 .route("/v1/responses", post(responses_api::responses))
4581 .route("/metrics", get(get_metrics))
4582 .route("/yield/metrics", get(yield_metrics))
4583 .with_state(state.clone());
4584 let app = apply_body_limit(app);
4587 let app = app.layer(middleware::from_fn_with_state(
4590 state,
4591 authenticate_inference_before_body,
4592 ));
4593 let app = if ttft::enabled() {
4594 app.layer(middleware::from_fn(ttft_request_start))
4595 } else {
4596 app
4597 };
4598
4599 let listener = tokio::net::TcpListener::bind(bind_addr).await?;
4600 eprintln!("[server] listening on http://{bind_addr}");
4601 drop(drain_shutdown_rx);
4602 health::sd_notify("READY=1\nSTATUS=serving");
4606 let inflight = inflight_handle;
4613 let signal_admin_shutdown = drain_shutdown_tx.clone();
4614 let serve_result = axum::serve(listener, app)
4615 .with_graceful_shutdown(async move {
4616 let mut sigterm =
4617 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
4618 Ok(s) => s,
4619 Err(err) => {
4620 eprintln!("[server] WARN: no SIGTERM handler ({err}); drain disabled");
4621 std::future::pending::<()>().await;
4622 unreachable!()
4623 }
4624 };
4625 sigterm.recv().await;
4626 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
4627 let _ = signal_admin_shutdown.send(true);
4628 health::sd_notify(&format!(
4632 "STOPPING=1\nSTATUS=draining\nEXTEND_TIMEOUT_USEC={}",
4633 (drain_deadline_s() + 5) * 1_000_000
4634 ));
4635 let n: usize = inflight
4636 .iter()
4637 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4638 .sum();
4639 eprintln!(
4640 "[server] SIGTERM: draining ({n} in flight, deadline {}s)",
4641 drain_deadline_s()
4642 );
4643 let deadline = std::time::Duration::from_secs(drain_deadline_s());
4644 let t0 = std::time::Instant::now();
4645 loop {
4646 let n: usize = inflight
4647 .iter()
4648 .map(|c| c.load(std::sync::atomic::Ordering::SeqCst))
4649 .sum();
4650 if n == 0 {
4651 eprintln!(
4652 "[server] drain complete in {:.1}s; exiting",
4653 t0.elapsed().as_secs_f64()
4654 );
4655 break;
4656 }
4657 if t0.elapsed() >= deadline {
4658 eprintln!(
4659 "[server] drain deadline ({}s) hit with {n} in flight; exiting",
4660 drain_deadline_s()
4661 );
4662 if let Some(metering) = drain_metering.as_ref() {
4669 metering.drain_kill();
4670 }
4671 break;
4672 }
4673 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
4674 }
4675 })
4676 .await;
4677 let _ = drain_shutdown_tx.send(true);
4680 serve_result?;
4681 if let Some(h) = bg_handle {
4685 h.shutdown();
4686 }
4687 worker_thread.join().map_err(|_| {
4696 std::io::Error::other("GPU worker thread panicked during graceful shutdown")
4697 })?;
4698 eprintln!("[server] GPU worker shutdown complete");
4699 Ok(())
4700}
4701
4702fn validate_model_path(path: &str) -> Result<(), String> {
4708 let p = std::path::Path::new(path);
4709 if !p.exists() {
4710 return Err(format!("model path {path:?} does not exist"));
4711 }
4712 if p.is_file() {
4713 return Ok(()); }
4715 if p.join("manifest.json").exists() {
4716 return Ok(()); }
4718 let has_st =
4719 p.join("model.safetensors").exists() || p.join("model.safetensors.index.json").exists();
4720 if !has_st {
4721 return Err(format!(
4722 "model dir {path:?} is not a servable checkpoint: want model.safetensors or \
4723 model.safetensors.index.json + config.json (HF safetensors dir), or \
4724 manifest.json (memra repack dir)"
4725 ));
4726 }
4727 if !p.join("config.json").exists() {
4728 return Err(format!(
4729 "model dir {path:?} has safetensors weights but no config.json"
4730 ));
4731 }
4732 Ok(())
4733}
4734
4735fn parse_models_config() -> Vec<(String, String, Option<String>)> {
4743 if let Ok(spec) = std::env::var("MEMRA_MODELS") {
4744 let mut out = Vec::new();
4745 for entry in spec.split(',').filter(|s| !s.trim().is_empty()) {
4746 if let Some((name, path)) = entry.split_once('=') {
4747 let (mpath, dpath) = match path.trim().split_once('+') {
4750 Some((m, d)) => (m.trim(), Some(d.trim())),
4751 None => (path.trim(), None),
4752 };
4753 let resolve = |p: &str| {
4754 memra_gguf::hf::resolve_arg(p).unwrap_or_else(|err| {
4755 eprintln!("[server] FATAL: model {name:?}: {err}");
4756 std::process::exit(1);
4757 })
4758 };
4759 let mpath = resolve(mpath);
4760 if let Err(err) = validate_model_path(&mpath) {
4761 eprintln!("[server] FATAL: model {name:?}: {err}");
4762 std::process::exit(1);
4763 }
4764 let dpath = dpath.map(|d| {
4774 let d = resolve(d);
4775 let p = std::path::Path::new(&d);
4776 if !p.exists() {
4777 eprintln!(
4778 "[server] FATAL: model {name:?}: drafter path {d:?} does not \
4779 exist (MEMRA_MODELS '+draft' attach). Refusing to start \
4780 rather than serving plain decode under a config that asked \
4781 for speculative decoding."
4782 );
4783 std::process::exit(1);
4784 }
4785 if !p.is_file() {
4786 eprintln!(
4787 "[server] FATAL: model {name:?}: drafter path {d:?} is not a \
4788 file — a '+draft' attach must be a NextN/MTP GGUF file."
4789 );
4790 std::process::exit(1);
4791 }
4792 d
4793 });
4794 out.push((name.trim().to_string(), mpath, dpath));
4795 } else {
4796 eprintln!(
4797 "[server] WARN: bad MEMRA_MODELS entry {entry:?} (want name=/path[+/draft]); skipping"
4798 );
4799 }
4800 }
4801 if !out.is_empty() {
4802 return out;
4803 }
4804 }
4805 vec![
4807 (
4808 "main".into(),
4809 "/data/ai-ml/hf-models/qwen36-27b-nvfp4-mtp/Qwen3.6-27B-NVFP4-Q4_K_M-mtp.gguf".into(),
4810 None,
4811 ),
4812 (
4813 "judge".into(),
4814 "/data/ai-ml/hf-models/qwen35-9b-nvfp4-gguf/Qwen3.5-9B-NVFP4-MTP-GGUF.gguf".into(),
4815 None,
4816 ),
4817 ]
4818}
4819
4820fn load_budget_tokenizers(
4821 models: &[(String, String, Option<String>)],
4822) -> Result<Arc<HashMap<String, Arc<Tokenizer>>>, String> {
4823 let mut tokenizers = HashMap::new();
4824 for (alias, path, _) in models {
4825 let path = std::path::Path::new(path);
4826 let tokenizer = if path.is_dir() {
4827 let tokenizer_dir = if path.join("manifest.json").exists() {
4828 let repack = memra_gguf::source::Hy3RepackSource::open(path).map_err(|err| {
4829 format!("model {alias:?}: open repack tokenizer source: {err}")
4830 })?;
4831 repack
4832 .source_dir()
4833 .filter(|source| source.join("tokenizer.json").exists())
4834 .unwrap_or(path)
4835 .to_path_buf()
4836 } else {
4837 path.to_path_buf()
4838 };
4839 Tokenizer::from_hf_dir(&tokenizer_dir)
4840 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4841 } else {
4842 let gguf = memra_gguf::GgufFile::open(path)
4843 .map_err(|err| format!("model {alias:?}: open reservation tokenizer: {err}"))?;
4844 Tokenizer::from_gguf(&gguf)
4845 .map_err(|err| format!("model {alias:?}: reservation tokenizer: {err}"))?
4846 };
4847 tokenizers.insert(alias.clone(), Arc::new(tokenizer));
4848 }
4849 Ok(Arc::new(tokenizers))
4850}
4851
4852fn health_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4854 let s = st.health.snapshot();
4855 let mut v = json!({
4856 "status": status,
4857 "models": *st.models,
4858 "worker": {
4859 "phase": health::phase_name(s.phase),
4860 "beat_age_ms": s.beat_age_ms,
4861 "tick_max_ms": s.tick_max_ms,
4862 "stall_threshold_ms": s.stall_threshold_ms,
4863 "generation": s.generation,
4864 "xid_warnings": s.xid_warns,
4865 },
4866 });
4867 if let Some(d) = detail {
4868 v["detail"] = json!(d);
4869 }
4870 v
4871}
4872
4873fn readiness_payload(st: &AppState, status: &str, detail: Option<&str>) -> serde_json::Value {
4877 let mut v = health_payload(st, status, detail);
4878 v["peer_probe_integrity"] = json!(st.health.peer_probe_integrity().detail());
4879 v
4880}
4881
4882async fn auth_check() -> impl IntoResponse {
4886 StatusCode::NO_CONTENT
4887}
4888
4889async fn health_live(State(st): State<AppState>) -> impl IntoResponse {
4905 if draining() {
4906 return (StatusCode::OK, Json(health_payload(&st, "draining", None))).into_response();
4909 }
4910 match st.health.live() {
4911 Ok(()) => (StatusCode::OK, Json(health_payload(&st, "ok", None))).into_response(),
4912 Err(why) => retry_contract_response(
4913 (
4914 StatusCode::SERVICE_UNAVAILABLE,
4915 Json(health_payload(&st, "unhealthy", Some(&why))),
4916 )
4917 .into_response(),
4918 Some(worker::WORKER_RESPAWN_BACKOFF_BASE_S),
4919 ),
4920 }
4921}
4922
4923async fn health_ready(State(st): State<AppState>) -> impl IntoResponse {
4934 let is_draining = draining();
4935 match st.health.ready(is_draining) {
4936 Ok(()) => (StatusCode::OK, Json(readiness_payload(&st, "ready", None))).into_response(),
4937 Err(why) => retry_contract_response(
4938 (
4939 StatusCode::SERVICE_UNAVAILABLE,
4940 Json(readiness_payload(&st, "not_ready", Some(&why))),
4941 )
4942 .into_response(),
4943 Some(if is_draining {
4944 drain_deadline_s()
4945 } else {
4946 worker::WORKER_RESPAWN_BACKOFF_BASE_S
4947 }),
4948 ),
4949 }
4950}
4951
4952#[derive(Clone, Copy)]
4953struct DualPpMetricsSnapshot {
4954 stage_ns: [u64; 4],
4955 stage_samples: [usize; 4],
4956 dropped_timing_samples: usize,
4957 overlaps: usize,
4958 slot_pairs: usize,
4959 slot_uses: [usize; 2],
4960 slot_collisions: usize,
4961}
4962
4963impl DualPpMetricsSnapshot {
4964 fn current() -> Self {
4965 let (stage_ns, stage_samples) = memra_engine::pp::dual_pp_timing_snapshot();
4966 let (slot_pairs, slot_uses, slot_collisions) = memra_engine::pp::dual_pp_slot_snapshot();
4967 Self {
4968 stage_ns,
4969 stage_samples,
4970 dropped_timing_samples: memra_engine::pp::dual_pp_timing_dropped(),
4971 overlaps: memra_engine::pp::dual_pp_overlaps(),
4972 slot_pairs,
4973 slot_uses,
4974 slot_collisions,
4975 }
4976 }
4977
4978 fn populated(self) -> bool {
4979 self.stage_samples.iter().any(|&n| n > 0)
4980 || self.dropped_timing_samples > 0
4981 || self.slot_pairs > 0
4982 || self.slot_collisions > 0
4983 }
4984}
4985
4986fn insert_dual_pp_metrics(
4987 body: &mut serde_json::Value,
4988 metrics_scope: &MetricsScope,
4989 snapshot: impl FnOnce() -> DualPpMetricsSnapshot,
4990) {
4991 if !metrics_scope.operator() {
4994 return;
4995 }
4996 let snapshot = snapshot();
4997 if !snapshot.populated() {
4998 return;
4999 }
5000 let timings: serde_json::Map<String, serde_json::Value> = memra_engine::pp::DUAL_PP_STAGE_NAMES
5001 .iter()
5002 .enumerate()
5003 .map(|(i, name)| {
5004 let total_ms = snapshot.stage_ns[i] as f64 / 1_000_000.0;
5005 (
5006 name.to_string(),
5007 json!({
5008 "samples": snapshot.stage_samples[i],
5009 "total_ms": total_ms,
5010 "mean_ms": if snapshot.stage_samples[i] > 0 {
5011 total_ms / snapshot.stage_samples[i] as f64
5012 } else { 0.0 },
5013 }),
5014 )
5015 })
5016 .collect();
5017 body["dual_pp"] = json!({
5018 "overlaps": snapshot.overlaps,
5019 "slot_pairs": snapshot.slot_pairs,
5020 "slot_uses": snapshot.slot_uses,
5021 "slot_collisions": snapshot.slot_collisions,
5022 "cuda_event_spans": timings,
5023 "dropped_timing_samples": snapshot.dropped_timing_samples,
5024 });
5025}
5026
5027fn insert_spec_acceptance_metrics(
5028 body: &mut serde_json::Value,
5029 metrics_scope: &MetricsScope,
5030 snapshot: impl FnOnce() -> HashMap<String, memra_engine::spec::SpecTelemetry>,
5031) {
5032 if !metrics_scope.operator() {
5035 return;
5036 }
5037 let snapshot = snapshot();
5038 if snapshot.is_empty() {
5039 return;
5040 }
5041
5042 let mut tau = serde_json::Map::new();
5043 let mut by_position = serde_json::Map::new();
5044 for (model, telemetry) in snapshot {
5045 if telemetry.rounds == 0 {
5046 continue;
5047 }
5048 let n_pos = telemetry
5049 .pos_drafted
5050 .iter()
5051 .rposition(|&n| n > 0)
5052 .map_or(0, |position| position + 1);
5053 tau.insert(model.clone(), json!(telemetry.tau()));
5054 by_position.insert(
5055 model,
5056 json!({
5057 "window_seconds": worker::SPEC_METRICS_WINDOW_S,
5058 "rounds": telemetry.rounds,
5059 "offered": telemetry.pos_drafted[..n_pos].to_vec(),
5060 "accepted": telemetry.pos_accepted[..n_pos].to_vec(),
5061 "accept_rate": (0..n_pos).map(|position| {
5062 let offered = telemetry.pos_drafted[position];
5063 if offered > 0 {
5064 telemetry.pos_accepted[position] as f64 / offered as f64
5065 } else {
5066 0.0
5067 }
5068 }).collect::<Vec<f64>>(),
5069 }),
5070 );
5071 }
5072 if !tau.is_empty() {
5073 body["spec_tau"] = serde_json::Value::Object(tau);
5074 body["spec_accept_by_position"] = serde_json::Value::Object(by_position);
5075 }
5076}
5077
5078fn insert_peer_probe_metrics(
5079 body: &mut serde_json::Value,
5080 metrics_scope: &MetricsScope,
5081 snapshot: impl FnOnce() -> memra_engine::pp::PeerProbeMetrics,
5082) {
5083 if !metrics_scope.operator() {
5086 return;
5087 }
5088 let snapshot = snapshot();
5089 body["peer_probe_bypassed"] = json!(snapshot.bypassed);
5090 body["peer_probe_boundary_copies"] = json!(snapshot.boundary_copies);
5091 body["peer_probe_runtime_reprobes"] = json!(snapshot.runtime_probes);
5092 body["peer_probe_runtime_failures"] = json!(snapshot.runtime_failures);
5093 body["peer_probe_deferred_total"] = json!(snapshot.deferred_total);
5094 body["peer_probe_integrity_degraded"] = json!(snapshot.integrity_degraded);
5095 body["peer_probe_degraded_to_host_bounce"] = json!(snapshot.degraded_to_host_bounce);
5096}
5097
5098async fn get_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5100 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5101 Ok(scope) => scope,
5102 Err(response) => return response,
5103 };
5104 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5105 let mut body = if metrics_scope.process_wide() {
5109 json!({
5110 "admitted": m.admitted,
5111 "completed": m.completed,
5112 "tokens_out": m.tokens_out,
5113 "step_p50_ms": m.step_p50_ms,
5114 "step_p99_ms": m.step_p99_ms,
5115 "prompt_tokens_in": m.prompt_tokens_in,
5117 "cached_tokens_in": m.cached_tokens_in,
5118 "computed_tokens_in": m.prompt_tokens_in.saturating_sub(m.cached_tokens_in),
5121 "admission_session_defers": m.admission_session_defers,
5124 "admission_vram_defers": m.admission_vram_defers,
5125 "step_oom_parks": m.step_oom_parks,
5126 "continuation_pool_hits": m.continuation_pool_hits,
5127 "continuation_pool_evictions": m.continuation_pool_evictions,
5128 "plain_affinity_rewinds": m.plain_affinity_rewinds,
5129 "served_dspark": m.served_dspark,
5130 "served_spec": m.served_spec,
5131 "served_plain": m.served_plain,
5132 "spec_pool_hits": m.spec_pool_hits,
5133 "spec_pool_misses": m.spec_pool_misses,
5134 "spec_pool_affinity_rewinds": m.spec_pool_affinity_rewinds,
5135 "spec_pool_evictions": m.spec_pool_evictions,
5136 "spec_pool_sampler_refusals": m.spec_pool_sampler_refusals,
5139 })
5140 } else {
5141 json!({})
5142 };
5143 if metrics_scope.operator() {
5147 if let Some(budget_health) = st.metering.as_ref().and_then(|m| m.limits_health()) {
5148 body["budget_source_reload_failed"] = json!(budget_health.source_reload_failed);
5149 body["budget_source_reload_consecutive"] =
5150 json!(budget_health.source_reload_consecutive);
5151 body["budget_source_available"] = json!(budget_health.source_available);
5152 }
5153 body["cache_hit_token_ratio"] = json!(if m.prompt_tokens_in > 0 {
5155 m.cached_tokens_in as f64 / m.prompt_tokens_in as f64
5156 } else {
5157 0.0
5158 });
5159 body["prefix_cache_hits"] = json!(m.prefix_hits);
5160 body["prefix_cache_misses"] = json!(m.prefix_misses);
5161 body["prefix_cache_inserts"] = json!(m.prefix_inserts);
5162 body["prefix_cache_evictions"] = json!(m.prefix_evictions);
5163 body["prefix_cache_skips_budget"] = json!(m.prefix_skips_budget);
5164 body["prefix_cache_skips_pinned"] = json!(m.prefix_skips_pinned);
5165 body["prefix_cache_hit_tokens"] = json!(m.prefix_hit_tokens);
5166 body["lcp_histogram"] = json!({
5169 "edges": worker::LCP_HIST_EDGES.to_vec(),
5170 "counts": m.lcp_hist.to_vec(),
5171 });
5172 let idle_s = darklane::ValleySignal::new(st.health.clone()).idle_seconds();
5176 body["prefix_cache_entries"] = json!(m.prefix_entries);
5177 body["prefix_cache_bytes"] = json!(m.prefix_bytes);
5178 body["active_sessions"] = json!(m.active_sessions);
5179 body["queued_requests"] = json!(m.queued_requests);
5180 body["continuation_pool_entries"] = json!(m.continuation_pool_entries);
5181 body["spec_pool_entries"] = json!(m.spec_pool_entries);
5182 body["cuda_driver_free_bytes"] = json!(m.cuda_driver_free_bytes);
5183 body["cuda_pool_reserved_bytes"] = json!(m.cuda_pool_reserved_bytes);
5184 body["cuda_pool_used_bytes"] = json!(m.cuda_pool_used_bytes);
5185 body["cuda_pool_cached_bytes"] = json!(m.cuda_pool_cached_bytes);
5186 if !m.constraint_compiler_fail_closed.is_empty() {
5187 body["constraint_compiler_fail_closed"] = serde_json::Value::Object(
5188 m.constraint_compiler_fail_closed
5189 .iter()
5190 .map(|(model, gauge)| {
5191 let value = u8::from(gauge.load(std::sync::atomic::Ordering::Acquire));
5192 (model.clone(), json!(value))
5193 })
5194 .collect(),
5195 );
5196 }
5197 body["serve_idle_seconds"] = json!((idle_s * 1000.0).round() / 1000.0);
5198 }
5199 if !m.ns_tokens.is_empty() {
5204 let tenants: serde_json::Map<String, serde_json::Value> = m
5205 .ns_tokens
5206 .iter()
5207 .filter(|(ns, _)| metrics_scope.includes(ns))
5208 .map(|(ns, [p, c])| {
5209 (
5210 ns.clone(),
5211 json!({
5212 "prompt_tokens_in": p,
5213 "cached_tokens_in": c,
5214 "cache_hit_token_ratio": if *p > 0 { *c as f64 / *p as f64 } else { 0.0 },
5215 }),
5216 )
5217 })
5218 .collect();
5219 if !tenants.is_empty() {
5220 body["tenants"] = serde_json::Value::Object(tenants);
5221 }
5222 }
5223 let adsd_suspect_total: serde_json::Map<String, serde_json::Value> = m
5224 .adsd_suspect_total
5225 .iter()
5226 .filter(|(tenant, _)| metrics_scope.includes(tenant))
5227 .map(|(tenant, total)| (tenant.clone(), json!(total)))
5228 .collect();
5229 if !adsd_suspect_total.is_empty() {
5230 body["adsd_suspect_total"] = serde_json::Value::Object(adsd_suspect_total);
5231 }
5232 if metrics_scope.operator() {
5234 if let Some((bg, mode)) = &st.bg {
5235 body["bg"] = bg.to_json(mode);
5236 }
5237 }
5238 if metrics_scope.operator() {
5245 let spec: serde_json::Map<String, serde_json::Value> = m
5246 .spec
5247 .iter()
5248 .map(|(model, t)| {
5249 let n_pos = t
5250 .pos_drafted
5251 .iter()
5252 .rposition(|&d| d > 0)
5253 .map_or(0, |p| p + 1);
5254 (
5255 model.clone(),
5256 json!({
5257 "rounds": t.rounds,
5258 "drafted": t.drafted,
5259 "accepted": t.accepted,
5260 "acceptance_rate": if t.drafted > 0 {
5261 t.accepted as f64 / t.drafted as f64 } else { 0.0 },
5262 "tokens_per_round": if t.rounds > 0 {
5263 (t.accepted + t.rounds) as f64 / t.rounds as f64 } else { 0.0 },
5264 "pos_drafted": t.pos_drafted[..n_pos].to_vec(),
5265 "pos_accepted": t.pos_accepted[..n_pos].to_vec(),
5266 "accept_rate_per_pos": (0..n_pos).map(|j| if t.pos_drafted[j] > 0 {
5267 t.pos_accepted[j] as f64 / t.pos_drafted[j] as f64 } else { 0.0 })
5268 .collect::<Vec<f64>>(),
5269 }),
5270 )
5271 })
5272 .collect();
5273 if !spec.is_empty() {
5274 body["spec"] = serde_json::Value::Object(spec);
5275 }
5276 }
5277 insert_spec_acceptance_metrics(&mut body, &metrics_scope, || m.spec_window.clone());
5278 insert_dual_pp_metrics(&mut body, &metrics_scope, DualPpMetricsSnapshot::current);
5279 insert_peer_probe_metrics(
5280 &mut body,
5281 &metrics_scope,
5282 memra_engine::pp::peer_probe_metrics,
5283 );
5284 Json(body).into_response()
5285}
5286
5287#[derive(Debug, Default, Deserialize)]
5288struct ModelsQuery {
5289 #[serde(default)]
5290 schema: Option<String>,
5291}
5292
5293fn models_openai_body(models: &[String]) -> serde_json::Value {
5294 let data: Vec<_> = models
5295 .iter()
5296 .map(|m| json!({ "id": m, "object": "model" }))
5297 .collect();
5298 json!({ "object": "list", "data": data })
5299}
5300
5301fn declared_surface(metadata: Option<&OpenRouterModelMetadata>) -> &'static str {
5306 match metadata.and_then(|m| m.surface.as_deref()) {
5307 Some("embedding") => "embedding",
5308 Some("rerank") => "rerank",
5309 _ => "chat",
5310 }
5311}
5312
5313fn openrouter_supported_parameters(
5314 caps: Option<&ModelCaps>,
5315 max_output_length: Option<u64>,
5316 is_chat: bool,
5317) -> serde_json::Value {
5318 let mut parameters = serde_json::Map::new();
5319 if !is_chat {
5327 return serde_json::Value::Object(parameters);
5328 }
5329 for name in [
5330 "temperature",
5331 "top_p",
5332 "min_p",
5333 "frequency_penalty",
5334 "presence_penalty",
5335 "repetition_penalty",
5336 "stop",
5337 ] {
5338 parameters.insert(name.into(), json!({ "type": "unknown" }));
5339 }
5340 parameters.insert("top_k".into(), json!({ "type": "integer", "min": 0 }));
5341 parameters.insert(
5342 "seed".into(),
5343 json!({ "type": "integer", "min": 0, "max": JSON_SAFE_INTEGER_MAX }),
5344 );
5345 let mut max_tokens = json!({ "type": "integer", "min": 1, "unit": "token" });
5346 if let Some(max) = max_output_length {
5347 max_tokens["max"] = json!(max);
5348 }
5349 parameters.insert("max_tokens".into(), max_tokens);
5350 parameters.insert("json_mode".into(), json!({ "type": "boolean" }));
5351 parameters.insert("structured_outputs".into(), json!({ "type": "boolean" }));
5352 if is_chat && caps.is_some_and(|c| c.tools_branch) {
5353 parameters.insert("tools".into(), json!({ "type": "boolean" }));
5354 parameters.insert(
5355 "tool_choice".into(),
5356 json!({ "type": "enum", "values": ["auto", "none"] }),
5357 );
5358 }
5359 if is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5360 parameters.insert("reasoning".into(), json!({ "type": "boolean" }));
5361 }
5362 serde_json::Value::Object(parameters)
5363}
5364
5365fn model_entry_openrouter(
5366 name: &str,
5367 caps: Option<&ModelCaps>,
5368 metadata: Option<&OpenRouterModelMetadata>,
5369) -> serde_json::Value {
5370 let empty = OpenRouterModelMetadata::default();
5371 let metadata = metadata.unwrap_or(&empty);
5372 let context_length = caps
5373 .map(|c| c.context_length as u64)
5374 .filter(|&v| v > 0 && v <= JSON_SAFE_INTEGER_MAX);
5375 let tokenizer = caps
5376 .map(|c| c.tokenizer.as_str())
5377 .filter(|tokenizer| !tokenizer.is_empty());
5378
5379 let mut input = serde_json::Map::new();
5380 input.insert("type".into(), json!("text"));
5381 let mut supported_inputs = serde_json::Map::new();
5382 if let Some(value) = context_length {
5383 supported_inputs.insert(
5384 "max_context_length".into(),
5385 json!({ "value": value, "unit": "token" }),
5386 );
5387 }
5388 if let Some(value) = metadata.max_prompt_length {
5389 supported_inputs.insert(
5390 "max_prompt_length".into(),
5391 json!({ "value": value, "unit": "token" }),
5392 );
5393 }
5394 if !supported_inputs.is_empty() {
5395 input.insert(
5396 "supported_inputs".into(),
5397 serde_json::Value::Object(supported_inputs),
5398 );
5399 }
5400 let mut input_pricing = Vec::new();
5401 for (kind, cost) in [
5402 ("prompt", metadata.pricing.prompt.as_deref()),
5403 ("cached_prompt", metadata.pricing.cached_prompt.as_deref()),
5404 ("cache_write", metadata.pricing.cache_write.as_deref()),
5405 ] {
5406 if let Some(cost) = cost {
5407 input_pricing.push(json!({
5408 "type": kind,
5409 "unit": "token",
5410 "cost_usd": cost,
5411 }));
5412 }
5413 }
5414 if !input_pricing.is_empty() {
5415 input.insert("pricing".into(), serde_json::Value::Array(input_pricing));
5416 }
5417 let mut input_capacity = Vec::new();
5418 for (kind, value) in [
5419 ("prompt", metadata.capacity.prompt_tpm),
5420 ("cached_prompt", metadata.capacity.cached_prompt_tpm),
5421 ] {
5422 if let Some(value) = value {
5423 input_capacity.push(json!({
5424 "type": kind,
5425 "unit": "token",
5426 "per": "minute",
5427 "value": value,
5428 }));
5429 }
5430 }
5431 if !input_capacity.is_empty() {
5432 input.insert("capacity".into(), serde_json::Value::Array(input_capacity));
5433 }
5434
5435 let or_surface = declared_surface(Some(metadata));
5436 let or_is_chat = or_surface == "chat";
5437 let mut output = serde_json::Map::new();
5438 output.insert(
5446 "type".into(),
5447 json!(match or_surface {
5448 "embedding" => "embeddings",
5449 "rerank" => "rerank",
5450 _ => "text",
5451 }),
5452 );
5453 output.insert(
5454 "supported_parameters".into(),
5455 openrouter_supported_parameters(caps, metadata.max_output_length, or_is_chat),
5456 );
5457 if or_is_chat {
5461 output.insert("streaming".into(), json!(true));
5462 }
5463 if let Some(value) = metadata.max_output_length
5466 && or_is_chat
5467 {
5468 output.insert(
5469 "max_length".into(),
5470 json!({ "value": value, "unit": "token" }),
5471 );
5472 }
5473 let mut output_pricing = Vec::new();
5474 for (kind, cost) in [
5475 ("completion", metadata.pricing.completion.as_deref()),
5476 (
5477 "internal_reasoning",
5478 metadata.pricing.internal_reasoning.as_deref(),
5479 ),
5480 ] {
5481 if let Some(cost) = cost {
5482 output_pricing.push(json!({
5483 "type": kind,
5484 "unit": "token",
5485 "cost_usd": cost,
5486 }));
5487 }
5488 }
5489 if !output_pricing.is_empty() {
5490 output.insert("pricing".into(), serde_json::Value::Array(output_pricing));
5491 }
5492 let mut output_capacity = Vec::new();
5493 if let Some(value) = metadata.capacity.completion_tpm {
5494 output_capacity.push(json!({
5495 "type": "completion",
5496 "unit": "token",
5497 "per": "minute",
5498 "value": value,
5499 }));
5500 }
5501 if let Some(value) = metadata.capacity.concurrency {
5502 output_capacity.push(json!({
5503 "type": "concurrency",
5504 "unit": "request",
5505 "value": value,
5506 }));
5507 }
5508 if !output_capacity.is_empty() {
5509 output.insert("capacity".into(), serde_json::Value::Array(output_capacity));
5510 }
5511
5512 let mut entry = serde_json::Map::new();
5513 entry.insert("schema_version".into(), json!(OPENROUTER_SCHEMA_VERSION));
5514 entry.insert("id".into(), json!(name));
5515 entry.insert("name".into(), json!(name));
5516 if let Some(value) = metadata.hugging_face_id.as_deref() {
5517 entry.insert("hugging_face_id".into(), json!(value));
5518 }
5519 if let Some(value) = metadata.created {
5520 entry.insert("created".into(), json!(value));
5521 }
5522 if let Some(value) = metadata.quantization.as_deref() {
5523 entry.insert("quantization".into(), json!(value));
5524 }
5525 if let Some(value) = tokenizer {
5526 entry.insert("tokenizer".into(), json!(value));
5527 }
5528 if let Some(value) = metadata.description.as_deref() {
5529 entry.insert("description".into(), json!(value));
5530 }
5531 let mut input_modalities = vec![serde_json::Value::Object(input)];
5532 for m in &metadata.input_modalities {
5533 let mut extra = serde_json::Map::new();
5534 extra.insert("type".into(), json!(m));
5535 if let Some(cost) = metadata.pricing.prompt.as_deref() {
5536 extra.insert(
5538 "pricing".into(),
5539 json!([{ "type": "prompt", "unit": "token", "cost_usd": cost }]),
5540 );
5541 }
5542 input_modalities.push(serde_json::Value::Object(extra));
5543 }
5544 entry.insert(
5545 "input_modalities".into(),
5546 serde_json::Value::Array(input_modalities),
5547 );
5548 entry.insert(
5549 "output_modalities".into(),
5550 serde_json::Value::Array(vec![serde_json::Value::Object(output)]),
5551 );
5552 if let Some(cost) = metadata.pricing.request.as_deref() {
5553 entry.insert(
5554 "pricing".into(),
5555 json!([{ "type": "request", "unit": "request", "cost_usd": cost }]),
5556 );
5557 }
5558 if let Some(value) = metadata.capacity.request_rpm {
5559 entry.insert(
5560 "capacity".into(),
5561 json!([{
5562 "type": "request",
5563 "unit": "request",
5564 "per": "minute",
5565 "value": value,
5566 }]),
5567 );
5568 }
5569 if let Some(value) = metadata.is_ready {
5570 entry.insert("is_ready".into(), json!(value));
5571 }
5572 if let Some(value) = metadata.is_free {
5573 entry.insert("is_free".into(), json!(value));
5574 }
5575 if let Some(value) = metadata.discount_to_user {
5576 entry.insert("discount_to_user".into(), json!(value));
5577 }
5578 if let Some(value) = metadata.openrouter_slug.as_deref() {
5579 entry.insert("openrouter".into(), json!({ "slug": value }));
5580 }
5581 if !metadata.datacenters.is_empty() {
5582 entry.insert("datacenters".into(), json!(metadata.datacenters));
5583 }
5584 let mut compliance = serde_json::Map::new();
5585 if let Some(value) = metadata.zdr {
5586 compliance.insert("zdr".into(), json!(value));
5587 }
5588 if let Some(value) = metadata.hipaa {
5589 compliance.insert("hipaa".into(), json!(value));
5590 }
5591 if !compliance.is_empty() {
5592 entry.insert("compliance".into(), serde_json::Value::Object(compliance));
5593 }
5594 serde_json::Value::Object(entry)
5595}
5596
5597fn models_openrouter_body(st: &AppState) -> serde_json::Value {
5598 let data: Vec<_> = st
5599 .models
5600 .iter()
5601 .map(|model| {
5602 model_entry_openrouter(model, st.caps.get(model), st.openrouter_metadata.get(model))
5603 })
5604 .collect();
5605 json!({ "data": data })
5606}
5607
5608fn model_entry_openmodels(
5609 name: &str,
5610 caps: Option<&ModelCaps>,
5611 metadata: Option<&OpenRouterModelMetadata>,
5612) -> Result<serde_json::Value, String> {
5613 let metadata = metadata.ok_or_else(|| {
5614 format!("OpenModels feed requires MEMRA_MODEL_METADATA for model {name:?}")
5615 })?;
5616 let context_length = caps
5617 .map(|c| c.context_length as u64)
5618 .filter(|&value| value > 0 && value <= JSON_SAFE_INTEGER_MAX)
5619 .ok_or_else(|| format!("OpenModels feed requires context_length for model {name:?}"))?;
5620 let created = metadata
5621 .created
5622 .ok_or_else(|| format!("OpenModels feed requires created for model {name:?}"))?;
5623 let max_output_length = metadata
5624 .max_output_length
5625 .ok_or_else(|| format!("OpenModels feed requires max_output_length for model {name:?}"))?;
5626 let prompt = metadata
5627 .pricing
5628 .prompt
5629 .as_deref()
5630 .ok_or_else(|| format!("OpenModels feed requires pricing.prompt for model {name:?}"))?;
5631 let completion =
5632 metadata.pricing.completion.as_deref().ok_or_else(|| {
5633 format!("OpenModels feed requires pricing.completion for model {name:?}")
5634 })?;
5635 let input_cache_read = metadata.pricing.cached_prompt.as_deref().ok_or_else(|| {
5636 format!("OpenModels feed requires pricing.cached_prompt for model {name:?}")
5637 })?;
5638 let is_ready = metadata
5639 .is_ready
5640 .ok_or_else(|| format!("OpenModels feed requires is_ready for model {name:?}"))?;
5641 let is_free = metadata
5642 .is_free
5643 .ok_or_else(|| format!("OpenModels feed requires is_free for model {name:?}"))?;
5644 let discount_to_user = metadata
5645 .discount_to_user
5646 .ok_or_else(|| format!("OpenModels feed requires discount_to_user for model {name:?}"))?;
5647
5648 let mut pricing = serde_json::Map::new();
5649 pricing.insert("prompt".into(), json!(prompt));
5650 pricing.insert("completion".into(), json!(completion));
5651 pricing.insert("input_cache_read".into(), json!(input_cache_read));
5652 if let Some(value) = metadata.pricing.request.as_deref() {
5653 pricing.insert("request".into(), json!(value));
5654 }
5655
5656 let om_surface = declared_surface(Some(metadata));
5657 let om_is_chat = om_surface == "chat";
5658 let mut supported_features = Vec::new();
5659 if om_is_chat && caps.is_some_and(|c| c.tools_branch) {
5660 supported_features.push("tool_calling");
5661 }
5662 if om_is_chat && caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think) {
5663 supported_features.push("reasoning");
5664 }
5665
5666 let mut entry = serde_json::Map::new();
5667 entry.insert("id".into(), json!(name));
5668 entry.insert("name".into(), json!(name));
5669 entry.insert("created".into(), json!(created));
5670 entry.insert("input_modalities".into(), json!(["text"]));
5671 entry.insert(
5672 "output_modalities".into(),
5673 json!(match om_surface {
5674 "embedding" => ["embeddings"],
5675 "rerank" => ["rerank"],
5676 _ => ["text"],
5677 }),
5678 );
5679 entry.insert("context_length".into(), json!(context_length));
5680 entry.insert("max_output_length".into(), json!(max_output_length));
5681 entry.insert("currency".into(), json!("USD"));
5684 entry.insert("pricing".into(), serde_json::Value::Object(pricing));
5685 entry.insert("supported_features".into(), json!(supported_features));
5686 entry.insert("is_ready".into(), json!(is_ready));
5687 entry.insert("is_free".into(), json!(is_free));
5688 entry.insert("discount_to_user".into(), json!(discount_to_user));
5689 Ok(serde_json::Value::Object(entry))
5690}
5691
5692fn models_openmodels_body(st: &AppState) -> Result<serde_json::Value, String> {
5693 let data: Result<Vec<_>, _> = st
5694 .models
5695 .iter()
5696 .map(|model| {
5697 model_entry_openmodels(model, st.caps.get(model), st.openrouter_metadata.get(model))
5698 })
5699 .collect();
5700 Ok(json!({ "data": data? }))
5701}
5702
5703async fn list_models(State(st): State<AppState>, Query(query): Query<ModelsQuery>) -> Response {
5704 match query.schema.as_deref() {
5705 None | Some("openai") => Json(models_openai_body(st.models.as_ref())).into_response(),
5706 Some("openrouter") => Json(models_openrouter_body(&st)).into_response(),
5707 Some("openmodels") => match models_openmodels_body(&st) {
5708 Ok(body) => Json(body).into_response(),
5709 Err(error) => bad_request(&error, Some("schema")),
5710 },
5711 Some(schema) => bad_request(
5712 &format!(
5713 "unsupported models schema {schema:?}; expected openai, openrouter, or openmodels"
5714 ),
5715 Some("schema"),
5716 ),
5717 }
5718}
5719
5720fn model_entry_v1(
5728 name: &str,
5729 caps: Option<&ModelCaps>,
5730 metadata: Option<&OpenRouterModelMetadata>,
5731) -> serde_json::Value {
5732 let ctx = caps.map(|c| c.context_length).filter(|&c| c > 0);
5733 let thinking = caps.is_some_and(|c| c.qwen_think || c.effort_levels || c.gemma_think || c.dsv4);
5737 let is_dsv4 = caps.is_some_and(|c| c.dsv4);
5740 let per_1m = |v: Option<&str>| match v.and_then(per_million_price) {
5741 Some(p) => json!(p),
5742 None => serde_json::Value::Null,
5743 };
5744 let owned_by = metadata
5745 .and_then(|m| m.owned_by.as_deref())
5746 .unwrap_or_else(|| name.split('/').next().unwrap_or(name));
5747 let mut input_modalities = vec!["text"];
5748 if let Some(meta) = metadata {
5749 input_modalities.extend(meta.input_modalities.iter().map(String::as_str));
5750 }
5751 let lifecycle = metadata.and_then(|m| m.lifecycle.as_ref());
5752 let reliability = metadata.and_then(|m| m.reliability.as_ref());
5753 let surface = declared_surface(metadata);
5759 let (model_type, endpoints, output_modalities) = match surface {
5760 "embedding" => ("embedding", vec!["embeddings"], vec!["embeddings"]),
5764 "rerank" => ("rerank", vec!["rerank"], vec!["rerank"]),
5765 _ => ("chat", vec!["chat/completions"], vec!["text"]),
5766 };
5767 let is_chat = surface == "chat";
5768 json!({
5769 "id": name,
5770 "name": name,
5771 "object": "model",
5772 "owned_by": owned_by,
5773 "type": model_type,
5774 "context_length": ctx,
5775 "max_output_tokens": if is_chat { metadata.and_then(|m| m.max_output_length) } else { None },
5778 "endpoints": endpoints,
5779 "input_modalities": input_modalities,
5780 "output_modalities": output_modalities,
5781 "capabilities": {
5782 "streaming": is_chat,
5785 "tools": is_chat && caps.is_some_and(|c| c.tools_branch),
5786 "structured_output": is_chat && !is_dsv4,
5787 "reasoning": is_chat && thinking,
5788 "prompt_caching": is_chat && !is_dsv4,
5789 },
5790 "pricing": {
5791 "currency": "USD",
5792 "unit": "per_1m_tokens",
5793 "input": per_1m(metadata.and_then(|m| m.pricing.prompt.as_deref())),
5794 "output": per_1m(metadata.and_then(|m| m.pricing.completion.as_deref())),
5795 "cached_input": per_1m(metadata.and_then(|m| m.pricing.cached_prompt.as_deref())),
5796 "cache_write": per_1m(metadata.and_then(|m| m.pricing.cache_write.as_deref())),
5797 "minimum_request": metadata
5799 .and_then(|m| m.pricing.request.as_deref())
5800 .unwrap_or("0"),
5801 },
5802 "lifecycle": {
5803 "status": lifecycle.and_then(|l| l.status.as_deref()).unwrap_or("active"),
5804 "deprecation_at": lifecycle.and_then(|l| l.deprecation_at.as_deref()),
5805 "retirement_at": lifecycle.and_then(|l| l.retirement_at.as_deref()),
5806 "replacement_model_id": lifecycle.and_then(|l| l.replacement_model_id.as_deref()),
5807 },
5808 "reliability": {
5809 "first_token_timeout_seconds":
5810 reliability.and_then(|r| r.first_token_timeout_seconds).unwrap_or(120),
5811 "completion_timeout_seconds":
5812 reliability.and_then(|r| r.completion_timeout_seconds).unwrap_or(900),
5813 "stream_idle_timeout_seconds":
5814 reliability.and_then(|r| r.stream_idle_timeout_seconds).unwrap_or(60),
5815 "capacity_scope":
5816 reliability.and_then(|r| r.capacity_scope.as_deref()).unwrap_or("model_region"),
5817 },
5818 })
5819}
5820
5821async fn list_models_v1(State(st): State<AppState>) -> impl IntoResponse {
5824 let data: Vec<_> = st
5825 .models
5826 .iter()
5827 .map(|m| model_entry_v1(m, st.caps.get(m), st.openrouter_metadata.get(m)))
5828 .collect();
5829 let mut body = json!({
5830 "object": "list",
5831 "contract_version": "2.0",
5832 "data": data,
5833 });
5834 if let Some(provider) = st.provider_metadata.as_ref() {
5839 body["provider"] = json!({
5840 "id": provider.id,
5841 "status_url": provider.status_url,
5842 "support_contact": provider.support_contact,
5843 "incident_contact": provider.incident_contact,
5844 "regions": provider.regions,
5845 "request_id_header": "x-request-id",
5846 "error_contract": {
5847 "rate_limit_status": 429,
5848 "overload_status": 503,
5849 "retry_after_header": "Retry-After",
5850 "account_quota_error_codes": ["insufficient_balance"],
5851 },
5852 });
5853 }
5854 Json(body)
5855}
5856
5857async fn yield_metrics(State(st): State<AppState>, headers: HeaderMap) -> Response {
5860 let metrics_scope = match authorize_metrics(&st.api_auth, &st.metrics_auth, &headers) {
5861 Ok(scope) => scope,
5862 Err(response) => return response,
5863 };
5864 if !metrics_scope.process_wide() {
5865 return error_response(
5866 StatusCode::FORBIDDEN,
5867 "completion api keys do not authorize process-wide yield metrics; configure \
5868 MEMRA_METRICS_TOKEN",
5869 "authentication_error",
5870 None,
5871 );
5872 }
5873 let m = st.metrics.lock().map(|m| m.clone()).unwrap_or_default();
5874 let lane = |i: usize| {
5875 json!({
5876 "admitted": m.lane_admitted[i], "shed": m.lane_shed[i],
5877 "completed": m.lane_completed[i], "tokens_out": m.lane_tokens[i],
5878 })
5879 };
5880 let mut body = json!({
5881 "lanes": {
5882 "interactive": lane(0), "judge": lane(1), "harvest": lane(2),
5883 },
5884 "interactive_step_ms": { "p50": m.step_p50_ms, "p99": m.step_p99_ms },
5885 });
5886 if metrics_scope.operator() {
5887 body["batch_size_last"] = json!(m.batch_size_last);
5888 }
5889 Json(body).into_response()
5890}
5891
5892async fn peek_admission(
5905 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5906) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, (Response, &'static str)> {
5907 match rx.recv().await {
5908 Some(Event::Error(e)) => {
5913 let error_code = engine_error_code(e.class);
5914 Err((engine_error_response(&e), error_code))
5915 }
5916 first => {
5917 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
5918 if let Some(ev) = first {
5919 let _ = tx2.send(ev);
5920 }
5921 tokio::spawn(forward_events(rx, tx2));
5922 Ok(rx2)
5923 }
5924 }
5925}
5926
5927async fn forward_events(
5938 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5939 tx2: tokio::sync::mpsc::UnboundedSender<Event>,
5940) {
5941 loop {
5942 tokio::select! {
5943 biased;
5944 () = tx2.closed() => break,
5945 ev = rx.recv() => match ev {
5946 Some(ev) => {
5947 if tx2.send(ev).is_err() {
5948 break;
5949 }
5950 }
5951 None => break,
5952 },
5953 }
5954 }
5955}
5956
5957async fn peek_first_token(
5972 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
5973 deadline: RequestDeadline,
5974) -> Result<tokio::sync::mpsc::UnboundedReceiver<Event>, ()> {
5975 let mut buffered: Vec<Event> = Vec::new();
5976 loop {
5977 match tokio::time::timeout_at(deadline.at, rx.recv()).await {
5978 Err(_) => return Err(()), Ok(None) => break, Ok(Some(ev)) => {
5981 let first_delivery = matches!(
5982 ev,
5983 Event::Token { .. } | Event::Done { .. } | Event::Error(_)
5984 );
5985 buffered.push(ev);
5986 if first_delivery {
5987 break;
5988 }
5989 }
5990 }
5991 }
5992 let (tx2, rx2) = tokio::sync::mpsc::unbounded_channel();
5993 for ev in buffered {
5994 let _ = tx2.send(ev);
5995 }
5996 tokio::spawn(forward_events(rx, tx2));
5997 Ok(rx2)
5998}
5999
6000#[cfg(test)]
6002fn build_request(
6006 req: &CompletionReq,
6007 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6008 lane: lanes::Lane,
6009 affinity: Option<String>,
6010) -> Request {
6011 build_request_with_trace(req, tx, lane, affinity, None, &SamplingDefaults::default())
6012}
6013
6014fn build_request_with_trace(
6015 req: &CompletionReq,
6016 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6017 lane: lanes::Lane,
6018 affinity: Option<String>,
6019 ttft: Option<Arc<ttft::Trace>>,
6020 sampling_defaults: &SamplingDefaults,
6021) -> Request {
6022 let params = GenParams {
6023 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6024 max_ctx: req.max_ctx,
6025 eos: Vec::new(), };
6027 let sampler_cfg = resolve_sampler_config(req.into(), sampling_defaults);
6032 Request {
6033 model: req.model.clone(),
6034 prompt_ids: req.prompt_ids.clone(),
6035 prompt_text: req.prompt.clone(),
6036 chat: req.chat,
6037 chat_turns: Vec::new(),
6038 tools_json: Vec::new(),
6039 tools_struct: Vec::new(),
6040 think: ThinkMode::Default,
6041 reasoning_effort: None, params,
6043 sampler_cfg,
6044 stop_strings: req.stop.clone().into_vec(),
6045 trace_id: req.trace_id.clone(),
6046 max_prompt_tokens: None,
6047 cache_ns: cache_namespace(&req.cache_salt),
6048 affinity,
6049 lane,
6050 grammar: None, prepared_constraint: None,
6052 constraint_ready: None,
6053 oom_retries: 0, spec_k_replay: None,
6055 prepared_prompt: None,
6056 capture: None, images: Vec::new(), gemma_images: Vec::new(),
6059 step_images: Vec::new(),
6060 vision_memory: None,
6061 ttft,
6062 tx,
6063 }
6064}
6065
6066struct ChatPlan {
6069 request: Request,
6070 parser: Option<ToolStreamParser>,
6073 pending_images: Vec<PendingVisionUnit>,
6076 pending_gemma: Vec<PendingGemmaImage>,
6077 pending_step: Vec<PendingStepImage>,
6078 vision_memory: Option<VisionMemoryPermit>,
6082}
6083
6084pub(crate) fn request_has_vision(req: &ChatCompletionReq) -> bool {
6085 req.messages.iter().any(|message| {
6086 message.content.as_array().is_some_and(|parts| {
6087 parts.iter().any(|part| {
6088 matches!(
6089 part.get("type").and_then(serde_json::Value::as_str),
6090 Some("image_url" | "video_url")
6091 )
6092 })
6093 })
6094 })
6095}
6096
6097fn planned_vision_bytes(plan: &ChatPlan) -> Result<usize, String> {
6098 let mut total = 0usize;
6099 let mut add = |bytes: usize| {
6100 total = total.checked_add(bytes).ok_or_else(|| {
6101 "vision patch memory reservation overflowed while planning".to_string()
6102 })?;
6103 Ok::<(), String>(())
6104 };
6105 for unit in &plan.pending_images {
6106 let bytes = match unit {
6107 PendingVisionUnit::Still { gh, gw, .. } => gh
6108 .checked_mul(*gw)
6109 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6110 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6111 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?,
6112 PendingVisionUnit::Video { groups, .. } => {
6113 groups.iter().try_fold(0usize, |total, group| {
6114 let bytes = group
6115 .gh
6116 .checked_mul(group.gw)
6117 .and_then(|n| n.checked_mul(memra_engine::vision::V_PATCH_IN))
6118 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6119 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6120 total.checked_add(bytes).ok_or_else(|| {
6121 "vision patch memory reservation overflowed while planning".to_string()
6122 })
6123 })?
6124 }
6125 };
6126 add(bytes)?;
6127 }
6128 for unit in &plan.pending_gemma {
6129 let bytes = unit
6130 .gw
6131 .checked_mul(unit.gh)
6132 .and_then(|n| n.checked_mul(memra_engine::vision_gemma::GV_PATCH_IN))
6133 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6134 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6135 add(bytes)?;
6136 }
6137 for unit in &plan.pending_step {
6138 use memra_engine::vision_step::{SV_GRID_MAIN, SV_GRID_TILE, SV_PATCH_IN};
6139 let patches = unit
6141 .plan
6142 .n_tiles
6143 .checked_mul(SV_GRID_TILE * SV_GRID_TILE)
6144 .and_then(|n| n.checked_add(SV_GRID_MAIN * SV_GRID_MAIN))
6145 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6146 let bytes = patches
6147 .checked_mul(SV_PATCH_IN)
6148 .and_then(|n| n.checked_mul(std::mem::size_of::<f32>()))
6149 .ok_or_else(|| "vision patch memory reservation overflowed".to_string())?;
6150 add(bytes)?;
6151 }
6152 Ok(total)
6153}
6154
6155pub(crate) fn reserve_vision_memory(
6156 plan: &ChatPlan,
6157) -> Result<Option<VisionMemoryPermit>, VisionMemoryError> {
6158 let bytes = planned_vision_bytes(plan).map_err(VisionMemoryError::Request)?;
6159 try_reserve_vision_memory(bytes)
6160}
6161
6162#[cfg(test)]
6163fn build_chat_request(
6164 req: ChatCompletionReq,
6165 caps: Option<&ModelCaps>,
6166 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6167 lane: lanes::Lane,
6168 affinity: Option<String>,
6169) -> Result<ChatPlan, String> {
6170 let defaults = ModelSamplingDefaults::resolve(None, caps);
6173 build_chat_request_with_trace(req, caps, tx, lane, affinity, None, None, &defaults)
6174}
6175
6176#[allow(clippy::too_many_arguments)]
6190fn build_chat_request_with_trace(
6191 req: ChatCompletionReq,
6192 caps: Option<&ModelCaps>,
6193 tx: tokio::sync::mpsc::UnboundedSender<Event>,
6194 lane: lanes::Lane,
6195 affinity: Option<String>,
6196 ttft: Option<Arc<ttft::Trace>>,
6197 default_effort: Option<&str>,
6198 sampling_defaults: &ModelSamplingDefaults,
6199) -> Result<ChatPlan, String> {
6200 let client_sampling: ClientSampling = (&req).into();
6203 let tool_choice = parse_tool_choice(&req.tool_choice)?;
6204 if let Some(c) = caps {
6209 if !c.chat_ok {
6210 return Err(format!(
6211 "model {:?} has no chat template (checkpoint carries neither \
6212 tokenizer_config.json chat_template nor chat_template.jinja) — \
6213 /v1/chat/completions unavailable; use /v1/completions with a raw prompt",
6214 req.model
6215 ));
6216 }
6217 }
6218 let vllm_switch = resolve_vllm_think_switch(req.enable_thinking, &req.chat_template_kwargs)?;
6219 let (mut think, effort_level, think_client_explicit) = parse_think(
6220 &req.reasoning_effort,
6221 &req.reasoning,
6222 vllm_switch,
6223 req.include_reasoning,
6224 default_effort,
6225 caps.is_some_and(|c| c.dsv4),
6226 )?;
6227 let level_template = caps
6231 .map(|c| c.effort_levels || c.dsv4 || c.qwen_effort)
6232 .unwrap_or(false);
6233 if think_client_explicit && think == ThinkMode::NoThink {
6254 if let Some(c) = caps {
6255 if c.qwen_think && !c.think_switch && !c.dsv4 {
6256 return Err(format!(
6257 "model {:?} cannot disable reasoning: its chat template opens a think \
6258 tail unconditionally and carries no enable_thinking switch, so \
6259 reasoning_effort/enable_thinking cannot turn it off on this model",
6260 req.model
6261 ));
6262 }
6263 }
6264 }
6265 let reasoning_effort = if level_template { effort_level } else { None };
6298 let grammar = constrained::parse_response_format(req.response_format.as_ref())?;
6301 if grammar.is_some() {
6307 if let Some(c) = caps {
6308 if c.qwen_think && think != ThinkMode::NoThink {
6309 if c.think_switch {
6310 think = ThinkMode::NoThink;
6311 } else {
6312 return Err(
6313 "response_format requires disabling the model's think tail, \
6314 but this chat template has no enable_thinking switch"
6315 .into(),
6316 );
6317 }
6318 }
6319 }
6320 }
6321
6322 let sampler_cfg = resolve_sampler_config(client_sampling, sampling_defaults.for_mode(think));
6329
6330 let (tools_json, tools_struct, schemas) =
6333 if !req.tools.is_empty() && tool_choice == ToolChoice::Auto {
6334 prepare_tools(&req.tools)?
6335 } else {
6336 (Vec::new(), Vec::new(), HashMap::new())
6337 };
6338
6339 let mut turns: Vec<TmplTurn> = Vec::with_capacity(req.messages.len());
6340 let mut images: Vec<PendingVisionUnit> = Vec::new();
6341 let mut gemma_images: Vec<PendingGemmaImage> = Vec::new();
6342 let mut step_images: Vec<PendingStepImage> = Vec::new();
6343 let mut next_video = 0usize;
6344 for msg in &req.messages {
6345 let content = content_to_text_vision(
6346 &msg.content,
6347 &mut images,
6348 &mut gemma_images,
6349 &mut step_images,
6350 &mut next_video,
6351 )
6352 .map_err(|e| format!("{} message: {e}", msg.role))?;
6353 let tool_calls = msg
6354 .tool_calls
6355 .iter()
6356 .map(render_req_tool_call)
6357 .collect::<Result<Vec<_>, _>>()?;
6358 if !tool_calls.is_empty() && msg.role != "assistant" {
6359 return Err("tool_calls are only valid on assistant messages".into());
6360 }
6361 let role = if msg.role == "developer" {
6364 "system".to_string()
6365 } else {
6366 msg.role.clone()
6367 };
6368 turns.push(TmplTurn {
6369 role,
6370 content,
6371 tool_calls,
6372 reasoning: msg.reasoning.clone().filter(|r| !r.is_empty()),
6374 tool_call_id: msg.tool_call_id.clone(),
6375 tool_name: msg.name.clone(),
6376 tool_responses: Vec::new(),
6377 task: None,
6381 tools: Vec::new(),
6382 });
6383 }
6384
6385 let has_tool_features = !tools_json.is_empty()
6388 || turns
6389 .iter()
6390 .any(|t| t.role == "tool" || !t.tool_calls.is_empty());
6391 if has_tool_features && !caps.map(|c| c.tools_branch).unwrap_or(false) {
6392 return Err(format!(
6393 "model {:?} chat template has no tools branch",
6394 req.model
6395 ));
6396 }
6397
6398 let think_open = caps
6401 .map(|c| c.qwen_think && !(think == ThinkMode::NoThink && c.think_switch))
6402 .unwrap_or(false);
6403 let gemma_tools = !tools_json.is_empty() && caps.map(|c| c.gemma_think).unwrap_or(false);
6421 let is_dsv4 = caps.map(|c| c.dsv4).unwrap_or(false);
6427 let dsv4_think_open = is_dsv4 && think != ThinkMode::NoThink;
6428 let dsv4_tools = is_dsv4 && !tools_struct.is_empty();
6429 let parser = if is_dsv4 && (dsv4_tools || dsv4_think_open) {
6430 Some(ToolStreamParser::dsv4(dsv4_think_open))
6431 } else if gemma_tools {
6432 Some(ToolStreamParser::gemma_tools())
6433 } else if !tools_json.is_empty() {
6434 Some(ToolStreamParser::new(schemas, think_open))
6435 } else if think_open {
6436 Some(ToolStreamParser::reasoning_only())
6437 } else if caps.map(|c| c.gemma_think).unwrap_or(false) {
6438 Some(ToolStreamParser::gemma_thought())
6446 } else {
6447 None
6448 };
6449
6450 Ok(ChatPlan {
6451 request: Request {
6452 model: req.model,
6453 prompt_ids: Vec::new(),
6454 prompt_text: String::new(),
6455 chat: false,
6456 chat_turns: turns,
6457 tools_json,
6458 tools_struct,
6459 think,
6460 reasoning_effort,
6461 params: GenParams {
6462 max_new: req.max_tokens.unwrap_or(worker::MAX_NEW_CTX_BOUNDED),
6463 max_ctx: req.max_ctx,
6464 eos: Vec::new(),
6465 },
6466 sampler_cfg,
6467 stop_strings: {
6468 let mut stops = req.stop.into_vec();
6473 if gemma_tools {
6474 stops.push("<tool_call|>".to_string());
6475 }
6476 if dsv4_tools {
6481 stops.push("</\u{ff5c}DSML\u{ff5c}tool_calls>".to_string());
6482 }
6483 stops
6484 },
6485 trace_id: None,
6486 max_prompt_tokens: None,
6487 cache_ns: cache_namespace(&req.cache_salt),
6488 affinity,
6489 lane,
6490 grammar,
6491 prepared_constraint: None,
6492 constraint_ready: None,
6493 oom_retries: 0, spec_k_replay: None,
6495 prepared_prompt: None,
6496 images: Vec::new(),
6501 gemma_images: Vec::new(),
6502 step_images: Vec::new(),
6503 capture: None, vision_memory: None,
6505 ttft,
6506 tx,
6507 },
6508 parser,
6509 pending_images: images,
6510 pending_gemma: gemma_images,
6511 pending_step: step_images,
6512 vision_memory: None,
6513 })
6514}
6515
6516fn decode_pending_vision(plan: &mut ChatPlan) -> Result<(), String> {
6522 for (i, unit) in plan.pending_images.drain(..).enumerate() {
6523 match unit {
6524 PendingVisionUnit::Still { bytes, gh, gw } => {
6525 let prep = memra_engine::vision_pre::prep_image_bytes(&bytes)
6526 .map_err(|e| format!("image {}: {e}", i + 1))?;
6527 if (prep.gh, prep.gw) != (gh, gw) {
6528 return Err(format!(
6529 "image {}: decoded grid {}x{} differs from its header-planned grid {gh}x{gw} — refusing (pad runs already rendered)",
6530 i + 1,
6531 prep.gh,
6532 prep.gw
6533 ));
6534 }
6535 plan.request
6536 .images
6537 .push(memra_engine::vision_pre::VisionUnit { prep, video: None });
6538 }
6539 PendingVisionUnit::Video {
6540 bytes,
6541 groups,
6542 video,
6543 } => {
6544 let prepared = memra_engine::vision_pre::prep_video_gif(&bytes)
6545 .map_err(|e| format!("video {}: {e}", i + 1))?;
6546 if prepared.groups.len() != groups.len() {
6547 return Err(format!(
6548 "video {}: decoded {} groups differ from its header-planned {} groups",
6549 i + 1,
6550 prepared.groups.len(),
6551 groups.len()
6552 ));
6553 }
6554 for ((group, prep), timestamp) in
6555 groups.iter().zip(prepared.groups).zip(prepared.timestamps)
6556 {
6557 if (prep.gh, prep.gw) != (group.gh, group.gw) {
6558 return Err(format!(
6559 "video {}: decoded grid {}x{} differs from its header-planned grid {}x{}",
6560 i + 1,
6561 prep.gh,
6562 prep.gw,
6563 group.gh,
6564 group.gw
6565 ));
6566 }
6567 if (timestamp - group.timestamp).abs() > 0.001 {
6568 return Err(format!(
6569 "video {}: decoded timestamp {timestamp:.3} differs from its header-planned timestamp {:.3}",
6570 i + 1,
6571 group.timestamp
6572 ));
6573 }
6574 plan.request
6575 .images
6576 .push(memra_engine::vision_pre::VisionUnit {
6577 prep,
6578 video: Some(video),
6579 });
6580 }
6581 }
6582 }
6583 }
6584 for (i, unit) in plan.pending_gemma.drain(..).enumerate() {
6585 let (patches, gw, gh) = memra_engine::vision_gemma::gemma_prep_image(&unit.bytes)
6586 .map_err(|e| format!("image {}: {e}", i + 1))?;
6587 if (gw, gh) != (unit.gw, unit.gh) {
6588 return Err(format!(
6589 "image {}: decoded grid {gw}x{gh} differs from its header-planned grid {}x{} — refusing (pad runs already rendered)",
6590 i + 1,
6591 unit.gw,
6592 unit.gh
6593 ));
6594 }
6595 plan.request
6596 .gemma_images
6597 .push(memra_engine::vision_gemma::GemmaVisionUnit { patches, gw, gh });
6598 }
6599 for (i, unit) in plan.pending_step.drain(..).enumerate() {
6600 let prepped = memra_engine::vision_step::step_prep_image(&unit.bytes)
6601 .map_err(|e| format!("image {}: {e}", i + 1))?;
6602 if prepped.tiles.len() != unit.plan.n_tiles
6603 || prepped.newline_mask != unit.plan.newline_mask
6604 {
6605 return Err(format!(
6606 "image {}: decoded tiling ({} tiles) differs from its header-planned tiling \
6607 ({} tiles) — refusing (pad runs already rendered)",
6608 i + 1,
6609 prepped.tiles.len(),
6610 unit.plan.n_tiles
6611 ));
6612 }
6613 plan.request.step_images.push(prepped);
6614 }
6615 Ok(())
6616}
6617
6618fn bearer_token(headers: &HeaderMap) -> Option<&str> {
6626 headers
6627 .get("authorization")
6628 .and_then(|value| value.to_str().ok())
6629 .and_then(|value| value.strip_prefix("Bearer "))
6630}
6631
6632fn authentication_error(why: auth::AuthDenied) -> Response {
6633 match why {
6634 auth::AuthDenied::Unknown => error_response(
6635 StatusCode::UNAUTHORIZED,
6636 "invalid api key",
6637 "authentication_error",
6638 None,
6639 ),
6640 auth::AuthDenied::Disabled => error_response(
6641 StatusCode::FORBIDDEN,
6642 "api key is disabled",
6643 "authentication_error",
6644 None,
6645 ),
6646 }
6647}
6648
6649fn authenticate(api_auth: &ApiAuth, headers: &HeaderMap) -> Result<auth::TenantCtx, Response> {
6650 auth::authenticate_with(
6651 api_auth.keyring,
6652 api_auth.single_key.as_deref(),
6653 bearer_token(headers),
6654 )
6655 .map_err(authentication_error)
6656}
6657
6658#[derive(Debug, Clone, PartialEq, Eq)]
6659enum MetricsScope {
6660 All,
6661 CompletionDomain,
6662 Tenant(String),
6663}
6664
6665impl MetricsScope {
6666 fn operator(&self) -> bool {
6667 matches!(self, MetricsScope::All)
6668 }
6669
6670 fn process_wide(&self) -> bool {
6671 matches!(self, MetricsScope::All | MetricsScope::CompletionDomain)
6672 }
6673
6674 fn includes(&self, tenant_row: &str) -> bool {
6675 match self {
6676 MetricsScope::All | MetricsScope::CompletionDomain => true,
6677 MetricsScope::Tenant(tenant) => tenant == tenant_row,
6678 }
6679 }
6680}
6681
6682fn authorize_metrics(
6683 api_auth: &ApiAuth,
6684 metrics_auth: &MetricsAuth,
6685 headers: &HeaderMap,
6686) -> Result<MetricsScope, Response> {
6687 if !metrics_auth.required {
6688 return Ok(MetricsScope::All);
6689 }
6690 let Some(candidate) = bearer_token(headers) else {
6691 return Err(authentication_error(auth::AuthDenied::Unknown));
6692 };
6693 if let Some(token) = metrics_auth.token.as_deref() {
6694 if auth::constant_time_secret_eq(token, candidate) {
6695 return Ok(MetricsScope::All);
6696 }
6697 if api_auth.configured() {
6698 return match auth::authenticate_with(
6699 api_auth.keyring,
6700 api_auth.single_key.as_deref(),
6701 Some(candidate),
6702 ) {
6703 Ok(_) => Err(error_response(
6704 StatusCode::FORBIDDEN,
6705 "completion api keys do not authorize metrics while \
6706 MEMRA_METRICS_TOKEN is configured",
6707 "authentication_error",
6708 None,
6709 )),
6710 Err(why) => Err(authentication_error(why)),
6711 };
6712 }
6713 return Err(authentication_error(auth::AuthDenied::Unknown));
6714 }
6715 if api_auth.configured() {
6716 let tenant = authenticate(api_auth, headers)?;
6717 return Ok(if api_auth.keyring.is_some() {
6718 MetricsScope::Tenant(format!("t:{}", tenant.tenant))
6719 } else {
6720 MetricsScope::CompletionDomain
6724 });
6725 }
6726 Err(authentication_error(auth::AuthDenied::Unknown))
6727}
6728
6729fn lane_for_tenant(
6735 headers: &axum::http::HeaderMap,
6736 tenant: &auth::TenantCtx,
6737) -> Result<lanes::Lane, Response> {
6738 let requested = match headers.get("x-lane").map(|v| v.to_str().unwrap_or("?")) {
6739 None => None,
6740 Some(v) => Some(lanes::Lane::parse(v).ok_or_else(|| {
6745 error_response_coded(
6746 StatusCode::BAD_REQUEST,
6747 &format!("unknown x-lane {v:?}; expected one of interactive, judge, harvest"),
6748 "invalid_request_error",
6749 Some("x-lane"),
6750 Some("invalid_lane"),
6751 )
6752 })?),
6753 };
6754 match tenant.lane_class {
6755 auth::LaneClass::Interactive => Ok(requested.unwrap_or(lanes::Lane::Interactive)),
6756 auth::LaneClass::Batch => match requested {
6757 None => Ok(lanes::Lane::Harvest),
6758 Some(lanes::Lane::Interactive) => Err(error_response(
6759 StatusCode::FORBIDDEN,
6760 "this api key is batch-class: x-lane interactive is not permitted \
6761 (use judge or harvest)",
6762 "authentication_error",
6763 Some("x-lane"),
6764 )),
6765 Some(l) => Ok(l),
6766 },
6767 }
6768}
6769
6770fn tenant_namespace(
6774 tenant: &auth::TenantCtx,
6775 cache_salt: &Option<String>,
6776) -> Result<String, &'static str> {
6777 let keyring_configured = auth::global().is_some();
6778 let raw = validate_cache_namespace(cache_salt, keyring_configured)?;
6779 if keyring_configured {
6780 Ok(auth::scope_namespace(&tenant.tenant, &raw))
6781 } else {
6782 Ok(raw)
6783 }
6784}
6785
6786fn meter_admit(env: &Envelope, tenant: &auth::TenantCtx, model: &str, lane: lanes::Lane) {
6791 eprintln!(
6792 "[meter] admit id={} tenant={} lane={} model={:?}",
6793 env.id,
6794 tenant.tenant,
6795 lane.as_str(),
6796 model
6797 );
6798}
6799
6800fn apply_model_request_limits(
6801 request: &mut Request,
6802 metadata: Option<&OpenRouterModelMetadata>,
6803 caps: Option<&ModelCaps>,
6804) -> Result<(), (String, &'static str)> {
6805 let Some(metadata) = metadata else {
6806 return Ok(());
6807 };
6808 let max_prompt = metadata
6809 .max_prompt_length
6810 .map(usize::try_from)
6811 .transpose()
6812 .map_err(|_| {
6813 (
6814 "configured model prompt limit does not fit this platform".into(),
6815 "model",
6816 )
6817 })?;
6818 let max_output = metadata
6819 .max_output_length
6820 .map(usize::try_from)
6821 .transpose()
6822 .map_err(|_| {
6823 (
6824 "configured model output limit does not fit this platform".into(),
6825 "model",
6826 )
6827 })?;
6828
6829 request.max_prompt_tokens = max_prompt;
6830 if let Some(max_output) = max_output {
6831 if request.params.max_new == worker::MAX_NEW_CTX_BOUNDED {
6832 request.params.max_new = metadata
6833 .default_output_length
6834 .map(usize::try_from)
6835 .transpose()
6836 .map_err(|_| {
6837 (
6838 "configured default output length does not fit this platform".into(),
6839 "model",
6840 )
6841 })?
6842 .unwrap_or(max_output);
6843 } else if request.params.max_new > max_output {
6844 return Err((
6845 format!(
6846 "max_tokens {} exceeds configured model maximum {max_output}",
6847 request.params.max_new
6848 ),
6849 "max_tokens",
6850 ));
6851 }
6852 }
6853
6854 if let (Some(max_prompt), Some(max_output), Some(requested_ctx)) =
6858 (max_prompt, max_output, request.params.max_ctx)
6859 {
6860 let operational_ctx = max_prompt
6861 .checked_add(max_output)
6862 .and_then(|value| value.checked_add(8))
6863 .ok_or_else(|| {
6864 (
6865 "configured model context envelope overflowed".into(),
6866 "model",
6867 )
6868 })?;
6869 let operational_ctx = caps
6870 .map(|caps| caps.context_length)
6871 .filter(|&context| context > 0)
6872 .map_or(operational_ctx, |context| operational_ctx.min(context));
6873 if requested_ctx > operational_ctx {
6874 return Err((
6875 format!(
6876 "max_ctx {requested_ctx} exceeds configured model envelope {operational_ctx}"
6877 ),
6878 "max_ctx",
6879 ));
6880 }
6881 }
6882 Ok(())
6883}
6884
6885#[allow(clippy::too_many_arguments)]
6886fn start_request_receipt(
6887 st: &AppState,
6888 env: &Envelope,
6889 tenant: &auth::TenantCtx,
6890 model: &str,
6891 route: &'static str,
6892 lane: lanes::Lane,
6893 stream: bool,
6894 budget_permit: Option<metering::Permit>,
6895) -> Option<Box<dyn metering::Receipt>> {
6896 st.metering.as_ref().map(|accounting| {
6897 accounting.open(
6898 &metering::RequestMeta {
6899 request_id: &env.id,
6900 tenant: &tenant.tenant,
6901 principal: tenant.key_prefix.as_deref(),
6902 model,
6903 route,
6904 lane: lane.as_str(),
6905 stream,
6906 },
6907 budget_permit,
6908 )
6909 })
6910}
6911
6912fn arm_capture(
6918 mut receipt: Option<Box<dyn metering::Receipt>>,
6919 prompt: impl FnOnce() -> serde_json::Value,
6920) -> Option<Box<dyn metering::Receipt>> {
6921 if let Some(receipt) = receipt.as_mut()
6922 && receipt.wants_capture()
6923 {
6924 receipt.arm_capture(prompt());
6925 }
6926 receipt
6927}
6928
6929fn capture_chat_messages(messages: &[ChatMessage]) -> serde_json::Value {
6933 serde_json::Value::Array(
6934 messages
6935 .iter()
6936 .map(|message| {
6937 let mut row = json!({ "role": message.role, "content": message.content });
6938 if !message.tool_calls.is_empty() {
6939 row["tool_calls"] = serde_json::Value::Array(
6940 message
6941 .tool_calls
6942 .iter()
6943 .map(|call| {
6944 json!({
6945 "id": call.id,
6946 "function": {
6947 "name": call.function.name,
6948 "arguments": call.function.arguments,
6949 },
6950 })
6951 })
6952 .collect(),
6953 );
6954 }
6955 row
6956 })
6957 .collect(),
6958 )
6959}
6960
6961enum BudgetRejection {
6962 Invalid(String),
6963 Insufficient,
6964 Unenrolled,
6965 PrincipalCapped,
6968 Unavailable(String),
6969}
6970
6971impl BudgetRejection {
6972 fn into_response(self) -> (Response, &'static str) {
6973 match self {
6974 Self::Invalid(message) => (bad_request(&message, Some("prompt")), "invalid_request"),
6975 Self::Insufficient => (
6976 error_response_coded(
6977 StatusCode::PAYMENT_REQUIRED,
6978 "tenant prepaid balance is insufficient for this request",
6979 "insufficient_balance",
6980 None,
6981 Some("insufficient_balance"),
6982 ),
6983 "insufficient_balance",
6984 ),
6985 Self::Unenrolled => (
6986 error_response_coded(
6987 StatusCode::PAYMENT_REQUIRED,
6988 "tenant is not enrolled for prepaid billing",
6989 "tenant_not_enrolled",
6990 None,
6991 Some("tenant_not_enrolled"),
6992 ),
6993 "tenant_not_enrolled",
6994 ),
6995 Self::PrincipalCapped => (
6996 error_response_coded(
6997 StatusCode::PAYMENT_REQUIRED,
6998 "this API key's spend cap is reached; raise or clear the key's cap to continue",
6999 "key_spend_cap_reached",
7000 None,
7001 Some("key_spend_cap_reached"),
7002 ),
7003 "key_spend_cap_reached",
7004 ),
7005 Self::Unavailable(err) => {
7006 eprintln!("[budget] ERROR: admission unavailable: {err}");
7007 (
7008 error_response_coded(
7009 StatusCode::SERVICE_UNAVAILABLE,
7010 "tenant budget accounting is unavailable",
7011 "server_error",
7012 None,
7013 Some("tenant_budget_unavailable"),
7014 ),
7015 "tenant_budget_unavailable",
7016 )
7017 }
7018 }
7019 }
7020}
7021
7022fn prepare_budget_prompt(
7023 request: &mut Request,
7024 tokenizer: Option<&Tokenizer>,
7025) -> Result<usize, String> {
7026 if let Some(error) = worker::prompt_source_limit_error(request) {
7027 return Err(error);
7028 }
7029 if request.prepared_prompt.is_none() {
7030 if let Some(trace) = request.ttft.as_ref() {
7031 trace.mark_tokenize_start();
7032 }
7033 let prompt = if !request.prompt_ids.is_empty() {
7034 request.prompt_ids.clone()
7035 } else if !request.chat_turns.is_empty() {
7036 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7037 let plain = worker::plain_chat_render_path(
7044 &request.tools_json,
7045 &request.think,
7046 request.reasoning_effort.as_deref(),
7047 &request.chat_turns,
7048 tokenizer.has_qwen_effort_ladder(),
7049 );
7050 let rendered = if plain {
7051 let messages: Vec<_> = request
7052 .chat_turns
7053 .iter()
7054 .map(|turn| (turn.role.as_str(), turn.content.as_str()))
7055 .collect();
7056 tokenizer.apply_chat_template(&messages, true)
7057 } else {
7058 tokenizer
7059 .apply_chat_template_tools_ex(
7060 &request.chat_turns,
7061 true,
7062 &request.tools_json,
7063 &request.tools_struct,
7064 request.think,
7065 request.reasoning_effort.as_deref(),
7066 )
7067 .map_err(|err| format!("chat template: {err}"))?
7068 };
7069 tokenizer.encode(&rendered, true)
7070 } else if request.chat {
7071 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7072 let rendered =
7073 tokenizer.apply_chat_template(&[("user", request.prompt_text.as_str())], true);
7074 tokenizer.encode(&rendered, true)
7075 } else {
7076 let tokenizer = tokenizer.ok_or("reservation tokenizer is unavailable")?;
7077 tokenizer.encode(&request.prompt_text, true)
7078 };
7079 if prompt.is_empty() {
7080 return Err("empty prompt after tokenization".into());
7081 }
7082 if let Some(trace) = request.ttft.as_ref() {
7083 trace.mark_tokenize_end(prompt.len());
7084 }
7085 request.prepared_prompt = Some(prompt);
7086 }
7087 let prompt_tokens = request
7088 .prepared_prompt
7089 .as_ref()
7090 .expect("budget prompt was prepared")
7091 .len();
7092 if let Some(limit) = request.max_prompt_tokens
7093 && prompt_tokens > limit
7094 {
7095 return Err(format!(
7096 "prompt ({prompt_tokens} tok) exceeds configured model maximum ({limit})"
7097 ));
7098 }
7099 Ok(prompt_tokens)
7100}
7101
7102fn budget_completion_bound(
7103 request: &Request,
7104 prompt_tokens: usize,
7105 caps: Option<&ModelCaps>,
7106) -> Result<usize, String> {
7107 let max_new = request.params.max_new;
7108 let requested_ctx = match (request.params.max_ctx, max_new) {
7109 (Some(cap), _) => cap,
7110 (None, worker::MAX_NEW_CTX_BOUNDED) => {
7111 let server_ctx = std::env::var("MEMRA_CTX")
7112 .ok()
7113 .and_then(|value| value.parse().ok())
7114 .unwrap_or(8192usize);
7115 let mut cap = server_ctx;
7116 if prompt_tokens.saturating_add(16) > cap {
7117 cap = prompt_tokens.saturating_add(server_ctx);
7118 }
7119 cap
7120 }
7121 (None, max_new) => prompt_tokens
7122 .checked_add(max_new)
7123 .and_then(|value| value.checked_add(8))
7124 .ok_or_else(|| "request context bound overflowed".to_string())?,
7125 };
7126 let ctx_cap = caps
7127 .map(|caps| caps.context_length)
7128 .filter(|&context| context > 0)
7129 .map_or(requested_ctx, |context| requested_ctx.min(context));
7130 if prompt_tokens >= ctx_cap {
7131 return Err(format!(
7132 "prompt ({prompt_tokens} tok) >= context cap ({ctx_cap})"
7133 ));
7134 }
7135 Ok(max_new.min(ctx_cap - prompt_tokens))
7136}
7137
7138fn admit_tenant_budget(
7139 st: &AppState,
7140 tenant: &auth::TenantCtx,
7141 request: &mut Request,
7142) -> Result<Option<metering::Permit>, BudgetRejection> {
7143 let Some(accounting) = st.metering.as_ref().filter(|m| m.enforces_limits()) else {
7144 return Ok(None);
7145 };
7146 match accounting.is_limited(&tenant.tenant) {
7147 Ok(false) => return Err(BudgetRejection::Unenrolled),
7148 Ok(true) => {}
7149 Err(metering::AdmitError::Unavailable(err)) => {
7150 return Err(BudgetRejection::Unavailable(err));
7151 }
7152 Err(other) => {
7153 return Err(BudgetRejection::Unavailable(format!(
7154 "unexpected budget enrollment result: {other:?}"
7155 )));
7156 }
7157 }
7158 let tokenizer = st
7159 .budget_tokenizers
7160 .as_ref()
7161 .and_then(|tokenizers| tokenizers.get(&request.model))
7162 .map(Arc::as_ref);
7163 if request.prompt_ids.is_empty() && tokenizer.is_none() {
7164 return Err(BudgetRejection::Unavailable(format!(
7165 "no reservation tokenizer for model {:?}",
7166 request.model
7167 )));
7168 }
7169 let prompt_tokens =
7170 prepare_budget_prompt(request, tokenizer).map_err(BudgetRejection::Invalid)?;
7171 let completion_tokens =
7172 budget_completion_bound(request, prompt_tokens, st.caps.get(&request.model))
7173 .map_err(BudgetRejection::Invalid)?;
7174 let prompt_tokens = u64::try_from(prompt_tokens)
7175 .map_err(|_| BudgetRejection::Unavailable("prompt token count exceeds u64".into()))?;
7176 let completion_tokens = u64::try_from(completion_tokens)
7177 .map_err(|_| BudgetRejection::Unavailable("completion token bound exceeds u64".into()))?;
7178 match accounting.reserve(
7179 &tenant.tenant,
7180 tenant.key_prefix.as_deref(),
7181 &request.model,
7182 prompt_tokens,
7183 completion_tokens,
7184 ) {
7185 Ok(permit) => Ok(permit),
7186 Err(metering::AdmitError::Insufficient) => Err(BudgetRejection::Insufficient),
7187 Err(metering::AdmitError::PrincipalCapped) => Err(BudgetRejection::PrincipalCapped),
7188 Err(metering::AdmitError::Blocked) => Err(BudgetRejection::Insufficient),
7192 Err(metering::AdmitError::Unenrolled) => Err(BudgetRejection::Unenrolled),
7193 Err(metering::AdmitError::Unavailable(err)) => Err(BudgetRejection::Unavailable(err)),
7194 }
7195}
7196
7197fn request_ledger_error_response() -> Response {
7198 error_response_coded(
7199 StatusCode::INTERNAL_SERVER_ERROR,
7200 "request completion could not be committed to the billing ledger",
7201 "server_error",
7202 None,
7203 Some("request_ledger_unavailable"),
7204 )
7205}
7206
7207fn request_ledger_error_body() -> serde_json::Value {
7208 error_body(
7209 "request completion could not be committed to the billing ledger",
7210 "server_error",
7211 None,
7212 Some("request_ledger_unavailable"),
7213 )
7214}
7215
7216fn ledger_rejected(
7217 mut receipt: Option<Box<dyn metering::Receipt>>,
7218 response: Response,
7219 error_code: &str,
7220 request_id: &str,
7221) -> Response {
7222 let status = response.status().as_u16();
7223 if let Some(receipt) = receipt.as_mut()
7224 && let Err(err) = receipt.reject(status, error_code)
7225 {
7226 eprintln!("[ledger] ERROR: request {request_id} rejection receipt failed: {err}");
7227 return with_request_id(request_id, request_ledger_error_response());
7228 }
7229 with_request_id(request_id, response)
7230}
7231
7232fn ledger_unbilled(
7237 mut receipt: Option<Box<dyn metering::Receipt>>,
7238 response: Response,
7239 outcome: &'static str,
7240 error_code: &str,
7241 request_id: &str,
7242) -> Response {
7243 let status = response.status().as_u16();
7244 if let Some(receipt) = receipt.as_mut()
7245 && let Err(err) = receipt.settle_unbilled(outcome, status, error_code)
7246 {
7247 eprintln!("[ledger] ERROR: request {request_id} {outcome} receipt failed: {err}");
7248 return with_request_id(request_id, request_ledger_error_response());
7249 }
7250 with_request_id(request_id, response)
7251}
7252
7253fn engine_error_code(class: worker::ErrClass) -> &'static str {
7254 use worker::ErrClass as C;
7255 match class {
7256 C::InvalidRequest => "invalid_request",
7257 C::ContextLength => "context_length_exceeded",
7258 C::ModelNotFound => "model_not_found",
7259 C::RateLimit => "rate_limit_exceeded",
7260 C::Overloaded => "overloaded",
7261 C::Engine => "engine_error",
7262 }
7263}
7264
7265fn model_not_found_response(models: &[String], requested: &str) -> Response {
7285 error_response_coded(
7286 StatusCode::BAD_REQUEST,
7287 &format!("unknown model {requested:?}; loaded: {models:?}"),
7288 "invalid_request_error",
7289 Some("model"),
7290 Some("model_not_found"),
7291 )
7292}
7293
7294fn validate_prompt_ids(ids: &[u32], caps: Option<&ModelCaps>) -> Result<(), String> {
7303 let Some(n_vocab) = caps.map(|c| c.n_vocab).filter(|&n| n > 0) else {
7304 return Ok(());
7305 };
7306 if let Some((pos, &id)) = ids
7307 .iter()
7308 .enumerate()
7309 .find(|&(_, &id)| id as usize >= n_vocab)
7310 {
7311 return Err(format!(
7312 "prompt_ids[{pos}] = {id} is out of vocabulary (model vocab size {n_vocab})"
7313 ));
7314 }
7315 Ok(())
7316}
7317
7318#[cfg(test)]
7319mod prompt_ids_tests {
7320 use super::*;
7321
7322 #[test]
7323 fn prompt_ids_are_bounded_by_the_model_vocab_at_intake() {
7324 let caps = ModelCaps {
7325 n_vocab: 8,
7326 ..Default::default()
7327 };
7328 assert!(validate_prompt_ids(&[0, 3, 7], Some(&caps)).is_ok());
7330 assert!(validate_prompt_ids(&[], Some(&caps)).is_ok());
7331 let err = validate_prompt_ids(&[1, 8, 2], Some(&caps)).unwrap_err();
7333 assert!(err.contains("prompt_ids[1] = 8"), "{err}");
7334 assert!(err.contains("vocab size 8"), "{err}");
7335 let err = validate_prompt_ids(&[u32::MAX], Some(&caps)).unwrap_err();
7336 assert!(err.contains("4294967295"), "{err}");
7337 let unknown = ModelCaps::default();
7339 assert!(validate_prompt_ids(&[u32::MAX], Some(&unknown)).is_ok());
7340 assert!(validate_prompt_ids(&[u32::MAX], None).is_ok());
7341 }
7342}
7343
7344fn canonical_model_id(models: &[String], requested: &str) -> Option<String> {
7345 if models.iter().any(|m| m == requested) {
7346 return Some(requested.to_string());
7347 }
7348 if requested.is_empty() || requested.contains('/') {
7349 return None;
7350 }
7351 let mut matches = models.iter().filter(|m| {
7352 m.rsplit('/')
7353 .next()
7354 .is_some_and(|suffix| suffix == requested)
7355 });
7356 match (matches.next(), matches.next()) {
7357 (Some(only), None) => Some(only.clone()),
7358 _ => None,
7359 }
7360}
7361
7362async fn completions(
7363 State(st): State<AppState>,
7364 headers: axum::http::HeaderMap,
7365 trace: Option<Extension<TtftRequestTrace>>,
7366 Json(mut req): Json<CompletionReq>,
7367) -> Response {
7368 let env = Envelope::new(false);
7369 match canonical_model_id(&st.models, &req.model) {
7370 Some(canonical) => req.model = canonical,
7371 None => {
7372 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7373 }
7374 }
7375 let ttft = trace.and_then(|Extension(trace)| trace.0);
7378 if let Some(trace) = ttft.as_ref() {
7379 trace.mark_parsed();
7380 trace.bind_request(&env.id, &req.model);
7381 }
7382 let tenant = match authenticate(&st.api_auth, &headers) {
7383 Ok(t) => t,
7384 Err(resp) => return with_request_id(&env.id, resp),
7385 };
7386 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7387 Ok(ns) => ns,
7388 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7389 };
7390 if let Err((msg, param)) = reject_unsupported(&[
7392 (
7393 "logit_bias",
7394 req.logit_bias.is_some(),
7395 " (device-side sampling has no bias hook yet)",
7396 ),
7397 ("logprobs", req.logprobs.is_some(), ""),
7398 (
7399 "n",
7400 req.n.is_some_and(|n| n != 1),
7401 " for n != 1 (single choice only)",
7402 ),
7403 (
7404 "best_of",
7405 req.best_of.is_some_and(|n| n != 1),
7406 " (single choice only)",
7407 ),
7408 ]) {
7409 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7410 }
7411 if let Err(msg) = validate_prompt_ids(&req.prompt_ids, st.caps.get(&req.model)) {
7414 return with_request_id(&env.id, bad_request(&msg, Some("prompt_ids")));
7415 }
7416 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7419 Ok(ms) => RequestDeadline::starting_now(ms),
7420 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7421 };
7422 let lane = match lane_for_tenant(&headers, &tenant) {
7423 Ok(l) => l,
7424 Err(resp) => return resp,
7425 };
7426 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7427 let model = req.model.clone();
7428 let stream = req.stream;
7429 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7430 let mut request = build_request_with_trace(
7431 &req,
7432 tx,
7433 lane,
7434 affinity,
7435 ttft.clone(),
7436 st.sampling_defaults(&model).for_mode(ThinkMode::Default),
7440 );
7441 request.cache_ns = cache_ns;
7442 if let Err((message, param)) = apply_model_request_limits(
7443 &mut request,
7444 st.openrouter_metadata.get(&model),
7445 st.caps.get(&model),
7446 ) {
7447 return with_request_id(&env.id, bad_request(&message, Some(param)));
7448 }
7449 if let Err(msg) = nonstream_deadline_gate(
7454 &request,
7455 req.stream,
7456 deadline,
7457 req.max_tokens.is_some(),
7458 st.budget_tokenizers
7459 .as_ref()
7460 .and_then(|t| t.get(&req.model))
7461 .map(Arc::as_ref),
7462 ) {
7463 return with_request_id(
7464 &env.id,
7465 error_response_coded(
7466 StatusCode::BAD_REQUEST,
7467 &msg,
7468 "invalid_request_error",
7469 Some("max_tokens"),
7470 Some("nonstream_deadline_infeasible"),
7471 ),
7472 );
7473 }
7474 if draining() {
7477 let receipt = start_request_receipt(
7478 &st,
7479 &env,
7480 &tenant,
7481 &req.model,
7482 "/v1/completions",
7483 lane,
7484 req.stream,
7485 None,
7486 );
7487 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7488 }
7489 let budget_permit = match admit_tenant_budget(&st, &tenant, &mut request) {
7490 Ok(permit) => permit,
7491 Err(rejection) => {
7492 let (response, error_code) = rejection.into_response();
7493 let receipt = start_request_receipt(
7494 &st,
7495 &env,
7496 &tenant,
7497 &req.model,
7498 "/v1/completions",
7499 lane,
7500 req.stream,
7501 None,
7502 );
7503 return ledger_rejected(receipt, response, error_code, &env.id);
7504 }
7505 };
7506 let receipt = start_request_receipt(
7507 &st,
7508 &env,
7509 &tenant,
7510 &req.model,
7511 "/v1/completions",
7512 lane,
7513 req.stream,
7514 budget_permit,
7515 );
7516 let receipt = arm_capture(receipt, || json!({ "prompt": req.prompt }));
7517 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7520 Ok(slot) => slot,
7521 Err(resp) => {
7522 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7523 }
7524 };
7525 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7528 Ok(guard) => guard,
7529 Err((resp, outcome)) => {
7530 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7531 }
7532 };
7533 meter_admit(&env, &tenant, &model, lane);
7534 let stop_strings = request.stop_strings.clone();
7535
7536 if let Some(trace) = ttft.as_ref() {
7541 trace.mark_submitted();
7542 }
7543 if st.cmd_tx.send(Cmd::Generate(Box::new(request))).is_err() {
7544 drop(pending_admit);
7545 return ledger_rejected(
7546 receipt,
7547 rl.attach(worker_unavailable_response()),
7548 "worker_unavailable",
7549 &env.id,
7550 );
7551 }
7552 pending_admit.commit();
7553 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7557 Ok(Ok(rx)) => rx,
7558 Ok(Err((resp, error_code))) => {
7559 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7560 }
7561 Err(_) => {
7562 return ledger_unbilled(
7563 receipt,
7564 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7565 "deadline_exceeded",
7566 "deadline_exceeded",
7567 &env.id,
7568 );
7569 }
7570 };
7571
7572 let resp = if stream {
7573 let rx = match peek_first_token(rx, deadline).await {
7577 Ok(rx) => rx,
7578 Err(()) => {
7579 return ledger_unbilled(
7580 receipt,
7581 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7582 "deadline_exceeded",
7583 "deadline_exceeded",
7584 &env.id,
7585 );
7586 }
7587 };
7588 sse_response_with_receipt(
7589 rx,
7590 model,
7591 false,
7592 None,
7593 env.clone(),
7594 stop_strings,
7595 Some(guard),
7596 receipt,
7597 )
7598 .into_response()
7599 } else {
7600 let mut receipt = receipt;
7606 let resp = blocking_response_with_receipt(
7607 rx,
7608 model,
7609 false,
7610 stop_strings,
7611 None,
7612 env.clone(),
7613 &mut receipt,
7614 Some(deadline),
7615 )
7616 .await;
7617 drop(guard); resp.into_response()
7619 };
7620 rl.attach(with_request_id(&env.id, resp))
7621}
7622
7623async fn chat_completions(
7624 State(st): State<AppState>,
7625 headers: axum::http::HeaderMap,
7626 trace: Option<Extension<TtftRequestTrace>>,
7627 Json(mut req): Json<ChatCompletionReq>,
7628) -> Response {
7629 let env = Envelope::new(true);
7630 match canonical_model_id(&st.models, &req.model) {
7635 Some(canonical) => req.model = canonical,
7636 None => {
7637 return with_request_id(&env.id, model_not_found_response(&st.models, &req.model));
7638 }
7639 }
7640 let ttft = trace.and_then(|Extension(trace)| trace.0);
7641 if let Some(trace) = ttft.as_ref() {
7642 trace.mark_parsed();
7643 trace.bind_request(&env.id, &req.model);
7644 }
7645 let tenant = match authenticate(&st.api_auth, &headers) {
7646 Ok(t) => t,
7647 Err(resp) => return with_request_id(&env.id, resp),
7648 };
7649 let cache_ns = match tenant_namespace(&tenant, &req.cache_salt) {
7650 Ok(ns) => ns,
7651 Err(msg) => return with_request_id(&env.id, bad_request(msg, Some("cache_salt"))),
7652 };
7653 if req.messages.is_empty()
7654 || req.messages.iter().any(|message| {
7655 !matches!(
7656 message.role.as_str(),
7657 "system" | "developer" | "user" | "assistant" | "tool"
7658 )
7659 })
7660 {
7661 return with_request_id(
7662 &env.id,
7663 bad_request(
7664 "messages must use system/developer/user/assistant/tool roles",
7665 Some("messages"),
7666 ),
7667 );
7668 }
7669 if let Err((msg, param)) = reject_unsupported(&[
7674 (
7675 "logit_bias",
7676 req.logit_bias.is_some(),
7677 " (device-side sampling has no bias hook yet)",
7678 ),
7679 (
7680 "logprobs",
7681 req.logprobs
7682 .as_ref()
7683 .is_some_and(|v| v.as_bool() != Some(false)),
7684 "",
7685 ),
7686 ("top_logprobs", req.top_logprobs.is_some(), ""),
7687 (
7688 "n",
7689 req.n.is_some_and(|n| n != 1),
7690 " for n != 1 (single choice only)",
7691 ),
7692 ]) {
7693 return with_request_id(&env.id, bad_request(&msg, Some(¶m)));
7694 }
7695 let deadline = match parse_timeout_ms(req.timeout_ms.as_ref()) {
7698 Ok(ms) => RequestDeadline::starting_now(ms),
7699 Err(msg) => return with_request_id(&env.id, bad_request(&msg, Some("timeout_ms"))),
7700 };
7701 let lane = match lane_for_tenant(&headers, &tenant) {
7702 Ok(l) => l,
7703 Err(resp) => return resp,
7704 };
7705 let model = req.model.clone();
7706 let stream = req.stream;
7707 let capture_prompt = st
7710 .metering
7711 .as_ref()
7712 .filter(|m| m.captures(&tenant.tenant))
7713 .map(|_| capture_chat_messages(&req.messages));
7714 let declared_max_tokens = req.max_tokens.is_some();
7718 let vision_preprocess_permit = if request_has_vision(&req) {
7722 match VISION_PREPROCESS_SEMAPHORE.acquire().await {
7723 Ok(permit) => Some(permit),
7724 Err(_) => {
7725 return with_request_id(
7726 &env.id,
7727 error_response(
7728 StatusCode::SERVICE_UNAVAILABLE,
7729 "vision preprocessing is unavailable",
7730 "server_error",
7731 None,
7732 ),
7733 );
7734 }
7735 }
7736 } else {
7737 None
7738 };
7739 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
7740 let affinity = affinity_key(&req.session_id, &req.user, &headers);
7741 let mut plan = match build_chat_request_with_trace(
7742 req,
7743 st.caps.get(&model),
7744 tx,
7745 lane,
7746 affinity,
7747 ttft.clone(),
7748 st.openrouter_metadata
7749 .get(&model)
7750 .and_then(|m| m.default_reasoning_effort.as_deref()),
7751 &st.sampling_defaults(&model),
7752 ) {
7753 Ok(plan) => plan,
7754 Err(err) => {
7755 return with_request_id(&env.id, bad_request(&err, None));
7756 }
7757 };
7758 plan.request.cache_ns = cache_ns;
7759 if let Err((message, param)) = apply_model_request_limits(
7760 &mut plan.request,
7761 st.openrouter_metadata.get(&model),
7762 st.caps.get(&model),
7763 ) {
7764 return with_request_id(&env.id, bad_request(&message, Some(param)));
7765 }
7766 if let Err(msg) = nonstream_deadline_gate(
7769 &plan.request,
7770 stream,
7771 deadline,
7772 declared_max_tokens,
7773 st.budget_tokenizers
7774 .as_ref()
7775 .and_then(|t| t.get(&model))
7776 .map(Arc::as_ref),
7777 ) {
7778 return with_request_id(
7779 &env.id,
7780 error_response_coded(
7781 StatusCode::BAD_REQUEST,
7782 &msg,
7783 "invalid_request_error",
7784 Some("max_tokens"),
7785 Some("nonstream_deadline_infeasible"),
7786 ),
7787 );
7788 }
7789 plan.vision_memory = match reserve_vision_memory(&plan) {
7790 Ok(permit) => permit,
7791 Err(err) => {
7792 return with_request_id(&env.id, vision_memory_error_response(err, Some("messages")));
7793 }
7794 };
7795 if draining() {
7798 let receipt = start_request_receipt(
7799 &st,
7800 &env,
7801 &tenant,
7802 &model,
7803 "/v1/chat/completions",
7804 lane,
7805 stream,
7806 None,
7807 );
7808 return ledger_rejected(receipt, drain_response(), "draining", &env.id);
7809 }
7810 let budget_permit = match admit_tenant_budget(&st, &tenant, &mut plan.request) {
7811 Ok(permit) => permit,
7812 Err(rejection) => {
7813 let (response, error_code) = rejection.into_response();
7814 let receipt = start_request_receipt(
7815 &st,
7816 &env,
7817 &tenant,
7818 &model,
7819 "/v1/chat/completions",
7820 lane,
7821 stream,
7822 None,
7823 );
7824 return ledger_rejected(receipt, response, error_code, &env.id);
7825 }
7826 };
7827 let receipt = start_request_receipt(
7828 &st,
7829 &env,
7830 &tenant,
7831 &model,
7832 "/v1/chat/completions",
7833 lane,
7834 stream,
7835 budget_permit,
7836 );
7837 let receipt = if let Some(prompt) = capture_prompt {
7838 arm_capture(receipt, move || prompt)
7839 } else {
7840 receipt
7841 };
7842 let (guard, rl) = match acquire_request_slot(&st, lane, &tenant, &env) {
7846 Ok(slot) => slot,
7847 Err(resp) => {
7848 return ledger_rejected(receipt, resp, "rate_limit_exceeded", &env.id);
7849 }
7850 };
7851 let pending_admit = match reserve_pending_admit(&st, lane, &rl, deadline) {
7854 Ok(guard) => guard,
7855 Err((resp, outcome)) => {
7856 return ledger_unbilled(receipt, rl.attach(resp), outcome, outcome, &env.id);
7857 }
7858 };
7859 if let Err(err) = decode_pending_vision(&mut plan) {
7864 return ledger_rejected(
7865 receipt,
7866 rl.attach(bad_request(&err, Some("messages"))),
7867 "invalid_request_error",
7868 &env.id,
7869 );
7870 }
7871 plan.request.vision_memory = plan.vision_memory.take();
7872 drop(vision_preprocess_permit);
7873 let constraint_ready = if plan.request.grammar.is_some() {
7874 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
7875 plan.request.constraint_ready = Some(ready_tx);
7876 Some(ready_rx)
7877 } else {
7878 None
7879 };
7880 meter_admit(&env, &tenant, &model, lane);
7881 let stop_strings = plan.request.stop_strings.clone();
7882 if let Some(trace) = ttft.as_ref() {
7884 trace.mark_submitted();
7885 }
7886 if st
7887 .cmd_tx
7888 .send(Cmd::Generate(Box::new(plan.request)))
7889 .is_err()
7890 {
7891 drop(pending_admit);
7892 return ledger_rejected(
7893 receipt,
7894 rl.attach(worker_unavailable_response()),
7895 "worker_unavailable",
7896 &env.id,
7897 );
7898 }
7899 pending_admit.commit();
7900 if let Some(ready) = constraint_ready {
7906 let bound = constrained::CONSTRAINT_COMPILE_TIMEOUT.min(deadline.remaining());
7907 match tokio::time::timeout(bound, ready).await {
7908 Ok(Ok(Ok(()))) => {}
7909 Ok(Ok(Err(err))) => {
7910 return ledger_rejected(
7911 receipt,
7912 rl.attach(engine_error_response(&err)),
7913 engine_error_code(err.class),
7914 &env.id,
7915 );
7916 }
7917 Ok(Err(_)) => {
7918 return ledger_rejected(
7919 receipt,
7920 rl.attach(worker_unavailable_response()),
7921 "worker_unavailable",
7922 &env.id,
7923 );
7924 }
7925 Err(_) if deadline.remaining().is_zero() => {
7926 return ledger_unbilled(
7927 receipt,
7928 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7929 "deadline_exceeded",
7930 "deadline_exceeded",
7931 &env.id,
7932 );
7933 }
7934 Err(_) => {
7935 return ledger_rejected(
7936 receipt,
7937 rl.attach(engine_error_response(&worker::constraint_timeout_error())),
7938 "constraint_compile_timeout",
7939 &env.id,
7940 );
7941 }
7942 }
7943 }
7944 let rx = match tokio::time::timeout_at(deadline.at, peek_admission(rx)).await {
7946 Ok(Ok(rx)) => rx,
7947 Ok(Err((resp, error_code))) => {
7948 return ledger_rejected(receipt, rl.attach(resp), error_code, &env.id);
7949 }
7950 Err(_) => {
7951 return ledger_unbilled(
7952 receipt,
7953 rl.attach(deadline_exceeded_response(deadline.ms, stream)),
7954 "deadline_exceeded",
7955 "deadline_exceeded",
7956 &env.id,
7957 );
7958 }
7959 };
7960 let resp = if stream {
7961 let rx = match peek_first_token(rx, deadline).await {
7963 Ok(rx) => rx,
7964 Err(()) => {
7965 return ledger_unbilled(
7966 receipt,
7967 rl.attach(deadline_exceeded_response(deadline.ms, true)),
7968 "deadline_exceeded",
7969 "deadline_exceeded",
7970 &env.id,
7971 );
7972 }
7973 };
7974 sse_response_with_receipt(
7975 rx,
7976 model,
7977 true,
7978 plan.parser,
7979 env.clone(),
7980 stop_strings,
7981 Some(guard),
7982 receipt,
7983 )
7984 .into_response()
7985 } else {
7986 let mut receipt = receipt;
7989 let resp = blocking_response_with_receipt(
7990 rx,
7991 model,
7992 true,
7993 stop_strings,
7994 plan.parser,
7995 env.clone(),
7996 &mut receipt,
7997 Some(deadline),
7998 )
7999 .await;
8000 drop(guard); resp.into_response()
8002 };
8003 rl.attach(with_request_id(&env.id, resp))
8004}
8005
8006#[cfg(test)]
8015fn sse_response(
8016 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8017 model: String,
8018 chat: bool,
8019 parser: Option<ToolStreamParser>,
8020 env: Envelope,
8021 stop_strings: Vec<String>,
8022 guard: Option<InflightGuard>,
8023) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8024 sse_response_with_receipt(rx, model, chat, parser, env, stop_strings, guard, None)
8025}
8026
8027fn sse_response_with_receipt(
8028 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8029 model: String,
8030 chat: bool,
8031 mut parser: Option<ToolStreamParser>,
8032 env: Envelope,
8033 stop_strings: Vec<String>,
8034 guard: Option<InflightGuard>,
8035 mut receipt: Option<Box<dyn metering::Receipt>>,
8036) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
8037 let mut scrub = (!stop_strings.is_empty() && (chat || openai_compat()))
8041 .then(|| StopScrubber::new(stop_strings));
8042 let stream = async_stream::stream! {
8043 let _guard = guard;
8046 let mut call_index: usize = 0;
8047 let mut role_sent = false;
8050 macro_rules! chat_chunk {
8051 ($delta:expr, $finish:expr) => {{
8052 let mut delta = $delta;
8053 if chat && !role_sent {
8054 role_sent = true;
8055 delta["role"] = json!("assistant");
8056 }
8057 env.stamp(json!({ "object": "chat.completion.chunk", "model": model,
8058 "choices": [{ "index": 0, "delta": delta,
8059 "finish_reason": $finish }] }))
8060 .to_string()
8061 }};
8062 }
8063 macro_rules! piece_chunks {
8065 ($piece:expr) => {{
8066 let mut payloads: Vec<String> = Vec::new();
8067 match $piece {
8068 Piece::Content(text) => {
8069 let text = match scrub.as_mut() {
8070 Some(sc) => sc.push(&text),
8071 None => text,
8072 };
8073 if !text.is_empty() {
8074 payloads.push(chat_chunk!(json!({ "content": text }),
8075 serde_json::Value::Null));
8076 }
8077 }
8078 Piece::Reasoning(text) => payloads.push(
8082 chat_chunk!(json!({ "reasoning": text }), serde_json::Value::Null)),
8083 Piece::Call(call) => {
8084 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8085 "index": call_index, "id": call.id, "type": "function",
8086 "function": { "name": call.name, "arguments": "" } }] }),
8087 serde_json::Value::Null));
8088 payloads.push(chat_chunk!(json!({ "tool_calls": [{
8089 "index": call_index,
8090 "function": { "arguments": call.arguments } }] }),
8091 serde_json::Value::Null));
8092 call_index += 1;
8093 }
8094 }
8095 payloads
8096 }};
8097 }
8098 let mut terminal = false;
8102 while let Some(ev) = rx.recv().await {
8103 match ev {
8104 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8106 if let Some(receipt) = receipt.as_mut()
8107 && let Err(err) = receipt.record_prompt_usage(
8108 n_prompt as u64,
8109 n_cached as u64,
8110 )
8111 {
8112 eprintln!(
8113 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8114 env.id
8115 );
8116 let _ = receipt.reject(500, "request_ledger_unavailable");
8119 let payload = request_ledger_error_body().to_string();
8120 if chat || openai_compat() {
8121 yield Ok(SseEvent::default().data(payload));
8122 yield Ok(SseEvent::default().data("[DONE]"));
8123 } else {
8124 yield Ok(SseEvent::default().event("error").data(payload));
8125 }
8126 terminal = true;
8127 break;
8128 }
8129 }
8130 Event::Token { id, text } => {
8131 if let Some(receipt) = receipt.as_mut()
8132 && let Err(err) = receipt.record_completion_token()
8133 {
8134 eprintln!(
8135 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8136 env.id
8137 );
8138 let _ = receipt.reject(500, "request_ledger_unavailable");
8139 let payload = request_ledger_error_body().to_string();
8140 if chat || openai_compat() {
8141 yield Ok(SseEvent::default().data(payload));
8142 yield Ok(SseEvent::default().data("[DONE]"));
8143 } else {
8144 yield Ok(SseEvent::default().event("error").data(payload));
8145 }
8146 terminal = true;
8147 break;
8148 }
8149 if let Some(receipt) = receipt.as_mut() {
8152 receipt.capture_completion_delta(&text);
8153 }
8154 if let Some(p) = parser.as_mut() {
8155 for piece in p.push(&text) {
8156 for payload in piece_chunks!(piece) {
8157 yield Ok(SseEvent::default().data(payload));
8158 }
8159 }
8160 continue;
8161 }
8162 let text = match scrub.as_mut() {
8163 Some(sc) => sc.push(&text),
8164 None => text,
8165 };
8166 if text.is_empty() && scrub.is_some() {
8167 continue; }
8169 let payload = if chat {
8170 chat_chunk!(json!({ "content": text }), serde_json::Value::Null)
8171 } else if openai_compat() {
8172 env.stamp(json!({ "object": "text_completion", "model": model,
8173 "choices": [{ "index": 0, "text": text, "finish_reason": null }] }))
8174 .to_string()
8175 } else {
8176 json!({ "model": model, "id": id, "text": text }).to_string()
8177 };
8178 yield Ok(SseEvent::default().data(payload));
8179 }
8180 Event::TokenSnapshot(_) => {}
8184 Event::Done { stop_reason, n_tokens, n_prompt, n_cached, elapsed_s, spec } => {
8185 let mut finish = stop_reason_to_finish(&stop_reason);
8186 if let Some(p) = parser.as_mut() {
8187 for piece in p.finish() {
8188 for payload in piece_chunks!(piece) {
8189 yield Ok(SseEvent::default().data(payload));
8190 }
8191 }
8192 if p.n_calls() > 0 { finish = "tool_calls"; }
8193 }
8194 if let Some(sc) = scrub.as_mut() {
8196 let tail = sc.finish();
8197 if !tail.is_empty() {
8198 let payload = if chat {
8199 chat_chunk!(json!({ "content": tail }),
8200 serde_json::Value::Null)
8201 } else {
8202 env.stamp(json!({ "object": "text_completion",
8203 "model": model,
8204 "choices": [{ "index": 0, "text": tail,
8205 "finish_reason": null }] })).to_string()
8206 };
8207 yield Ok(SseEvent::default().data(payload));
8208 }
8209 }
8210 if let Some(receipt) = receipt.as_mut()
8211 && let Err(err) = receipt.complete(
8212 metering::UsageCounts {
8213 prompt_tokens: n_prompt as u64,
8214 cached_prompt_tokens: n_cached as u64,
8215 completion_tokens: n_tokens as u64,
8216 },
8217 elapsed_s,
8218 )
8219 {
8220 eprintln!(
8221 "[ledger] ERROR: request {} completion receipt failed: {err}",
8222 env.id
8223 );
8224 let _ = receipt.reject(500, "request_ledger_unavailable");
8228 let payload = request_ledger_error_body().to_string();
8229 if chat || openai_compat() {
8230 yield Ok(SseEvent::default().data(payload));
8231 yield Ok(SseEvent::default().data("[DONE]"));
8232 } else {
8233 yield Ok(SseEvent::default().event("error").data(payload));
8234 }
8235 terminal = true;
8236 break;
8237 }
8238 if chat || openai_compat() {
8239 let usage = usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec);
8240 let fin = if chat {
8241 let mut v = env.stamp(json!({
8242 "object": "chat.completion.chunk", "model": model,
8243 "choices": [{ "index": 0, "delta": {},
8244 "finish_reason": finish }],
8245 "usage": usage }));
8246 if !role_sent {
8248 v["choices"][0]["delta"]["role"] = json!("assistant");
8249 }
8250 v
8251 } else {
8252 env.stamp(json!({ "object": "text_completion", "model": model,
8253 "choices": [{ "index": 0, "text": "",
8254 "finish_reason": finish }],
8255 "usage": usage }))
8256 }.to_string();
8257 yield Ok(SseEvent::default().data(fin));
8258 yield Ok(SseEvent::default().data("[DONE]"));
8259 } else {
8260 let payload = json!({
8261 "stop_reason": stop_reason, "n_tokens": n_tokens,
8262 "prompt_tokens": n_prompt, "cached_tokens": n_cached,
8263 "elapsed_s": elapsed_s
8264 }).to_string();
8265 yield Ok(SseEvent::default().event("done").data(payload));
8266 }
8267 terminal = true;
8268 break;
8269 }
8270 Event::Error(err) => {
8271 let ledger_error = if let Some(receipt) = receipt.as_mut() {
8281 receipt
8282 .reject(class_http(err.class).0.as_u16(), engine_error_code(err.class))
8283 .err()
8284 } else {
8285 None
8286 };
8287 if let Some(ref ledger_error) = ledger_error {
8288 eprintln!(
8289 "[ledger] ERROR: request {} failure receipt failed: {ledger_error}",
8290 env.id
8291 );
8292 }
8293 let payload = if ledger_error.is_some() {
8294 request_ledger_error_body().to_string()
8295 } else {
8296 engine_error_body(&err).to_string()
8297 };
8298 if chat || openai_compat() {
8299 yield Ok(SseEvent::default().data(payload));
8302 yield Ok(SseEvent::default().data("[DONE]"));
8303 } else {
8304 yield Ok(SseEvent::default().event("error").data(payload));
8307 }
8308 terminal = true;
8309 break;
8310 }
8311 }
8312 }
8313 if !terminal {
8314 let e = worker::EngineError::overloaded(
8320 "worker closed the stream without completing (worker restart in progress)",
8321 );
8322 if let Some(receipt) = receipt.as_mut()
8323 && let Err(ledger_err) = receipt.reject(
8324 class_http(e.class).0.as_u16(),
8325 engine_error_code(e.class),
8326 )
8327 {
8328 eprintln!(
8329 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8330 env.id
8331 );
8332 }
8333 let payload = engine_error_body(&e).to_string();
8334 if chat || openai_compat() {
8335 yield Ok(SseEvent::default().data(payload));
8336 yield Ok(SseEvent::default().data("[DONE]"));
8337 } else {
8338 yield Ok(SseEvent::default().event("error").data(payload));
8339 }
8340 }
8341 };
8342 Sse::new(stream).keep_alive(
8343 axum::response::sse::KeepAlive::new().interval(std::time::Duration::from_secs(5)),
8346 )
8347}
8348
8349fn truncate_at_stop(text: &mut String, stop_strings: &[String]) {
8351 if let Some(offset) = stop_strings.iter().filter_map(|stop| text.find(stop)).min() {
8352 text.truncate(offset);
8353 }
8354}
8355
8356fn partial_stop_suffix(s: &str, tag: &str) -> usize {
8359 let mut best = 0;
8360 for (k, _) in tag.char_indices().skip(1) {
8361 if k <= s.len() && s.ends_with(&tag[..k]) {
8362 best = k;
8363 }
8364 }
8365 best
8366}
8367
8368struct StopScrubber {
8374 stops: Vec<String>,
8375 buf: String,
8376 done: bool,
8377}
8378
8379impl StopScrubber {
8380 fn new(stops: Vec<String>) -> Self {
8381 Self {
8382 stops,
8383 buf: String::new(),
8384 done: false,
8385 }
8386 }
8387
8388 fn push(&mut self, text: &str) -> String {
8390 if self.done {
8391 return String::new();
8392 }
8393 self.buf.push_str(text);
8394 if let Some(i) = self
8395 .stops
8396 .iter()
8397 .filter_map(|s| self.buf.find(s.as_str()))
8398 .min()
8399 {
8400 self.done = true;
8401 let out = self.buf[..i].to_string();
8402 self.buf.clear();
8403 return out;
8404 }
8405 let keep = self
8406 .stops
8407 .iter()
8408 .map(|s| partial_stop_suffix(&self.buf, s))
8409 .max()
8410 .unwrap_or(0);
8411 let emit_to = self.buf.len() - keep;
8412 let out = self.buf[..emit_to].to_string();
8413 self.buf.drain(..emit_to);
8414 out
8415 }
8416
8417 fn finish(&mut self) -> String {
8419 if self.done {
8420 self.buf.clear();
8421 return String::new();
8422 }
8423 std::mem::take(&mut self.buf)
8424 }
8425}
8426
8427#[cfg(test)]
8428async fn blocking_response(
8429 rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8430 model: String,
8431 chat: bool,
8432 stop_strings: Vec<String>,
8433 parser: Option<ToolStreamParser>,
8434 env: Envelope,
8435) -> Response {
8436 blocking_response_with_receipt(rx, model, chat, stop_strings, parser, env, &mut None, None)
8437 .await
8438}
8439
8440struct BlockingPayload<'a> {
8444 env: &'a Envelope,
8445 model: String,
8446 chat: bool,
8447 finish: &'static str,
8448 text: String,
8449 reasoning: String,
8450 calls: Vec<ParsedToolCall>,
8451 tokens: Vec<u32>,
8452 stop_reason: String,
8453 n_prompt: usize,
8454 n_tokens: usize,
8455 n_cached: usize,
8456 elapsed_s: f64,
8457 spec: Option<worker::SpecUsage>,
8458 deadline_error: Option<serde_json::Value>,
8464}
8465
8466fn blocking_payload(p: BlockingPayload<'_>) -> Response {
8467 let BlockingPayload {
8468 env,
8469 model,
8470 chat,
8471 finish,
8472 text,
8473 reasoning,
8474 calls,
8475 tokens,
8476 stop_reason,
8477 n_prompt,
8478 n_tokens,
8479 n_cached,
8480 elapsed_s,
8481 spec,
8482 deadline_error,
8483 } = p;
8484 if chat {
8485 let content = if !calls.is_empty() && text.is_empty() {
8487 serde_json::Value::Null
8488 } else {
8489 serde_json::Value::String(text)
8490 };
8491 let mut message = json!({ "role": "assistant", "content": content });
8492 if !reasoning.is_empty() {
8495 message["reasoning"] = json!(reasoning);
8496 message["reasoning_details"] = json!([{
8497 "type": "reasoning.text", "text": reasoning }]);
8498 }
8499 if !calls.is_empty() {
8500 message["tool_calls"] =
8501 serde_json::Value::Array(calls.iter().map(tool_call_json).collect());
8502 }
8503 let mut body = json!({
8504 "object": "chat.completion", "model": model,
8505 "choices": [{ "index": 0,
8506 "message": message,
8507 "finish_reason": finish }],
8508 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8509 });
8510 if let Some(err) = deadline_error {
8511 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8512 body["error"] = err;
8513 }
8514 return Json(env.stamp(body)).into_response();
8515 }
8516 if openai_compat() {
8517 let mut body = json!({
8518 "object": "text_completion", "model": model,
8519 "choices": [{ "index": 0, "text": text,
8520 "finish_reason": finish }],
8521 "usage": usage_json(n_prompt, n_tokens, n_cached, elapsed_s, spec)
8522 });
8523 if let Some(err) = deadline_error {
8524 body["choices"][0]["native_finish_reason"] = json!("deadline_exceeded");
8525 body["error"] = err;
8526 }
8527 return Json(env.stamp(body)).into_response();
8528 }
8529 Json(CompletionResp {
8530 model,
8531 text,
8532 tokens,
8533 stop_reason,
8534 error: deadline_error,
8535 n_tokens,
8536 prompt_tokens: n_prompt,
8537 cached_tokens: n_cached,
8538 elapsed_s,
8539 })
8540 .into_response()
8541}
8542
8543async fn blocking_response_with_receipt(
8568 mut rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
8569 model: String,
8570 chat: bool,
8571 stop_strings: Vec<String>,
8572 mut parser: Option<ToolStreamParser>,
8573 env: Envelope,
8574 receipt: &mut Option<Box<dyn metering::Receipt>>,
8575 deadline: Option<RequestDeadline>,
8576) -> Response {
8577 let mut text = String::new();
8578 let mut reasoning = String::new();
8579 let mut tokens: Vec<u32> = Vec::new();
8580 let mut calls: Vec<ParsedToolCall> = Vec::new();
8581 let consume = |pieces: Vec<Piece>,
8582 text: &mut String,
8583 reasoning: &mut String,
8584 calls: &mut Vec<ParsedToolCall>| {
8585 for piece in pieces {
8586 match piece {
8587 Piece::Content(t) => text.push_str(&t),
8588 Piece::Reasoning(t) => reasoning.push_str(&t),
8589 Piece::Call(c) => calls.push(c),
8590 }
8591 }
8592 };
8593 let started = std::time::Instant::now();
8595 let mut seen_prompt: usize = 0;
8596 let mut seen_cached: usize = 0;
8597 let mut seen_tokens: usize = 0;
8598 loop {
8599 let ev = match deadline {
8600 Some(d) => tokio::select! {
8601 biased;
8602 ev = rx.recv() => ev,
8603 () = tokio::time::sleep_until(d.at) => {
8604 drop(rx);
8607 if seen_tokens == 0 {
8608 if let Some(receipt) = receipt.as_mut()
8612 && let Err(err) = receipt.settle_unbilled(
8613 "deadline_exceeded",
8614 StatusCode::REQUEST_TIMEOUT.as_u16(),
8615 "deadline_exceeded",
8616 )
8617 {
8618 eprintln!(
8619 "[ledger] ERROR: request {} deadline receipt failed: {err}",
8620 env.id
8621 );
8622 return request_ledger_error_response();
8623 }
8624 return deadline_exceeded_response(d.ms, false);
8625 }
8626 if let Some(p) = parser.as_mut() {
8627 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8628 }
8629 truncate_at_stop(&mut text, &stop_strings);
8630 let elapsed_s = started.elapsed().as_secs_f64();
8631 if let Some(receipt) = receipt.as_mut()
8634 && let Err(err) = receipt.complete_deadline_partial(
8635 metering::UsageCounts {
8636 prompt_tokens: seen_prompt as u64,
8637 cached_prompt_tokens: seen_cached as u64,
8638 completion_tokens: seen_tokens as u64,
8639 },
8640 elapsed_s,
8641 )
8642 {
8643 eprintln!(
8644 "[ledger] ERROR: request {} partial-deadline receipt failed: {err}",
8645 env.id
8646 );
8647 let _ = receipt.reject(500, "request_ledger_unavailable");
8648 return request_ledger_error_response();
8649 }
8650 eprintln!(
8651 "[deadline] request {} delivered PARTIAL: {} tokens in {:.1}s of a \
8652 {} ms deadline (prompt {}); non-streaming caller advised to stream",
8653 env.id, seen_tokens, elapsed_s, d.ms, seen_prompt
8654 );
8655 let err_obj = json!({
8656 "message": format!(
8657 "deadline of {} ms (timeout_ms; default {}) elapsed mid-generation; \
8658 the {} tokens produced before the cut are delivered above and are \
8659 billed. Set \"stream\": true for work this long — a stream's \
8660 deadline bounds only the time to first token — or lower max_tokens.",
8661 d.ms, TIMEOUT_MS_DEFAULT, seen_tokens
8662 ),
8663 "code": "deadline_exceeded",
8664 "metadata": { "error_type": "timeout", "provider_name": "memra" }
8665 });
8666 return blocking_payload(BlockingPayload {
8667 env: &env,
8668 model,
8669 chat,
8670 finish: "error",
8671 text,
8672 reasoning,
8673 calls,
8674 tokens,
8675 stop_reason: "Deadline".to_string(),
8676 n_prompt: seen_prompt,
8677 n_tokens: seen_tokens,
8678 n_cached: seen_cached,
8679 elapsed_s,
8680 spec: None,
8681 deadline_error: Some(err_obj),
8682 });
8683 }
8684 },
8685 None => rx.recv().await,
8686 };
8687 let Some(ev) = ev else { break };
8688 match ev {
8689 Event::PromptCapture { .. } => {} Event::PromptUsage { n_prompt, n_cached } => {
8691 if let Some(receipt) = receipt.as_mut()
8692 && let Err(err) = receipt.record_prompt_usage(n_prompt as u64, n_cached as u64)
8693 {
8694 eprintln!(
8695 "[ledger] ERROR: request {} partial prompt receipt failed: {err}",
8696 env.id
8697 );
8698 let _ = receipt.reject(500, "request_ledger_unavailable");
8701 return request_ledger_error_response();
8702 }
8703 seen_prompt = n_prompt;
8704 seen_cached = n_cached;
8705 }
8706 Event::Token { id, text: delta } => {
8707 if let Some(receipt) = receipt.as_mut()
8708 && let Err(err) = receipt.record_completion_token()
8709 {
8710 eprintln!(
8711 "[ledger] ERROR: request {} partial completion receipt failed: {err}",
8712 env.id
8713 );
8714 let _ = receipt.reject(500, "request_ledger_unavailable");
8715 return request_ledger_error_response();
8716 }
8717 if let Some(receipt) = receipt.as_mut() {
8719 receipt.capture_completion_delta(&delta);
8720 }
8721 tokens.push(id);
8722 seen_tokens += 1;
8723 match parser.as_mut() {
8724 Some(p) => consume(p.push(&delta), &mut text, &mut reasoning, &mut calls),
8725 None => text.push_str(&delta),
8726 }
8727 }
8728 Event::TokenSnapshot(ids) => tokens = ids,
8729 Event::Done {
8730 stop_reason,
8731 n_tokens,
8732 n_prompt,
8733 n_cached,
8734 elapsed_s,
8735 spec,
8736 } => {
8737 if let Some(p) = parser.as_mut() {
8738 consume(p.finish(), &mut text, &mut reasoning, &mut calls);
8739 }
8740 truncate_at_stop(&mut text, &stop_strings);
8741 let finish = if calls.is_empty() {
8742 stop_reason_to_finish(&stop_reason)
8743 } else {
8744 "tool_calls"
8745 };
8746 if let Some(receipt) = receipt.as_mut()
8747 && let Err(err) = receipt.complete(
8748 metering::UsageCounts {
8749 prompt_tokens: n_prompt as u64,
8750 cached_prompt_tokens: n_cached as u64,
8751 completion_tokens: n_tokens as u64,
8752 },
8753 elapsed_s,
8754 )
8755 {
8756 eprintln!(
8757 "[ledger] ERROR: request {} completion receipt failed: {err}",
8758 env.id
8759 );
8760 let _ = receipt.reject(500, "request_ledger_unavailable");
8763 return request_ledger_error_response();
8764 }
8765 return blocking_payload(BlockingPayload {
8766 env: &env,
8767 model,
8768 chat,
8769 finish,
8770 text,
8771 reasoning,
8772 calls,
8773 tokens,
8774 stop_reason,
8775 n_prompt,
8776 n_tokens,
8777 n_cached,
8778 elapsed_s,
8779 spec,
8780 deadline_error: None,
8781 });
8782 }
8783 Event::Error(err) => {
8784 if let Some(receipt) = receipt.as_mut()
8788 && let Err(ledger_err) = receipt.reject(
8789 class_http(err.class).0.as_u16(),
8790 engine_error_code(err.class),
8791 )
8792 {
8793 eprintln!(
8794 "[ledger] ERROR: request {} failure receipt failed: {ledger_err}",
8795 env.id
8796 );
8797 return request_ledger_error_response();
8798 }
8799 return engine_error_response(&err);
8800 }
8801 }
8802 }
8803 let e = worker::EngineError::overloaded(
8808 "worker closed the stream without completing (worker restart in progress)",
8809 );
8810 if let Some(receipt) = receipt.as_mut()
8811 && let Err(ledger_err) =
8812 receipt.reject(class_http(e.class).0.as_u16(), engine_error_code(e.class))
8813 {
8814 eprintln!(
8815 "[ledger] ERROR: request {} closed-stream receipt failed: {ledger_err}",
8816 env.id
8817 );
8818 return request_ledger_error_response();
8819 }
8820 engine_error_response(&e)
8821}
8822
8823#[cfg(test)]
8824mod tests {
8825 use super::*;
8826
8827 #[derive(Debug, Clone, PartialEq)]
8835 enum MeterEvent {
8836 Reserve {
8837 tenant: String,
8838 principal: Option<String>,
8839 model: String,
8840 },
8841 Open {
8842 request_id: String,
8843 tenant: String,
8844 model: String,
8845 route: &'static str,
8846 stream: bool,
8847 with_permit: bool,
8848 },
8849 PromptUsage {
8850 prompt: u64,
8851 cached: u64,
8852 },
8853 Token,
8854 CapturePrompt(serde_json::Value),
8855 CaptureDelta(String),
8856 Complete {
8857 prompt: u64,
8858 cached: u64,
8859 completion: u64,
8860 },
8861 DeadlinePartial {
8862 prompt: u64,
8863 cached: u64,
8864 completion: u64,
8865 },
8866 Reject {
8867 status: u16,
8868 code: String,
8869 },
8870 Unbilled {
8871 outcome: &'static str,
8872 status: u16,
8873 code: String,
8874 },
8875 Dropped {
8878 prompt: u64,
8879 cached: u64,
8880 completion: u64,
8881 },
8882 }
8883
8884 enum ReserveScript {
8887 Admit { with_permit: bool },
8888 Insufficient,
8889 Blocked,
8890 Unenrolled,
8891 PrincipalCapped,
8892 }
8893
8894 struct MockMetering {
8895 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
8896 limits: bool,
8897 limited: bool,
8898 reserve_script: std::sync::Mutex<std::collections::VecDeque<ReserveScript>>,
8899 captures: bool,
8900 }
8901
8902 impl MockMetering {
8903 fn admit_all() -> Arc<Self> {
8904 Arc::new(MockMetering {
8905 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8906 limits: false,
8907 limited: true,
8908 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
8909 captures: false,
8910 })
8911 }
8912
8913 fn with_limits(script: Vec<ReserveScript>) -> Arc<Self> {
8914 Arc::new(MockMetering {
8915 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8916 limits: true,
8917 limited: true,
8918 reserve_script: std::sync::Mutex::new(script.into()),
8919 captures: false,
8920 })
8921 }
8922
8923 fn capturing() -> Arc<Self> {
8924 Arc::new(MockMetering {
8925 events: Arc::new(std::sync::Mutex::new(Vec::new())),
8926 limits: false,
8927 limited: true,
8928 reserve_script: std::sync::Mutex::new(std::collections::VecDeque::new()),
8929 captures: true,
8930 })
8931 }
8932
8933 fn events(&self) -> Vec<MeterEvent> {
8934 self.events.lock().unwrap().clone()
8935 }
8936 }
8937
8938 impl metering::Metering for MockMetering {
8939 fn enforces_limits(&self) -> bool {
8940 self.limits
8941 }
8942
8943 fn is_limited(&self, _tenant: &str) -> Result<bool, metering::AdmitError> {
8944 Ok(self.limited)
8945 }
8946
8947 fn reserve(
8948 &self,
8949 tenant: &str,
8950 principal: Option<&str>,
8951 model: &str,
8952 _prompt_tokens: u64,
8953 _completion_bound: u64,
8954 ) -> Result<Option<metering::Permit>, metering::AdmitError> {
8955 self.events.lock().unwrap().push(MeterEvent::Reserve {
8956 tenant: tenant.into(),
8957 principal: principal.map(str::to_owned),
8958 model: model.into(),
8959 });
8960 match self.reserve_script.lock().unwrap().pop_front() {
8961 None | Some(ReserveScript::Admit { with_permit: false }) => Ok(None),
8962 Some(ReserveScript::Admit { with_permit: true }) => {
8963 Ok(Some(Box::new(()) as metering::Permit))
8964 }
8965 Some(ReserveScript::Insufficient) => Err(metering::AdmitError::Insufficient),
8966 Some(ReserveScript::Blocked) => Err(metering::AdmitError::Blocked),
8967 Some(ReserveScript::Unenrolled) => Err(metering::AdmitError::Unenrolled),
8968 Some(ReserveScript::PrincipalCapped) => Err(metering::AdmitError::PrincipalCapped),
8969 }
8970 }
8971
8972 fn open(
8973 &self,
8974 meta: &metering::RequestMeta<'_>,
8975 permit: Option<metering::Permit>,
8976 ) -> Box<dyn metering::Receipt> {
8977 self.events.lock().unwrap().push(MeterEvent::Open {
8978 request_id: meta.request_id.into(),
8979 tenant: meta.tenant.into(),
8980 model: meta.model.into(),
8981 route: meta.route,
8982 stream: meta.stream,
8983 with_permit: permit.is_some(),
8984 });
8985 Box::new(MockReceipt {
8986 events: self.events.clone(),
8987 wants_capture: self.captures,
8988 prompt: 0,
8989 cached: 0,
8990 completion: 0,
8991 finalized: false,
8992 })
8993 }
8994
8995 fn captures(&self, _tenant: &str) -> bool {
8996 self.captures
8997 }
8998
8999 fn limits_health(&self) -> Option<metering::LimitsHealth> {
9000 self.limits.then_some(metering::LimitsHealth {
9001 source_reload_failed: 0,
9002 source_reload_consecutive: 0,
9003 source_available: true,
9004 })
9005 }
9006 }
9007
9008 struct MockReceipt {
9009 events: Arc<std::sync::Mutex<Vec<MeterEvent>>>,
9010 wants_capture: bool,
9011 prompt: u64,
9012 cached: u64,
9013 completion: u64,
9014 finalized: bool,
9015 }
9016
9017 impl metering::Receipt for MockReceipt {
9018 fn wants_capture(&self) -> bool {
9019 self.wants_capture
9020 }
9021
9022 fn arm_capture(&mut self, prompt: serde_json::Value) {
9023 self.events
9024 .lock()
9025 .unwrap()
9026 .push(MeterEvent::CapturePrompt(prompt));
9027 }
9028
9029 fn capture_completion_delta(&mut self, text: &str) {
9030 if self.wants_capture {
9031 self.events
9032 .lock()
9033 .unwrap()
9034 .push(MeterEvent::CaptureDelta(text.into()));
9035 }
9036 }
9037
9038 fn record_prompt_usage(&mut self, prompt: u64, cached: u64) -> Result<(), String> {
9039 self.prompt = prompt;
9040 self.cached = cached;
9041 self.events
9042 .lock()
9043 .unwrap()
9044 .push(MeterEvent::PromptUsage { prompt, cached });
9045 Ok(())
9046 }
9047
9048 fn record_completion_token(&mut self) -> Result<(), String> {
9049 self.completion += 1;
9050 self.events.lock().unwrap().push(MeterEvent::Token);
9051 Ok(())
9052 }
9053
9054 fn complete(
9055 &mut self,
9056 usage: metering::UsageCounts,
9057 _worker_elapsed_s: f64,
9058 ) -> Result<(), String> {
9059 self.finalized = true;
9060 self.events.lock().unwrap().push(MeterEvent::Complete {
9061 prompt: usage.prompt_tokens,
9062 cached: usage.cached_prompt_tokens,
9063 completion: usage.completion_tokens,
9064 });
9065 Ok(())
9066 }
9067
9068 fn complete_deadline_partial(
9069 &mut self,
9070 usage: metering::UsageCounts,
9071 _worker_elapsed_s: f64,
9072 ) -> Result<(), String> {
9073 self.finalized = true;
9074 self.events
9075 .lock()
9076 .unwrap()
9077 .push(MeterEvent::DeadlinePartial {
9078 prompt: usage.prompt_tokens,
9079 cached: usage.cached_prompt_tokens,
9080 completion: usage.completion_tokens,
9081 });
9082 Ok(())
9083 }
9084
9085 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
9086 self.finalized = true;
9087 self.events.lock().unwrap().push(MeterEvent::Reject {
9088 status,
9089 code: error_code.into(),
9090 });
9091 Ok(())
9092 }
9093
9094 fn settle_unbilled(
9095 &mut self,
9096 outcome: &'static str,
9097 status: u16,
9098 error_code: &str,
9099 ) -> Result<(), String> {
9100 self.finalized = true;
9101 self.events.lock().unwrap().push(MeterEvent::Unbilled {
9102 outcome,
9103 status,
9104 code: error_code.into(),
9105 });
9106 Ok(())
9107 }
9108 }
9109
9110 impl Drop for MockReceipt {
9111 fn drop(&mut self) {
9112 if !self.finalized {
9113 self.events.lock().unwrap().push(MeterEvent::Dropped {
9114 prompt: self.prompt,
9115 cached: self.cached,
9116 completion: self.completion,
9117 });
9118 }
9119 }
9120 }
9121
9122 static GATE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
9127
9128 fn gate_request(max_new: usize, prompt_ids: usize) -> worker::Request {
9131 let req: CompletionReq = serde_json::from_value(json!({
9132 "model": "qwen/qwen3.8-27b",
9133 "prompt_ids": vec![7u32; prompt_ids],
9134 }))
9135 .unwrap();
9136 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9137 let mut request = build_request(&req, tx, lanes::Lane::Interactive, None);
9138 request.params.max_new = max_new;
9139 request
9140 }
9141
9142 #[test]
9149 fn the_feasibility_gate_boundary_matches_the_measured_ladder() {
9150 let prompt = 30_278u64;
9151 let deadline_ms = TIMEOUT_MS_DEFAULT;
9152 let margin = |max_new: u64| {
9153 let prefill_ms = prompt * 1_000 / PREFILL_FLOOR_TOK_S;
9154 let decode_ms = max_new * 1_000 / DECODE_FLOOR_TOK_S;
9155 (prefill_ms + decode_ms) <= deadline_ms * DEADLINE_INFEASIBLE_MARGIN_PCT / 100
9156 };
9157 for allowed in [64u64, 2048, 4096, 5120, 6144] {
9158 assert!(margin(allowed), "{allowed} measured OK and must be allowed");
9159 }
9160 for refused in [8192u64, 16384, 262_144] {
9161 assert!(
9162 !margin(refused),
9163 "{refused} measured as a 408 and must be refused"
9164 );
9165 }
9166 }
9167
9168 #[test]
9169 fn the_gate_names_a_max_tokens_that_actually_fits() {
9170 let fits = deadline_fitting_max_tokens(30_278, TIMEOUT_MS_DEFAULT).unwrap();
9173 assert!(
9174 fits > 0 && fits < 7_800,
9175 "advice {fits} must fit the measured ceiling"
9176 );
9177 assert_eq!(
9179 deadline_fitting_max_tokens(400_000, TIMEOUT_MS_DEFAULT),
9180 None
9181 );
9182 }
9183
9184 #[test]
9185 fn streaming_is_never_gated_and_the_gate_can_be_switched_off() {
9186 let req = gate_request(262_144, 30_000);
9187 let deadline = RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT);
9188 let err = nonstream_deadline_gate(&req, false, deadline, true, None).unwrap_err();
9190 assert!(
9191 err.contains("stream"),
9192 "message must name the streaming alternative: {err}"
9193 );
9194 assert!(
9195 err.contains("max_tokens"),
9196 "message must name the knob: {err}"
9197 );
9198 assert!(nonstream_deadline_gate(&req, true, deadline, true, None).is_ok());
9200 let _l = GATE_ENV_LOCK.lock().unwrap(); for off in ["0", "off", "false"] {
9207 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", off) };
9208 assert!(
9209 nonstream_deadline_gate(&req, false, deadline, true, None).is_ok(),
9210 "MEMRA_NONSTREAM_DEADLINE_GATE={off} must disable the gate"
9211 );
9212 }
9213 unsafe { std::env::set_var("MEMRA_NONSTREAM_DEADLINE_GATE", "1") };
9214 assert!(nonstream_deadline_gate(&req, false, deadline, true, None).is_err());
9215 unsafe { std::env::remove_var("MEMRA_NONSTREAM_DEADLINE_GATE") };
9216 assert!(
9217 nonstream_deadline_gate(&req, false, deadline, true, None).is_err(),
9218 "unset means ON (the documented default)"
9219 );
9220 }
9221
9222 #[test]
9229 fn the_feasibility_gate_is_wired_on_every_surface_not_just_the_two_i_remembered() {
9230 let strip = |src: &str| -> String {
9237 src.lines()
9238 .map(|line| match line.find("//") {
9239 Some(i) => line[..i].to_string(),
9240 None => line.to_string(),
9241 })
9242 .collect::<Vec<_>>()
9243 .join("\n")
9244 };
9245 fn body<'a>(src: &'a str, signature: &str) -> &'a str {
9247 let start = src
9248 .find(signature)
9249 .unwrap_or_else(|| panic!("{signature} not found — did the handler get renamed?"));
9250 let rest = &src[start + signature.len()..];
9251 let end = rest.find("\nasync fn ").unwrap_or(rest.len());
9252 let end = rest[..end].find("\npub(crate) async fn ").unwrap_or(end);
9253 &rest[..end]
9254 }
9255 let main_src = strip(include_str!("lib.rs"));
9256 let surfaces_src = strip(include_str!("surfaces.rs"));
9257 for (surface, src, signature) in [
9258 ("/v1/completions", &main_src, "async fn completions("),
9259 (
9260 "/v1/chat/completions",
9261 &main_src,
9262 "async fn chat_completions(",
9263 ),
9264 (
9265 "/v1/messages + /v1/responses (shared admission)",
9266 &surfaces_src,
9267 "pub(crate) async fn admit_translated(",
9268 ),
9269 ] {
9270 let handler = body(src, signature);
9271 assert!(
9272 handler.contains("nonstream_deadline_gate("),
9273 "{surface} must CALL the feasibility gate inside {signature}"
9274 );
9275 let limits = handler
9278 .find("apply_model_request_limits(")
9279 .unwrap_or_else(|| panic!("{surface}: no apply_model_request_limits call"));
9280 let gate = handler.find("nonstream_deadline_gate(").unwrap();
9281 assert!(
9282 limits < gate,
9283 "{surface}: the gate must run after apply_model_request_limits"
9284 );
9285 }
9286 }
9287
9288 #[test]
9292 fn the_native_shape_carries_the_deadline_error_and_omits_it_otherwise() {
9293 let err = json!({"code": "deadline_exceeded",
9294 "metadata": {"error_type": "timeout"}});
9295 let cut = CompletionResp {
9296 model: "m".into(),
9297 text: "partial".into(),
9298 tokens: vec![1, 2],
9299 stop_reason: "Deadline".into(),
9300 error: Some(err.clone()),
9301 n_tokens: 2,
9302 prompt_tokens: 9,
9303 cached_tokens: 0,
9304 elapsed_s: 1.0,
9305 };
9306 let v = serde_json::to_value(&cut).unwrap();
9307 assert_eq!(v["stop_reason"], "Deadline");
9308 assert_eq!(v["error"]["code"], "deadline_exceeded");
9309 assert_eq!(v["error"]["metadata"]["error_type"], "timeout");
9310 let whole = CompletionResp {
9312 error: None,
9313 stop_reason: "Eos".into(),
9314 ..cut
9315 };
9316 let v = serde_json::to_value(&whole).unwrap();
9317 assert!(
9318 v.get("error").is_none(),
9319 "a complete response must not grow an error key: {v}"
9320 );
9321 }
9322
9323 #[test]
9324 fn a_ctx_bounded_request_is_not_gated_because_context_is_its_only_limit() {
9325 let _l = GATE_ENV_LOCK.lock().unwrap();
9326 let req = gate_request(worker::MAX_NEW_CTX_BOUNDED, 30_000);
9330 assert!(
9331 nonstream_deadline_gate(
9332 &req,
9333 false,
9334 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9335 false,
9336 None,
9337 )
9338 .is_ok(),
9339 "an omitted max_tokens is never gated — context is its only limit"
9340 );
9341 let resolved = gate_request(32_768, 30_000);
9346 assert!(
9347 nonstream_deadline_gate(
9348 &resolved,
9349 false,
9350 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9351 false,
9352 None,
9353 )
9354 .is_ok(),
9355 "a resolved-but-undeclared cap is not the caller's number to be refused over"
9356 );
9357 assert!(
9359 nonstream_deadline_gate(
9360 &resolved,
9361 false,
9362 RequestDeadline::starting_now(TIMEOUT_MS_DEFAULT),
9363 true,
9364 None,
9365 )
9366 .is_err()
9367 );
9368 }
9369
9370 #[test]
9371 fn the_prompt_estimate_is_exact_for_ids_and_a_proxy_otherwise() {
9372 let req = gate_request(64, 1234);
9373 assert_eq!(prompt_tokens_estimate(&req, None), 1234, "ids are exact");
9374 let mut text = gate_request(64, 0);
9375 text.prompt_ids.clear();
9376 text.prompt_text = "x".repeat(6_000);
9377 assert_eq!(
9378 prompt_tokens_estimate(&text, None),
9379 1_000,
9380 "the fallback under-counts on purpose (bytes/6): an over-count refuses work \
9381 that would have succeeded"
9382 );
9383 }
9384
9385 #[test]
9386 fn vision_memory_reservation_is_bounded_and_released() {
9387 let permit = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES).unwrap();
9388 let Err(capacity) = try_reserve_vision_memory(1) else {
9389 panic!("a full process vision budget admitted another request");
9390 };
9391 assert!(matches!(capacity, VisionMemoryError::Capacity(_)));
9392 let response = vision_memory_error_response(capacity, Some("messages"));
9393 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
9394 assert_eq!(response.headers()["retry-after"], "5");
9395 assert_eq!(response.headers()["retry-after-ms"], "5000");
9396 drop(permit);
9397 assert!(try_reserve_vision_memory(1).is_ok());
9398 let Err(request) = try_reserve_vision_memory(MAX_VISION_PATCH_BYTES + 1) else {
9399 panic!("an over-limit vision request was admitted");
9400 };
9401 assert!(matches!(request, VisionMemoryError::Request(_)));
9402 let response = vision_memory_error_response(request, Some("messages"));
9403 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
9404 assert_eq!(response.headers()["x-should-retry"], "false");
9405 let _ = try_reserve_vision_memory(1);
9406 }
9407
9408 #[test]
9409 fn header_auth_gate_covers_only_inference_dialects() {
9410 for path in [
9411 "/v1/auth/check",
9412 "/v1/completions",
9413 "/v1/chat/completions",
9414 "/v1/messages",
9415 "/v1/responses",
9416 "/v1/embeddings",
9417 "/v1/rerank",
9418 ] {
9419 assert!(protected_inference_path(path), "{path}");
9420 }
9421 for path in ["/health", "/readyz", "/models", "/v1/models", "/metrics"] {
9422 assert!(!protected_inference_path(path), "{path}");
9423 }
9424 }
9425 #[tokio::test]
9432 async fn served_completion_capture_is_byte_exact_and_armed_receipts_only() {
9433 use crate::metering::Metering as _;
9434 let prompt = json!([{ "role": "user", "content": "capture me — exactly" }]);
9435
9436 let drive = |receipt: Option<Box<dyn metering::Receipt>>| async {
9437 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
9438 tx.send(Event::PromptUsage {
9439 n_prompt: 7,
9440 n_cached: 0,
9441 })
9442 .unwrap();
9443 tx.send(Event::Token {
9444 id: 1,
9445 text: "Hel".into(),
9446 })
9447 .unwrap();
9448 tx.send(Event::Token {
9449 id: 2,
9450 text: "lo".into(),
9451 })
9452 .unwrap();
9453 tx.send(Event::Done {
9454 stop_reason: "eos".into(),
9455 n_tokens: 2,
9456 n_prompt: 7,
9457 n_cached: 0,
9458 elapsed_s: 0.05,
9459 spec: None,
9460 })
9461 .unwrap();
9462 drop(tx);
9463 let mut receipt = receipt;
9464 blocking_response_with_receipt(
9465 rx,
9466 "m".into(),
9467 true,
9468 Vec::new(),
9469 None,
9470 Envelope::new(true),
9471 &mut receipt,
9472 None,
9473 )
9474 .await
9475 };
9476
9477 let plain = MockMetering::admit_all();
9479 let receipt = plain.open(
9480 &metering::RequestMeta {
9481 request_id: "cap-unmarked",
9482 tenant: "unmarked",
9483 principal: None,
9484 model: "m",
9485 route: "/v1/chat/completions",
9486 lane: "interactive",
9487 stream: false,
9488 },
9489 None,
9490 );
9491 let response = drive(Some(receipt)).await;
9492 assert_eq!(response.status(), StatusCode::OK);
9493 assert!(
9494 !plain.events().iter().any(|e| matches!(
9495 e,
9496 MeterEvent::CaptureDelta(_) | MeterEvent::CapturePrompt(_)
9497 )),
9498 "an unarmed receipt must see no capture traffic: {:?}",
9499 plain.events()
9500 );
9501
9502 let capturing = MockMetering::capturing();
9505 let mut receipt = capturing.open(
9506 &metering::RequestMeta {
9507 request_id: "cap-marked",
9508 tenant: "marked",
9509 principal: None,
9510 model: "m",
9511 route: "/v1/chat/completions",
9512 lane: "interactive",
9513 stream: false,
9514 },
9515 None,
9516 );
9517 assert!(receipt.wants_capture());
9518 receipt.arm_capture(prompt.clone());
9519 let response = drive(Some(receipt)).await;
9520 assert_eq!(response.status(), StatusCode::OK);
9521 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
9522 .await
9523 .unwrap();
9524 let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
9525 assert_eq!(body["choices"][0]["message"]["content"], "Hello");
9526
9527 let events = capturing.events();
9528 assert!(
9529 events.contains(&MeterEvent::CapturePrompt(prompt.clone())),
9530 "prompt must arm byte-exact: {events:?}"
9531 );
9532 let completion: String = events
9533 .iter()
9534 .filter_map(|e| match e {
9535 MeterEvent::CaptureDelta(text) => Some(text.as_str()),
9536 _ => None,
9537 })
9538 .collect();
9539 assert_eq!(
9540 completion, "Hello",
9541 "the deltas must reassemble the served completion byte-exact: {events:?}"
9542 );
9543 assert!(
9544 events.contains(&MeterEvent::Complete {
9545 prompt: 7,
9546 cached: 0,
9547 completion: 2,
9548 }),
9549 "worker-truth usage settles alongside the capture: {events:?}"
9550 );
9551 }
9552
9553 fn tool_caps() -> ModelCaps {
9554 ModelCaps {
9555 tools_branch: true,
9556 qwen_think: true,
9557 think_switch: true,
9558 chat_ok: true,
9559 ..Default::default()
9560 }
9561 }
9562
9563 fn ladder_caps() -> ModelCaps {
9568 ModelCaps {
9569 qwen_effort: true,
9570 ..tool_caps()
9571 }
9572 }
9573
9574 fn gemma_tool_caps() -> ModelCaps {
9575 ModelCaps {
9576 tools_branch: true,
9577 gemma_think: true,
9578 chat_ok: true,
9579 instruct_type: Some("gemma".into()),
9580 ..Default::default()
9581 }
9582 }
9583
9584 fn gemma_template(kind: &str) -> String {
9585 let file = match kind {
9586 "qat" => "qat-trunk-template.jinja",
9587 _ => "official-tooluse-template.jinja",
9588 };
9589 let path = format!(
9590 "{}/../../research/gemma4-tools-20260817/{file}",
9591 env!("CARGO_MANIFEST_DIR")
9592 );
9593 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9594 }
9595
9596 fn render_fixture(request: &serde_json::Value, template: &str) -> String {
9601 let tools_arr = request
9602 .get("tools")
9603 .and_then(|t| t.as_array())
9604 .cloned()
9605 .unwrap_or_default();
9606 let (tools_json, tools_struct, _schemas) = if tools_arr.is_empty() {
9607 (Vec::new(), Vec::new(), HashMap::new())
9608 } else {
9609 prepare_tools(&tools_arr).unwrap()
9610 };
9611 let effort = request
9612 .get("reasoning_effort")
9613 .and_then(|v| v.as_str())
9614 .map(String::from);
9615 let (think, _lvl, _explicit) =
9616 parse_think(&effort, &None, None, None, None, false).unwrap();
9617
9618 let mut turns: Vec<TmplTurn> = Vec::new();
9619 for msg in request["messages"].as_array().unwrap() {
9620 let role = msg["role"].as_str().unwrap();
9621 let role = if role == "developer" { "system" } else { role };
9622 let content =
9623 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9624 let tool_calls = msg
9625 .get("tool_calls")
9626 .and_then(|a| a.as_array())
9627 .map(|a| {
9628 a.iter()
9629 .map(|tc| {
9630 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9631 render_req_tool_call(&rtc).unwrap()
9632 })
9633 .collect()
9634 })
9635 .unwrap_or_default();
9636 let tool_responses = msg
9637 .get("tool_responses")
9638 .and_then(|a| a.as_array())
9639 .map(|a| {
9640 a.iter()
9641 .map(|tr| {
9642 (
9643 tr.get("name").and_then(|n| n.as_str()).unwrap().to_string(),
9644 json_to_val(&tr["response"]),
9645 )
9646 })
9647 .collect()
9648 })
9649 .unwrap_or_default();
9650 turns.push(TmplTurn {
9651 role: role.to_string(),
9652 content,
9653 tool_calls,
9654 reasoning: msg
9655 .get("reasoning")
9656 .and_then(|r| r.as_str())
9657 .map(String::from)
9658 .filter(|s| !s.is_empty()),
9659 tool_call_id: msg
9660 .get("tool_call_id")
9661 .and_then(|s| s.as_str())
9662 .map(String::from),
9663 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9664 tool_responses,
9665 task: None,
9666 tools: Vec::new(),
9667 });
9668 }
9669 chat::apply_chat_template_tools_ex(
9670 Some(template),
9671 &turns,
9672 true,
9673 &tools_json,
9674 &tools_struct,
9675 think,
9676 None,
9677 None,
9678 )
9679 .unwrap()
9680 }
9681
9682 #[test]
9686 fn gemma4_tools_fixtures_match_the_official_jinja() {
9687 let dir = format!(
9688 "{}/../../research/gemma4-tools-20260817/fixtures",
9689 env!("CARGO_MANIFEST_DIR")
9690 );
9691 let mut entries: Vec<_> = std::fs::read_dir(&dir)
9692 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
9693 .map(|e| e.unwrap().path())
9694 .filter(|p| p.is_dir())
9695 .collect();
9696 entries.sort();
9697 assert!(
9698 entries.len() >= 14,
9699 "expected >=14 fixtures, found {}",
9700 entries.len()
9701 );
9702 let (mut official, mut qat) = (0u32, 0u32);
9703 for d in entries {
9704 let input: serde_json::Value =
9705 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
9706 .unwrap();
9707 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
9708 let kind = input
9709 .get("template")
9710 .and_then(|t| t.as_str())
9711 .unwrap_or("official");
9712 match kind {
9713 "qat" => qat += 1,
9714 _ => official += 1,
9715 }
9716 let tmpl = gemma_template(kind);
9717 let got = render_fixture(&input["request"], &tmpl);
9718 assert_eq!(
9719 got, expected,
9720 "fixture {:?} diverged from the jinja oracle",
9721 d
9722 );
9723 }
9724 assert!(
9725 official >= 12 && qat >= 2,
9726 "coverage: {official} official, {qat} qat"
9727 );
9728 }
9729
9730 #[test]
9736 fn gemma4_tools_flow_through_build_chat_request() {
9737 let tmpl = gemma_template("official");
9738 for name in [
9739 "01-system-tools-basic",
9740 "04-single-call-cycle",
9741 "07-multi-cycle-agentic",
9742 ] {
9743 let path = format!(
9744 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/input.json",
9745 env!("CARGO_MANIFEST_DIR")
9746 );
9747 let input: serde_json::Value =
9748 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
9749 let expected_path = format!(
9750 "{}/../../research/gemma4-tools-20260817/fixtures/{name}/expected.txt",
9751 env!("CARGO_MANIFEST_DIR")
9752 );
9753 let expected = std::fs::read_to_string(&expected_path).unwrap();
9754 let req: ChatCompletionReq = serde_json::from_value(input["request"].clone()).unwrap();
9755 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
9756 let plan = build_chat_request(
9757 req,
9758 Some(&gemma_tool_caps()),
9759 tx,
9760 lanes::Lane::Interactive,
9761 None,
9762 )
9763 .unwrap();
9764 let got = chat::apply_chat_template_tools_ex(
9765 Some(&tmpl),
9766 &plan.request.chat_turns,
9767 true,
9768 &plan.request.tools_json,
9769 &plan.request.tools_struct,
9770 plan.request.think,
9771 plan.request.reasoning_effort.as_deref(),
9772 None,
9773 )
9774 .unwrap();
9775 assert_eq!(got, expected, "pipeline render diverged for {name}");
9776 }
9777 }
9778
9779 fn dsv4_sentinel() -> String {
9787 let path = format!(
9788 "{}/../../research/dsv4-template-20260818/dsv4-chat-template.sentinel.jinja",
9789 env!("CARGO_MANIFEST_DIR")
9790 );
9791 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
9792 }
9793
9794 fn dsv4_turn(msg: &serde_json::Value) -> TmplTurn {
9799 let role = msg["role"].as_str().unwrap().to_string();
9800 let content =
9801 content_to_text(msg.get("content").unwrap_or(&serde_json::Value::Null)).unwrap();
9802 let reasoning = msg
9803 .get("reasoning")
9804 .or_else(|| msg.get("reasoning_content"))
9805 .and_then(|r| r.as_str())
9806 .map(String::from)
9807 .filter(|s| !s.is_empty());
9808 let tool_calls = msg
9809 .get("tool_calls")
9810 .and_then(|a| a.as_array())
9811 .map(|a| {
9812 a.iter()
9813 .map(|tc| {
9814 let rtc: ReqToolCall = serde_json::from_value(tc.clone()).unwrap();
9815 render_req_tool_call(&rtc).unwrap()
9816 })
9817 .collect()
9818 })
9819 .unwrap_or_default();
9820 let tools = msg
9821 .get("tools")
9822 .and_then(|a| a.as_array())
9823 .map(|a| {
9824 a.iter()
9825 .filter_map(|t| t.get("function").map(json_to_val))
9826 .collect()
9827 })
9828 .unwrap_or_default();
9829 TmplTurn {
9830 role,
9831 content,
9832 tool_calls,
9833 reasoning,
9834 tool_call_id: msg
9835 .get("tool_call_id")
9836 .and_then(|s| s.as_str())
9837 .map(String::from),
9838 tool_name: msg.get("name").and_then(|s| s.as_str()).map(String::from),
9839 tool_responses: Vec::new(),
9840 task: msg.get("task").and_then(|s| s.as_str()).map(String::from),
9841 tools,
9842 }
9843 }
9844
9845 fn dsv4_req_tools(v: Option<&serde_json::Value>) -> Vec<chat::Val> {
9846 v.and_then(|t| t.as_array())
9847 .map(|a| {
9848 a.iter()
9849 .filter_map(|t| t.get("function").map(json_to_val))
9850 .collect()
9851 })
9852 .unwrap_or_default()
9853 }
9854
9855 fn dsv4_run_fixture_dir(subdir: &str, encoding: chat::Dsv4Encoding, min_fixtures: usize) {
9859 let dir = format!(
9860 "{}/../../research/dsv4-template-20260818/{subdir}",
9861 env!("CARGO_MANIFEST_DIR")
9862 );
9863 let tmpl = dsv4_sentinel();
9864 let mut entries: Vec<_> = std::fs::read_dir(&dir)
9865 .unwrap_or_else(|e| panic!("read fixtures dir {dir}: {e}"))
9866 .map(|e| e.unwrap().path())
9867 .filter(|p| p.is_dir())
9868 .collect();
9869 entries.sort();
9870 assert!(
9871 entries.len() >= min_fixtures,
9872 "expected >={min_fixtures} fixtures, found {}",
9873 entries.len()
9874 );
9875 for d in &entries {
9876 let input: serde_json::Value =
9877 serde_json::from_str(&std::fs::read_to_string(d.join("input.json")).unwrap())
9878 .unwrap();
9879 let expected = std::fs::read_to_string(d.join("expected.txt")).unwrap();
9880 let turns: Vec<TmplTurn> = input["turns"]
9881 .as_array()
9882 .unwrap()
9883 .iter()
9884 .map(dsv4_turn)
9885 .collect();
9886 let think = match input["think"].as_str().unwrap() {
9887 "chat" => ThinkMode::NoThink,
9888 _ => ThinkMode::Think,
9889 };
9890 let effort = input
9891 .get("reasoning_effort")
9892 .and_then(|v| v.as_str())
9893 .map(String::from);
9894 let req_tools = dsv4_req_tools(input.get("req_tools"));
9895 let agp = input["add_generation_prompt"].as_bool().unwrap_or(true);
9896 let got = chat::apply_chat_template_tools_ex(
9897 Some(&tmpl),
9898 &turns,
9899 agp,
9900 &[],
9901 &req_tools,
9902 think,
9903 effort.as_deref(),
9904 Some(encoding),
9905 )
9906 .unwrap();
9907 assert_eq!(got, expected, "fixture {:?} diverged from the oracle", d);
9908 }
9909 }
9910
9911 #[test]
9912 fn dsv4_template_fixtures_match_the_oracle() {
9913 dsv4_run_fixture_dir("fixtures", chat::Dsv4Encoding::Preview, 20);
9914 }
9915
9916 #[test]
9922 fn dsv4_0731_fixtures_match_the_oracle() {
9923 dsv4_run_fixture_dir("fixtures-0731", chat::Dsv4Encoding::V0731, 40);
9924 }
9925
9926 #[test]
9927 fn dsv4_artifact_fixtures_are_byte_identical() {
9928 let base = format!(
9932 "{}/../../research/dsv4-template-20260818/ref/artifact-encoding/tests",
9933 env!("CARGO_MANIFEST_DIR")
9934 );
9935 let tmpl = dsv4_sentinel();
9936 for (n, think) in [
9937 (1u32, ThinkMode::Think),
9938 (2, ThinkMode::Think),
9939 (3, ThinkMode::Think),
9940 (4, ThinkMode::NoThink),
9941 ] {
9942 let td: serde_json::Value = serde_json::from_str(
9943 &std::fs::read_to_string(format!("{base}/test_input_{n}.json")).unwrap(),
9944 )
9945 .unwrap();
9946 let (messages, tools) = if td.is_object() {
9947 (td["messages"].clone(), td.get("tools").cloned())
9948 } else {
9949 (td.clone(), None)
9950 };
9951 let mut turns: Vec<TmplTurn> = Vec::new();
9952 for (i, msg) in messages.as_array().unwrap().iter().enumerate() {
9953 let mut t = dsv4_turn(msg);
9954 if i == 0 {
9955 if let Some(tl) = &tools {
9956 t.tools = tl
9957 .as_array()
9958 .unwrap()
9959 .iter()
9960 .filter_map(|x| x.get("function").map(json_to_val))
9961 .collect();
9962 }
9963 }
9964 turns.push(t);
9965 }
9966 let expected = std::fs::read_to_string(format!("{base}/test_output_{n}.txt")).unwrap();
9967 for encoding in [chat::Dsv4Encoding::Preview, chat::Dsv4Encoding::V0731] {
9971 let got = chat::apply_chat_template_tools_ex(
9972 Some(&tmpl),
9973 &turns,
9974 true,
9975 &[],
9976 &[],
9977 think,
9978 None,
9979 Some(encoding),
9980 )
9981 .unwrap();
9982 assert_eq!(
9983 got, expected,
9984 "artifact fixture {n} diverged from the oracle under {encoding:?}"
9985 );
9986 }
9987 }
9988 }
9989
9990 #[test]
9991 fn dsv4_default_thinkmode_renders_thinking() {
9992 let tmpl = dsv4_sentinel();
9995 let turns = vec![TmplTurn {
9996 role: "user".into(),
9997 content: "Hi".into(),
9998 ..Default::default()
9999 }];
10000 let dflt = chat::apply_chat_template_tools_ex(
10001 Some(&tmpl),
10002 &turns,
10003 true,
10004 &[],
10005 &[],
10006 ThinkMode::Default,
10007 None,
10008 None,
10009 )
10010 .unwrap();
10011 let think = chat::apply_chat_template_tools_ex(
10012 Some(&tmpl),
10013 &turns,
10014 true,
10015 &[],
10016 &[],
10017 ThinkMode::Think,
10018 None,
10019 None,
10020 )
10021 .unwrap();
10022 assert_eq!(dflt, think);
10023 assert!(
10024 dflt.ends_with("<\u{ff5c}Assistant\u{ff5c}><think>"),
10025 "{dflt:?}"
10026 );
10027 let chat_mode = chat::apply_chat_template_tools_ex(
10028 Some(&tmpl),
10029 &turns,
10030 true,
10031 &[],
10032 &[],
10033 ThinkMode::NoThink,
10034 None,
10035 None,
10036 )
10037 .unwrap();
10038 assert!(
10039 chat_mode.ends_with("<\u{ff5c}Assistant\u{ff5c}></think>"),
10040 "{chat_mode:?}"
10041 );
10042 }
10043
10044 fn dsv4_run_tokenization_crosscheck(subdir: &str) {
10049 let base = format!(
10050 "{}/../../research/dsv4-template-20260818",
10051 env!("CARGO_MANIFEST_DIR")
10052 );
10053 let refdir = std::path::Path::new(&base).join("ref");
10054 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
10055 .expect("load dsv4 tokenizer from ref dir");
10056 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
10057 let banked: serde_json::Value = serde_json::from_str(
10058 &std::fs::read_to_string(format!("{base}/{subdir}/tokenization-crosscheck.json"))
10059 .unwrap(),
10060 )
10061 .unwrap();
10062 let obj = banked.as_object().unwrap();
10063 assert!(obj.len() >= 3, "expected >=3 cross-check fixtures");
10064 for (name, ids_v) in obj {
10065 let rendered =
10066 std::fs::read_to_string(format!("{base}/{subdir}/{name}/expected.txt")).unwrap();
10067 let want: Vec<u32> = ids_v
10068 .as_array()
10069 .unwrap()
10070 .iter()
10071 .map(|v| v.as_u64().unwrap() as u32)
10072 .collect();
10073 let got = tok.encode(&rendered, true);
10074 assert_eq!(got, want, "tokenization diverged for {name}");
10075 }
10076 }
10077
10078 #[test]
10079 fn dsv4_tokenization_crosscheck_matches_official_ids() {
10080 dsv4_run_tokenization_crosscheck("fixtures");
10081 }
10082
10083 #[test]
10087 fn dsv4_0731_tokenization_crosscheck_matches_official_ids() {
10088 dsv4_run_tokenization_crosscheck("fixtures-0731");
10089 }
10090
10091 #[test]
10092 fn dsv4_tool_result_long_runs_render_tokenize_roundtrip() {
10093 let base = format!(
10103 "{}/../../research/dsv4-template-20260818",
10104 env!("CARGO_MANIFEST_DIR")
10105 );
10106 let refdir = std::path::Path::new(&base).join("ref");
10107 let tok = memra_tokenizer::Tokenizer::from_hf_dir(&refdir)
10108 .expect("load dsv4 tokenizer from ref dir");
10109 assert_eq!(tok.pre(), "deepseek-v3", "pre-tokenizer family detection");
10110 let tmpl = dsv4_sentinel();
10111 let req_tools = dsv4_req_tools(Some(&serde_json::json!([
10112 {"type": "function", "function": {
10113 "name": "get_data",
10114 "description": "Fetch a blob",
10115 "parameters": {"type": "object", "properties": {"key": {"type": "string"}},
10116 "required": ["key"]}
10117 }}
10118 ])));
10119
10120 let cases: Vec<(&str, String)> = vec![
10121 ("ascii-letter-131k", "Z".repeat(131_072)), ("ascii-letter-1m", "Z".repeat(1_048_576)),
10123 ("space-131k", " ".repeat(131_072)),
10124 ("digit-131k", "7".repeat(131_072)),
10125 (
10126 "mixed-runs",
10127 format!(
10128 "{}{}{}{}",
10129 "Z".repeat(65_536),
10130 " ".repeat(65_536),
10131 "7".repeat(65_536),
10132 "\n".repeat(65_536)
10133 ),
10134 ),
10135 ("cjk-64k", "中".repeat(65_536)),
10136 ("accented-letter-64k", "é".repeat(65_536)),
10137 ];
10138 for (name, blob) in &cases {
10139 let msgs = serde_json::json!([
10140 {"role": "system", "content": "You are a tool-using assistant."},
10141 {"role": "user", "content": "Fetch the blob."},
10142 {"role": "assistant", "reasoning": "Use get_data.", "content": "",
10143 "tool_calls": [{"id": "call_001", "type": "function",
10144 "function": {"name": "get_data",
10145 "arguments": "{\"key\": \"blob\"}"}}]},
10146 {"role": "tool", "tool_call_id": "call_001", "content": blob}
10147 ]);
10148 let turns: Vec<TmplTurn> = msgs.as_array().unwrap().iter().map(dsv4_turn).collect();
10149 let rendered = chat::apply_chat_template_tools_ex(
10150 Some(&tmpl),
10151 &turns,
10152 true,
10153 &[],
10154 &req_tools,
10155 ThinkMode::Think,
10156 None,
10157 None,
10158 )
10159 .unwrap_or_else(|e| panic!("{name}: render failed: {e}"));
10160 assert!(
10161 rendered.contains(blob.as_str()),
10162 "{name}: tool result missing from render"
10163 );
10164 let t0 = std::time::Instant::now();
10165 let ids = tok.encode(&rendered, true);
10166 let encode_dt = t0.elapsed();
10167 assert!(!ids.is_empty(), "{name}: empty encode");
10168 let back = tok.decode(&ids);
10169 assert_eq!(back, rendered, "{name}: decode(encode(x)) != x");
10170 assert!(
10174 encode_dt < std::time::Duration::from_secs(60),
10175 "{name}: encode took {encode_dt:?}"
10176 );
10177 if *name == "ascii-letter-131k" {
10181 if let Ok(dir) = std::env::var("DSV4_LONGRUN_DUMP_DIR") {
10182 std::fs::write(format!("{dir}/rendered-131k.txt"), &rendered).unwrap();
10183 let csv: Vec<String> = ids.iter().map(|i| i.to_string()).collect();
10184 std::fs::write(format!("{dir}/memra-ids-131k.csv"), csv.join(",")).unwrap();
10185 }
10186 }
10187 }
10188 }
10189
10190 #[test]
10191 fn models_v1_entry_advertises_thinking_support() {
10192 let step_caps = ModelCaps {
10195 effort_levels: true,
10196 ..tool_caps()
10197 };
10198 let entry = model_entry_v1("stepfun/step-3.7-flash", Some(&step_caps), None);
10199 assert_eq!(entry["capabilities"]["reasoning"], true);
10200 assert_eq!(entry["capabilities"]["tools"], true);
10201
10202 let plain = ModelCaps {
10204 chat_ok: true,
10205 ..Default::default()
10206 };
10207 let entry = model_entry_v1("plain", Some(&plain), None);
10208 assert_eq!(entry["capabilities"]["reasoning"], false);
10209 assert_eq!(entry["capabilities"]["tools"], false);
10210 let entry = model_entry_v1("unknown", None, None);
10212 assert_eq!(entry["capabilities"]["reasoning"], false);
10213 assert_eq!(entry["capabilities"]["streaming"], true);
10214 }
10215
10216 #[test]
10217 fn chat_request_preserves_turns_and_openai_stop_forms() {
10218 let payload = serde_json::json!({
10219 "model": "plain_quant",
10220 "messages": [
10221 {"role": "system", "content": "rules"},
10222 {"role": "developer", "content": "dev rules"},
10223 {"role": "user", "content": "task"},
10224 {"role": "assistant", "content": "work"}
10225 ],
10226 "max_tokens": 64,
10227 "temperature": 0.0,
10228 "stop": "<stop>"
10229 });
10230 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10231 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10232 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10233 let request = plan.request;
10234 assert!(
10235 plan.parser.is_none(),
10236 "no tools -> no parser (isolation contract)"
10237 );
10238 assert!(request.tools_json.is_empty());
10239 assert_eq!(request.think, ThinkMode::Default);
10240 assert_eq!(request.model, "plain_quant");
10241 assert_eq!(request.params.max_new, 64);
10242 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10244 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}]
10245 }))
10246 .unwrap();
10247 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10248 let plan = build_chat_request(req, None, tx, lanes::Lane::Interactive, None).unwrap();
10249 assert_eq!(plan.request.params.max_new, worker::MAX_NEW_CTX_BOUNDED);
10250 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10252 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10253 "max_completion_tokens": 7
10254 }))
10255 .unwrap();
10256 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10257 assert_eq!(
10258 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
10259 .unwrap()
10260 .request
10261 .params
10262 .max_new,
10263 7
10264 );
10265 let req: CompletionReq = serde_json::from_value(serde_json::json!({
10267 "model": "plain_quant", "prompt": "task"
10268 }))
10269 .unwrap();
10270 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10271 assert_eq!(
10272 build_request(&req, tx, lanes::Lane::Interactive, None)
10273 .params
10274 .max_new,
10275 worker::MAX_NEW_CTX_BOUNDED
10276 );
10277 let turns: Vec<(String, String)> = request
10278 .chat_turns
10279 .iter()
10280 .map(|t| (t.role.clone(), t.content.clone()))
10281 .collect();
10282 assert_eq!(
10283 turns,
10284 vec![
10285 ("system".into(), "rules".into()),
10286 ("system".into(), "dev rules".into()), ("user".into(), "task".into()),
10288 ("assistant".into(), "work".into()),
10289 ]
10290 );
10291 assert!(request.chat_turns.iter().all(|t| t.tool_calls.is_empty()));
10292 assert_eq!(request.stop_strings, vec!["<stop>"]);
10293
10294 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10295 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10296 "stop": ["a", "b"]
10297 }))
10298 .unwrap();
10299 assert_eq!(req.stop.into_vec(), vec!["a", "b"]);
10300
10301 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10305 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10306 "stop": ["", "real", ""]
10307 }))
10308 .unwrap();
10309 assert_eq!(req.stop.into_vec(), vec!["real"]);
10310 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10311 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10312 "stop": ""
10313 }))
10314 .unwrap();
10315 assert!(req.stop.into_vec().is_empty());
10316
10317 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
10318 "model": "plain_quant", "messages": [{"role": "user", "content": "task"}],
10319 "stop": null
10320 }))
10321 .unwrap();
10322 assert!(req.stop.into_vec().is_empty());
10323 }
10324
10325 #[tokio::test]
10326 async fn chat_response_has_openai_message_shape() {
10327 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10328 tx.send(Event::Token {
10329 id: 1,
10330 text: "hello".into(),
10331 })
10332 .unwrap();
10333 tx.send(Event::Done {
10334 stop_reason: "Eos".into(),
10335 n_tokens: 1,
10336 n_prompt: 42,
10337 n_cached: 30,
10338 elapsed_s: 0.5,
10339 spec: None,
10340 })
10341 .unwrap();
10342 drop(tx);
10343 let response = blocking_response(
10344 rx,
10345 "plain_quant".into(),
10346 true,
10347 Vec::new(),
10348 None,
10349 Envelope::new(true),
10350 )
10351 .await;
10352 assert_eq!(response.status(), StatusCode::OK);
10353 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10354 .await
10355 .unwrap();
10356 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10357 assert_eq!(payload["object"], "chat.completion");
10358 assert!(payload["id"].as_str().unwrap().starts_with("chatcmpl-"));
10360 assert!(payload["created"].as_u64().unwrap() > 1_700_000_000);
10361 assert!(
10362 payload["system_fingerprint"]
10363 .as_str()
10364 .unwrap()
10365 .starts_with("memra-")
10366 );
10367 assert_eq!(payload["choices"][0]["message"]["role"], "assistant");
10368 assert_eq!(payload["choices"][0]["message"]["content"], "hello");
10369 assert_eq!(payload["choices"][0]["finish_reason"], "stop");
10370 assert_eq!(payload["usage"]["prompt_tokens"], 42);
10372 assert_eq!(payload["usage"]["completion_tokens"], 1);
10373 assert_eq!(payload["usage"]["total_tokens"], 43);
10374 assert_eq!(
10375 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
10376 30
10377 );
10378 assert!(payload["usage"].get("spec").is_none());
10381 }
10382
10383 #[tokio::test]
10384 async fn native_response_uses_terminal_token_snapshot_for_coalesced_events() {
10385 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10386 tx.send(Event::Token {
10388 id: 4,
10389 text: "hello".into(),
10390 })
10391 .unwrap();
10392 tx.send(Event::TokenSnapshot(vec![1, 2, 3, 4])).unwrap();
10393 tx.send(Event::Done {
10394 stop_reason: "MaxNew".into(),
10395 n_tokens: 4,
10396 n_prompt: 2,
10397 n_cached: 0,
10398 elapsed_s: 0.5,
10399 spec: None,
10400 })
10401 .unwrap();
10402 drop(tx);
10403
10404 let response = blocking_response(
10405 rx,
10406 "plain_quant".into(),
10407 false,
10408 Vec::new(),
10409 None,
10410 Envelope::new(false),
10411 )
10412 .await;
10413 assert_eq!(response.status(), StatusCode::OK);
10414 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10415 .await
10416 .unwrap();
10417 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10418 assert_eq!(payload["text"], "hello");
10419 assert_eq!(payload["tokens"], serde_json::json!([1, 2, 3, 4]));
10420 assert_eq!(payload["n_tokens"], 4);
10421 }
10422
10423 #[tokio::test]
10426 async fn chat_usage_carries_spec_acceptance_summary() {
10427 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
10428 tx.send(Event::Token {
10429 id: 1,
10430 text: "hello".into(),
10431 })
10432 .unwrap();
10433 tx.send(Event::Done {
10434 stop_reason: "Eos".into(),
10435 n_tokens: 1,
10436 n_prompt: 42,
10437 n_cached: 0,
10438 elapsed_s: 0.5,
10439 spec: Some(worker::SpecUsage {
10440 rounds: 10,
10441 drafted: 30,
10442 accepted: 21,
10443 }),
10444 })
10445 .unwrap();
10446 drop(tx);
10447 let response = blocking_response(
10448 rx,
10449 "plain_quant".into(),
10450 true,
10451 Vec::new(),
10452 None,
10453 Envelope::new(true),
10454 )
10455 .await;
10456 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
10457 .await
10458 .unwrap();
10459 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
10460 let sp = &payload["usage"]["spec"];
10461 assert_eq!(sp["rounds"], 10);
10462 assert_eq!(sp["drafted"], 30);
10463 assert_eq!(sp["accepted"], 21);
10464 assert!((sp["acceptance_rate"].as_f64().unwrap() - 0.7).abs() < 1e-9);
10465 assert_eq!(payload["usage"]["total_tokens"], 43);
10467 }
10468
10469 fn weather_request(extra: serde_json::Value) -> ChatCompletionReq {
10470 let mut payload = serde_json::json!({
10471 "model": "m",
10472 "messages": [{"role": "user", "content": "Weather in Paris?"}],
10473 "tools": [{"type": "function", "function": {
10474 "name": "get_weather",
10475 "description": "Get current weather",
10476 "parameters": {"type": "object",
10477 "properties": {"city": {"type": "string"},
10478 "days": {"type": "integer"}},
10479 "required": ["city"]}}}],
10480 });
10481 if let Some(obj) = extra.as_object() {
10482 for (k, v) in obj {
10483 payload[k] = v.clone();
10484 }
10485 }
10486 serde_json::from_value(payload).unwrap()
10487 }
10488
10489 #[test]
10490 fn vision_decode_is_deferred_and_grid_pinned() {
10491 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10496 let req: ChatCompletionReq = serde_json::from_value(json!({
10497 "model": "m", "messages": [{"role": "user", "content": "hi"}],
10498 }))
10499 .unwrap();
10500 let mut plan = build_chat_request(
10501 req,
10502 Some(&ModelCaps {
10503 chat_ok: true,
10504 ..Default::default()
10505 }),
10506 tx,
10507 lanes::Lane::Interactive,
10508 None,
10509 )
10510 .unwrap();
10511 let bmp = |w: i32, h: i32, with_pixels: bool| -> Vec<u8> {
10515 let mut b = Vec::new();
10516 b.extend_from_slice(b"BM");
10517 b.extend_from_slice(&54u32.to_le_bytes());
10518 b.extend_from_slice(&0u32.to_le_bytes());
10519 b.extend_from_slice(&54u32.to_le_bytes());
10520 b.extend_from_slice(&40u32.to_le_bytes());
10521 b.extend_from_slice(&w.to_le_bytes());
10522 b.extend_from_slice(&h.to_le_bytes());
10523 b.extend_from_slice(&1u16.to_le_bytes());
10524 b.extend_from_slice(&24u16.to_le_bytes());
10525 b.extend_from_slice(&[0u8; 24]);
10526 if with_pixels {
10527 b.extend(std::iter::repeat_n(0x7fu8, (w * h * 3) as usize));
10528 }
10529 b
10530 };
10531 let bytes = bmp(64, 64, true);
10532 let (gh, gw) = memra_engine::vision_pre::plan_image_bytes(&bytes).unwrap();
10533 plan.pending_images.push(PendingVisionUnit::Still {
10534 bytes: bytes.clone(),
10535 gh,
10536 gw,
10537 });
10538 decode_pending_vision(&mut plan).unwrap();
10539 assert_eq!(plan.request.images.len(), 1);
10540 assert_eq!(
10541 (
10542 plan.request.images[0].prep.gh,
10543 plan.request.images[0].prep.gw
10544 ),
10545 (gh, gw),
10546 "decoded grid must equal the header-planned grid the pad run was rendered from"
10547 );
10548 plan.request.images.clear();
10550 plan.pending_images.push(PendingVisionUnit::Still {
10551 bytes,
10552 gh: gh + 2,
10553 gw,
10554 });
10555 let err = decode_pending_vision(&mut plan).unwrap_err();
10556 assert!(err.contains("header-planned"), "got: {err}");
10557 let bomb = bmp(16_000, 16_000, false);
10560 plan.pending_images.clear();
10561 plan.pending_images.push(PendingVisionUnit::Still {
10562 bytes: bomb,
10563 gh: 2,
10564 gw: 2,
10565 });
10566 let err = decode_pending_vision(&mut plan).unwrap_err();
10567 assert!(err.contains("exceeds the decode budget"), "got: {err}");
10568 }
10569
10570 #[test]
10571 fn tools_request_renders_client_key_order_and_arms_parser() {
10572 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10573 let plan = build_chat_request(
10574 weather_request(json!({})),
10575 Some(&tool_caps()),
10576 tx,
10577 lanes::Lane::Interactive,
10578 None,
10579 )
10580 .unwrap();
10581 assert!(plan.parser.is_some());
10582 assert_eq!(plan.request.tools_json.len(), 1);
10583 assert_eq!(
10585 plan.request.tools_json[0],
10586 "{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \
10587 \"description\": \"Get current weather\", \"parameters\": {\"type\": \"object\", \
10588 \"properties\": {\"city\": {\"type\": \"string\"}, \"days\": {\"type\": \
10589 \"integer\"}}, \"required\": [\"city\"]}}}"
10590 );
10591 }
10592
10593 #[test]
10594 fn tool_choice_none_strips_tools_and_parser() {
10595 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10596 let plan = build_chat_request(
10597 weather_request(json!({"tool_choice": "none"})),
10598 Some(&tool_caps()),
10599 tx,
10600 lanes::Lane::Interactive,
10601 None,
10602 )
10603 .unwrap();
10604 let mut p = plan
10607 .parser
10608 .expect("think-open chat arms the reasoning splitter");
10609 let pieces = p.push("x</think>\n\n<tool_call> stays prose");
10610 assert_eq!(
10611 pieces,
10612 vec![
10613 Piece::Reasoning("x".into()),
10614 Piece::Content("<tool_call> stays prose".into()),
10615 ]
10616 );
10617 assert!(plan.request.tools_json.is_empty());
10618 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10620 assert!(
10621 build_chat_request(
10622 weather_request(json!({"tool_choice": "required"})),
10623 Some(&tool_caps()),
10624 tx,
10625 lanes::Lane::Interactive,
10626 None
10627 )
10628 .is_err()
10629 );
10630 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10631 assert!(
10632 build_chat_request(
10633 weather_request(json!({"tool_choice":
10634 {"type": "function", "function": {"name": "get_weather"}}})),
10635 Some(&tool_caps()),
10636 tx,
10637 lanes::Lane::Interactive,
10638 None
10639 )
10640 .is_err()
10641 );
10642 }
10643
10644 #[test]
10645 fn model_plan_accepts_st_dir_and_rejects_bogus_dir() {
10646 let root = std::env::temp_dir().join(format!("memra_plan_test_{}", std::process::id()));
10647 let _ = std::fs::remove_dir_all(&root);
10648
10649 let st = root.join("st_single");
10651 std::fs::create_dir_all(&st).unwrap();
10652 std::fs::write(st.join("config.json"), "{}").unwrap();
10653 std::fs::write(st.join("model.safetensors"), b"x").unwrap();
10654 assert!(validate_model_path(st.to_str().unwrap()).is_ok());
10655
10656 let sh = root.join("st_sharded");
10658 std::fs::create_dir_all(&sh).unwrap();
10659 std::fs::write(sh.join("config.json"), "{}").unwrap();
10660 std::fs::write(sh.join("model.safetensors.index.json"), "{}").unwrap();
10661 assert!(validate_model_path(sh.to_str().unwrap()).is_ok());
10662
10663 let rp = root.join("repack");
10665 std::fs::create_dir_all(&rp).unwrap();
10666 std::fs::write(rp.join("manifest.json"), "{}").unwrap();
10667 assert!(validate_model_path(rp.to_str().unwrap()).is_ok());
10668
10669 let bogus = root.join("bogus");
10671 std::fs::create_dir_all(&bogus).unwrap();
10672 let err = validate_model_path(bogus.to_str().unwrap()).unwrap_err();
10673 assert!(
10674 err.contains("model.safetensors"),
10675 "error should say what is missing: {err}"
10676 );
10677 assert!(
10678 err.contains("manifest.json"),
10679 "error should mention the repack form: {err}"
10680 );
10681
10682 let nc = root.join("no_config");
10684 std::fs::create_dir_all(&nc).unwrap();
10685 std::fs::write(nc.join("model.safetensors"), b"x").unwrap();
10686 let err = validate_model_path(nc.to_str().unwrap()).unwrap_err();
10687 assert!(
10688 err.contains("config.json"),
10689 "error should name config.json: {err}"
10690 );
10691
10692 let err = validate_model_path(root.join("nowhere").to_str().unwrap()).unwrap_err();
10694 assert!(err.contains("does not exist"), "{err}");
10695
10696 let f = root.join("model.gguf");
10698 std::fs::write(&f, b"g").unwrap();
10699 assert!(validate_model_path(f.to_str().unwrap()).is_ok());
10700
10701 let _ = std::fs::remove_dir_all(&root);
10702 }
10703
10704 #[test]
10705 fn chat_on_templateless_dir_checkpoint_is_rejected_with_clear_message() {
10706 let caps = ModelCaps {
10709 tools_branch: false,
10710 qwen_think: false,
10711 think_switch: false,
10712 chat_ok: false,
10713 ..Default::default()
10714 };
10715 let payload = serde_json::json!({
10716 "model": "st_model",
10717 "messages": [{"role": "user", "content": "hello"}],
10718 });
10719 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
10720 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10721 let err = match build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None) {
10722 Err(e) => e,
10723 Ok(_) => panic!("templateless dir checkpoint must reject chat"),
10724 };
10725 assert!(
10726 err.contains("no chat template"),
10727 "message should name the cause: {err}"
10728 );
10729 assert!(
10730 err.contains("/v1/completions"),
10731 "message should point at the raw-prompt escape hatch: {err}"
10732 );
10733 }
10734
10735 #[test]
10736 fn tools_on_model_without_tools_branch_is_rejected() {
10737 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10738 let caps = ModelCaps {
10739 chat_ok: true,
10740 ..Default::default()
10741 };
10742 assert!(
10743 build_chat_request(
10744 weather_request(json!({})),
10745 Some(&caps),
10746 tx,
10747 lanes::Lane::Interactive,
10748 None
10749 )
10750 .is_err()
10751 );
10752 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10753 assert!(
10754 build_chat_request(
10755 weather_request(json!({})),
10756 None,
10757 tx,
10758 lanes::Lane::Interactive,
10759 None
10760 )
10761 .is_err()
10762 );
10763 }
10764
10765 #[test]
10766 fn reasoning_effort_maps_to_think_switch() {
10767 for (extra, want) in [
10773 (json!({}), ThinkMode::Default),
10774 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
10775 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
10776 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
10777 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
10778 (json!({"reasoning_effort": "medium"}), ThinkMode::Think),
10779 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
10780 (json!({"reasoning": {"effort": "low"}}), ThinkMode::Think),
10781 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
10782 (json!({"reasoning_effort": "xhigh"}), ThinkMode::Think),
10786 (json!({"reasoning_effort": "max"}), ThinkMode::Think),
10787 (json!({"reasoning_effort": "ultra"}), ThinkMode::Think),
10788 (
10792 json!({"reasoning": {"enabled": true, "effort": "none"}}),
10793 ThinkMode::Think,
10794 ),
10795 (
10796 json!({"reasoning": {"enabled": false, "effort": "high"}}),
10797 ThinkMode::NoThink,
10798 ),
10799 ] {
10800 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10801 let plan = build_chat_request(
10802 weather_request(extra.clone()),
10803 Some(&ladder_caps()),
10808 tx,
10809 lanes::Lane::Interactive,
10810 None,
10811 )
10812 .unwrap();
10813 assert_eq!(plan.request.think, want, "extra={extra}");
10814 }
10815 for extra in [
10819 json!({"reasoning_effort": "extreme"}),
10820 json!({"reasoning": {"effort": "banana"}}),
10821 json!({"reasoning": {"enabled": false, "effort": "banana"}}),
10822 json!({"reasoning": {"enabled": true, "effort": ""}}),
10823 ] {
10824 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10825 assert!(
10826 build_chat_request(
10827 weather_request(extra.clone()),
10828 Some(&tool_caps()),
10829 tx,
10830 lanes::Lane::Interactive,
10831 None
10832 )
10833 .is_err(),
10834 "extra={extra} must be rejected by the one allowlist"
10835 );
10836 }
10837 for (raw, want) in [
10840 ("none", Some("none")),
10841 ("minimal", Some("minimal")),
10842 ("low", Some("low")),
10843 ("medium", Some("medium")),
10844 ("high", Some("high")),
10845 ("xhigh", Some("high")),
10846 ("max", Some("high")),
10847 ("ultra", Some("high")),
10848 ("banana", None),
10849 ("", None),
10850 ("HIGH", None),
10851 ] {
10852 assert_eq!(canonical_effort(raw), want, "canonical_effort({raw:?})");
10853 }
10854 for (raw, want) in [
10857 ("none", Some("none")),
10858 ("minimal", Some("minimal")),
10859 ("low", Some("low")),
10860 ("medium", Some("medium")),
10861 ("high", Some("high")),
10862 ("xhigh", Some("max")),
10863 ("max", Some("max")),
10864 ("ultra", Some("max")),
10865 ("banana", None),
10866 ("", None),
10867 ("MAX", None),
10868 ] {
10869 assert_eq!(
10870 canonical_effort_for(raw, true),
10871 want,
10872 "canonical_effort_for({raw:?}, dsv4)"
10873 );
10874 }
10875 }
10876
10877 #[test]
10878 fn dsv4_reasoning_effort_max_survives_canonicalization() {
10879 let dsv4_caps = ModelCaps {
10885 chat_ok: true,
10886 dsv4: true,
10887 ..Default::default()
10888 };
10889 let build = |caps: &ModelCaps, effort: &str| {
10890 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10891 let req: ChatCompletionReq = serde_json::from_value(json!({
10892 "model": "m",
10893 "messages": [{"role": "user", "content": "hi"}],
10894 "reasoning_effort": effort,
10895 }))
10896 .unwrap();
10897 build_chat_request(req, Some(caps), tx, lanes::Lane::Interactive, None)
10898 };
10899 for raw in ["max", "xhigh", "ultra"] {
10900 let plan = build(&dsv4_caps, raw).unwrap();
10901 assert_eq!(
10902 plan.request.reasoning_effort.as_deref(),
10903 Some("max"),
10904 "dsv4 {raw:?} must reach the renderer as the max rung"
10905 );
10906 assert_eq!(plan.request.think, chat::ThinkMode::Think);
10907 }
10908 let plan = build(&dsv4_caps, "high").unwrap();
10910 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
10911 let step_caps = ModelCaps {
10913 chat_ok: true,
10914 effort_levels: true,
10915 ..Default::default()
10916 };
10917 let plan = build(&step_caps, "max").unwrap();
10918 assert_eq!(plan.request.reasoning_effort.as_deref(), Some("high"));
10919 }
10920
10921 #[test]
10922 fn default_reasoning_effort_flips_only_the_unset_request() {
10923 let build = |extra: serde_json::Value, default_effort: Option<&str>| {
10930 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
10931 build_chat_request_with_trace(
10932 weather_request(extra),
10933 Some(&ladder_caps()),
10934 tx,
10935 lanes::Lane::Interactive,
10936 None,
10937 None,
10938 default_effort,
10939 &ModelSamplingDefaults::default(),
10940 )
10941 .unwrap()
10942 };
10943 for (extra, want) in [
10944 (json!({}), ThinkMode::Think),
10946 (json!({"reasoning": {"exclude": true}}), ThinkMode::NoThink),
10950 (json!({"include_reasoning": false}), ThinkMode::NoThink),
10951 (json!({"reasoning": {"exclude": false}}), ThinkMode::Think),
10953 (json!({"include_reasoning": true}), ThinkMode::Think),
10954 (json!({"reasoning_effort": "none"}), ThinkMode::NoThink),
10956 (json!({"reasoning_effort": "minimal"}), ThinkMode::NoThink),
10957 (json!({"reasoning": {"enabled": false}}), ThinkMode::NoThink),
10958 (json!({"reasoning_effort": "low"}), ThinkMode::Think),
10960 (json!({"reasoning_effort": "high"}), ThinkMode::Think),
10961 (json!({"reasoning": {"enabled": true}}), ThinkMode::Think),
10962 ] {
10963 let plan = build(extra.clone(), Some("high"));
10964 assert_eq!(plan.request.think, want, "extra={extra}");
10965 }
10966 assert_eq!(
10968 build(json!({}), Some("none")).request.think,
10969 ThinkMode::NoThink
10970 );
10971 assert_eq!(
10972 build(json!({"reasoning_effort": "high"}), Some("none"))
10973 .request
10974 .think,
10975 ThinkMode::Think
10976 );
10977 assert_eq!(build(json!({}), None).request.think, ThinkMode::Default);
10981 }
10982
10983 const SWITCHED_QWEN_TMPL: &str = "<tools> ... add_generation_prompt ... \
10989 {%- if enable_thinking is defined and enable_thinking is false %}'<think>\\n\\n</think>\\n\\n'\
10990 {%- else %}'<think>\\n'{%- endif %}";
10991
10992 #[test]
10993 fn vllm_enable_thinking_switch_is_wired_not_ignored() {
10994 let build = |extra: serde_json::Value| {
11000 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11001 build_chat_request(
11002 weather_request(extra),
11003 Some(&tool_caps()),
11004 tx,
11005 lanes::Lane::Interactive,
11006 None,
11007 )
11008 };
11009 for (extra, want) in [
11010 (json!({"enable_thinking": false}), ThinkMode::NoThink),
11011 (json!({"enable_thinking": true}), ThinkMode::Think),
11012 (
11013 json!({"chat_template_kwargs": {"enable_thinking": false}}),
11014 ThinkMode::NoThink,
11015 ),
11016 (
11017 json!({"chat_template_kwargs": {"enable_thinking": true}}),
11018 ThinkMode::Think,
11019 ),
11020 (
11023 json!({"enable_thinking": false, "reasoning_effort": "high"}),
11024 ThinkMode::NoThink,
11025 ),
11026 (
11028 json!({"enable_thinking": false,
11029 "chat_template_kwargs": {"enable_thinking": false}}),
11030 ThinkMode::NoThink,
11031 ),
11032 ] {
11033 let plan = build(extra.clone()).unwrap_or_else(|e| {
11034 panic!("{extra} must be accepted and honored, got 400: {e}");
11035 });
11036 assert_eq!(
11037 plan.request.think, want,
11038 "{extra} was ACCEPTED AND IGNORED — the banned silent-accept class"
11039 );
11040 }
11041 let render = |extra: serde_json::Value| -> String {
11044 let plan = build(extra).unwrap();
11045 chat::apply_chat_template_tools_ex(
11046 Some(SWITCHED_QWEN_TMPL),
11047 &plan.request.chat_turns,
11048 true,
11049 &plan.request.tools_json,
11050 &plan.request.tools_struct,
11051 plan.request.think,
11052 plan.request.reasoning_effort.as_deref(),
11053 None,
11054 )
11055 .unwrap()
11056 };
11057 let off = render(json!({"enable_thinking": false}));
11058 assert!(
11059 off.ends_with("<|im_start|>assistant\n<think>\n\n</think>\n\n"),
11060 "enable_thinking:false must render the CLOSED think pair: {off:?}"
11061 );
11062 let on = render(json!({}));
11063 assert!(
11064 on.ends_with("<|im_start|>assistant\n<think>\n"),
11065 "an unset request must still render the template's OPEN think tail: {on:?}"
11066 );
11067 assert_eq!(
11068 off,
11069 render(json!({"chat_template_kwargs": {"enable_thinking": false}})),
11070 "both vLLM spellings must render byte-identically"
11071 );
11072 assert_eq!(
11073 off,
11074 render(json!({"reasoning_effort": "none"})),
11075 "the vLLM spelling must render byte-identically to the OpenAI spelling"
11076 );
11077 }
11078
11079 #[test]
11080 fn unknown_chat_template_kwarg_refuses_by_name() {
11081 let build = |extra: serde_json::Value| {
11084 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11085 build_chat_request(
11086 weather_request(extra),
11087 Some(&tool_caps()),
11088 tx,
11089 lanes::Lane::Interactive,
11090 None,
11091 )
11092 };
11093 let refusal = |extra: serde_json::Value, why: &str| -> String {
11094 build(extra).err().unwrap_or_else(|| panic!("{why}"))
11095 };
11096 let err = refusal(
11097 json!({"chat_template_kwargs": {"add_generation_prompt": false}}),
11098 "an unimplementable template kwarg must not be accepted",
11099 );
11100 assert!(
11101 err.contains("add_generation_prompt") && err.contains("enable_thinking"),
11102 "the refusal must name the offending key AND the supported one: {err}"
11103 );
11104 let err = refusal(
11105 json!({"chat_template_kwargs": "enable_thinking=false"}),
11106 "a non-object chat_template_kwargs must not be accepted",
11107 );
11108 assert!(
11109 err.contains("must be an object"),
11110 "refusal must say what shape is expected: {err}"
11111 );
11112 let err = refusal(
11113 json!({"chat_template_kwargs": {"enable_thinking": "false"}}),
11114 "a stringly-typed switch must not be accepted",
11115 );
11116 assert!(
11117 err.contains("true or false"),
11118 "refusal must name the expected type: {err}"
11119 );
11120 let plan = build(json!({"chat_template_kwargs": null}))
11122 .expect("null chat_template_kwargs is the unset case");
11123 assert_eq!(plan.request.think, ThinkMode::Default);
11124 }
11125
11126 const Q38_TMPL: &str =
11141 include_str!("../../../research/reasoning-schema-20260823/qwen38-27b.chat_template.jinja");
11142
11143 fn render_with(
11146 tmpl: &str,
11147 caps: &ModelCaps,
11148 extra: serde_json::Value,
11149 default_effort: Option<&str>,
11150 ) -> Result<String, String> {
11151 let mut payload = serde_json::json!({
11152 "model": "m",
11153 "messages": [{"role": "user", "content": "hi"}],
11154 });
11155 if let Some(obj) = extra.as_object() {
11156 for (k, v) in obj {
11157 payload[k] = v.clone();
11158 }
11159 }
11160 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11161 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11162 let plan = build_chat_request_with_trace(
11163 req,
11164 Some(caps),
11165 tx,
11166 lanes::Lane::Interactive,
11167 None,
11168 None,
11169 default_effort,
11170 &ModelSamplingDefaults::default(),
11171 )?;
11172 Ok(chat::apply_chat_template_tools_ex(
11173 Some(tmpl),
11174 &plan.request.chat_turns,
11175 true,
11176 &plan.request.tools_json,
11177 &plan.request.tools_struct,
11178 plan.request.think,
11179 plan.request.reasoning_effort.as_deref(),
11180 None,
11181 )
11182 .unwrap())
11183 }
11184
11185 #[test]
11186 fn qwen38_effort_ladder_reaches_prompt_bytes_through_the_whole_api() {
11187 let r = |extra: serde_json::Value| render_with(Q38_TMPL, &ladder_caps(), extra, None);
11194 let xhigh = "Reasoning effort is set to xhigh.";
11195 let low = "Reasoning effort is set to low.";
11196 assert!(r(json!({"reasoning_effort": "low"})).unwrap().contains(low));
11198 assert!(
11199 r(json!({"reasoning_effort": "high"}))
11200 .unwrap()
11201 .contains(xhigh)
11202 );
11203 let medium = r(json!({"reasoning_effort": "medium"})).unwrap();
11206 assert!(!medium.contains("Reasoning effort is set to"), "{medium:?}");
11207 let low_p = r(json!({"reasoning_effort": "low"})).unwrap();
11209 let high_p = r(json!({"reasoning_effort": "high"})).unwrap();
11210 assert_ne!(low_p, high_p);
11211 assert_ne!(low_p, medium);
11212 assert_ne!(high_p, medium);
11213 for alias in ["xhigh", "max", "ultra"] {
11216 assert_eq!(r(json!({"reasoning_effort": alias})).unwrap(), high_p);
11217 }
11218 assert_eq!(r(json!({})).unwrap(), high_p);
11221 assert_eq!(
11224 render_with(Q38_TMPL, &ladder_caps(), json!({}), Some("medium")).unwrap(),
11225 medium
11226 );
11227 let off = r(json!({"reasoning_effort": "none"})).unwrap();
11230 assert!(off.ends_with("<think>\n\n</think>\n\n"), "{off:?}");
11231 assert!(!off.contains("Reasoning effort is set to"), "{off:?}");
11232 }
11233
11234 #[test]
11235 fn the_effort_sentence_is_measurable_on_the_deployed_binary_without_a_deploy() {
11236 const LOW_SENTENCE: &str = "Reasoning effort is set to low. Keep your thinking brief and \
11249focused, moving directly to the conclusion without unnecessary elaboration.";
11250 let expected = format!(
11251 "<|im_start|>system\n{LOW_SENTENCE}<|im_end|>\n\
11252 <|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n<think>\n"
11253 );
11254 let after_fix = render_with(
11256 Q38_TMPL,
11257 &ladder_caps(),
11258 json!({"reasoning_effort": "low"}),
11259 None,
11260 )
11261 .unwrap();
11262 assert_eq!(
11263 after_fix, expected,
11264 "the shipped prompt for reasoning_effort:\"low\""
11265 );
11266 const ORNITH_TMPL: &str = include_str!(
11269 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11270 );
11271 let on_deployed_binary = render_with(
11272 ORNITH_TMPL,
11273 &tool_caps(),
11274 json!({"messages": [{"role": "system", "content": LOW_SENTENCE},
11275 {"role": "user", "content": "hi"}]}),
11276 None,
11277 )
11278 .unwrap();
11279 assert_eq!(
11280 on_deployed_binary, expected,
11281 "the live cell's system-message stand-in must render the SAME bytes as the post-fix \
11282 level, or its reasoning-volume numbers do not describe the shipped prompt"
11283 );
11284 let ladderless_unset = render_with(ORNITH_TMPL, &tool_caps(), json!({}), None).unwrap();
11287 assert!(
11288 !ladderless_unset.contains("Reasoning effort is set to"),
11289 "pre-lane q38 injected no effort instruction at any level: {ladderless_unset:?}"
11290 );
11291 assert_eq!(
11292 ladderless_unset,
11293 render_with(
11294 Q38_TMPL,
11295 &ladder_caps(),
11296 json!({"reasoning_effort": "medium"}),
11297 None
11298 )
11299 .unwrap(),
11300 "medium is the vendor's zero-steering rung and therefore the pre-lane byte baseline"
11301 );
11302 }
11303
11304 #[test]
11305 fn include_reasoning_false_stops_reasoning_it_does_not_hide_it() {
11306 let off = render_with(
11313 Q38_TMPL,
11314 &ladder_caps(),
11315 json!({"reasoning_effort": "none"}),
11316 None,
11317 )
11318 .unwrap();
11319 for extra in [
11320 json!({"include_reasoning": false}),
11321 json!({"reasoning": {"exclude": true}}),
11322 ] {
11323 let got = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap();
11324 assert!(
11325 got.ends_with("<think>\n\n</think>\n\n"),
11326 "{extra} must render the CLOSED think pair, not a hidden reasoning block: {got:?}"
11327 );
11328 assert_eq!(got, off, "{extra} must be byte-identical to reasoning-off");
11329 }
11330 for extra in [
11335 json!({"enable_thinking": true, "include_reasoning": false}),
11336 json!({"reasoning": {"enabled": true}, "include_reasoning": false}),
11337 json!({"reasoning": {"enabled": true, "exclude": true}}),
11338 ] {
11339 let e = render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None)
11340 .err()
11341 .unwrap_or_else(|| panic!("{extra} must be refused as contradictory"));
11342 assert!(e.contains("contradictory"), "{extra}: {e}");
11343 assert!(
11344 e.contains("include_reasoning") || e.contains("exclude"),
11345 "{extra}: the refusal must name the suppression field the caller sent: {e}"
11346 );
11347 }
11348 let dflt = render_with(Q38_TMPL, &ladder_caps(), json!({}), None).unwrap();
11351 for extra in [
11352 json!({"include_reasoning": true}),
11353 json!({"reasoning": {"exclude": false}}),
11354 ] {
11355 assert_eq!(
11356 render_with(Q38_TMPL, &ladder_caps(), extra.clone(), None).unwrap(),
11357 dflt,
11358 "{extra} must not perturb the model's default"
11359 );
11360 }
11361 let switchless = ModelCaps {
11365 think_switch: false,
11366 ..tool_caps()
11367 };
11368 let err = render_with(
11369 Q38_TMPL,
11370 &switchless,
11371 json!({"include_reasoning": false}),
11372 None,
11373 )
11374 .expect_err("include_reasoning:false must not silently bill for hidden reasoning");
11375 assert!(err.contains("cannot disable reasoning"), "{err}");
11376 }
11377
11378 #[test]
11379 fn the_reasoning_object_refuses_every_key_it_cannot_honour() {
11380 let build = |extra: serde_json::Value| {
11381 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11382 build_chat_request(
11383 weather_request(extra),
11384 Some(&ladder_caps()),
11385 tx,
11386 lanes::Lane::Interactive,
11387 None,
11388 )
11389 };
11390 let err = |extra: serde_json::Value, why: &str| -> String {
11391 build(extra).err().unwrap_or_else(|| panic!("{why}"))
11392 };
11393 let e = err(
11397 json!({"reasoning": {"max_tokens": 1024}}),
11398 "reasoning.max_tokens must not be accepted-and-ignored",
11399 );
11400 assert!(e.contains("reasoning.max_tokens"), "{e}");
11401 assert!(e.contains("ONE output budget"), "{e}");
11402 for extra in [
11407 json!({"reasoning": {"max_tokens": null}}),
11408 json!({"reasoning": {"banana": null}}),
11409 ] {
11410 let e = err(
11411 extra.clone(),
11412 "a null-valued unhonourable key must still refuse",
11413 );
11414 assert!(
11415 e.contains("max_tokens") || e.contains("banana"),
11416 "{extra}: {e}"
11417 );
11418 }
11419 let e = err(
11421 json!({"reasoning": {"budget": 5}}),
11422 "an unknown reasoning key must not be accepted",
11423 );
11424 assert!(
11425 e.contains("reasoning.budget") && e.contains("enabled"),
11426 "{e}"
11427 );
11428 for (extra, want) in [
11432 (json!({"reasoning": {"enabled": "false"}}), "true or false"),
11433 (json!({"reasoning": {"exclude": 1}}), "true or false"),
11434 (json!({"reasoning": {"effort": 3}}), "must be a string"),
11435 ] {
11436 let e = err(
11437 extra.clone(),
11438 "a wrong-typed reasoning key must not be ignored",
11439 );
11440 assert!(e.contains(want), "{extra}: {e}");
11441 }
11442 for extra in [
11447 json!({"reasoning": {"enabled": true}}),
11448 json!({"reasoning": {"effort": "low"}}),
11449 json!({"reasoning": {"exclude": false}}),
11450 json!({"reasoning": null}),
11451 json!({"reasoning": {"effort": null}}),
11452 json!({"reasoning": {"enabled": null, "exclude": null}}),
11453 ] {
11454 build(extra.clone()).unwrap_or_else(|e| panic!("{extra} must be served: {e}"));
11455 }
11456 }
11457
11458 #[test]
11459 fn a_graded_level_on_a_binary_model_translates_to_reasoning_on() {
11460 const ORNITH_TMPL: &str = include_str!(
11468 "../../../research/reasoning-schema-20260823/ornith15.chat_template.jinja"
11469 );
11470 let explicit_on = render_with(
11474 ORNITH_TMPL,
11475 &tool_caps(),
11476 json!({"reasoning": {"enabled": true}}),
11477 None,
11478 )
11479 .unwrap();
11480 assert!(explicit_on.ends_with("<think>\n"), "{explicit_on:?}");
11481 for extra in [
11482 json!({"reasoning_effort": "low"}),
11483 json!({"reasoning_effort": "medium"}),
11484 json!({"reasoning_effort": "high"}),
11485 json!({"reasoning_effort": "xhigh"}),
11487 json!({"reasoning": {"effort": "xhigh"}}),
11488 ] {
11489 let got = render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11490 .unwrap_or_else(|e| panic!("{extra} must TRANSLATE to reasoning-on, got 400: {e}"));
11491 assert_eq!(
11492 got, explicit_on,
11493 "{extra} must render byte-identical to reasoning:{{enabled:true}} — the \
11494 documented translation, not a decorative accept"
11495 );
11496 }
11497 for extra in [
11499 json!({}),
11500 json!({"reasoning_effort": "none"}),
11501 json!({"reasoning_effort": "minimal"}),
11502 json!({"enable_thinking": false}),
11503 ] {
11504 render_with(ORNITH_TMPL, &tool_caps(), extra.clone(), None)
11505 .unwrap_or_else(|e| panic!("{extra} must still be served: {e}"));
11506 }
11507 let minimal = render_with(
11510 ORNITH_TMPL,
11511 &tool_caps(),
11512 json!({"reasoning_effort": "minimal"}),
11513 None,
11514 )
11515 .unwrap();
11516 assert!(
11517 minimal.ends_with("<think>\n\n</think>\n\n"),
11518 "minimal must close the think pair (OFF), not clamp to a reasoning level: {minimal:?}"
11519 );
11520 let ladder_low = render_with(
11523 Q38_TMPL,
11524 &ladder_caps(),
11525 json!({"reasoning_effort": "low"}),
11526 None,
11527 )
11528 .unwrap();
11529 assert!(
11530 ladder_low.contains("Reasoning effort is set to low."),
11531 "{ladder_low:?}"
11532 );
11533 assert_ne!(
11534 ladder_low,
11535 render_with(
11536 Q38_TMPL,
11537 &ladder_caps(),
11538 json!({"reasoning_effort": "high"}),
11539 None
11540 )
11541 .unwrap(),
11542 "the ladder model's rungs stay distinct prompts"
11543 );
11544 }
11545
11546 #[test]
11547 fn one_semantic_reasoning_request_renders_identical_bytes_on_all_three_surfaces() {
11548 let render_chat = |body: serde_json::Value| -> Result<String, String> {
11559 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11560 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11561 let plan = build_chat_request(
11562 req,
11563 Some(&ladder_caps()),
11564 tx,
11565 lanes::Lane::Interactive,
11566 None,
11567 )?;
11568 Ok(chat::apply_chat_template_tools_ex(
11569 Some(Q38_TMPL),
11570 &plan.request.chat_turns,
11571 true,
11572 &plan.request.tools_json,
11573 &plan.request.tools_struct,
11574 plan.request.think,
11575 plan.request.reasoning_effort.as_deref(),
11576 None,
11577 )
11578 .unwrap())
11579 };
11580 for (intent, chat_body, responses_body, messages_body) in [
11585 (
11586 "reasoning OFF",
11587 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11588 "reasoning_effort": "none"}),
11589 json!({"model": "m", "input": "hi", "reasoning": {"effort": "none"}}),
11590 json!({"model": "m", "max_tokens": 16,
11591 "messages": [{"role": "user", "content": "hi"}],
11592 "thinking": {"type": "disabled"}}),
11593 ),
11594 (
11595 "reasoning ON at the top rung",
11596 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11597 "reasoning_effort": "xhigh"}),
11598 json!({"model": "m", "input": "hi", "reasoning": {"effort": "xhigh"}}),
11599 json!({"model": "m", "max_tokens": 16,
11600 "messages": [{"role": "user", "content": "hi"}],
11601 "output_config": {"effort": "xhigh"}}),
11602 ),
11603 (
11604 "reasoning ON at the bottom rung",
11605 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}],
11606 "reasoning_effort": "low"}),
11607 json!({"model": "m", "input": "hi", "reasoning": {"effort": "low"}}),
11608 json!({"model": "m", "max_tokens": 16,
11609 "messages": [{"role": "user", "content": "hi"}],
11610 "output_config": {"effort": "low"}}),
11611 ),
11612 (
11613 "the model's own default",
11614 json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}),
11615 json!({"model": "m", "input": "hi"}),
11616 json!({"model": "m", "max_tokens": 16,
11617 "messages": [{"role": "user", "content": "hi"}]}),
11618 ),
11619 ] {
11620 let chat = render_chat(chat_body).unwrap_or_else(|e| panic!("{intent} on chat: {e}"));
11621 let via_responses = responses_api::translate(&responses_body)
11622 .unwrap_or_else(|e| panic!("{intent} on /v1/responses: {e:?}"));
11623 let via_messages = anthropic::translate(&messages_body)
11624 .unwrap_or_else(|e| panic!("{intent} on /v1/messages: {e}"));
11625 for (surface, translated) in [
11626 ("/v1/responses", via_responses),
11627 ("/v1/messages", via_messages),
11628 ] {
11629 let got = render_chat(translated)
11630 .unwrap_or_else(|e| panic!("{intent} via {surface}: {e}"));
11631 assert_eq!(
11632 got, chat,
11633 "{intent}: {surface} rendered DIFFERENT prompt bytes than \
11634 /v1/chat/completions — the parameter is honoured on one format and not \
11635 the other"
11636 );
11637 }
11638 }
11639 let switchless = ModelCaps {
11642 think_switch: false,
11643 ..ladder_caps()
11644 };
11645 let render_switchless = |body: serde_json::Value| -> Result<String, String> {
11646 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
11647 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11648 let plan =
11649 build_chat_request(req, Some(&switchless), tx, lanes::Lane::Interactive, None)?;
11650 Ok(format!("{:?}", plan.request.think))
11651 };
11652 for (surface, body) in [
11653 (
11654 "/v1/responses",
11655 responses_api::translate(&json!({
11656 "model": "m", "input": "hi", "reasoning": {"effort": "none"}}))
11657 .unwrap(),
11658 ),
11659 (
11660 "/v1/messages",
11661 anthropic::translate(&json!({
11662 "model": "m", "max_tokens": 16,
11663 "messages": [{"role": "user", "content": "hi"}],
11664 "thinking": {"type": "disabled"}}))
11665 .unwrap(),
11666 ),
11667 ] {
11668 let err = render_switchless(body)
11669 .err()
11670 .unwrap_or_else(|| panic!("{surface} must refuse an unhonourable off-request"));
11671 assert!(err.contains("cannot disable reasoning"), "{surface}: {err}");
11672 }
11673 }
11674
11675 #[test]
11676 fn preserve_thinking_true_is_the_implemented_default_and_false_refuses() {
11677 let build = |extra: serde_json::Value| {
11685 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11686 build_chat_request(
11687 weather_request(extra),
11688 Some(&ladder_caps()),
11689 tx,
11690 lanes::Lane::Interactive,
11691 None,
11692 )
11693 };
11694 build(json!({"chat_template_kwargs": {"preserve_thinking": true}}))
11695 .expect("preserve_thinking:true is the vendor default the renderer implements");
11696 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": false}}))
11697 .err()
11698 .expect("preserve_thinking:false (the strip arm) must refuse");
11699 assert!(e.contains("preserve_thinking"), "{e}");
11700 assert!(e.contains("strip"), "{e}");
11701 assert_eq!(
11704 build(json!({"chat_template_kwargs": {"enable_thinking": false}}))
11705 .unwrap()
11706 .request
11707 .think,
11708 ThinkMode::NoThink
11709 );
11710 let e = build(json!({"chat_template_kwargs": {"preserve_thinking": "false"}}))
11712 .err()
11713 .expect("a stringly-typed preserve_thinking must not be accepted");
11714 assert!(e.contains("true or false"), "{e}");
11715 }
11716
11717 #[test]
11718 fn dsv4_is_exempt_from_the_switchless_off_refusal() {
11719 let dsv4_caps = ModelCaps {
11725 qwen_think: true,
11726 think_switch: false,
11727 dsv4: true,
11728 ..tool_caps()
11729 };
11730 for extra in [
11731 json!({"reasoning_effort": "none"}),
11732 json!({"reasoning": {"enabled": false}}),
11733 json!({"enable_thinking": false}),
11734 json!({"include_reasoning": false}),
11735 ] {
11736 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11737 let plan = build_chat_request(
11738 weather_request(extra.clone()),
11739 Some(&dsv4_caps),
11740 tx,
11741 lanes::Lane::Interactive,
11742 None,
11743 )
11744 .unwrap_or_else(|e| panic!("{extra} must be served on dsv4: {e}"));
11745 assert_eq!(plan.request.think, ThinkMode::NoThink, "extra={extra}");
11746 }
11747 }
11748
11749 #[test]
11750 fn contradictory_think_switches_refuse_instead_of_picking_one() {
11751 let build = |extra: serde_json::Value| {
11754 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11755 build_chat_request(
11756 weather_request(extra),
11757 Some(&tool_caps()),
11758 tx,
11759 lanes::Lane::Interactive,
11760 None,
11761 )
11762 };
11763 for extra in [
11764 json!({"enable_thinking": true, "reasoning": {"enabled": false}}),
11765 json!({"enable_thinking": false, "reasoning": {"enabled": true}}),
11766 json!({"enable_thinking": false, "chat_template_kwargs": {"enable_thinking": true}}),
11767 ] {
11768 match build(extra.clone()) {
11769 Err(err) => assert!(
11770 err.contains("contradictory"),
11771 "the refusal must say the switches contradict: {err}"
11772 ),
11773 Ok(plan) => panic!(
11774 "{extra} must be rejected as contradictory; it silently resolved to {:?}",
11775 plan.request.think
11776 ),
11777 }
11778 }
11779 for extra in [
11781 json!({"enable_thinking": false, "reasoning": {"enabled": false}}),
11782 json!({"enable_thinking": true, "reasoning": {"enabled": true}}),
11783 json!({"enable_thinking": false, "reasoning": {"effort": "high"}}),
11784 ] {
11785 build(extra.clone())
11786 .unwrap_or_else(|e| panic!("{extra} is not a contradiction, but got 400: {e}"));
11787 }
11788 }
11789
11790 #[test]
11791 fn explicit_reasoning_off_on_a_switchless_template_refuses_loudly() {
11792 let switchless = ModelCaps {
11797 tools_branch: true,
11798 qwen_think: true,
11799 think_switch: false,
11800 chat_ok: true,
11801 ..Default::default()
11802 };
11803 let build = |extra: serde_json::Value, caps: &ModelCaps, default_effort: Option<&str>| {
11804 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11805 build_chat_request_with_trace(
11806 weather_request(extra),
11807 Some(caps),
11808 tx,
11809 lanes::Lane::Interactive,
11810 None,
11811 None,
11812 default_effort,
11813 &ModelSamplingDefaults::default(),
11814 )
11815 };
11816 for extra in [
11817 json!({"reasoning_effort": "none"}),
11818 json!({"reasoning_effort": "minimal"}),
11819 json!({"reasoning": {"enabled": false}}),
11820 json!({"enable_thinking": false}),
11821 json!({"chat_template_kwargs": {"enable_thinking": false}}),
11822 ] {
11823 let err = build(extra.clone(), &switchless, None)
11824 .err()
11825 .unwrap_or_else(|| {
11826 panic!(
11827 "{extra} on a switchless think template must not be accepted-and-ignored"
11828 )
11829 });
11830 assert!(
11831 err.contains("cannot disable reasoning"),
11832 "the refusal must say the model cannot disable reasoning: {err}"
11833 );
11834 }
11835 for (extra, default_effort) in [
11839 (json!({}), None),
11840 (json!({"reasoning_effort": "high"}), None),
11843 (json!({"reasoning": {"enabled": true}}), None),
11844 (json!({"enable_thinking": true}), None),
11845 (json!({}), Some("none")),
11846 (json!({}), Some("minimal")),
11847 (json!({}), Some("high")),
11848 ] {
11849 build(extra.clone(), &switchless, default_effort).unwrap_or_else(|e| {
11850 panic!("{extra} (default={default_effort:?}) must still be served: {e}")
11851 });
11852 }
11853 assert_eq!(
11856 build(json!({"enable_thinking": false}), &tool_caps(), None)
11857 .unwrap()
11858 .request
11859 .think,
11860 ThinkMode::NoThink
11861 );
11862 }
11863
11864 #[test]
11865 fn gemma4_default_think_on_renders_byte_identical_to_explicit_think_on() {
11866 let gemma_caps = ModelCaps {
11873 tools_branch: true,
11874 chat_ok: true,
11875 gemma_think: true,
11876 instruct_type: Some("gemma".into()),
11877 ..Default::default()
11878 };
11879 let render =
11880 |tmpl: &str, extra: serde_json::Value, default_effort: Option<&str>| -> String {
11881 let mut payload = serde_json::json!({
11882 "model": "google/gemma-4-31b-it",
11883 "messages": [{"role": "user", "content": "Weather in Paris?"}],
11884 });
11885 if let Some(obj) = extra.as_object() {
11886 for (k, v) in obj {
11887 payload[k] = v.clone();
11888 }
11889 }
11890 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
11891 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
11892 let plan = build_chat_request_with_trace(
11893 req,
11894 Some(&gemma_caps),
11895 tx,
11896 lanes::Lane::Interactive,
11897 None,
11898 None,
11899 default_effort,
11900 &ModelSamplingDefaults::default(),
11901 )
11902 .unwrap();
11903 chat::apply_chat_template_tools_ex(
11904 Some(tmpl),
11905 &plan.request.chat_turns,
11906 true,
11907 &plan.request.tools_json,
11908 &plan.request.tools_struct,
11909 plan.request.think,
11910 plan.request.reasoning_effort.as_deref(),
11911 None, )
11913 .unwrap()
11914 };
11915 let official = gemma_template("official");
11916 let unset_with_knob = render(&official, json!({}), Some("high"));
11917 let explicit_on = render(&official, json!({"reasoning_effort": "high"}), None);
11918 assert_eq!(
11919 unset_with_knob, explicit_on,
11920 "knob render must be byte-identical to the explicit think-on render"
11921 );
11922 assert!(
11923 unset_with_knob.starts_with("<|turn>system\n<|think|>\n"),
11924 "think-on injects the <|think|> system token: {unset_with_knob:?}"
11925 );
11926 assert!(
11927 unset_with_knob.ends_with("<|turn>model\n"),
11928 "think-on generation turn is OPEN: {unset_with_knob:?}"
11929 );
11930 let explicit_off_with_knob =
11934 render(&official, json!({"reasoning_effort": "none"}), Some("high"));
11935 let explicit_off = render(&official, json!({"reasoning_effort": "none"}), None);
11936 assert_eq!(explicit_off_with_knob, explicit_off);
11937 assert!(
11938 !explicit_off_with_knob.contains("<|think|>")
11939 && explicit_off_with_knob.ends_with("<|turn>model\n"),
11940 "explicit off keeps the official template's thinking-off bytes: \
11941 {explicit_off_with_knob:?}"
11942 );
11943 let unset_no_knob = render(&official, json!({}), None);
11945 assert_eq!(
11946 unset_no_knob, explicit_off,
11947 "knobless unset stays the template's own thinking-off default"
11948 );
11949 assert_ne!(unset_no_knob, unset_with_knob);
11950 let qat = gemma_template("qat");
11953 assert!(
11954 render(&qat, json!({}), None).ends_with("<|turn>model\n<|channel>thought\n<channel|>"),
11955 "QAT knobless unset keeps the closed-channel default"
11956 );
11957 assert_eq!(
11958 render(&qat, json!({}), Some("high")),
11959 render(&qat, json!({"reasoning_effort": "high"}), None),
11960 "QAT knob render must equal the explicit think-on render"
11961 );
11962 }
11963
11964 #[test]
11965 fn default_reasoning_effort_is_validated_at_metadata_load() {
11966 let parsed = OpenRouterMetadataFile::from_toml(
11968 r#"
11969[models.g]
11970default_reasoning_effort = "high"
11971"#,
11972 )
11973 .unwrap();
11974 assert_eq!(
11975 parsed.get("g").unwrap().default_reasoning_effort.as_deref(),
11976 Some("high")
11977 );
11978 let err = OpenRouterMetadataFile::from_toml(
11979 r#"
11980[models.g]
11981default_reasoning_effort = "always"
11982"#,
11983 )
11984 .unwrap_err();
11985 assert!(err.contains("default_reasoning_effort"), "{err}");
11986 }
11987
11988 #[test]
11989 fn reasoning_effort_maps_to_effort_level_on_step35_class_templates() {
11990 let effort_caps = ModelCaps {
12002 effort_levels: true,
12003 think_switch: false,
12004 ..tool_caps()
12005 };
12006 for (extra, want) in [
12007 (json!({}), None),
12008 (json!({"reasoning_effort": "low"}), Some("low")),
12009 (json!({"reasoning_effort": "medium"}), Some("medium")),
12010 (json!({"reasoning_effort": "high"}), Some("high")),
12011 (json!({"reasoning": {"effort": "high"}}), Some("high")),
12012 (json!({"reasoning_effort": "xhigh"}), Some("high")),
12014 (json!({"reasoning": {"effort": "max"}}), Some("high")),
12015 ] {
12016 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12017 let plan = build_chat_request(
12018 weather_request(extra.clone()),
12019 Some(&effort_caps),
12020 tx,
12021 lanes::Lane::Interactive,
12022 None,
12023 )
12024 .unwrap();
12025 assert_eq!(
12026 plan.request.reasoning_effort.as_deref(),
12027 want,
12028 "extra={extra}"
12029 );
12030 }
12031 for extra in [
12038 json!({"reasoning_effort": "none"}),
12039 json!({"reasoning_effort": "minimal"}),
12040 json!({"reasoning": {"enabled": false}}),
12041 json!({"enable_thinking": false}),
12042 json!({"include_reasoning": false}),
12043 ] {
12044 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12045 let err = build_chat_request(
12046 weather_request(extra.clone()),
12047 Some(&effort_caps),
12048 tx,
12049 lanes::Lane::Interactive,
12050 None,
12051 )
12052 .err()
12053 .unwrap_or_else(|| panic!("{extra} must not be clamped to a reasoning level"));
12054 assert!(
12055 err.contains("cannot disable reasoning"),
12056 "extra={extra}: {err}"
12057 );
12058 }
12059 for extra in [
12066 json!({"reasoning_effort": "high"}),
12067 json!({"reasoning": {"effort": "low"}}),
12068 ] {
12069 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12070 let plan = build_chat_request(
12071 weather_request(extra.clone()),
12072 Some(&tool_caps()),
12073 tx,
12074 lanes::Lane::Interactive,
12075 None,
12076 )
12077 .unwrap_or_else(|e| panic!("{extra} must translate, not refuse: {e}"));
12078 assert_eq!(plan.request.think, ThinkMode::Think, "extra={extra}");
12079 assert_eq!(plan.request.reasoning_effort, None, "extra={extra}");
12080 }
12081 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12083 let plan = build_chat_request(
12084 weather_request(json!({})),
12085 Some(&tool_caps()),
12086 tx,
12087 lanes::Lane::Interactive,
12088 None,
12089 )
12090 .unwrap();
12091 assert_eq!(plan.request.reasoning_effort, None);
12092 }
12093
12094 #[test]
12095 fn assistant_history_tool_calls_and_tool_role_render_into_turns() {
12096 let payload = serde_json::json!({
12097 "model": "m",
12098 "messages": [
12099 {"role": "user", "content": "Weather in Paris?"},
12100 {"role": "assistant", "content": null, "tool_calls": [
12101 {"id": "call_x", "type": "function", "function": {
12102 "name": "get_weather",
12103 "arguments": "{\"city\": \"Paris\", \"days\": 3}"}}]},
12104 {"role": "tool", "tool_call_id": "call_x", "content": "{\"temp_c\": 21}"}
12105 ],
12106 });
12107 let req: ChatCompletionReq = serde_json::from_value(payload).unwrap();
12108 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12109 let plan = build_chat_request(req, Some(&tool_caps()), tx, lanes::Lane::Interactive, None)
12110 .unwrap();
12111 let turns = &plan.request.chat_turns;
12112 assert_eq!(turns[1].tool_calls.len(), 1);
12113 assert_eq!(turns[1].tool_calls[0].name, "get_weather");
12114 assert_eq!(
12115 turns[1].tool_calls[0].params,
12116 vec![("city".into(), "Paris".into()), ("days".into(), "3".into())]
12117 );
12118 assert_eq!(turns[2].role, "tool");
12119 assert_eq!(turns[2].content, "{\"temp_c\": 21}");
12120 let mut p = plan
12123 .parser
12124 .expect("think-open chat arms the reasoning splitter");
12125 let pieces = p.push("thought</think>\n\nanswer <tool_call> is prose here");
12126 assert_eq!(
12127 pieces,
12128 vec![
12129 Piece::Reasoning("thought".into()),
12130 Piece::Content("answer <tool_call> is prose here".into()),
12131 ]
12132 );
12133 }
12134
12135 #[tokio::test]
12136 async fn blocking_tools_response_carries_tool_calls_and_finish_reason() {
12137 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12138 tx.send(Event::Token {
12139 id: 1,
12140 text: "plan</think>\n\n".into(),
12141 })
12142 .unwrap();
12143 tx.send(Event::Token {
12144 id: 2,
12145 text: "<tool_call>\n<function=get_weather>\n\
12146<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>"
12147 .into(),
12148 })
12149 .unwrap();
12150 tx.send(Event::Done {
12151 stop_reason: "Eos".into(),
12152 n_tokens: 2,
12153 n_prompt: 40,
12154 n_cached: 0,
12155 elapsed_s: 0.5,
12156 spec: None,
12157 })
12158 .unwrap();
12159 drop(tx);
12160 let parser = ToolStreamParser::new(HashMap::new(), true);
12161 let response = blocking_response(
12162 rx,
12163 "m".into(),
12164 true,
12165 Vec::new(),
12166 Some(parser),
12167 Envelope::new(true),
12168 )
12169 .await;
12170 assert_eq!(response.status(), StatusCode::OK);
12171 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12172 .await
12173 .unwrap();
12174 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12175 assert_eq!(payload["choices"][0]["finish_reason"], "tool_calls");
12176 assert_eq!(
12179 payload["choices"][0]["message"]["content"],
12180 serde_json::Value::Null
12181 );
12182 assert_eq!(payload["choices"][0]["message"]["reasoning"], "plan");
12183 assert_eq!(
12184 payload["choices"][0]["message"]["reasoning_details"][0]["text"],
12185 "plan"
12186 );
12187 let call = &payload["choices"][0]["message"]["tool_calls"][0];
12188 assert_eq!(call["type"], "function");
12189 assert_eq!(call["function"]["name"], "get_weather");
12190 assert_eq!(call["function"]["arguments"], "{\"city\":\"Paris\"}");
12191 assert_eq!(payload["usage"]["prompt_tokens"], 40);
12194 assert_eq!(payload["usage"]["completion_tokens"], 2);
12195 assert_eq!(payload["usage"]["total_tokens"], 42);
12196 assert_eq!(
12197 payload["usage"]["prompt_tokens_details"]["cached_tokens"],
12198 0
12199 );
12200 }
12201
12202 #[test]
12203 fn cache_salt_plumbs_to_the_worker_namespace() {
12204 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12206 "model": "m", "prompt": "task", "cache_salt": "tenant-a"
12207 }))
12208 .unwrap();
12209 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12210 assert_eq!(
12211 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12212 "tenant-a"
12213 );
12214
12215 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12216 "model": "m", "messages": [{"role": "user", "content": "task"}],
12217 "cache_salt": "tenant-b"
12218 }))
12219 .unwrap();
12220 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12221 assert_eq!(
12222 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12223 .unwrap()
12224 .request
12225 .cache_ns,
12226 "tenant-b"
12227 );
12228
12229 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12231 "model": "m", "prompt": "task"
12232 }))
12233 .unwrap();
12234 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12235 assert_eq!(
12236 build_request(&req, tx, lanes::Lane::Interactive, None).cache_ns,
12237 ""
12238 );
12239 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12240 "model": "m", "messages": [{"role": "user", "content": "task"}]
12241 }))
12242 .unwrap();
12243 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12244 assert_eq!(
12245 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
12246 .unwrap()
12247 .request
12248 .cache_ns,
12249 ""
12250 );
12251 }
12252
12253 #[test]
12254 fn cache_salt_validation_rejects_oversized_value() {
12255 let salt = Some("a".repeat(CACHE_SALT_MAX_BYTES + 1));
12256 assert_eq!(
12257 validate_cache_namespace(&salt, false),
12258 Err("cache_salt must be at most 64 bytes")
12259 );
12260 }
12261
12262 #[test]
12263 fn cache_salt_validation_rejects_reserved_open_namespace() {
12264 let salt = Some("t:acme\u{1f}private".to_string());
12265 assert_eq!(
12266 validate_cache_namespace(&salt, false),
12267 Err("cache_salt must not use the reserved t: prefix without a keyring")
12268 );
12269 }
12270
12271 #[test]
12272 fn cache_salt_validation_accepts_normal_value() {
12273 let salt = Some("tenant-A_7.c2VjcmV0LXNjb3Bl+/=".to_string());
12274 assert_eq!(
12275 validate_cache_namespace(&salt, false).unwrap(),
12276 salt.unwrap()
12277 );
12278 assert_eq!(validate_cache_namespace(&None, false).unwrap(), "");
12279 let max = Some("a".repeat(CACHE_SALT_MAX_BYTES));
12280 assert_eq!(validate_cache_namespace(&max, false).unwrap(), max.unwrap());
12281 }
12282
12283 #[test]
12284 fn cache_salt_validation_rejects_unsupported_characters() {
12285 let salt = Some("tenant salt".to_string());
12286 assert_eq!(
12287 validate_cache_namespace(&salt, false),
12288 Err("cache_salt contains unsupported characters")
12289 );
12290 }
12291
12292 #[test]
12293 fn affinity_key_honors_both_client_conventions_in_priority_order() {
12294 use axum::http::HeaderMap;
12295 let hdr = |v: &str| {
12296 let mut h = HeaderMap::new();
12297 h.insert("x-session-id", v.parse().unwrap());
12298 h
12299 };
12300 let empty = HeaderMap::new();
12301 let s = |v: &str| Some(v.to_string());
12302 assert_eq!(affinity_key(&s("explicit"), &None, &empty), s("explicit"));
12304 assert_eq!(
12305 affinity_key(&None, &s("openai-user"), &empty),
12306 s("openai-user")
12307 );
12308 assert_eq!(affinity_key(&None, &None, &hdr("hdr-id")), s("hdr-id"));
12309 assert_eq!(affinity_key(&s("a"), &s("b"), &hdr("c")), s("a"));
12312 assert_eq!(affinity_key(&None, &s("b"), &hdr("c")), s("b"));
12313 assert_eq!(affinity_key(&s(" "), &s(""), &hdr(" ")), None);
12316 assert_eq!(affinity_key(&s(""), &s("real"), &empty), s("real"));
12317 assert_eq!(affinity_key(&s(" padded "), &None, &empty), s("padded"));
12319 assert_eq!(affinity_key(&None, &None, &empty), None);
12321 }
12322
12323 #[test]
12324 fn affinity_key_plumbs_to_the_worker_request_on_both_bodies() {
12325 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12326 "model": "m", "prompt": "task", "session_id": "conv-1"
12327 }))
12328 .unwrap();
12329 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12330 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12331 assert_eq!(
12332 build_request(&req, tx, lanes::Lane::Interactive, key)
12333 .affinity
12334 .as_deref(),
12335 Some("conv-1")
12336 );
12337 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
12339 "model": "m", "messages": [{"role": "user", "content": "task"}],
12340 "user": "conv-2"
12341 }))
12342 .unwrap();
12343 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12344 let key = affinity_key(&req.session_id, &req.user, &axum::http::HeaderMap::new());
12345 assert_eq!(
12346 build_chat_request(req, None, tx, lanes::Lane::Interactive, key)
12347 .unwrap()
12348 .request
12349 .affinity
12350 .as_deref(),
12351 Some("conv-2")
12352 );
12353 let req: CompletionReq = serde_json::from_value(serde_json::json!({
12355 "model": "m", "prompt": "task"
12356 }))
12357 .unwrap();
12358 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
12359 assert!(
12360 build_request(&req, tx, lanes::Lane::Interactive, None)
12361 .affinity
12362 .is_none()
12363 );
12364 }
12365
12366 async fn sse_data_lines(resp: Response) -> Vec<String> {
12368 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12369 .await
12370 .unwrap();
12371 String::from_utf8(bytes.to_vec())
12372 .unwrap()
12373 .lines()
12374 .filter_map(|l| l.strip_prefix("data: ").map(str::to_string))
12375 .collect()
12376 }
12377
12378 #[tokio::test]
12379 async fn chat_returns_reasoning_text_when_on_and_no_field_when_off() {
12380 let feed = |think: bool| {
12386 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12387 let body = if think {
12388 "a plan</think>\n\nanswer"
12389 } else {
12390 "answer"
12391 };
12392 tx.send(Event::Token {
12393 id: 1,
12394 text: body.into(),
12395 })
12396 .unwrap();
12397 tx.send(Event::Done {
12398 stop_reason: "Eos".into(),
12399 n_tokens: 3,
12400 n_prompt: 10,
12401 n_cached: 0,
12402 elapsed_s: 0.1,
12403 spec: None,
12404 })
12405 .unwrap();
12406 drop(tx);
12407 rx
12408 };
12409 let resp = blocking_response(
12411 feed(true),
12412 "m".into(),
12413 true,
12414 Vec::new(),
12415 Some(ToolStreamParser::reasoning_only()),
12416 Envelope::new(true),
12417 )
12418 .await;
12419 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12420 .await
12421 .unwrap();
12422 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12423 assert_eq!(v["choices"][0]["message"]["reasoning"], "a plan");
12424 assert_eq!(
12425 v["choices"][0]["message"]["reasoning_details"][0]["text"],
12426 "a plan"
12427 );
12428 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12429 let resp = blocking_response(
12432 feed(false),
12433 "m".into(),
12434 true,
12435 Vec::new(),
12436 None,
12437 Envelope::new(true),
12438 )
12439 .await;
12440 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12441 .await
12442 .unwrap();
12443 let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12444 assert!(
12445 v["choices"][0]["message"].get("reasoning").is_none(),
12446 "a reasoning-off response must carry no reasoning field: {v}"
12447 );
12448 assert_eq!(v["choices"][0]["message"]["content"], "answer");
12449 let resp = sse_response(
12451 feed(true),
12452 "m".into(),
12453 true,
12454 Some(ToolStreamParser::reasoning_only()),
12455 Envelope::new(true),
12456 Vec::new(),
12457 None,
12458 )
12459 .into_response();
12460 let lines = sse_data_lines(resp).await;
12461 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12462 .iter()
12463 .map(|l| serde_json::from_str(l).unwrap())
12464 .collect();
12465 let reasoning: String = chunks
12466 .iter()
12467 .filter_map(|c| c["choices"][0]["delta"]["reasoning"].as_str())
12468 .collect();
12469 assert_eq!(
12470 reasoning, "a plan",
12471 "think text must stream as delta.reasoning"
12472 );
12473 let content: String = chunks
12474 .iter()
12475 .filter_map(|c| c["choices"][0]["delta"]["content"].as_str())
12476 .collect();
12477 assert_eq!(content, "answer", "content must exclude the think segment");
12478 let resp = sse_response(
12480 feed(false),
12481 "m".into(),
12482 true,
12483 None,
12484 Envelope::new(true),
12485 Vec::new(),
12486 None,
12487 )
12488 .into_response();
12489 let lines = sse_data_lines(resp).await;
12490 for l in &lines[..lines.len() - 1] {
12491 let c: serde_json::Value = serde_json::from_str(l).unwrap();
12492 assert!(
12493 c["choices"][0]["delta"].get("reasoning").is_none(),
12494 "a reasoning-off stream must carry no reasoning deltas: {c}"
12495 );
12496 }
12497 }
12498
12499 #[tokio::test]
12500 async fn stream_chunks_carry_envelope_and_first_delta_role() {
12501 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12502 tx.send(Event::Token {
12503 id: 1,
12504 text: "he".into(),
12505 })
12506 .unwrap();
12507 tx.send(Event::Token {
12508 id: 2,
12509 text: "llo".into(),
12510 })
12511 .unwrap();
12512 tx.send(Event::Done {
12513 stop_reason: "Eos".into(),
12514 n_tokens: 2,
12515 n_prompt: 10,
12516 n_cached: 0,
12517 elapsed_s: 0.1,
12518 spec: None,
12519 })
12520 .unwrap();
12521 drop(tx);
12522 let resp = sse_response(
12523 rx,
12524 "m".into(),
12525 true,
12526 None,
12527 Envelope::new(true),
12528 Vec::new(),
12529 None,
12530 )
12531 .into_response();
12532 let lines = sse_data_lines(resp).await;
12533 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12534 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12535 .iter()
12536 .map(|l| serde_json::from_str(l).unwrap())
12537 .collect();
12538 let id = chunks[0]["id"].as_str().unwrap().to_string();
12540 assert!(id.starts_with("chatcmpl-"));
12541 for c in &chunks {
12542 assert_eq!(c["id"], id.as_str());
12543 assert!(c["created"].as_u64().unwrap() > 1_700_000_000);
12544 assert!(
12545 c["system_fingerprint"]
12546 .as_str()
12547 .unwrap()
12548 .starts_with("memra-")
12549 );
12550 assert_eq!(c["object"], "chat.completion.chunk");
12551 }
12552 assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant");
12554 assert_eq!(chunks[0]["choices"][0]["delta"]["content"], "he");
12555 assert!(chunks[1]["choices"][0]["delta"].get("role").is_none());
12556 let fin = chunks.last().unwrap();
12558 assert_eq!(fin["choices"][0]["finish_reason"], "stop");
12559 assert_eq!(fin["usage"]["prompt_tokens"], 10);
12560 }
12561
12562 #[tokio::test]
12563 async fn stream_token_events_equal_usage_on_every_finish_path() {
12564 for (stop_reason, expected_finish) in [
12565 ("Eos", "stop"),
12566 ("Callback", "stop"),
12567 ("MaxNew", "length"),
12568 ("ContextFull", "length"),
12569 ] {
12570 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12571 tx.send(Event::Token {
12574 id: 248_046,
12575 text: String::new(),
12576 })
12577 .unwrap();
12578 tx.send(Event::Done {
12579 stop_reason: stop_reason.into(),
12580 n_tokens: 1,
12581 n_prompt: 8,
12582 n_cached: 8,
12583 elapsed_s: 0.1,
12584 spec: None,
12585 })
12586 .unwrap();
12587 drop(tx);
12588
12589 let resp = sse_response(
12590 rx,
12591 "m".into(),
12592 true,
12593 None,
12594 Envelope::new(true),
12595 Vec::new(),
12596 None,
12597 )
12598 .into_response();
12599 let lines = sse_data_lines(resp).await;
12600 assert_eq!(lines.last().map(String::as_str), Some("[DONE]"));
12601 let chunks: Vec<serde_json::Value> = lines[..lines.len() - 1]
12602 .iter()
12603 .map(|line| serde_json::from_str(line).unwrap())
12604 .collect();
12605 let token_events = chunks
12606 .iter()
12607 .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null())
12608 .count();
12609 let terminal = chunks.last().unwrap();
12610 assert_eq!(token_events, 1, "{stop_reason} SSE token count");
12611 assert_eq!(terminal["usage"]["completion_tokens"], token_events);
12612 assert_eq!(terminal["choices"][0]["finish_reason"], expected_finish);
12613 }
12614 }
12615
12616 #[tokio::test]
12617 async fn stream_excludes_stop_text_like_non_stream_does() {
12618 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12622 tx.send(Event::Token {
12623 id: 1,
12624 text: "answer\nPro".into(),
12625 })
12626 .unwrap();
12627 tx.send(Event::Token {
12628 id: 2,
12629 text: "blem: leaked prompt".into(),
12630 })
12631 .unwrap();
12632 tx.send(Event::Done {
12633 stop_reason: "Callback".into(),
12634 n_tokens: 2,
12635 n_prompt: 8,
12636 n_cached: 0,
12637 elapsed_s: 0.1,
12638 spec: None,
12639 })
12640 .unwrap();
12641 drop(tx);
12642 let resp = sse_response(
12643 rx,
12644 "m".into(),
12645 true,
12646 None,
12647 Envelope::new(true),
12648 vec!["Problem:".into()],
12649 None,
12650 )
12651 .into_response();
12652 let lines = sse_data_lines(resp).await;
12653 let content: String = lines
12654 .iter()
12655 .filter(|l| *l != "[DONE]")
12656 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12657 .filter_map(|c| {
12658 c["choices"][0]["delta"]["content"]
12659 .as_str()
12660 .map(str::to_string)
12661 })
12662 .collect();
12663 assert_eq!(content, "answer\n");
12664
12665 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12667 tx.send(Event::Token {
12668 id: 1,
12669 text: "ends in Pro".into(),
12670 })
12671 .unwrap();
12672 tx.send(Event::Done {
12673 stop_reason: "Eos".into(),
12674 n_tokens: 1,
12675 n_prompt: 8,
12676 n_cached: 0,
12677 elapsed_s: 0.1,
12678 spec: None,
12679 })
12680 .unwrap();
12681 drop(tx);
12682 let resp = sse_response(
12683 rx,
12684 "m".into(),
12685 true,
12686 None,
12687 Envelope::new(true),
12688 vec!["Problem:".into()],
12689 None,
12690 )
12691 .into_response();
12692 let lines = sse_data_lines(resp).await;
12693 let content: String = lines
12694 .iter()
12695 .filter(|l| *l != "[DONE]")
12696 .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
12697 .filter_map(|c| {
12698 c["choices"][0]["delta"]["content"]
12699 .as_str()
12700 .map(str::to_string)
12701 })
12702 .collect();
12703 assert_eq!(content, "ends in Pro");
12704 }
12705
12706 #[tokio::test]
12707 async fn stream_worker_error_is_a_data_chunk_not_a_named_event() {
12708 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12709 tx.send(Event::Error(worker::EngineError::engine("boom")))
12710 .unwrap();
12711 drop(tx);
12712 let resp = sse_response(
12713 rx,
12714 "m".into(),
12715 true,
12716 None,
12717 Envelope::new(true),
12718 Vec::new(),
12719 None,
12720 )
12721 .into_response();
12722 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12723 .await
12724 .unwrap();
12725 let body = String::from_utf8(bytes.to_vec()).unwrap();
12726 assert!(
12728 !body.contains("event: error"),
12729 "named SSE event leaked: {body}"
12730 );
12731 let lines: Vec<&str> = body
12732 .lines()
12733 .filter_map(|l| l.strip_prefix("data: "))
12734 .collect();
12735 let err: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
12736 assert_eq!(err["error"]["message"], "boom");
12737 assert_eq!(err["error"]["type"], "server_error");
12738 assert_eq!(err["error"]["code"], "engine_error");
12739 assert_eq!(lines.last(), Some(&"[DONE]"));
12740 }
12741
12742 #[test]
12743 fn ttft_sse_marker_ignores_keepalive_comments() {
12744 assert!(!is_sse_data_frame(b": keep-alive\n\n"));
12745 assert!(is_sse_data_frame(b"data: {\"choices\":[]}\n\n"));
12746 assert!(is_sse_data_frame(
12747 b"event: error\ndata: {\"error\":\"failed\"}\n\n"
12748 ));
12749 }
12750
12751 #[tokio::test]
12752 async fn error_bodies_use_the_openai_object_shape() {
12753 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
12754 tx.send(Event::Error(worker::EngineError::model_not_found(
12755 "unknown model \"x\"",
12756 )))
12757 .unwrap();
12758 drop(tx);
12759 let response =
12760 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
12761 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
12762 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
12763 .await
12764 .unwrap();
12765 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
12766 assert_eq!(payload["error"]["message"], "unknown model \"x\"");
12768 assert_eq!(payload["error"]["type"], "invalid_request_error");
12769 assert_eq!(payload["error"]["param"], "model");
12770 assert_eq!(payload["error"]["code"], "model_not_found");
12771 }
12772
12773 fn retry_after(resp: &Response) -> Option<String> {
12780 resp.headers()
12781 .get(axum::http::header::RETRY_AFTER)
12782 .and_then(|v| v.to_str().ok())
12783 .map(str::to_string)
12784 }
12785
12786 async fn body_value(resp: Response) -> serde_json::Value {
12789 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
12790 .await
12791 .expect("body");
12792 serde_json::from_slice(&bytes).expect("json body")
12793 }
12794
12795 #[test]
12796 fn timeout_ms_parses_clamps_nothing_and_names_every_refusal() {
12797 assert_eq!(parse_timeout_ms(None).unwrap(), TIMEOUT_MS_DEFAULT);
12799 assert_eq!(
12800 parse_timeout_ms(Some(&serde_json::Value::Null)).unwrap(),
12801 TIMEOUT_MS_DEFAULT
12802 );
12803 for ms in [TIMEOUT_MS_MIN, 5_000, 45_000, TIMEOUT_MS_MAX] {
12807 assert_eq!(parse_timeout_ms(Some(&json!(ms))).unwrap(), ms);
12808 }
12809 for bad in [0u64, TIMEOUT_MS_MIN - 1, TIMEOUT_MS_MAX + 1, 600_000] {
12811 let err = parse_timeout_ms(Some(&json!(bad))).expect_err("out of range must refuse");
12812 assert!(err.contains("timeout_ms"), "{err}");
12813 assert!(
12814 err.contains(&TIMEOUT_MS_MIN.to_string())
12815 && err.contains(&TIMEOUT_MS_MAX.to_string()),
12816 "the message must state the range: {err}"
12817 );
12818 assert!(
12819 err.contains("stream"),
12820 "the message must point at streaming for longer work: {err}"
12821 );
12822 }
12823 for bad in [json!("30s"), json!(1.5), json!(true), json!({}), json!([])] {
12825 let err = parse_timeout_ms(Some(&bad)).expect_err("bad type must refuse");
12826 assert!(
12827 err.contains("timeout_ms") && err.contains("stream"),
12828 "{err}"
12829 );
12830 }
12831 assert!(parse_timeout_ms(Some(&json!(-1))).is_err());
12833 }
12834
12835 #[tokio::test]
12838 async fn a_bad_timeout_ms_is_the_same_named_400_on_every_surface() {
12839 let _l = DRAIN_LOCK.lock().unwrap();
12840 let st = fake_worker_state();
12841
12842 let comp = completions(
12843 State(st.clone()),
12844 HeaderMap::new(),
12845 None,
12846 Json(
12847 serde_json::from_value(json!({
12848 "model": "m", "prompt": "t", "timeout_ms": 90_001}))
12849 .unwrap(),
12850 ),
12851 )
12852 .await;
12853 assert_eq!(comp.status(), StatusCode::BAD_REQUEST);
12854 let chat = chat_completions(
12855 State(st.clone()),
12856 HeaderMap::new(),
12857 None,
12858 Json(
12859 serde_json::from_value(json!({
12860 "model": "m", "messages": [{"role": "user", "content": "t"}],
12861 "timeout_ms": 90_001}))
12862 .unwrap(),
12863 ),
12864 )
12865 .await;
12866 assert_eq!(chat.status(), StatusCode::BAD_REQUEST);
12867 let resp_api = responses_api::responses(
12868 State(st.clone()),
12869 HeaderMap::new(),
12870 None,
12871 axum::body::Bytes::from(
12872 json!({"model": "m", "input": "t", "timeout_ms": 90_001}).to_string(),
12873 ),
12874 )
12875 .await;
12876 assert_eq!(resp_api.status(), StatusCode::BAD_REQUEST);
12877 let msgs = anthropic::messages(
12878 State(st.clone()),
12879 HeaderMap::new(),
12880 None,
12881 axum::body::Bytes::from(
12882 json!({"model": "m", "max_tokens": 16,
12883 "messages": [{"role": "user", "content": "t"}],
12884 "timeout_ms": 90_001})
12885 .to_string(),
12886 ),
12887 )
12888 .await;
12889 assert_eq!(msgs.status(), StatusCode::BAD_REQUEST);
12890
12891 for (surface, resp) in [
12893 ("/v1/completions", comp),
12894 ("/v1/chat/completions", chat),
12895 ("/v1/responses", resp_api),
12896 ] {
12897 let body = body_value(resp).await;
12898 assert_eq!(body["error"]["type"], "invalid_request_error", "{surface}");
12899 assert_eq!(body["error"]["param"], "timeout_ms", "{surface}");
12900 let m = body["error"]["message"].as_str().unwrap();
12901 assert!(
12902 m.contains("90000") && m.contains("stream"),
12903 "{surface}: {m}"
12904 );
12905 }
12906 let body = body_value(msgs).await;
12908 assert_eq!(body["error"]["type"], "invalid_request_error");
12909 let m = body["error"]["message"].as_str().unwrap();
12910 assert!(m.contains("timeout_ms") && m.contains("stream"), "{m}");
12911 }
12912
12913 #[tokio::test]
12916 async fn a_non_integer_timeout_ms_is_a_named_400() {
12917 let _l = DRAIN_LOCK.lock().unwrap();
12918 let st = fake_worker_state();
12919 let resp = chat_completions(
12920 State(st),
12921 HeaderMap::new(),
12922 None,
12923 Json(
12924 serde_json::from_value(json!({
12925 "model": "m", "messages": [{"role": "user", "content": "t"}],
12926 "timeout_ms": "30s"}))
12927 .unwrap(),
12928 ),
12929 )
12930 .await;
12931 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
12932 let body = body_value(resp).await;
12933 assert_eq!(body["error"]["param"], "timeout_ms");
12934 }
12935
12936 #[tokio::test]
12942 async fn a_missed_non_stream_deadline_delivers_the_partial_bills_it_and_cancels_generation() {
12943 let _l = DRAIN_LOCK.lock().unwrap();
12944 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
12949 let cancel_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));
12950 let worker_cancel = cancel_seen.clone();
12951 let health = health::WorkerHealth::new();
12952 let h = health.clone();
12953 std::thread::spawn(move || {
12954 h.mark_ready();
12955 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
12956 worker::release_pending_admit();
12957 worker::release_admission_reservation(req.lane);
12958 let _ = req.tx.send(Event::PromptUsage {
12959 n_prompt: 1,
12960 n_cached: 0,
12961 });
12962 let _ = req.tx.send(Event::Token {
12963 id: 1,
12964 text: "partial".into(),
12965 });
12966 for _ in 0..5_000 {
12969 if req.tx.is_closed() {
12970 worker_cancel.store(true, std::sync::atomic::Ordering::SeqCst);
12971 break;
12972 }
12973 std::thread::sleep(std::time::Duration::from_millis(1));
12974 }
12975 }
12976 });
12977 for _ in 0..2_000 {
12978 if health.live().is_ok() {
12979 break;
12980 }
12981 std::thread::sleep(std::time::Duration::from_millis(1));
12982 }
12983 let mut st = fake_worker_state();
12984 st.cmd_tx = cmd_tx;
12985 st.health = health;
12986 let mock = MockMetering::admit_all();
12987 st.metering = Some(mock.clone());
12988
12989 let resp = chat_completions(
12990 State(st),
12991 HeaderMap::new(),
12992 None,
12993 Json(
12994 serde_json::from_value(json!({
12995 "model": "m", "messages": [{"role": "user", "content": "t"}],
12996 "timeout_ms": 1_000}))
12997 .unwrap(),
12998 ),
12999 )
13000 .await;
13001
13002 assert_eq!(resp.status(), StatusCode::OK);
13007 let body = body_value(resp).await;
13008 assert!(
13009 body["choices"][0]["message"]["content"]
13010 .as_str()
13011 .unwrap()
13012 .contains("partial"),
13013 "the tokens generated before the cut must be delivered: {body}"
13014 );
13015 assert_eq!(body["choices"][0]["finish_reason"], "error");
13019 assert_eq!(
13020 body["choices"][0]["native_finish_reason"],
13021 "deadline_exceeded"
13022 );
13023 assert_eq!(body["error"]["code"], "deadline_exceeded");
13024 assert_eq!(body["error"]["metadata"]["error_type"], "timeout");
13025 let message = body["error"]["message"].as_str().unwrap();
13026 assert!(
13027 message.contains("1000") && message.contains("stream"),
13028 "the partial must name the deadline and the streaming alternative: {message}"
13029 );
13030 assert_eq!(body["usage"]["completion_tokens"], 1);
13031
13032 let mut cancelled = false;
13037 for _ in 0..500 {
13038 if cancel_seen.load(std::sync::atomic::Ordering::SeqCst) {
13039 cancelled = true;
13040 break;
13041 }
13042 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
13043 }
13044 assert!(
13045 cancelled,
13046 "the deadline must CANCEL generation (worker's event channel closed)"
13047 );
13048
13049 let events = mock.events();
13054 assert!(
13055 events.contains(&MeterEvent::DeadlinePartial {
13056 prompt: 1,
13057 cached: 0,
13058 completion: 1,
13059 }),
13060 "the partial must settle as a deadline-partial with worker-truth counts: {events:?}"
13061 );
13062 assert!(
13063 !events
13064 .iter()
13065 .any(|e| matches!(e, MeterEvent::Complete { .. })),
13066 "a deadline cut must stay distinguishable from a full answer: {events:?}"
13067 );
13068 }
13069
13070 #[tokio::test]
13074 async fn a_deadline_missed_before_any_token_is_still_408_and_unbilled() {
13075 let _l = DRAIN_LOCK.lock().unwrap();
13076 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13077 let health = health::WorkerHealth::new();
13078 let h = health.clone();
13079 std::thread::spawn(move || {
13080 h.mark_ready();
13081 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
13084 worker::release_pending_admit();
13085 worker::release_admission_reservation(req.lane);
13086 let _ = req.tx.send(Event::PromptUsage {
13087 n_prompt: 1,
13088 n_cached: 0,
13089 });
13090 for _ in 0..5_000 {
13091 if req.tx.is_closed() {
13092 break;
13093 }
13094 std::thread::sleep(std::time::Duration::from_millis(1));
13095 }
13096 }
13097 });
13098 for _ in 0..2_000 {
13099 if health.live().is_ok() {
13100 break;
13101 }
13102 std::thread::sleep(std::time::Duration::from_millis(1));
13103 }
13104 let mut st = fake_worker_state();
13105 st.cmd_tx = cmd_tx;
13106 st.health = health;
13107 let mock = MockMetering::admit_all();
13108 st.metering = Some(mock.clone());
13109 let resp = chat_completions(
13110 State(st),
13111 HeaderMap::new(),
13112 None,
13113 Json(
13114 serde_json::from_value(json!({
13115 "model": "m", "messages": [{"role": "user", "content": "t"}],
13116 "timeout_ms": 1_000}))
13117 .unwrap(),
13118 ),
13119 )
13120 .await;
13121 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
13122 assert!(resp.headers().get("x-should-retry").is_none());
13124 assert_eq!(retry_after(&resp), None);
13125 let body = body_value(resp).await;
13126 assert_eq!(body["error"]["code"], "deadline_exceeded");
13127 assert!(
13128 body["error"]["message"]
13129 .as_str()
13130 .unwrap()
13131 .contains("not billed"),
13132 "the zero-token 408 keeps the billing promise: {body}"
13133 );
13134 let events = mock.events();
13135 assert!(
13136 events.contains(&MeterEvent::Unbilled {
13137 outcome: "deadline_exceeded",
13138 status: 408,
13139 code: "deadline_exceeded".into(),
13140 }),
13141 "the named zero-debit census outcome, not the generic reject — every sibling \
13142 deadline path settles this one: {events:?}"
13143 );
13144 }
13145
13146 #[tokio::test]
13149 async fn a_stream_that_misses_ttft_is_a_preheader_408_and_not_billed() {
13150 let _l = DRAIN_LOCK.lock().unwrap();
13151 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13153 let health = health::WorkerHealth::new();
13154 let h = health.clone();
13155 std::thread::spawn(move || {
13156 h.mark_ready();
13157 while let Ok(Cmd::Generate(req)) = cmd_rx.recv() {
13158 worker::release_pending_admit();
13159 worker::release_admission_reservation(req.lane);
13160 let _ = req.tx.send(Event::PromptUsage {
13161 n_prompt: 1,
13162 n_cached: 0,
13163 });
13164 while !req.tx.is_closed() {
13165 std::thread::sleep(std::time::Duration::from_millis(1));
13166 }
13167 }
13168 });
13169 for _ in 0..2_000 {
13170 if health.live().is_ok() {
13171 break;
13172 }
13173 std::thread::sleep(std::time::Duration::from_millis(1));
13174 }
13175 let mut st = fake_worker_state();
13176 st.cmd_tx = cmd_tx;
13177 st.health = health;
13178 let mock = MockMetering::admit_all();
13179 st.metering = Some(mock.clone());
13180
13181 let resp = chat_completions(
13182 State(st),
13183 HeaderMap::new(),
13184 None,
13185 Json(
13186 serde_json::from_value(json!({
13187 "model": "m", "messages": [{"role": "user", "content": "t"}],
13188 "stream": true, "timeout_ms": 1_000}))
13189 .unwrap(),
13190 ),
13191 )
13192 .await;
13193 assert_eq!(resp.status(), StatusCode::REQUEST_TIMEOUT);
13196 let body = body_value(resp).await;
13197 assert_eq!(body["error"]["code"], "deadline_exceeded");
13198 assert!(
13199 body["error"]["message"]
13200 .as_str()
13201 .unwrap()
13202 .contains("first token"),
13203 "the streaming message must say the deadline bounded TIME TO FIRST TOKEN: {body}"
13204 );
13205 let events = mock.events();
13206 assert!(
13207 events.contains(&MeterEvent::Unbilled {
13208 outcome: "deadline_exceeded",
13209 status: 408,
13210 code: "deadline_exceeded".into(),
13211 }),
13212 "a TTFT miss must settle unbilled under the deadline outcome: {events:?}"
13213 );
13214 }
13215
13216 #[tokio::test]
13220 async fn a_stream_is_immune_to_the_deadline_after_its_first_token() {
13221 let _l = DRAIN_LOCK.lock().unwrap();
13222 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(400));
13225 let mock = MockMetering::admit_all();
13226 st.metering = Some(mock.clone());
13227 let resp = chat_completions(
13228 State(st),
13229 HeaderMap::new(),
13230 None,
13231 Json(
13232 serde_json::from_value(json!({
13233 "model": "m", "messages": [{"role": "user", "content": "t"}],
13234 "stream": true, "timeout_ms": 1_000}))
13235 .unwrap(),
13236 ),
13237 )
13238 .await;
13239 assert_eq!(
13240 resp.status(),
13241 StatusCode::OK,
13242 "TTFT was met — 200 is correct"
13243 );
13244 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13245 .await
13246 .expect("the stream must run to completion past the deadline");
13247 let text = String::from_utf8(bytes.to_vec()).unwrap();
13248 assert!(text.contains("[DONE]"), "stream did not complete: {text}");
13249 let events = mock.events();
13250 assert!(
13251 events
13252 .iter()
13253 .any(|e| matches!(e, MeterEvent::Complete { completion: 4, .. })),
13254 "a stream past its deadline after first token still settles as COMPLETE with \
13255 all four tokens: {events:?}"
13256 );
13257 }
13258
13259 #[test]
13262 fn the_queue_bound_sheds_with_429_retry_after_and_the_ratelimit_trio() {
13263 let st = fake_worker_state();
13264 let cap = lane_cap(lanes::Lane::Interactive);
13266 {
13267 let mut m = st.metrics.lock().unwrap();
13268 m.completed = 10;
13269 m.tokens_out = 1_000; m.step_p50_ms = 10.0; m.queued_requests = (cap * 4 + 1) as u64;
13272 }
13273 let rl = RateLimit {
13274 limit: cap,
13275 remaining: 0,
13276 reset_s: 1,
13277 };
13278 let (resp, outcome) = admission_backpressure(
13279 &st,
13280 lanes::Lane::Interactive,
13281 &rl,
13282 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13283 )
13284 .expect_err("a backlog past the bound must shed");
13285 assert_eq!(outcome, "shed_queue");
13286 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13287 assert!(
13288 retry_after(&resp).is_some(),
13289 "a shed must carry Retry-After so the router's spill can act on it"
13290 );
13291 let stamped = rl.attach(resp);
13293 for h in [
13294 "x-ratelimit-limit",
13295 "x-ratelimit-remaining",
13296 "x-ratelimit-reset",
13297 ] {
13298 assert!(stamped.headers().get(h).is_some(), "missing {h}");
13299 }
13300 }
13301
13302 #[test]
13306 fn admission_sheds_only_when_the_estimated_wait_cannot_fit_the_deadline() {
13307 let st = fake_worker_state();
13308 let cap = lane_cap(lanes::Lane::Interactive);
13309 {
13310 let mut m = st.metrics.lock().unwrap();
13311 m.completed = 10;
13312 m.tokens_out = 1_000;
13313 m.step_p50_ms = 10.0; m.queued_requests = cap as u64; }
13316 let rl = RateLimit {
13317 limit: cap,
13318 remaining: 0,
13319 reset_s: 1,
13320 };
13321 assert!(
13323 admission_backpressure(
13324 &st,
13325 lanes::Lane::Interactive,
13326 &rl,
13327 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13328 )
13329 .is_ok(),
13330 "a request whose deadline covers the estimate must be admitted"
13331 );
13332 let (resp, outcome) = admission_backpressure(
13334 &st,
13335 lanes::Lane::Interactive,
13336 &rl,
13337 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13338 )
13339 .expect_err("a deadline shorter than the estimated wait must shed");
13340 assert_eq!(outcome, "shed_deadline");
13341 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13342 assert!(retry_after(&resp).is_some());
13343 }
13344
13345 #[test]
13348 fn admission_backpressure_is_interactive_only_and_silent_with_free_slots() {
13349 let st = fake_worker_state();
13350 let cap = lane_cap(lanes::Lane::Interactive);
13351 {
13352 let mut m = st.metrics.lock().unwrap();
13353 m.completed = 10;
13354 m.tokens_out = 100_000; m.step_p50_ms = 100.0;
13356 m.queued_requests = (cap * 100) as u64;
13357 }
13358 let free = RateLimit {
13360 limit: cap,
13361 remaining: 1,
13362 reset_s: 0,
13363 };
13364 assert!(
13365 admission_backpressure(
13366 &st,
13367 lanes::Lane::Interactive,
13368 &free,
13369 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13370 )
13371 .is_ok()
13372 );
13373 let full = RateLimit {
13375 limit: cap,
13376 remaining: 0,
13377 reset_s: 5,
13378 };
13379 for lane in [lanes::Lane::Judge, lanes::Lane::Harvest] {
13380 assert!(
13381 admission_backpressure(
13382 &st,
13383 lane,
13384 &full,
13385 RequestDeadline::starting_now(TIMEOUT_MS_MIN)
13386 )
13387 .is_ok(),
13388 "{lane:?} must not be shed by the interactive gate"
13389 );
13390 }
13391 }
13392
13393 #[test]
13394 fn pending_admission_reservation_is_atomic_and_rolls_back_on_drop() {
13395 let st = fake_worker_state();
13396 let cap = lane_cap(lanes::Lane::Interactive);
13397 let bound = max_queue_depth(cap);
13398 assert!(bound > 0, "the queue bound must admit at least one request");
13399 let rl = RateLimit {
13400 limit: cap,
13401 remaining: 0,
13402 reset_s: 1,
13403 };
13404 let _ = worker::PENDING_ADMITS.fetch_update(
13405 std::sync::atomic::Ordering::AcqRel,
13406 std::sync::atomic::Ordering::Acquire,
13407 |_| Some(0),
13408 );
13409 let counter = &worker::ADMISSION_RESERVATIONS[lanes::Lane::Interactive.idx()];
13410 counter.store(bound - 1, std::sync::atomic::Ordering::Release);
13411 let guard = reserve_pending_admit(
13412 &st,
13413 lanes::Lane::Interactive,
13414 &rl,
13415 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13416 )
13417 .expect("the final queue slot should be reservable");
13418 assert_eq!(
13419 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13420 1
13421 );
13422 assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), bound);
13423 drop(guard);
13424 assert_eq!(
13425 worker::PENDING_ADMITS.load(std::sync::atomic::Ordering::Acquire),
13426 0
13427 );
13428 assert_eq!(
13429 counter.load(std::sync::atomic::Ordering::Acquire),
13430 bound - 1
13431 );
13432
13433 counter.store(bound, std::sync::atomic::Ordering::Release);
13434 let rejected = reserve_pending_admit(
13435 &st,
13436 lanes::Lane::Interactive,
13437 &rl,
13438 RequestDeadline::starting_now(TIMEOUT_MS_MAX),
13439 );
13440 assert!(matches!(rejected, Err((_, "shed_queue"))));
13441 counter.store(0, std::sync::atomic::Ordering::Release);
13442 }
13443
13444 #[test]
13445 fn admission_reservations_are_lane_scoped() {
13446 let st = fake_worker_state();
13447 let harvest = lanes::Lane::Harvest;
13448 let interactive = lanes::Lane::Interactive;
13449 let harvest_counter = &worker::ADMISSION_RESERVATIONS[harvest.idx()];
13450 let interactive_counter = &worker::ADMISSION_RESERVATIONS[interactive.idx()];
13451 harvest_counter.store(
13452 max_queue_depth(lane_cap(harvest)),
13453 std::sync::atomic::Ordering::Release,
13454 );
13455 interactive_counter.store(0, std::sync::atomic::Ordering::Release);
13456 let free = RateLimit {
13457 limit: lane_cap(interactive),
13458 remaining: 1,
13459 reset_s: 0,
13460 };
13461 let guard = reserve_pending_admit(
13462 &st,
13463 interactive,
13464 &free,
13465 RequestDeadline::starting_now(TIMEOUT_MS_MIN),
13466 )
13467 .expect("a full harvest queue must not consume interactive capacity");
13468 drop(guard);
13469 let harvest_rl = RateLimit {
13470 limit: lane_cap(harvest),
13471 remaining: 0,
13472 reset_s: 1,
13473 };
13474 assert!(matches!(
13475 reserve_pending_admit(
13476 &st,
13477 harvest,
13478 &harvest_rl,
13479 RequestDeadline::starting_now(TIMEOUT_MS_MAX)
13480 ),
13481 Err((_, "shed_queue"))
13482 ));
13483 harvest_counter.store(0, std::sync::atomic::Ordering::Release);
13484 }
13485
13486 #[test]
13487 fn taxonomy_maps_every_class_to_its_status_and_code() {
13488 use worker::{EngineError as E, ErrClass as C};
13489 let cases: Vec<(worker::EngineError, StatusCode, &str, &str)> = vec![
13490 (
13491 E::invalid_param("bad json", "response_format"),
13492 StatusCode::BAD_REQUEST,
13493 "invalid_request_error",
13494 "",
13495 ),
13496 (
13497 E::context_length("prompt (9000 tok) >= context cap (8192)"),
13498 StatusCode::BAD_REQUEST,
13499 "invalid_request_error",
13500 "context_length_exceeded",
13501 ),
13502 (
13503 E::model_not_found("unknown model \"nope\""),
13504 StatusCode::BAD_REQUEST,
13505 "invalid_request_error",
13506 "model_not_found",
13507 ),
13508 (
13509 E::rate_limit("lane judge is at capacity, retry"),
13510 StatusCode::TOO_MANY_REQUESTS,
13511 "rate_limit_error",
13512 "rate_limit_exceeded",
13513 ),
13514 (
13515 E::overloaded("no VRAM for a new session"),
13516 StatusCode::SERVICE_UNAVAILABLE,
13517 "server_error",
13518 "overloaded",
13519 ),
13520 (
13521 E::engine("graph step failed: launch error"),
13522 StatusCode::INTERNAL_SERVER_ERROR,
13523 "server_error",
13524 "engine_error",
13525 ),
13526 ];
13527 for (err, want_status, want_type, want_code) in cases {
13528 let (status, etype, code) = class_http(err.class);
13529 assert_eq!(status, want_status, "{:?}", err);
13530 assert_eq!(etype, want_type, "{:?}", err);
13531 if !want_code.is_empty() {
13532 assert_eq!(code, Some(want_code), "{:?}", err);
13533 }
13534 let body = engine_error_body(&err);
13536 assert_eq!(body["error"]["message"], err.message);
13537 assert_eq!(body["error"]["type"], want_type);
13538 }
13539 for c in [
13541 C::InvalidRequest,
13542 C::ContextLength,
13543 C::ModelNotFound,
13544 C::RateLimit,
13545 C::Overloaded,
13546 C::Engine,
13547 ] {
13548 let (s, t, _) = class_http(c);
13549 assert!(s.is_client_error() || s.is_server_error(), "{c:?} -> {s}");
13550 assert!(!t.is_empty());
13551 }
13552 }
13553
13554 #[test]
13555 fn a_cuda_oom_message_is_capacity_503_not_a_500() {
13556 let e = worker::EngineError::engine(
13561 "step error: DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")",
13562 );
13563 let resp = engine_error_response(&e);
13564 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13565 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13566 }
13567
13568 #[test]
13569 fn retry_headers_follow_the_sdk_contract() {
13570 for e in [
13574 worker::EngineError::rate_limit("shed"),
13575 worker::EngineError::overloaded("no VRAM"),
13576 ] {
13577 let resp = engine_error_response(&e);
13578 let ra = retry_after(&resp).expect("retryable class must carry Retry-After");
13579 let secs: u64 = ra
13580 .parse()
13581 .expect("Retry-After must be integer delay-seconds");
13582 assert!(
13583 secs > 0 && secs <= 60,
13584 "Retry-After {secs}s outside the honored window"
13585 );
13586 let ms = resp
13587 .headers()
13588 .get("retry-after-ms")
13589 .unwrap()
13590 .to_str()
13591 .unwrap();
13592 assert_eq!(
13593 ms.parse::<u64>().unwrap(),
13594 secs * 1000,
13595 "the two headers disagree"
13596 );
13597 assert!(
13598 resp.headers().get("x-should-retry").is_none(),
13599 "a retryable class must not say x-should-retry: false"
13600 );
13601 }
13602 }
13603
13604 #[tokio::test]
13605 async fn command_send_failure_obeys_the_retry_contract() {
13606 let _l = DRAIN_LOCK.lock().unwrap();
13607 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
13608 let mut st = fake_worker_state();
13609 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
13610 drop(cmd_rx);
13611 st.cmd_tx = cmd_tx;
13612
13613 let completion = completions(
13614 State(st.clone()),
13615 axum::http::HeaderMap::new(),
13616 None,
13617 Json(
13618 serde_json::from_value(serde_json::json!({
13619 "model": "m", "prompt": "test"
13620 }))
13621 .unwrap(),
13622 ),
13623 )
13624 .await;
13625 let chat = chat_completions(
13626 State(st),
13627 axum::http::HeaderMap::new(),
13628 None,
13629 Json(
13630 serde_json::from_value(serde_json::json!({
13631 "model": "m", "messages": [{"role": "user", "content": "test"}]
13632 }))
13633 .unwrap(),
13634 ),
13635 )
13636 .await;
13637
13638 for resp in [completion, chat] {
13639 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13640 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13641 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
13642 assert_ne!(
13643 resp.headers()
13644 .get("x-should-retry")
13645 .and_then(|v| v.to_str().ok()),
13646 Some("false")
13647 );
13648 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13649 .await
13650 .unwrap();
13651 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13652 assert_eq!(payload["error"]["type"], "server_error");
13653 assert_eq!(payload["error"]["code"], "overloaded");
13654 }
13655 }
13656
13657 #[test]
13658 fn unfixable_client_errors_say_x_should_retry_false() {
13659 for e in [
13662 worker::EngineError::model_not_found("unknown model \"x\""),
13663 worker::EngineError::context_length("prompt too long"),
13664 worker::EngineError::invalid_param("bad", "messages"),
13665 ] {
13666 let resp = engine_error_response(&e);
13667 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13668 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
13669 assert!(
13670 retry_after(&resp).is_none(),
13671 "a 400 must not promise a retry window"
13672 );
13673 }
13674 }
13675
13676 #[tokio::test]
13677 async fn a_closed_worker_channel_is_503_not_500() {
13678 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
13682 drop(tx);
13683 let resp =
13684 blocking_response(rx, "m".into(), true, Vec::new(), None, Envelope::new(true)).await;
13685 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
13686 assert_eq!(retry_after(&resp).as_deref(), Some("5"));
13687 }
13688
13689 #[tokio::test]
13690 async fn a_dark_lane_shed_is_429_with_an_openai_object_body() {
13691 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13694 tx.send(Event::Error(worker::EngineError::rate_limit(
13695 "lane judge shed: interactive p99 over budget, retry",
13696 )))
13697 .unwrap();
13698 let (resp, error_code) = peek_admission(rx)
13699 .await
13700 .expect_err("a shed must not be forwarded into the stream");
13701 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13702 assert_eq!(error_code, "rate_limit_exceeded");
13703 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
13704 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
13705 .await
13706 .unwrap();
13707 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
13708 assert!(
13709 payload["error"].is_object(),
13710 "bare-string error body: {payload}"
13711 );
13712 assert_eq!(payload["error"]["type"], "rate_limit_error");
13713 assert!(
13714 payload["error"]["message"]
13715 .as_str()
13716 .unwrap()
13717 .contains("shed")
13718 );
13719 }
13720
13721 #[tokio::test]
13722 async fn interactive_admission_error_is_a_preheader_429() {
13723 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13726 tx.send(Event::Error(worker::EngineError::rate_limit(
13727 "KV capacity unavailable",
13728 )))
13729 .unwrap();
13730 let (resp, error_code) = peek_admission(rx)
13731 .await
13732 .expect_err("admission error must stay pre-header");
13733 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
13734 assert_eq!(error_code, "rate_limit_exceeded");
13735 }
13736
13737 #[tokio::test]
13738 async fn admission_peek_preserves_context_error_for_the_ledger() {
13739 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13740 tx.send(Event::Error(worker::EngineError::context_length(
13741 "prompt exceeds configured model maximum",
13742 )))
13743 .unwrap();
13744 let (resp, error_code) = peek_admission(rx)
13745 .await
13746 .expect_err("context rejection must stay pre-header");
13747 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
13748 assert_eq!(error_code, "context_length_exceeded");
13749 }
13750
13751 #[tokio::test]
13752 async fn admission_peek_replays_prompt_usage_without_waiting_for_a_token() {
13753 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
13754 tx.send(Event::PromptUsage {
13755 n_prompt: 262_143,
13756 n_cached: 0,
13757 })
13758 .unwrap();
13759 let mut replay = peek_admission(rx).await.expect("successful admission");
13760 assert!(matches!(
13761 replay.recv().await,
13762 Some(Event::PromptUsage {
13763 n_prompt: 262_143,
13764 n_cached: 0
13765 }),
13766 ));
13767 }
13768
13769 #[test]
13770 fn penalties_plumb_from_http_to_sampler_config() {
13771 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
13774 "model": "m", "messages": [{"role": "user", "content": "task"}],
13775 "frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1
13776 }))
13777 .unwrap();
13778 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13779 let cfg = build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13780 .unwrap()
13781 .request
13782 .sampler_cfg;
13783 assert_eq!(cfg.penalty_freq, 0.5);
13784 assert_eq!(cfg.penalty_present, 0.25);
13785 assert_eq!(cfg.penalty_repeat, 1.1);
13786 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
13787
13788 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13789 "model": "m", "prompt": "task", "frequency_penalty": 1.5
13790 }))
13791 .unwrap();
13792 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13793 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13794 assert_eq!(cfg.penalty_freq, 1.5);
13795 assert_eq!(cfg.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
13796
13797 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13799 "model": "m", "prompt": "task"
13800 }))
13801 .unwrap();
13802 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13803 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13804 assert_eq!(cfg.penalty_last_n, 0);
13805 assert_eq!(cfg.penalty_repeat, 1.0);
13806 }
13807
13808 #[test]
13809 fn omitted_temperature_is_openai_default_not_greedy() {
13810 let chat_temp = |body: serde_json::Value| {
13826 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13827 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13828 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
13829 .unwrap()
13830 .request
13831 .sampler_cfg
13832 .temperature
13833 };
13834 let comp_temp = |body: serde_json::Value| {
13835 let req: CompletionReq = serde_json::from_value(body).unwrap();
13836 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13837 build_request(&req, tx, lanes::Lane::Interactive, None)
13838 .sampler_cfg
13839 .temperature
13840 };
13841
13842 assert_eq!(
13844 chat_temp(serde_json::json!({
13845 "model": "m", "messages": [{"role": "user", "content": "t"}]})),
13846 1.0,
13847 "omitted chat temperature must be the OpenAI 1.0 default, not 0.0/greedy"
13848 );
13849 assert_eq!(
13850 comp_temp(serde_json::json!({
13851 "model": "m", "prompt": "t"})),
13852 1.0,
13853 "omitted completions temperature must be the OpenAI 1.0 default"
13854 );
13855
13856 assert_eq!(
13858 chat_temp(serde_json::json!({
13859 "model": "m", "messages": [{"role": "user", "content": "t"}],
13860 "temperature": 0.0})),
13861 0.0,
13862 "explicit temperature 0 must stay greedy"
13863 );
13864 assert_eq!(
13865 comp_temp(serde_json::json!({
13866 "model": "m", "prompt": "t", "temperature": 0})),
13867 0.0,
13868 "explicit temperature 0 must stay greedy"
13869 );
13870 assert!(
13872 memra_engine::sampler::Sampler::new(sampler_config(
13873 0.0,
13874 0,
13875 1.0,
13876 0.0,
13877 0.0,
13878 0.0,
13879 1.0,
13880 Some(0)
13881 ))
13882 .is_greedy()
13883 );
13884 assert!(
13885 !memra_engine::sampler::Sampler::new(sampler_config(
13886 1.0,
13887 0,
13888 1.0,
13889 0.0,
13890 0.0,
13891 0.0,
13892 1.0,
13893 Some(0)
13894 ))
13895 .is_greedy()
13896 );
13897
13898 assert_eq!(
13900 chat_temp(serde_json::json!({
13901 "model": "m", "messages": [{"role": "user", "content": "t"}],
13902 "temperature": 0.7})),
13903 0.7
13904 );
13905
13906 let req: CompletionReq = serde_json::from_value(serde_json::json!({
13910 "model": "m", "prompt": "t"}))
13911 .unwrap();
13912 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13913 let cfg = build_request(&req, tx, lanes::Lane::Interactive, None).sampler_cfg;
13914 assert_eq!(cfg.top_p, 1.0, "omitted top_p = OpenAI 1.0 = disabled");
13915 assert_eq!(cfg.top_k, 0, "omitted top_k = disabled");
13916 assert_eq!(cfg.min_p, 0.0, "omitted min_p = disabled");
13917 assert_eq!(cfg.penalty_last_n, 0, "omitted penalties = window off");
13918 assert!(
13923 memra_engine::sampler::Sampler::new(cfg).is_spec_sampling(),
13924 "the omitted-temperature default must ride sampled spec's pure-temp regime"
13925 );
13926 }
13927
13928 #[test]
13929 fn step35_chat_uses_published_sampling_defaults_only_when_omitted() {
13930 let caps = ModelCaps {
13931 chat_temperature_default: Some(0.5),
13932 chat_top_p_default: Some(0.9),
13933 chat_ok: true,
13934 ..Default::default()
13935 };
13936 let cfg = |extra: serde_json::Value| {
13937 let mut body = serde_json::json!({
13938 "model": "step35",
13939 "messages": [{"role": "user", "content": "task"}]
13940 });
13941 body.as_object_mut()
13942 .unwrap()
13943 .extend(extra.as_object().unwrap().clone());
13944 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
13945 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
13946 build_chat_request(req, Some(&caps), tx, lanes::Lane::Interactive, None)
13947 .unwrap()
13948 .request
13949 .sampler_cfg
13950 };
13951
13952 let omitted = cfg(serde_json::json!({}));
13953 assert_eq!(omitted.temperature, 0.5);
13954 assert_eq!(omitted.top_p, 0.9);
13955
13956 let explicit_temp = cfg(serde_json::json!({"temperature": 0.7}));
13957 assert_eq!(explicit_temp.temperature, 0.7);
13958 assert_eq!(
13959 explicit_temp.top_p, 0.9,
13960 "omitting top_p must retain StepFun's nucleus default"
13961 );
13962
13963 let explicit = cfg(serde_json::json!({"temperature": 0.0, "top_p": 1.0}));
13964 assert_eq!(
13965 explicit.temperature, 0.0,
13966 "explicit greedy must remain authoritative"
13967 );
13968 assert_eq!(
13969 explicit.top_p, 1.0,
13970 "explicit untruncated sampling must remain authoritative"
13971 );
13972 }
13973
13974 fn qwen38_vendor_defaults() -> SamplingDefaults {
13978 SamplingDefaults {
13979 temperature: Some(1.0),
13980 top_p: Some(0.95),
13981 top_k: Some(20),
13982 min_p: Some(0.0),
13983 presence_penalty: Some(0.0),
13984 repetition_penalty: Some(1.0),
13985 frequency_penalty: None,
13986 }
13987 }
13988
13989 fn gemma4_vendor_defaults() -> SamplingDefaults {
13994 SamplingDefaults {
13995 temperature: Some(1.0),
13996 top_p: Some(0.95),
13997 top_k: Some(64),
13998 ..Default::default()
13999 }
14000 }
14001
14002 #[test]
14003 fn vendor_sampling_defaults_fill_only_the_omitted_fields() {
14004 let d = ModelSamplingDefaults::single(gemma4_vendor_defaults());
14009 let chat = |extra: serde_json::Value| {
14010 let mut body = serde_json::json!({
14011 "model": "google/gemma-4-31b-it",
14012 "messages": [{"role": "user", "content": "task"}],
14013 "seed": 7
14015 });
14016 body.as_object_mut()
14017 .unwrap()
14018 .extend(extra.as_object().unwrap().clone());
14019 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14020 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14021 build_chat_request_with_trace(
14022 req,
14023 Some(&ModelCaps {
14024 chat_ok: true,
14025 ..Default::default()
14026 }),
14027 tx,
14028 lanes::Lane::Interactive,
14029 None,
14030 None,
14031 None,
14032 &d,
14033 )
14034 .unwrap()
14035 .request
14036 .sampler_cfg
14037 };
14038
14039 let omitted = chat(serde_json::json!({}));
14041 assert_eq!(omitted.temperature, 1.0, "gemma-4 card temperature");
14042 assert_eq!(omitted.top_p, 0.95, "gemma-4 card top_p");
14043 assert_eq!(omitted.top_k, 64, "gemma-4 card top_k");
14044 assert_eq!(omitted.min_p, 0.0, "undeclared min_p stays API-standard");
14046 assert_eq!(omitted.penalty_repeat, 1.0);
14047 assert_eq!(omitted.penalty_freq, 0.0);
14048 assert_eq!(omitted.penalty_present, 0.0);
14049 assert_eq!(omitted.penalty_last_n, 0, "no penalty => no history window");
14050 assert!(
14051 !memra_engine::sampler::Sampler::new(omitted).is_greedy(),
14052 "the vendor default must NOT be greedy — that is the whole point of the lane"
14053 );
14054
14055 let greedy = chat(serde_json::json!({"temperature": 0}));
14058 assert_eq!(
14059 greedy.temperature, 0.0,
14060 "explicit temperature 0 stays greedy"
14061 );
14062 assert!(
14063 memra_engine::sampler::Sampler::new(greedy).is_greedy(),
14064 "an explicit temperature 0 must satisfy the greedy predicate that gates the \
14065 spec/graph exactness arms"
14066 );
14067
14068 let one_field = chat(serde_json::json!({"top_k": 3}));
14070 assert_eq!(one_field.top_k, 3, "explicit top_k wins");
14071 assert_eq!(
14072 one_field.temperature, 1.0,
14073 "omitting temperature still takes the vendor value"
14074 );
14075 assert_eq!(one_field.top_p, 0.95, "omitting top_p still takes vendor");
14076
14077 let disabled = chat(serde_json::json!({"top_k": 0, "top_p": 1.0}));
14080 assert_eq!(
14081 disabled.top_k, 0,
14082 "an explicit top_k 0 means KEEP ALL, not 'unset'"
14083 );
14084 assert_eq!(
14085 disabled.top_p, 1.0,
14086 "an explicit top_p 1.0 means untruncated"
14087 );
14088
14089 let penal = chat(serde_json::json!({"presence_penalty": 1.5}));
14091 assert_eq!(penal.penalty_present, 1.5);
14092 assert_eq!(penal.penalty_last_n, memra_engine::spec::PEN_WINDOW_MAX);
14093 }
14094
14095 #[test]
14096 fn vendor_sampling_defaults_are_identical_on_every_surface() {
14097 let d = qwen38_vendor_defaults();
14109 let md = ModelSamplingDefaults::single(d);
14110 let comp = |extra: serde_json::Value| {
14111 let mut body = serde_json::json!({
14112 "model": "qwen/qwen3.8-27b", "prompt": "task", "seed": 11 });
14113 body.as_object_mut()
14114 .unwrap()
14115 .extend(extra.as_object().unwrap().clone());
14116 let req: CompletionReq = serde_json::from_value(body).unwrap();
14117 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14118 build_request_with_trace(&req, tx, lanes::Lane::Interactive, None, None, &d).sampler_cfg
14119 };
14120 let chat = |extra: serde_json::Value| {
14121 let mut body = serde_json::json!({
14122 "model": "qwen/qwen3.8-27b",
14123 "messages": [{"role": "user", "content": "task"}],
14124 "seed": 11 });
14125 body.as_object_mut()
14126 .unwrap()
14127 .extend(extra.as_object().unwrap().clone());
14128 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14129 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14130 build_chat_request_with_trace(
14131 req,
14132 Some(&ModelCaps {
14133 chat_ok: true,
14134 ..Default::default()
14135 }),
14136 tx,
14137 lanes::Lane::Interactive,
14138 None,
14139 None,
14140 None,
14141 &md,
14142 )
14143 .unwrap()
14144 .request
14145 .sampler_cfg
14146 };
14147
14148 for extra in [
14149 serde_json::json!({}),
14150 serde_json::json!({"temperature": 0}),
14151 serde_json::json!({"temperature": 0.0}),
14152 serde_json::json!({"temperature": 0.7}),
14153 serde_json::json!({"top_p": 1.0}),
14154 serde_json::json!({"top_k": 0}),
14155 serde_json::json!({"min_p": 0.05}),
14156 serde_json::json!({"repetition_penalty": 1.1}),
14157 serde_json::json!({"frequency_penalty": 0.5}),
14158 serde_json::json!({"presence_penalty": 1.5}),
14159 serde_json::json!({
14160 "temperature": 0.3, "top_p": 0.5, "top_k": 7, "min_p": 0.02,
14161 "frequency_penalty": 0.1, "presence_penalty": 0.2,
14162 "repetition_penalty": 1.05 }),
14163 ] {
14164 let c = comp(extra.clone());
14165 let h = chat(extra.clone());
14166 assert_eq!(
14167 (
14168 c.temperature,
14169 c.top_p,
14170 c.top_k,
14171 c.min_p,
14172 c.penalty_repeat,
14173 c.penalty_freq,
14174 c.penalty_present,
14175 c.penalty_last_n,
14176 c.seed
14177 ),
14178 (
14179 h.temperature,
14180 h.top_p,
14181 h.top_k,
14182 h.min_p,
14183 h.penalty_repeat,
14184 h.penalty_freq,
14185 h.penalty_present,
14186 h.penalty_last_n,
14187 h.seed
14188 ),
14189 "/v1/completions and /v1/chat/completions disagree on {extra} — \
14190 standard-surface-law violation"
14191 );
14192 }
14193
14194 let omitted = comp(serde_json::json!({}));
14196 assert_eq!(
14197 omitted.temperature, 1.0,
14198 "qwen3.8 card thinking temperature"
14199 );
14200 assert_eq!(omitted.top_p, 0.95, "qwen3.8 card top_p");
14201 assert_eq!(omitted.top_k, 20, "qwen3.8 card top_k");
14202 assert!(
14204 memra_engine::sampler::Sampler::new(comp(serde_json::json!({"temperature": 0})))
14205 .is_greedy()
14206 );
14207 }
14208
14209 #[tokio::test]
14222 async fn same_omitted_request_resolves_identically_on_all_four_surfaces() {
14223 let _l = DRAIN_LOCK.lock().unwrap();
14224 let step_caps = ModelCaps {
14225 chat_ok: true,
14226 chat_temperature_default: Some(0.5),
14227 chat_top_p_default: Some(0.9),
14228 ..Default::default()
14229 };
14230 let (cfg_tx, cfg_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14231 let st = fake_worker_state_full(
14232 1,
14233 std::time::Duration::ZERO,
14234 HashMap::from([("m".to_string(), step_caps)]),
14235 Some(cfg_tx),
14236 );
14237 let fields = |saw: &WorkerSaw| {
14242 let c = &saw.sampler_cfg;
14243 (
14244 c.temperature,
14245 c.top_p,
14246 c.top_k,
14247 c.min_p,
14248 c.penalty_repeat,
14249 c.penalty_freq,
14250 c.penalty_present,
14251 c.penalty_last_n,
14252 )
14253 };
14254 let worker_saw = |surface: &str| {
14255 cfg_rx
14256 .recv_timeout(std::time::Duration::from_secs(10))
14257 .unwrap_or_else(|_| panic!("{surface}: request never reached the worker"))
14258 };
14259
14260 let resp = completions(
14261 State(st.clone()),
14262 axum::http::HeaderMap::new(),
14263 None,
14264 Json(serde_json::from_value(serde_json::json!({"model": "m", "prompt": "t"})).unwrap()),
14265 )
14266 .await;
14267 assert_eq!(
14268 resp.status(),
14269 StatusCode::OK,
14270 "/v1/completions rejected the omitted-sampling request"
14271 );
14272 let comp = worker_saw("/v1/completions");
14273
14274 let resp = chat_completions(
14275 State(st.clone()),
14276 axum::http::HeaderMap::new(),
14277 None,
14278 Json(
14279 serde_json::from_value(serde_json::json!({
14280 "model": "m", "messages": [{"role": "user", "content": "t"}]}))
14281 .unwrap(),
14282 ),
14283 )
14284 .await;
14285 assert_eq!(
14286 resp.status(),
14287 StatusCode::OK,
14288 "/v1/chat/completions rejected the omitted-sampling request"
14289 );
14290 let chat = worker_saw("/v1/chat/completions");
14291
14292 let resp = anthropic::messages(
14293 State(st.clone()),
14294 axum::http::HeaderMap::new(),
14295 None,
14296 axum::body::Bytes::from(
14297 serde_json::json!({
14298 "model": "m", "max_tokens": 16,
14299 "messages": [{"role": "user", "content": "t"}]})
14300 .to_string(),
14301 ),
14302 )
14303 .await;
14304 assert_eq!(
14305 resp.status(),
14306 StatusCode::OK,
14307 "/v1/messages rejected the omitted-sampling request"
14308 );
14309 let msg = worker_saw("/v1/messages");
14310
14311 let resp = responses_api::responses(
14312 State(st.clone()),
14313 axum::http::HeaderMap::new(),
14314 None,
14315 axum::body::Bytes::from(serde_json::json!({"model": "m", "input": "t"}).to_string()),
14316 )
14317 .await;
14318 assert_eq!(
14319 resp.status(),
14320 StatusCode::OK,
14321 "/v1/responses rejected the omitted-sampling request"
14322 );
14323 let rsp = worker_saw("/v1/responses");
14324
14325 for (surface, cfg) in [
14326 ("/v1/completions", &comp),
14327 ("/v1/messages", &msg),
14328 ("/v1/responses", &rsp),
14329 ] {
14330 assert_eq!(
14331 fields(cfg),
14332 fields(&chat),
14333 "{surface} resolved DIFFERENT effective sampling than /v1/chat/completions \
14334 for the same omitted-sampling request — standard-surface-law violation \
14335 (hermes d991b51699218285)"
14336 );
14337 }
14338 assert_eq!(
14341 (comp.sampler_cfg.temperature, comp.sampler_cfg.top_p),
14342 (0.5, 0.9),
14343 "an omitting client must get the model's vendor caps (Step-3.7: 0.5/0.9) on \
14344 EVERY surface, not the API-standard 1.0/1.0 (hermes d991b51699218285)"
14345 );
14346 }
14347
14348 #[tokio::test]
14358 async fn same_effort_value_resolves_identically_on_every_surface() {
14359 let _l = DRAIN_LOCK.lock().unwrap();
14360 let caps = ModelCaps {
14363 chat_ok: true,
14364 effort_levels: true,
14365 ..Default::default()
14366 };
14367 let (saw_tx, saw_rx) = std::sync::mpsc::channel::<WorkerSaw>();
14368 let st = fake_worker_state_full(
14369 1,
14370 std::time::Duration::ZERO,
14371 HashMap::from([("m".to_string(), caps)]),
14372 Some(saw_tx),
14373 );
14374 let send = |st: AppState, surface: &'static str, effort: &'static str| async move {
14375 match surface {
14376 "/v1/chat/completions" => {
14377 chat_completions(
14378 State(st),
14379 axum::http::HeaderMap::new(),
14380 None,
14381 Json(
14382 serde_json::from_value(serde_json::json!({
14383 "model": "m", "max_tokens": 8,
14384 "reasoning_effort": effort,
14385 "messages": [{"role": "user", "content": "t"}]}))
14386 .unwrap(),
14387 ),
14388 )
14389 .await
14390 }
14391 "/v1/responses" => {
14392 responses_api::responses(
14393 State(st),
14394 axum::http::HeaderMap::new(),
14395 None,
14396 axum::body::Bytes::from(
14397 serde_json::json!({
14398 "model": "m", "max_output_tokens": 8, "input": "t",
14399 "reasoning": {"effort": effort}})
14400 .to_string(),
14401 ),
14402 )
14403 .await
14404 }
14405 "/v1/messages" => {
14406 anthropic::messages(
14407 State(st),
14408 axum::http::HeaderMap::new(),
14409 None,
14410 axum::body::Bytes::from(
14411 serde_json::json!({
14412 "model": "m", "max_tokens": 8,
14413 "messages": [{"role": "user", "content": "t"}],
14414 "output_config": {"effort": effort}})
14415 .to_string(),
14416 ),
14417 )
14418 .await
14419 }
14420 other => panic!("unknown surface {other}"),
14421 }
14422 };
14423 const SURFACES: [&str; 3] = ["/v1/chat/completions", "/v1/responses", "/v1/messages"];
14424
14425 for (effort, want_think, want_level) in [
14428 ("none", ThinkMode::NoThink, Some("low")),
14429 ("minimal", ThinkMode::NoThink, Some("low")),
14430 ("low", ThinkMode::Think, Some("low")),
14431 ("medium", ThinkMode::Think, Some("medium")),
14432 ("high", ThinkMode::Think, Some("high")),
14433 ("xhigh", ThinkMode::Think, Some("high")),
14435 ] {
14436 for surface in SURFACES {
14437 let resp = send(st.clone(), surface, effort).await;
14438 assert_eq!(
14439 resp.status(),
14440 StatusCode::OK,
14441 "{surface} rejected effort {effort:?} — the surfaces' allowlists \
14442 diverged again (issue #31)"
14443 );
14444 let saw = saw_rx
14445 .recv_timeout(std::time::Duration::from_secs(10))
14446 .unwrap_or_else(|_| {
14447 panic!("{surface}: effort {effort:?} request never reached the worker")
14448 });
14449 assert_eq!(
14450 (saw.think, saw.reasoning_effort.as_deref()),
14451 (want_think, want_level),
14452 "{surface} resolved effort {effort:?} to a DIFFERENT worker-truth \
14453 reasoning surface — the parameter was dropped or remapped before \
14454 parse_think (issue #31 regression)"
14455 );
14456 }
14457 }
14458
14459 for effort in ["bogus", "banana", ""] {
14462 for surface in SURFACES {
14463 let resp = send(st.clone(), surface, effort).await;
14464 assert_eq!(
14465 resp.status(),
14466 StatusCode::BAD_REQUEST,
14467 "{surface} accepted effort {effort:?} — silent-accept regression \
14468 (issue #31: the value never reached parse_think's allowlist)"
14469 );
14470 let body = axum::body::to_bytes(resp.into_body(), 1 << 20)
14472 .await
14473 .unwrap();
14474 let v: serde_json::Value = serde_json::from_slice(&body)
14475 .unwrap_or_else(|_| panic!("{surface}: non-JSON 400 body for {effort:?}"));
14476 match surface {
14477 "/v1/messages" => {
14478 assert_eq!(v["type"], "error", "{surface} error envelope");
14479 assert_eq!(
14480 v["error"]["type"], "invalid_request_error",
14481 "{surface} error type"
14482 );
14483 }
14484 _ => {
14485 assert!(
14486 v["error"]["message"].is_string(),
14487 "{surface} OpenAI-shaped error body: {v}"
14488 );
14489 }
14490 }
14491 }
14492 }
14493
14494 let resp = anthropic::messages(
14498 State(st.clone()),
14499 axum::http::HeaderMap::new(),
14500 None,
14501 axum::body::Bytes::from(
14502 serde_json::json!({
14503 "model": "m", "max_tokens": 8,
14504 "messages": [{"role": "user", "content": "t"}],
14505 "thinking": {"type": "enabled"},
14506 "output_config": {"effort": "none"}})
14507 .to_string(),
14508 ),
14509 )
14510 .await;
14511 assert_eq!(resp.status(), StatusCode::OK);
14512 let saw = saw_rx
14513 .recv_timeout(std::time::Duration::from_secs(10))
14514 .expect("thinking+effort request never reached the worker");
14515 assert_eq!(
14516 saw.think,
14517 ThinkMode::Think,
14518 "thinking.type (the documented Anthropic lever) must win the switch over \
14519 output_config.effort"
14520 );
14521 let resp = anthropic::messages(
14522 State(st.clone()),
14523 axum::http::HeaderMap::new(),
14524 None,
14525 axum::body::Bytes::from(
14526 serde_json::json!({
14527 "model": "m", "max_tokens": 8,
14528 "messages": [{"role": "user", "content": "t"}],
14529 "thinking": {"type": "enabled"},
14530 "output_config": {"effort": "banana"}})
14531 .to_string(),
14532 ),
14533 )
14534 .await;
14535 assert_eq!(
14536 resp.status(),
14537 StatusCode::BAD_REQUEST,
14538 "an invalid effort must 400 even next to an explicit thinking.type — \
14539 precedence must not re-open the silent-accept hole"
14540 );
14541 }
14542
14543 #[test]
14544 fn vendor_sampling_defaults_are_boot_validated() {
14545 let parsed = OpenRouterMetadataFile::from_toml(
14548 r#"
14549[models.g]
14550default_temperature = 1.0
14551default_top_p = 0.95
14552default_top_k = 64
14553default_min_p = 0.0
14554default_presence_penalty = 0.0
14555default_frequency_penalty = 0.0
14556default_repetition_penalty = 1.0
14557"#,
14558 )
14559 .unwrap();
14560 let g = parsed.get("g").unwrap();
14561 assert_eq!(g.default_temperature, Some(1.0));
14562 assert_eq!(g.default_top_p, Some(0.95));
14563 assert_eq!(g.default_top_k, Some(64));
14564
14565 let err = OpenRouterMetadataFile::from_toml(
14569 r#"
14570[models.g]
14571default_temperature = 0.0
14572"#,
14573 )
14574 .unwrap_err();
14575 assert!(err.contains("default_temperature"), "{err}");
14576 assert!(
14577 err.contains("greedy"),
14578 "the refusal must say WHY a zero default is refused: {err}"
14579 );
14580
14581 for bad in [
14582 "default_temperature = 2.5",
14583 "default_temperature = -1.0",
14584 "default_top_p = 0.0",
14585 "default_top_p = 1.5",
14586 "default_min_p = 1.0",
14587 "default_min_p = -0.1",
14588 "default_presence_penalty = 3.0",
14589 "default_frequency_penalty = -2.5",
14590 "default_repetition_penalty = 0.0",
14591 ] {
14592 let err =
14593 OpenRouterMetadataFile::from_toml(&format!("[models.g]\n{bad}\n")).unwrap_err();
14594 let key = bad.split(' ').next().unwrap();
14595 assert!(err.contains(key), "{bad} must be refused by name: {err}");
14596 }
14597
14598 let err = OpenRouterMetadataFile::from_toml(
14602 r#"
14603[models.g]
14604default_temperture = 1.0
14605"#,
14606 )
14607 .unwrap_err();
14608 assert!(
14609 err.contains("unknown field"),
14610 "an unknown key must be fatal, which is what makes binary-first ordering \
14611 mandatory: {err}"
14612 );
14613 }
14614
14615 #[test]
14616 fn non_thinking_sampling_arm_is_boot_validated() {
14617 let parsed = OpenRouterMetadataFile::from_toml(
14621 r#"
14622[models.q]
14623default_temperature = 1.0
14624default_top_p = 0.95
14625default_top_k = 20
14626
14627[models.q.non_thinking_sampling]
14628temperature = 0.7
14629top_p = 0.8
14630top_k = 20
14631presence_penalty = 1.5
14632"#,
14633 )
14634 .unwrap();
14635 let arm = parsed
14636 .get("q")
14637 .unwrap()
14638 .non_thinking_sampling
14639 .as_ref()
14640 .unwrap();
14641 assert_eq!(arm.temperature, Some(0.7));
14642 assert_eq!(arm.top_p, Some(0.8));
14643 assert_eq!(arm.top_k, Some(20));
14644 assert_eq!(arm.presence_penalty, Some(1.5));
14645 assert_eq!(
14646 arm.min_p, None,
14647 "undeclared arm fields stay undeclared, never invented"
14648 );
14649
14650 let err = OpenRouterMetadataFile::from_toml(
14654 r#"
14655[models.q]
14656[models.q.non_thinking_sampling]
14657temperature = 0.0
14658"#,
14659 )
14660 .unwrap_err();
14661 assert!(err.contains("non_thinking_sampling.temperature"), "{err}");
14662 assert!(err.contains("greedy"), "{err}");
14663
14664 let err = OpenRouterMetadataFile::from_toml(
14667 r#"
14668[models.q]
14669[models.q.non_thinking_sampling]
14670"#,
14671 )
14672 .unwrap_err();
14673 assert!(err.contains("non_thinking_sampling"), "{err}");
14674 assert!(err.contains("declare"), "{err}");
14675
14676 for bad in [
14678 "temperature = 2.5",
14679 "top_p = 0.0",
14680 "top_p = 1.5",
14681 "min_p = 1.0",
14682 "presence_penalty = 3.0",
14683 "frequency_penalty = -2.5",
14684 "repetition_penalty = 0.0",
14685 ] {
14686 let err = OpenRouterMetadataFile::from_toml(&format!(
14687 "[models.q]\n[models.q.non_thinking_sampling]\n{bad}\n"
14688 ))
14689 .unwrap_err();
14690 let key = bad.split(' ').next().unwrap();
14691 assert!(
14692 err.contains(&format!("non_thinking_sampling.{key}")),
14693 "the refusal for {bad:?} must name the nested key: {err}"
14694 );
14695 }
14696
14697 let err = OpenRouterMetadataFile::from_toml(
14701 r#"
14702[models.q]
14703[models.q.non_thinking_sampling]
14704temperture = 0.7
14705"#,
14706 )
14707 .unwrap_err();
14708 assert!(err.contains("unknown field"), "{err}");
14709 }
14710
14711 fn qwen38_non_thinking_defaults() -> SamplingDefaults {
14716 SamplingDefaults {
14717 temperature: Some(0.7),
14718 top_p: Some(0.8),
14719 top_k: Some(20),
14720 presence_penalty: Some(1.5),
14721 ..Default::default()
14722 }
14723 }
14724
14725 fn qwen38_two_arm_defaults() -> ModelSamplingDefaults {
14726 ModelSamplingDefaults {
14727 thinking: qwen38_vendor_defaults(),
14728 non_thinking: Some(qwen38_non_thinking_defaults()),
14729 }
14730 }
14731
14732 fn qwen38_caps() -> ModelCaps {
14736 ModelCaps {
14737 chat_ok: true,
14738 qwen_think: true,
14739 think_switch: true,
14740 ..Default::default()
14741 }
14742 }
14743
14744 fn sampler_key(c: &SamplerConfig) -> (f32, f32, usize, f32, f32, f32, f32, usize, u64) {
14747 (
14748 c.temperature,
14749 c.top_p,
14750 c.top_k,
14751 c.min_p,
14752 c.penalty_present,
14753 c.penalty_freq,
14754 c.penalty_repeat,
14755 c.penalty_last_n,
14756 c.seed,
14757 )
14758 }
14759
14760 fn build_with_arms(
14761 defaults: &ModelSamplingDefaults,
14762 caps: &ModelCaps,
14763 default_effort: Option<&str>,
14764 extra: serde_json::Value,
14765 ) -> Request {
14766 let mut body = serde_json::json!({
14767 "model": "m",
14768 "messages": [{"role": "user", "content": "task"}],
14769 "seed": 3
14771 });
14772 body.as_object_mut()
14773 .unwrap()
14774 .extend(extra.as_object().unwrap().clone());
14775 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
14776 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
14777 build_chat_request_with_trace(
14778 req,
14779 Some(caps),
14780 tx,
14781 lanes::Lane::Interactive,
14782 None,
14783 None,
14784 default_effort,
14785 defaults,
14786 )
14787 .unwrap()
14788 .request
14789 }
14790
14791 #[test]
14792 fn resolved_thinking_mode_picks_the_vendor_sampling_arm() {
14793 let two_arm = qwen38_two_arm_defaults();
14798 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
14799 let caps = qwen38_caps();
14800
14801 let off_spellings = [
14803 serde_json::json!({"reasoning_effort": "none"}),
14804 serde_json::json!({"enable_thinking": false}),
14805 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
14806 serde_json::json!({"reasoning": {"enabled": false}}),
14807 ];
14808 for extra in &off_spellings {
14809 let r = build_with_arms(&two_arm, &caps, None, extra.clone());
14810 assert_eq!(r.think, ThinkMode::NoThink, "{extra}");
14811 let c = &r.sampler_cfg;
14812 assert_eq!(c.temperature, 0.7, "{extra}: non-thinking card temperature");
14813 assert_eq!(c.top_p, 0.8, "{extra}: non-thinking card top_p");
14814 assert_eq!(c.top_k, 20, "{extra}: non-thinking card top_k");
14815 assert_eq!(
14816 c.penalty_present, 1.5,
14817 "{extra}: non-thinking presence_penalty"
14818 );
14819 assert_eq!(
14820 c.penalty_last_n,
14821 memra_engine::spec::PEN_WINDOW_MAX,
14822 "{extra}: the arm's presence penalty uses the cross-path history window"
14823 );
14824 assert_eq!(
14825 c.min_p, 0.0,
14826 "{extra}: the arm recommends no min_p — API standard, never the other arm's"
14827 );
14828
14829 let s = build_with_arms(&single_arm, &caps, None, extra.clone());
14832 assert_eq!(s.think, ThinkMode::NoThink, "{extra}");
14833 assert_eq!(s.sampler_cfg.temperature, 1.0, "{extra}: single-arm model");
14834 assert_eq!(s.sampler_cfg.top_p, 0.95, "{extra}: single-arm model");
14835 assert_eq!(
14836 s.sampler_cfg.penalty_present, 0.0,
14837 "{extra}: single-arm model"
14838 );
14839 }
14840
14841 for extra in [
14844 serde_json::json!({}),
14845 serde_json::json!({"enable_thinking": true}),
14846 serde_json::json!({"reasoning_effort": "high"}),
14847 serde_json::json!({"reasoning": {"enabled": true}}),
14848 ] {
14849 for defaults in [&two_arm, &single_arm] {
14850 let c = build_with_arms(defaults, &caps, None, extra.clone()).sampler_cfg;
14851 assert_eq!(c.temperature, 1.0, "{extra}: thinking card temperature");
14852 assert_eq!(c.top_p, 0.95, "{extra}: thinking card top_p");
14853 assert_eq!(c.top_k, 20, "{extra}: thinking card top_k");
14854 assert_eq!(
14855 c.penalty_present, 0.0,
14856 "{extra}: thinking arm has no presence"
14857 );
14858 }
14859 }
14860
14861 let c = build_with_arms(&two_arm, &caps, Some("none"), serde_json::json!({})).sampler_cfg;
14864 assert_eq!(
14865 c.temperature, 0.7,
14866 "deployment-default off = non-thinking arm"
14867 );
14868 let c = build_with_arms(
14870 &two_arm,
14871 &caps,
14872 Some("none"),
14873 serde_json::json!({"enable_thinking": true}),
14874 )
14875 .sampler_cfg;
14876 assert_eq!(
14877 c.temperature, 1.0,
14878 "explicit ON beats the deployment default"
14879 );
14880
14881 let c = build_with_arms(
14883 &two_arm,
14884 &caps,
14885 None,
14886 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
14887 )
14888 .sampler_cfg;
14889 assert_eq!(c.temperature, 0.55, "explicit temperature survives the arm");
14890 assert_eq!(c.top_p, 0.8, "unset top_p still takes the non-thinking arm");
14891 let c = build_with_arms(
14892 &two_arm,
14893 &caps,
14894 None,
14895 serde_json::json!({
14896 "reasoning_effort": "none", "top_p": 0.99, "presence_penalty": 0.0}),
14897 )
14898 .sampler_cfg;
14899 assert_eq!(c.top_p, 0.99, "explicit top_p wins");
14900 assert_eq!(
14901 c.penalty_present, 0.0,
14902 "an explicit presence_penalty 0.0 wins over the arm's 1.5 — a disabling value \
14903 is a value, not an absence"
14904 );
14905 assert_eq!(
14906 c.penalty_last_n, 0,
14907 "all penalties off => no history window"
14908 );
14909 assert_eq!(c.top_k, 20, "unset top_k still takes the arm");
14910
14911 let c = build_with_arms(
14914 &two_arm,
14915 &caps,
14916 None,
14917 serde_json::json!({"enable_thinking": false, "temperature": 0}),
14918 )
14919 .sampler_cfg;
14920 assert!(
14921 memra_engine::sampler::Sampler::new(c).is_greedy(),
14922 "explicit temperature 0 must stay greedy on the non-thinking arm"
14923 );
14924
14925 let c = build_with_arms(
14928 &single_arm,
14929 &caps,
14930 None,
14931 serde_json::json!({"enable_thinking": false, "temperature": 0.55}),
14932 )
14933 .sampler_cfg;
14934 assert_eq!(c.temperature, 0.55);
14935 assert_eq!(
14936 c.top_p, 0.95,
14937 "single-arm model: unset top_p takes its one arm"
14938 );
14939 }
14940
14941 #[test]
14942 fn sampling_arms_never_blend_field_by_field() {
14943 let parsed = OpenRouterMetadataFile::from_toml(
14948 r#"
14949[models.m]
14950default_temperature = 1.0
14951default_min_p = 0.05
14952
14953[models.m.non_thinking_sampling]
14954temperature = 0.6
14955"#,
14956 )
14957 .unwrap();
14958 let caps = ModelCaps {
14959 chat_temperature_default: Some(0.5),
14960 chat_top_p_default: Some(0.9),
14961 ..Default::default()
14962 };
14963 let d = ModelSamplingDefaults::resolve(parsed.get("m"), Some(&caps));
14964 let client = ClientSampling {
14965 seed: Some(1),
14966 ..Default::default()
14967 };
14968
14969 let off = resolve_sampler_config(client, d.for_mode(ThinkMode::NoThink));
14970 assert_eq!(off.temperature, 0.6, "the arm's own field applies");
14971 assert_eq!(
14972 off.min_p, 0.0,
14973 "min_p undeclared on the arm = API standard, NOT the thinking arm's 0.05"
14974 );
14975 assert_eq!(
14976 off.top_p, 1.0,
14977 "top_p undeclared on the arm = API standard, NOT the arch cap's 0.9"
14978 );
14979
14980 for mode in [ThinkMode::Default, ThinkMode::Think] {
14982 let on = resolve_sampler_config(client, d.for_mode(mode));
14983 assert_eq!(on.temperature, 1.0);
14984 assert_eq!(on.min_p, 0.05);
14985 assert_eq!(on.top_p, 0.9, "primary arm keeps the arch-cap fallback");
14986 }
14987 }
14988
14989 #[test]
14990 fn single_arm_models_and_thinking_on_requests_match_the_pre_arm_law_exactly() {
14991 let caps = qwen38_caps();
15000 let single_arm = ModelSamplingDefaults::single(qwen38_vendor_defaults());
15001 let two_arm = qwen38_two_arm_defaults();
15002
15003 let bodies = [
15004 serde_json::json!({}),
15005 serde_json::json!({"enable_thinking": true}),
15006 serde_json::json!({"reasoning_effort": "high"}),
15007 serde_json::json!({"reasoning_effort": "none"}),
15008 serde_json::json!({"enable_thinking": false}),
15009 serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}),
15010 serde_json::json!({"temperature": 0.3, "top_p": 0.5}),
15011 serde_json::json!({"enable_thinking": false, "temperature": 0}),
15012 ];
15013 for extra in &bodies {
15014 let r = build_with_arms(&single_arm, &caps, None, extra.clone());
15016 let mut client = ClientSampling {
15017 seed: Some(3),
15018 ..Default::default()
15019 };
15020 if let Some(t) = extra.get("temperature").and_then(|v| v.as_f64()) {
15021 client.temperature = Some(t as f32);
15022 }
15023 if let Some(p) = extra.get("top_p").and_then(|v| v.as_f64()) {
15024 client.top_p = Some(p as f32);
15025 }
15026 let pre_arm = resolve_sampler_config(client, &qwen38_vendor_defaults());
15027 assert_eq!(
15028 sampler_key(&r.sampler_cfg),
15029 sampler_key(&pre_arm),
15030 "{extra}: single-arm model diverged from the pre-arm resolution law"
15031 );
15032
15033 if r.think != ThinkMode::NoThink {
15036 let t = build_with_arms(&two_arm, &caps, None, extra.clone());
15037 assert_eq!(t.think, r.think, "{extra}");
15038 assert_eq!(t.reasoning_effort, r.reasoning_effort, "{extra}");
15039 assert_eq!(
15040 sampler_key(&t.sampler_cfg),
15041 sampler_key(&r.sampler_cfg),
15042 "{extra}: a thinking-on request must not feel the non-thinking arm"
15043 );
15044 }
15045 }
15046 }
15047
15048 #[test]
15049 fn constraint_forced_nothink_takes_the_non_thinking_arm() {
15050 let r = build_with_arms(
15056 &qwen38_two_arm_defaults(),
15057 &qwen38_caps(),
15058 None,
15059 serde_json::json!({"response_format": {"type": "json_object"}}),
15060 );
15061 assert_eq!(
15062 r.think,
15063 ThinkMode::NoThink,
15064 "constraint forces the switch off"
15065 );
15066 assert_eq!(
15067 r.sampler_cfg.temperature, 0.7,
15068 "and the arm follows the real mode"
15069 );
15070 assert_eq!(r.sampler_cfg.penalty_present, 1.5);
15071 }
15072
15073 #[test]
15074 fn metadata_sampling_defaults_outrank_arch_caps_but_never_the_client() {
15075 let caps = ModelCaps {
15080 chat_temperature_default: Some(0.5),
15081 chat_top_p_default: Some(0.9),
15082 chat_ok: true,
15083 ..Default::default()
15084 };
15085 let metadata = OpenRouterModelMetadata {
15086 default_temperature: Some(1.0),
15087 default_top_p: Some(0.95),
15088 default_top_k: Some(64),
15089 ..Default::default()
15090 };
15091
15092 let caps_only = SamplingDefaults::resolve(None, Some(&caps));
15093 assert_eq!(caps_only.temperature, Some(0.5), "arch cap is the fallback");
15094 assert_eq!(caps_only.top_p, Some(0.9));
15095 assert_eq!(caps_only.top_k, None, "caps declare no top_k");
15096
15097 let both = SamplingDefaults::resolve(Some(&metadata), Some(&caps));
15098 assert_eq!(
15099 both.temperature,
15100 Some(1.0),
15101 "metadata outranks the arch cap"
15102 );
15103 assert_eq!(both.top_p, Some(0.95));
15104 assert_eq!(both.top_k, Some(64));
15105
15106 let partial = SamplingDefaults::resolve(
15108 Some(&OpenRouterModelMetadata {
15109 default_temperature: Some(0.7),
15110 ..Default::default()
15111 }),
15112 Some(&caps),
15113 );
15114 assert_eq!(partial.temperature, Some(0.7));
15115 assert_eq!(
15116 partial.top_p,
15117 Some(0.9),
15118 "an undeclared metadata field must fall through to the cap, not to 1.0"
15119 );
15120
15121 assert_eq!(
15123 SamplingDefaults::resolve(None, None),
15124 SamplingDefaults::default()
15125 );
15126 }
15127
15128 #[test]
15129 fn vendor_defaults_leave_the_pure_temp_sampled_spec_regime() {
15130 let resolved = |d: &SamplingDefaults| {
15143 resolve_sampler_config(
15144 ClientSampling {
15145 seed: Some(1),
15146 ..Default::default()
15147 },
15148 d,
15149 )
15150 };
15151
15152 assert!(
15154 memra_engine::sampler::Sampler::new(resolved(&SamplingDefaults::default()))
15155 .is_spec_sampling(),
15156 "the API-standard default must stay in the fast pure-temp regime"
15157 );
15158
15159 for (name, d) in [
15160 ("qwen/qwen3.8-27b", qwen38_vendor_defaults()),
15161 ("google/gemma-4-31b-it", gemma4_vendor_defaults()),
15162 ] {
15163 let sampler = memra_engine::sampler::Sampler::new(resolved(&d));
15164 assert!(
15165 !sampler.is_greedy(),
15166 "{name}: vendor default must not be greedy"
15167 );
15168 assert!(
15169 !sampler.is_spec_sampling(),
15170 "{name}: vendor top_p/top_k DO leave the pure-temp regime — if this ever \
15171 starts passing, either the vendor numbers changed or the in-graph draft \
15172 learned filters, and the perf note in docs/SERVING.md needs revisiting"
15173 );
15174 }
15175
15176 let opted_out = resolve_sampler_config(
15178 ClientSampling {
15179 top_p: Some(1.0),
15180 top_k: Some(0),
15181 seed: Some(1),
15182 ..Default::default()
15183 },
15184 &qwen38_vendor_defaults(),
15185 );
15186 assert!(
15187 memra_engine::sampler::Sampler::new(opted_out).is_spec_sampling(),
15188 "explicitly disabling the filters must restore the pure-temp regime"
15189 );
15190 }
15191
15192 #[test]
15193 fn omitted_seed_is_fresh_entropy_not_a_pinned_zero() {
15194 let comp_seed = |body: serde_json::Value| {
15201 let req: CompletionReq = serde_json::from_value(body).unwrap();
15202 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15203 build_request(&req, tx, lanes::Lane::Interactive, None)
15204 .sampler_cfg
15205 .seed
15206 };
15207 let chat_seed = |body: serde_json::Value| {
15208 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15209 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15210 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15211 .unwrap()
15212 .request
15213 .sampler_cfg
15214 .seed
15215 };
15216
15217 let a = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15220 let b = comp_seed(serde_json::json!({"model": "m", "prompt": "t"}));
15221 let c = chat_seed(serde_json::json!({
15222 "model": "m", "messages": [{"role": "user", "content": "t"}]}));
15223 assert_ne!(
15224 a, 0,
15225 "omitted seed must not be the pinned 0 that caused the loop"
15226 );
15227 assert_ne!(b, 0);
15228 assert_ne!(c, 0);
15229 assert_ne!(
15230 a, b,
15231 "two seed-omitting requests must get DIFFERENT streams"
15232 );
15233 assert_ne!(a, c);
15234
15235 assert_eq!(
15238 comp_seed(serde_json::json!({
15239 "model": "m", "prompt": "t", "seed": 0})),
15240 0,
15241 "explicit seed 0 must stay 0 — the determinism gates depend on it"
15242 );
15243 assert_eq!(
15244 comp_seed(serde_json::json!({
15245 "model": "m", "prompt": "t", "seed": 12345})),
15246 12345
15247 );
15248 assert_eq!(
15249 chat_seed(serde_json::json!({
15250 "model": "m", "messages": [{"role": "user", "content": "t"}],
15251 "seed": 777})),
15252 777
15253 );
15254 assert_eq!(
15256 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42})),
15257 comp_seed(serde_json::json!({"model": "m", "prompt": "t", "seed": 42}))
15258 );
15259
15260 let seeds: std::collections::HashSet<u64> = (0..256).map(|_| fresh_seed()).collect();
15263 assert_eq!(
15264 seeds.len(),
15265 256,
15266 "fresh_seed must not collide across rapid calls"
15267 );
15268 assert!(!seeds.contains(&0));
15269 }
15270
15271 #[test]
15272 fn response_format_builds_grammar_only_when_present() {
15273 let mk = |rf: Option<serde_json::Value>| {
15277 let mut body = serde_json::json!({
15278 "model": "m", "messages": [{"role": "user", "content": "t"}]});
15279 if let Some(rf) = rf {
15280 body["response_format"] = rf;
15281 }
15282 let req: ChatCompletionReq = serde_json::from_value(body).unwrap();
15283 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
15284 build_chat_request(req, None, tx, lanes::Lane::Interactive, None)
15285 };
15286 assert!(mk(None).unwrap().request.grammar.is_none());
15287 assert!(
15288 mk(Some(serde_json::json!({"type": "text"})))
15289 .unwrap()
15290 .request
15291 .grammar
15292 .is_none()
15293 );
15294 assert!(matches!(
15295 mk(Some(serde_json::json!({"type": "json_object"})))
15296 .unwrap()
15297 .request
15298 .grammar,
15299 Some(constrained::GrammarSpec::JsonObject)
15300 ));
15301 assert!(matches!(
15302 mk(Some(serde_json::json!({"type": "json_schema",
15303 "json_schema": {"schema": {"type": "object"}}})))
15304 .unwrap()
15305 .request
15306 .grammar,
15307 Some(constrained::GrammarSpec::JsonSchema(_))
15308 ));
15309 assert!(mk(Some(serde_json::json!({"type": "yaml"}))).is_err());
15311 }
15312
15313 #[test]
15314 fn unsupported_semantic_params_are_named_rejections() {
15315 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15317 "model": "m", "messages": [{"role": "user", "content": "t"}],
15318 "response_format": {"type": "json_object"}
15319 }))
15320 .unwrap();
15321 assert!(req.response_format.is_some());
15322 let req: ChatCompletionReq = serde_json::from_value(serde_json::json!({
15323 "model": "m", "messages": [{"role": "user", "content": "t"}],
15324 "response_format": {"type": "text"}, "logprobs": false, "n": 1,
15325 "user": "u-1", "stream_options": {"include_usage": true}
15326 }))
15327 .unwrap();
15328 assert_eq!(req.response_format.as_ref().unwrap()["type"], "text");
15330 assert_eq!(req.logprobs.as_ref().unwrap().as_bool(), Some(false));
15331 assert_eq!(req.n, Some(1));
15332 assert!(reject_unsupported(&[("logit_bias", false, "")]).is_ok());
15334 let (msg, param) = reject_unsupported(&[("logit_bias", true, " (why)")]).unwrap_err();
15335 assert_eq!(param, "logit_bias");
15336 assert_eq!(msg, "logit_bias is not supported (why)");
15337 }
15338
15339 #[test]
15340 fn completions_accept_openai_stop_forms() {
15341 for (value, expected) in [
15342 (serde_json::json!("Problem:"), vec!["Problem:"]),
15343 (
15344 serde_json::json!(["Question:", "Problem:"]),
15345 vec!["Question:", "Problem:"],
15346 ),
15347 (serde_json::Value::Null, Vec::<&str>::new()),
15348 ] {
15349 let req: CompletionReq = serde_json::from_value(serde_json::json!({
15350 "model": "plain_quant", "prompt": "task", "stop": value
15351 }))
15352 .unwrap();
15353 assert_eq!(req.stop.into_vec(), expected);
15354 }
15355 }
15356
15357 fn fake_worker_state() -> AppState {
15364 fake_worker_state_with_steps(1, std::time::Duration::ZERO)
15365 }
15366
15367 fn fake_worker_state_with_steps(steps: usize, step_delay: std::time::Duration) -> AppState {
15368 fake_worker_state_full(steps, step_delay, HashMap::new(), None)
15369 }
15370
15371 struct WorkerSaw {
15376 sampler_cfg: SamplerConfig,
15377 think: ThinkMode,
15378 reasoning_effort: Option<String>,
15379 }
15380
15381 fn fake_worker_state_full(
15388 steps: usize,
15389 step_delay: std::time::Duration,
15390 caps: HashMap<String, ModelCaps>,
15391 saw_tx: Option<std::sync::mpsc::Sender<WorkerSaw>>,
15392 ) -> AppState {
15393 let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<Cmd>();
15394 let health = health::WorkerHealth::new();
15395 let h = health.clone();
15396 std::thread::spawn(move || {
15397 h.mark_ready();
15398 while let Ok(Cmd::Generate(mut req)) = cmd_rx.recv() {
15399 if let Some(tx) = &saw_tx {
15400 let _ = tx.send(WorkerSaw {
15401 sampler_cfg: req.sampler_cfg.clone(),
15402 think: req.think,
15403 reasoning_effort: req.reasoning_effort.clone(),
15404 });
15405 }
15406 worker::release_pending_admit();
15410 worker::release_admission_reservation(req.lane);
15411 h.beat_busy();
15412 if let Some(ready) = req.constraint_ready.take() {
15413 let _ = ready.send(Ok(()));
15414 }
15415 let _ = req.tx.send(Event::PromptUsage {
15416 n_prompt: 1,
15417 n_cached: 0,
15418 });
15419 for step in 0..steps {
15420 h.beat_busy();
15421 let text = if steps == 1 { "ok" } else { "x" };
15422 let _ = req.tx.send(Event::Token {
15423 id: step as u32 + 1,
15424 text: text.into(),
15425 });
15426 if !step_delay.is_zero() {
15427 std::thread::sleep(step_delay);
15428 }
15429 }
15430 let _ = req.tx.send(Event::Done {
15431 stop_reason: "Eos".into(),
15432 n_tokens: steps,
15433 n_prompt: 1,
15434 n_cached: 0,
15435 elapsed_s: 0.01,
15436 spec: None,
15437 });
15438 h.set_phase(health::PHASE_IDLE);
15439 }
15440 });
15441 for _ in 0..2000 {
15444 if health.live().is_ok() {
15445 break;
15446 }
15447 std::thread::sleep(std::time::Duration::from_millis(1));
15448 }
15449 AppState {
15450 cmd_tx,
15451 models: Arc::new(vec!["m".into()]),
15452 caps: Arc::new(caps),
15453 openrouter_metadata: Arc::new(HashMap::new()),
15454 provider_metadata: Arc::new(None),
15455 metering: None,
15456
15457 budget_tokenizers: None,
15458 api_auth: ApiAuth::default(),
15459 metrics_auth: MetricsAuth::default(),
15460 metrics: SharedMetrics::default(),
15461 started: 1,
15462 inflight: Arc::new(Default::default()),
15463 tenant_inflight: Arc::new(Default::default()),
15464 health,
15465 bg: None,
15466 }
15467 }
15468
15469 #[tokio::test]
15470 async fn deep_schema_fails_while_normal_decode_keeps_stepping() {
15471 let _l = DRAIN_LOCK.lock().unwrap();
15472 let st = fake_worker_state_with_steps(64, std::time::Duration::from_millis(5));
15473 let normal_state = st.clone();
15474 let normal = tokio::spawn(async move {
15475 chat_completions(
15476 State(normal_state),
15477 axum::http::HeaderMap::new(),
15478 None,
15479 Json(
15480 serde_json::from_value(serde_json::json!({
15481 "model": "m",
15482 "messages": [{"role": "user", "content": "keep decoding"}],
15483 }))
15484 .unwrap(),
15485 ),
15486 )
15487 .await
15488 });
15489 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
15490
15491 let mut deep = serde_json::json!({"type": "string"});
15492 for _ in 0..(constrained::MAX_SCHEMA_DEPTH / 2 + 1) {
15493 deep = serde_json::json!({"allOf": [deep]});
15494 }
15495 let bad = chat_completions(
15496 State(st.clone()),
15497 axum::http::HeaderMap::new(),
15498 None,
15499 Json(
15500 serde_json::from_value(serde_json::json!({
15501 "model": "m",
15502 "messages": [{"role": "user", "content": "bad schema"}],
15503 "response_format": {
15504 "type": "json_schema",
15505 "json_schema": {"schema": deep},
15506 },
15507 }))
15508 .unwrap(),
15509 ),
15510 )
15511 .await;
15512 assert_eq!(bad.status(), StatusCode::BAD_REQUEST);
15513 assert_eq!(bad.headers().get("x-should-retry").unwrap(), "false");
15514 let bytes = axum::body::to_bytes(bad.into_body(), usize::MAX)
15515 .await
15516 .unwrap();
15517 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15518 assert!(
15519 payload["error"]["message"]
15520 .as_str()
15521 .unwrap()
15522 .contains("maximum nesting depth")
15523 );
15524 assert!(
15525 !normal.is_finished(),
15526 "bad schema stalled or replaced the normal decode"
15527 );
15528
15529 let normal_response = normal.await.unwrap();
15530 assert_eq!(normal_response.status(), StatusCode::OK);
15531 let snapshot = st.health.snapshot();
15532 assert!(
15533 st.health.live().is_ok(),
15534 "normal decode left health stalled"
15535 );
15536 assert!(snapshot.beat_age_ms < snapshot.stall_threshold_ms);
15537 }
15538
15539 #[tokio::test]
15540 async fn valid_response_format_preflight_preserves_generation() {
15541 let _l = DRAIN_LOCK.lock().unwrap();
15542 let response = chat_completions(
15543 State(fake_worker_state()),
15544 axum::http::HeaderMap::new(),
15545 None,
15546 Json(
15547 serde_json::from_value(serde_json::json!({
15548 "model": "m",
15549 "messages": [{"role": "user", "content": "valid schema"}],
15550 "response_format": {"type": "json_object"},
15551 }))
15552 .unwrap(),
15553 ),
15554 )
15555 .await;
15556 assert_eq!(response.status(), StatusCode::OK);
15557 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15558 .await
15559 .unwrap();
15560 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15561 assert_eq!(payload["choices"][0]["message"]["content"], "ok");
15562 }
15563
15564 #[tokio::test]
15565 async fn unknown_model_refuses_model_not_found_before_admission() {
15566 let _l = DRAIN_LOCK.lock().unwrap();
15567 let response = chat_completions(
15572 State(fake_worker_state()),
15573 axum::http::HeaderMap::new(),
15574 None,
15575 Json(
15576 serde_json::from_value(serde_json::json!({
15577 "model": "qwen/qwen3.8-27b-typo",
15578 "messages": [{"role": "user", "content": "hi"}],
15579 }))
15580 .unwrap(),
15581 ),
15582 )
15583 .await;
15584 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15585 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15586 .await
15587 .unwrap();
15588 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15589 assert_eq!(payload["error"]["code"], "model_not_found");
15590 assert_eq!(payload["error"]["type"], "invalid_request_error");
15591
15592 let response = completions(
15594 State(fake_worker_state()),
15595 axum::http::HeaderMap::new(),
15596 None,
15597 Json(
15598 serde_json::from_value(serde_json::json!({
15599 "model": "nope",
15600 "prompt": "hi",
15601 }))
15602 .unwrap(),
15603 ),
15604 )
15605 .await;
15606 assert_eq!(response.status(), StatusCode::BAD_REQUEST);
15607 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15608 .await
15609 .unwrap();
15610 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
15611 assert_eq!(payload["error"]["code"], "model_not_found");
15612 }
15613
15614 const METRICS_KEY_ACME: &str = "completion-acme-secret";
15615 const METRICS_KEY_BLUE: &str = "completion-blue-secret";
15616
15617 fn multi_key_metrics_state(metrics_token: Option<&str>) -> AppState {
15618 let spec = format!(
15619 "acme:{},blue:{}",
15620 auth::sha256_hex(METRICS_KEY_ACME),
15621 auth::sha256_hex(METRICS_KEY_BLUE),
15622 );
15623 let keyring = Box::leak(Box::new(auth::KeyStore::from_spec(&spec).unwrap()));
15624 let mut st = fake_worker_state();
15625 st.api_auth.keyring = Some(keyring);
15626 st.metrics_auth = MetricsAuth::new(
15627 true,
15628 st.api_auth.configured(),
15629 metrics_token.map(str::to_string),
15630 );
15631 {
15632 let mut metrics = st.metrics.lock().unwrap();
15633 metrics.admitted = 17;
15634 metrics.prompt_tokens_in = 400;
15635 metrics.cached_tokens_in = 60;
15636 metrics.prefix_hits = 2;
15637 metrics.prefix_misses = 3;
15638 metrics.prefix_inserts = 5;
15639 metrics.prefix_evictions = 7;
15640 metrics.prefix_skips_budget = 9;
15641 metrics.prefix_skips_pinned = 10;
15642 metrics.prefix_hit_tokens = 11;
15643 metrics.lcp_hist[4] = 13;
15644 metrics.ns_tokens.insert("t:acme".into(), [100, 40]);
15645 metrics.ns_tokens.insert("t:blue".into(), [300, 20]);
15646 metrics.adsd_suspect_total.insert("t:acme".into(), 1);
15647 metrics.adsd_suspect_total.insert("t:blue".into(), 2);
15648 metrics.prefix_entries = 29;
15649 metrics.prefix_bytes = 31;
15650 metrics.active_sessions = 3;
15651 metrics.queued_requests = 5;
15652 metrics.continuation_pool_entries = 7;
15653 metrics.spec_pool_entries = 11;
15654 metrics.cuda_driver_free_bytes = 13;
15655 metrics.cuda_pool_reserved_bytes = 17;
15656 metrics.cuda_pool_used_bytes = 19;
15657 metrics.cuda_pool_cached_bytes = 23;
15658 metrics.batch_size_last = 37;
15659 metrics.spec.insert(
15660 "m".into(),
15661 memra_engine::spec::SpecTelemetry {
15662 rounds: 2,
15663 drafted: 6,
15664 accepted: 4,
15665 ..Default::default()
15666 },
15667 );
15668 let mut spec_window = memra_engine::spec::SpecTelemetry {
15669 rounds: 4,
15670 drafted: 12,
15671 accepted: 6,
15672 ..Default::default()
15673 };
15674 spec_window.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
15675 spec_window.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
15676 metrics.spec_window.insert("m".into(), spec_window);
15677 metrics.constraint_compiler_fail_closed.insert(
15678 "m".into(),
15679 Arc::new(std::sync::atomic::AtomicBool::new(true)),
15680 );
15681 }
15682 st
15683 }
15684
15685 async fn metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15686 let mut headers = HeaderMap::new();
15687 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15688 let response = get_metrics(State(st), headers).await;
15689 assert_eq!(response.status(), StatusCode::OK);
15690 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15691 .await
15692 .unwrap();
15693 serde_json::from_slice(&bytes).unwrap()
15694 }
15695
15696 async fn yield_metrics_json(st: AppState, bearer: &str) -> serde_json::Value {
15697 let mut headers = HeaderMap::new();
15698 headers.insert("authorization", format!("Bearer {bearer}").parse().unwrap());
15699 let response = yield_metrics(State(st), headers).await;
15700 assert_eq!(response.status(), StatusCode::OK);
15701 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
15702 .await
15703 .unwrap();
15704 serde_json::from_slice(&bytes).unwrap()
15705 }
15706
15707 #[test]
15708 fn exposed_open_bind_is_refused_before_server_start() {
15709 assert!(validate_bind_security("127.0.0.1:8080", false, false).unwrap());
15710 assert!(validate_bind_security("[::1]:8080", false, false).unwrap());
15711
15712 let err = validate_bind_security("0.0.0.0:8000", false, false).unwrap_err();
15713 assert!(err.contains("refusing unauthenticated non-loopback bind"));
15714 assert!(err.contains("MEMRA_API_KEY"));
15715 assert!(err.contains("MEMRA_ALLOW_OPEN_BIND=1"));
15716 assert!(validate_bind_security("[::]:8000", false, false).is_err());
15717
15718 assert!(!validate_bind_security("0.0.0.0:8000", true, false).unwrap());
15719 assert!(!validate_bind_security("0.0.0.0:8000", false, true).unwrap());
15720 }
15721
15722 #[tokio::test]
15723 async fn keyed_metrics_require_and_accept_api_bearer() {
15724 let mut st = fake_worker_state();
15725 st.api_auth.single_key = Some(Arc::from("completion-secret"));
15726 st.metrics_auth = MetricsAuth::new(true, st.api_auth.configured(), None);
15727
15728 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
15729 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
15730 let response = yield_metrics(State(st.clone()), HeaderMap::new()).await;
15731 assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
15732
15733 let mut headers = HeaderMap::new();
15734 headers.insert("authorization", "Bearer completion-secret".parse().unwrap());
15735 assert_eq!(
15736 get_metrics(State(st.clone()), headers.clone())
15737 .await
15738 .status(),
15739 StatusCode::OK,
15740 );
15741 let body = metrics_json(st.clone(), "completion-secret").await;
15742 assert!(
15743 body.get("admitted").is_some(),
15744 "the legacy single-key domain keeps cumulative counters",
15745 );
15746 assert!(
15747 body.get("active_sessions").is_none(),
15748 "a static completion key is not an operator metrics principal",
15749 );
15750 assert_eq!(
15751 yield_metrics(State(st), headers).await.status(),
15752 StatusCode::OK
15753 );
15754 }
15755
15756 #[tokio::test]
15757 async fn keyring_metrics_bearer_sees_only_its_tenant_rows() {
15758 let st = multi_key_metrics_state(None);
15759 let body = metrics_json(st.clone(), METRICS_KEY_ACME).await;
15760 assert_eq!(
15761 body.as_object().unwrap().len(),
15762 2,
15763 "completion metrics must contain only tenant-scoped rows",
15764 );
15765 let tenants = body["tenants"].as_object().unwrap();
15766 assert_eq!(tenants.len(), 1);
15767 assert_eq!(tenants["t:acme"]["prompt_tokens_in"], 100);
15768 assert!(!tenants.contains_key("t:blue"));
15769 let adsd = body["adsd_suspect_total"].as_object().unwrap();
15770 assert_eq!(adsd.len(), 1);
15771 assert_eq!(adsd["t:acme"], 1);
15772 assert!(!adsd.contains_key("t:blue"));
15773
15774 let mut headers = HeaderMap::new();
15775 headers.insert(
15776 "authorization",
15777 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
15778 );
15779 assert_eq!(
15780 yield_metrics(State(st), headers).await.status(),
15781 StatusCode::FORBIDDEN,
15782 "the process-wide yield view requires an operator metrics token",
15783 );
15784 }
15785
15786 #[tokio::test]
15787 async fn tenant_metrics_hide_capacity_and_aggregate_spec() {
15788 let body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
15789 for operator_only in [
15790 "prefix_cache_entries",
15791 "prefix_cache_bytes",
15792 "prefix_cache_skips_budget",
15793 "prefix_cache_skips_pinned",
15794 "active_sessions",
15795 "queued_requests",
15796 "continuation_pool_entries",
15797 "spec_pool_entries",
15798 "cuda_driver_free_bytes",
15799 "cuda_pool_reserved_bytes",
15800 "cuda_pool_used_bytes",
15801 "cuda_pool_cached_bytes",
15802 "constraint_compiler_fail_closed",
15803 "serve_idle_seconds",
15804 "spec",
15805 "spec_tau",
15806 "spec_accept_by_position",
15807 "dual_pp",
15808 "peer_probe_bypassed",
15809 "peer_probe_boundary_copies",
15810 "peer_probe_runtime_reprobes",
15811 "peer_probe_runtime_failures",
15812 "peer_probe_deferred_total",
15813 "peer_probe_integrity_degraded",
15814 "peer_probe_degraded_to_host_bounce",
15815 ] {
15816 assert!(
15817 body.get(operator_only).is_none(),
15818 "tenant metrics must not expose operator field {operator_only}",
15819 );
15820 }
15821 }
15822
15823 #[test]
15824 fn populated_spec_acceptance_metrics_are_operator_only() {
15825 for scope in [
15826 MetricsScope::CompletionDomain,
15827 MetricsScope::Tenant("t:acme".into()),
15828 ] {
15829 let mut body = json!({});
15830 insert_spec_acceptance_metrics(&mut body, &scope, || {
15831 panic!("tenant scope evaluated the process-wide spec snapshot")
15832 });
15833 assert!(body.get("spec_tau").is_none(), "{scope:?} leaked spec tau");
15834 assert!(
15835 body.get("spec_accept_by_position").is_none(),
15836 "{scope:?} leaked the accept histogram"
15837 );
15838 }
15839
15840 let mut telemetry = memra_engine::spec::SpecTelemetry {
15841 rounds: 4,
15842 drafted: 12,
15843 accepted: 6,
15844 ..Default::default()
15845 };
15846 telemetry.pos_drafted[..3].copy_from_slice(&[4, 4, 4]);
15847 telemetry.pos_accepted[..3].copy_from_slice(&[3, 2, 1]);
15848 let mut body = json!({});
15849 insert_spec_acceptance_metrics(&mut body, &MetricsScope::All, || {
15850 HashMap::from([("model-a".to_string(), telemetry)])
15851 });
15852 assert_eq!(body["spec_tau"]["model-a"], 1.5);
15853 let histogram = &body["spec_accept_by_position"]["model-a"];
15854 assert_eq!(histogram["window_seconds"], worker::SPEC_METRICS_WINDOW_S);
15855 assert_eq!(histogram["rounds"], 4);
15856 assert_eq!(histogram["offered"], json!([4, 4, 4]));
15857 assert_eq!(histogram["accepted"], json!([3, 2, 1]));
15858 assert_eq!(histogram["accept_rate"], json!([0.75, 0.5, 0.25]));
15859 }
15860
15861 #[test]
15862 fn populated_dual_pp_metrics_are_operator_only() {
15863 let populated = DualPpMetricsSnapshot {
15864 stage_ns: [1_000_000, 2_000_000, 3_000_000, 4_000_000],
15865 stage_samples: [1, 1, 1, 1],
15866 dropped_timing_samples: 0,
15867 overlaps: 17,
15868 slot_pairs: 19,
15869 slot_uses: [19, 19],
15870 slot_collisions: 0,
15871 };
15872 for scope in [
15873 MetricsScope::CompletionDomain,
15874 MetricsScope::Tenant("t:acme".into()),
15875 ] {
15876 let mut body = json!({});
15877 insert_dual_pp_metrics(&mut body, &scope, || populated);
15878 assert!(
15879 body.get("dual_pp").is_none(),
15880 "{scope:?} leaked dual PP topology"
15881 );
15882 }
15883
15884 let mut body = json!({});
15885 insert_dual_pp_metrics(&mut body, &MetricsScope::All, || populated);
15886 assert_eq!(body["dual_pp"]["overlaps"], 17);
15887 assert_eq!(body["dual_pp"]["slot_pairs"], 19);
15888 assert_eq!(body["dual_pp"]["slot_uses"], json!([19, 19]));
15889 assert_eq!(body["dual_pp"]["slot_collisions"], 0);
15890 assert_eq!(
15891 body["dual_pp"]["cuda_event_spans"]["wave_a_stage0"]["mean_ms"],
15892 1.0
15893 );
15894 }
15895
15896 #[test]
15897 fn peer_probe_metrics_are_operator_only() {
15898 let populated = memra_engine::pp::PeerProbeMetrics {
15899 bypassed: 1,
15900 boundary_copies: 8_192,
15901 runtime_probes: 1,
15902 runtime_failures: 0,
15903 deferred_total: 4,
15904 integrity_degraded: true,
15905 degraded_to_host_bounce: true,
15906 };
15907 for scope in [
15908 MetricsScope::CompletionDomain,
15909 MetricsScope::Tenant("t:acme".into()),
15910 ] {
15911 let mut body = json!({});
15912 insert_peer_probe_metrics(&mut body, &scope, || populated);
15913 assert!(body.get("peer_probe_bypassed").is_none());
15914 }
15915
15916 let mut body = json!({});
15917 insert_peer_probe_metrics(&mut body, &MetricsScope::All, || populated);
15918 assert_eq!(body["peer_probe_bypassed"], 1);
15919 assert_eq!(body["peer_probe_boundary_copies"], 8_192);
15920 assert_eq!(body["peer_probe_runtime_reprobes"], 1);
15921 assert_eq!(body["peer_probe_runtime_failures"], 0);
15922 assert_eq!(body["peer_probe_deferred_total"], 4);
15923 assert_eq!(body["peer_probe_integrity_degraded"], true);
15924 assert_eq!(body["peer_probe_degraded_to_host_bounce"], true);
15925 }
15926
15927 #[tokio::test]
15928 async fn prefix_aggregate_metrics_are_operator_only_but_tenant_ratio_remains() {
15929 let tenant_body = metrics_json(multi_key_metrics_state(None), METRICS_KEY_ACME).await;
15930 for operator_only in [
15931 "lcp_histogram",
15932 "cache_hit_token_ratio",
15933 "prefix_cache_hits",
15934 "prefix_cache_misses",
15935 "prefix_cache_inserts",
15936 "prefix_cache_evictions",
15937 "prefix_cache_skips_budget",
15938 "prefix_cache_skips_pinned",
15939 "prefix_cache_hit_tokens",
15940 ] {
15941 assert!(
15942 tenant_body.get(operator_only).is_none(),
15943 "tenant metrics must not expose global prefix field {operator_only}",
15944 );
15945 }
15946 assert_eq!(tenant_body["tenants"].as_object().unwrap().len(), 1);
15947 assert_eq!(tenant_body["tenants"]["t:acme"]["prompt_tokens_in"], 100);
15948 assert_eq!(tenant_body["tenants"]["t:acme"]["cached_tokens_in"], 40);
15949 assert_eq!(
15950 tenant_body["tenants"]["t:acme"]["cache_hit_token_ratio"],
15951 0.4
15952 );
15953
15954 let operator_body = metrics_json(
15955 multi_key_metrics_state(Some("scrape-secret")),
15956 "scrape-secret",
15957 )
15958 .await;
15959 assert_eq!(operator_body["prefix_cache_hits"], 2);
15960 assert_eq!(operator_body["prefix_cache_misses"], 3);
15961 assert_eq!(operator_body["prefix_cache_inserts"], 5);
15962 assert_eq!(operator_body["prefix_cache_evictions"], 7);
15963 assert_eq!(operator_body["prefix_cache_skips_budget"], 9);
15964 assert_eq!(operator_body["prefix_cache_skips_pinned"], 10);
15965 assert_eq!(operator_body["prefix_cache_hit_tokens"], 11);
15966 assert_eq!(operator_body["cache_hit_token_ratio"], 0.15);
15967 assert_eq!(operator_body["lcp_histogram"]["counts"][4], 13);
15968 }
15969
15970 #[tokio::test]
15971 async fn configured_metrics_token_is_exclusive_and_sees_all_tenants() {
15972 let st = multi_key_metrics_state(Some("scrape-secret"));
15973 let mut completion_headers = HeaderMap::new();
15974 completion_headers.insert(
15975 "authorization",
15976 format!("Bearer {METRICS_KEY_ACME}").parse().unwrap(),
15977 );
15978 assert_eq!(
15979 get_metrics(State(st.clone()), completion_headers.clone())
15980 .await
15981 .status(),
15982 StatusCode::FORBIDDEN,
15983 );
15984 assert_eq!(
15985 yield_metrics(State(st.clone()), completion_headers)
15986 .await
15987 .status(),
15988 StatusCode::FORBIDDEN,
15989 );
15990
15991 let body = metrics_json(st.clone(), "scrape-secret").await;
15992 let tenants = body["tenants"].as_object().unwrap();
15993 assert_eq!(tenants.len(), 2);
15994 assert!(tenants.contains_key("t:acme"));
15995 assert!(tenants.contains_key("t:blue"));
15996 assert_eq!(body["adsd_suspect_total"]["t:acme"], 1);
15997 assert_eq!(body["adsd_suspect_total"]["t:blue"], 2);
15998 assert_eq!(body["active_sessions"], 3);
15999 assert_eq!(body["queued_requests"], 5);
16000 assert_eq!(body["prefix_cache_bytes"], 31);
16001 assert_eq!(body["cuda_driver_free_bytes"], 13);
16002 assert_eq!(body["constraint_compiler_fail_closed"]["m"], 1);
16003 assert_eq!(body["spec"]["m"]["drafted"], 6);
16004 assert_eq!(body["spec_tau"]["m"], 1.5);
16005 assert_eq!(
16006 body["spec_accept_by_position"]["m"]["accepted"],
16007 json!([3, 2, 1])
16008 );
16009 let yield_body = yield_metrics_json(st, "scrape-secret").await;
16010 assert_eq!(yield_body["batch_size_last"], 37);
16011 }
16012
16013 #[tokio::test]
16014 async fn metrics_token_protects_public_override_without_api_keys() {
16015 let mut st = fake_worker_state();
16016 st.metrics_auth = MetricsAuth::new(false, false, Some("scrape-secret".into()));
16017
16018 assert_eq!(
16019 get_metrics(State(st.clone()), HeaderMap::new())
16020 .await
16021 .status(),
16022 StatusCode::UNAUTHORIZED,
16023 );
16024 let mut headers = HeaderMap::new();
16025 headers.insert("authorization", "Bearer scrape-secret".parse().unwrap());
16026 assert_eq!(
16027 get_metrics(State(st.clone()), headers.clone())
16028 .await
16029 .status(),
16030 StatusCode::OK,
16031 );
16032 assert_eq!(
16033 yield_metrics(State(st), headers).await.status(),
16034 StatusCode::OK
16035 );
16036 }
16037
16038 #[tokio::test]
16039 async fn no_key_loopback_metrics_remain_open_for_development() {
16040 let mut st = fake_worker_state();
16041 st.metrics_auth = MetricsAuth::new(true, false, None);
16042 let response = get_metrics(State(st.clone()), HeaderMap::new()).await;
16043 assert_eq!(response.status(), StatusCode::OK);
16044 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16045 .await
16046 .unwrap();
16047 let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16048 assert!(
16049 body.get("active_sessions").is_some(),
16050 "no-key loopback development keeps full operator visibility",
16051 );
16052 assert_eq!(
16053 yield_metrics(State(st), HeaderMap::new()).await.status(),
16054 StatusCode::OK,
16055 );
16056 }
16057
16058 #[test]
16059 fn rate_limit_math_remaining_hits_zero_at_cap_and_reset_arms() {
16060 let metrics = SharedMetrics::default();
16061 let rl = RateLimit::compute(4, 1, &metrics);
16063 assert_eq!((rl.limit, rl.remaining, rl.reset_s), (4, 3, 0));
16064 let rl = RateLimit::compute(4, 3, &metrics);
16065 assert_eq!(rl.remaining, 1);
16066 let rl = RateLimit::compute(4, 4, &metrics);
16068 assert_eq!(rl.remaining, 0);
16069 assert!(rl.reset_s > 0, "reset must arm when no slots are free");
16070 assert_eq!(RateLimit::compute(4, 9, &metrics).remaining, 0);
16072 let m = worker::Metrics {
16074 completed: 2,
16075 tokens_out: 200,
16076 step_p50_ms: 20.0,
16077 ..Default::default()
16078 };
16079 assert_eq!(reset_estimate_s(&m), 2); }
16081
16082 #[test]
16083 fn inflight_guard_counts_up_and_frees_on_drop() {
16084 let counts: InflightCounts = Arc::new(Default::default());
16085 let tenants: TenantGauge = Arc::new(Default::default());
16086 let (g1, n1, t1) = InflightGuard::try_acquire(
16087 counts.clone(),
16088 lanes::Lane::Interactive,
16089 tenants.clone(),
16090 "acme",
16091 None,
16092 )
16093 .unwrap();
16094 let (g2, n2, t2) = InflightGuard::try_acquire(
16095 counts.clone(),
16096 lanes::Lane::Interactive,
16097 tenants.clone(),
16098 "acme",
16099 None,
16100 )
16101 .unwrap();
16102 assert_eq!((n1, n2), (1, 2));
16103 assert_eq!((t1, t2), (1, 2));
16105 let (gj, nj, tj) = InflightGuard::try_acquire(
16107 counts.clone(),
16108 lanes::Lane::Judge,
16109 tenants.clone(),
16110 "blue",
16111 None,
16112 )
16113 .unwrap();
16114 assert_eq!((nj, tj), (1, 1));
16115 drop(g1);
16116 drop(gj);
16117 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 1);
16118 assert_eq!(counts[1].load(std::sync::atomic::Ordering::SeqCst), 0);
16119 assert_eq!(tenants.lock().unwrap().get("acme"), Some(&1));
16120 assert!(tenants.lock().unwrap().get("blue").is_none());
16122 drop(g2);
16123 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
16124 assert!(tenants.lock().unwrap().is_empty());
16125 }
16126
16127 #[test]
16128 fn tenant_concurrency_cap_is_atomic_across_arrivals() {
16129 let counts: InflightCounts = Arc::new(Default::default());
16130 let tenants: TenantGauge = Arc::new(Default::default());
16131 let start = Arc::new(std::sync::Barrier::new(3));
16132 let attempted = Arc::new(std::sync::Barrier::new(3));
16133 let mut joins = Vec::new();
16134 for _ in 0..2 {
16135 let counts = counts.clone();
16136 let tenants = tenants.clone();
16137 let start = start.clone();
16138 let attempted = attempted.clone();
16139 joins.push(std::thread::spawn(move || {
16140 start.wait();
16141 let result = InflightGuard::try_acquire(
16142 counts,
16143 lanes::Lane::Interactive,
16144 tenants,
16145 "preview_001",
16146 Some(1),
16147 );
16148 let won = result.is_ok();
16149 attempted.wait(); drop(result);
16151 won
16152 }));
16153 }
16154 start.wait();
16155 attempted.wait();
16156 let wins = joins
16157 .into_iter()
16158 .map(|join| join.join().unwrap())
16159 .filter(|won| *won)
16160 .count();
16161 assert_eq!(wins, 1, "exactly one simultaneous request may pass cap=1");
16162 assert_eq!(counts[0].load(std::sync::atomic::Ordering::SeqCst), 0);
16163 assert!(tenants.lock().unwrap().is_empty());
16164 }
16165
16166 #[tokio::test]
16167 async fn tenant_concurrency_cap_rejects_before_worker_admission() {
16168 let st = fake_worker_state();
16169 let tenant = auth::TenantCtx {
16170 tenant: "preview_001".into(),
16171 lane_class: auth::LaneClass::Interactive,
16172 rate_limit: Some(1),
16173 key_prefix: None,
16174 };
16175 let first_env = Envelope::new(true);
16176 let (guard, first_rl) =
16177 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &first_env) {
16178 Ok(slot) => slot,
16179 Err(_) => panic!("the first request must acquire the tenant slot"),
16180 };
16181 assert_eq!((first_rl.limit, first_rl.remaining), (1, 0));
16182
16183 let second_env = Envelope::new(true);
16184 let response =
16185 match acquire_request_slot(&st, lanes::Lane::Interactive, &tenant, &second_env) {
16186 Err(response) => response,
16187 Ok(_) => panic!("the second request must be rejected at the tenant cap"),
16188 };
16189 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
16190 assert_eq!(response.headers()["retry-after"], "2");
16191 assert_eq!(response.headers()["retry-after-ms"], "2000");
16192 assert_eq!(response.headers()["x-ratelimit-limit"], "1");
16193 assert_eq!(response.headers()["x-ratelimit-remaining"], "0");
16194 assert_eq!(response.headers()["x-request-id"], second_env.id);
16195 assert_eq!(
16196 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16197 1,
16198 "rejected request must not consume a lane slot"
16199 );
16200 assert_eq!(
16201 st.tenant_inflight
16202 .lock()
16203 .unwrap()
16204 .get("preview_001")
16205 .copied(),
16206 Some(1),
16207 "rejected request must not increment the tenant gauge"
16208 );
16209 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
16210 .await
16211 .unwrap();
16212 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16213 assert_eq!(payload["error"]["type"], "rate_limit_error");
16214 assert_eq!(payload["error"]["code"], "rate_limit_exceeded");
16215 assert!(
16216 payload["error"]["message"]
16217 .as_str()
16218 .unwrap()
16219 .contains("concurrent request limit")
16220 );
16221
16222 drop(guard);
16223 let _ = InflightGuard::try_acquire(
16224 st.inflight.clone(),
16225 lanes::Lane::Interactive,
16226 st.tenant_inflight.clone(),
16227 "preview_001",
16228 Some(1),
16229 )
16230 .expect("slot must reopen after the in-flight request completes");
16231 }
16232
16233 #[test]
16234 fn tenant_rate_limit_override_is_min_with_global_cap() {
16235 let metrics = SharedMetrics::default();
16236 let unlimited = auth::TenantCtx::default_tenant();
16237 let capped = auth::TenantCtx {
16238 tenant: "acme".into(),
16239 lane_class: auth::LaneClass::Interactive,
16240 rate_limit: Some(2),
16241 key_prefix: None,
16242 };
16243 let global = lane_cap(lanes::Lane::Interactive);
16244 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &unlimited, 1);
16246 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16247 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 1);
16249 assert_eq!((rl.limit, rl.remaining), (2, 1));
16250 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 5, &metrics, &capped, 2);
16251 assert_eq!(rl.remaining, 0);
16252 assert!(rl.reset_s > 0, "reset must arm at the tenant cap too");
16253 let rl = RateLimit::at_admit(lanes::Lane::Interactive, global, &metrics, &capped, 0);
16257 assert_eq!(rl.remaining, 0);
16258 let wide = auth::TenantCtx {
16259 rate_limit: Some(global + 100),
16260 ..capped.clone()
16261 };
16262 let rl = RateLimit::at_admit(lanes::Lane::Interactive, 1, &metrics, &wide, 1);
16263 assert_eq!((rl.limit, rl.remaining), (global, global - 1));
16264 }
16265
16266 #[test]
16267 fn batch_class_keys_default_to_harvest_and_cannot_claim_interactive() {
16268 let batch = auth::TenantCtx {
16269 tenant: "bulk".into(),
16270 lane_class: auth::LaneClass::Batch,
16271 rate_limit: None,
16272 key_prefix: None,
16273 };
16274 let interactive = auth::TenantCtx::default_tenant();
16275 let hdr = |v: Option<&str>| {
16276 let mut h = axum::http::HeaderMap::new();
16277 if let Some(v) = v {
16278 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16279 }
16280 h
16281 };
16282 assert_eq!(
16284 lane_for_tenant(&hdr(None), &interactive).unwrap(),
16285 lanes::Lane::Interactive
16286 );
16287 assert_eq!(
16288 lane_for_tenant(&hdr(Some("judge")), &interactive).unwrap(),
16289 lanes::Lane::Judge
16290 );
16291 assert_eq!(
16293 lane_for_tenant(&hdr(None), &batch).unwrap(),
16294 lanes::Lane::Harvest
16295 );
16296 assert_eq!(
16297 lane_for_tenant(&hdr(Some("judge")), &batch).unwrap(),
16298 lanes::Lane::Judge
16299 );
16300 let resp = lane_for_tenant(&hdr(Some("interactive")), &batch).unwrap_err();
16301 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16302 let resp = lane_for_tenant(&hdr(Some("turbo")), &interactive).unwrap_err();
16304 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16305 }
16306
16307 #[tokio::test]
16308 async fn handler_layer_refusals_are_openai_objects_with_x_should_retry() {
16309 let hdr = |v: &str| {
16314 let mut h = axum::http::HeaderMap::new();
16315 h.insert("x-lane", axum::http::HeaderValue::from_str(v).unwrap());
16316 h
16317 };
16318 let body = |resp: Response| async move {
16319 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16320 .await
16321 .unwrap();
16322 serde_json::from_slice::<serde_json::Value>(&bytes).unwrap()
16323 };
16324
16325 let resp = lane_for_tenant(&hdr("turbo"), &auth::TenantCtx::default_tenant()).unwrap_err();
16326 assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
16327 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16328 let payload = body(resp).await;
16329 assert!(
16330 payload["error"].is_object(),
16331 "bare-string error body: {payload}"
16332 );
16333 assert_eq!(payload["error"]["type"], "invalid_request_error");
16334 assert_eq!(payload["error"]["param"], "x-lane");
16335 assert_eq!(payload["error"]["code"], "invalid_lane");
16336
16337 let batch = auth::TenantCtx {
16338 tenant: "bulk".into(),
16339 lane_class: auth::LaneClass::Batch,
16340 rate_limit: None,
16341 key_prefix: None,
16342 };
16343 let resp = lane_for_tenant(&hdr("interactive"), &batch).unwrap_err();
16344 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
16345 assert_eq!(resp.headers().get("x-should-retry").unwrap(), "false");
16346 let payload = body(resp).await;
16347 assert_eq!(payload["error"]["type"], "authentication_error");
16348 assert_eq!(payload["error"]["param"], "x-lane");
16349 }
16350
16351 static DRAIN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
16354
16355 #[tokio::test]
16356 async fn responses_carry_rate_limit_headers_and_slot_frees() {
16357 let _l = DRAIN_LOCK.lock().unwrap();
16358 let st = fake_worker_state();
16359 let resp = chat_completions(
16362 State(st.clone()),
16363 axum::http::HeaderMap::new(),
16364 None,
16365 Json(
16366 serde_json::from_value(serde_json::json!({
16367 "model": "m", "messages": [{"role": "user", "content": "t"}]
16368 }))
16369 .unwrap(),
16370 ),
16371 )
16372 .await;
16373 assert_eq!(resp.status(), StatusCode::OK);
16374 let h = resp.headers();
16375 let limit: usize = h["x-ratelimit-limit"].to_str().unwrap().parse().unwrap();
16376 let remaining: usize = h["x-ratelimit-remaining"]
16377 .to_str()
16378 .unwrap()
16379 .parse()
16380 .unwrap();
16381 assert_eq!(remaining, limit - 1);
16382 assert_eq!(h["x-ratelimit-reset"], "0");
16383 assert_eq!(
16384 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16385 0,
16386 "slot must free at completion"
16387 );
16388 let resp = completions(
16391 State(st.clone()),
16392 axum::http::HeaderMap::new(),
16393 None,
16394 Json(
16395 serde_json::from_value(serde_json::json!({
16396 "model": "m", "prompt": "t", "stream": true
16397 }))
16398 .unwrap(),
16399 ),
16400 )
16401 .await;
16402 assert_eq!(resp.status(), StatusCode::OK);
16403 assert!(resp.headers().contains_key("x-ratelimit-limit"));
16404 assert!(resp.headers().contains_key("x-ratelimit-remaining"));
16405 assert!(resp.headers().contains_key("x-ratelimit-reset"));
16406 assert_eq!(
16407 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16408 1,
16409 "stream in flight holds the slot"
16410 );
16411 let _ = axum::body::to_bytes(resp.into_body(), usize::MAX)
16412 .await
16413 .unwrap();
16414 assert_eq!(
16415 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16416 0,
16417 "slot must free when the stream completes"
16418 );
16419 }
16420
16421 #[tokio::test]
16422 async fn handlers_sync_worker_truth_usage_and_cost_before_terminal_response() {
16423 let _l = DRAIN_LOCK.lock().unwrap();
16424 let mut st = fake_worker_state();
16425 let mock = MockMetering::admit_all();
16426 st.metering = Some(mock.clone());
16427
16428 let nonstream = chat_completions(
16429 State(st.clone()),
16430 HeaderMap::new(),
16431 None,
16432 Json(
16433 serde_json::from_value(json!({
16434 "model": "m",
16435 "messages": [{"role": "user", "content": "t"}],
16436 }))
16437 .unwrap(),
16438 ),
16439 )
16440 .await;
16441 assert_eq!(nonstream.status(), StatusCode::OK);
16442 let nonstream_id = nonstream.headers()["x-request-id"]
16443 .to_str()
16444 .unwrap()
16445 .to_string();
16446
16447 let stream = completions(
16448 State(st),
16449 HeaderMap::new(),
16450 None,
16451 Json(
16452 serde_json::from_value(json!({
16453 "model": "m",
16454 "prompt": "t",
16455 "stream": true,
16456 }))
16457 .unwrap(),
16458 ),
16459 )
16460 .await;
16461 assert_eq!(stream.status(), StatusCode::OK);
16462 let stream_id = stream.headers()["x-request-id"]
16463 .to_str()
16464 .unwrap()
16465 .to_string();
16466 let _ = axum::body::to_bytes(stream.into_body(), usize::MAX)
16467 .await
16468 .unwrap();
16469
16470 let events = mock.events();
16474 let opened: Vec<&str> = events
16475 .iter()
16476 .filter_map(|e| match e {
16477 MeterEvent::Open { request_id, .. } => Some(request_id.as_str()),
16478 _ => None,
16479 })
16480 .collect();
16481 assert_eq!(opened, vec![nonstream_id.as_str(), stream_id.as_str()]);
16482 let completes = events
16483 .iter()
16484 .filter(|e| {
16485 matches!(
16486 e,
16487 MeterEvent::Complete {
16488 prompt: 1,
16489 cached: 0,
16490 completion: 1,
16491 }
16492 )
16493 })
16494 .count();
16495 assert_eq!(
16496 completes, 2,
16497 "both surfaces settle complete with worker-truth usage: {events:?}"
16498 );
16499 }
16500
16501 #[tokio::test]
16502 async fn completion_admission_supports_metered_blocked_and_paid_transitions() {
16503 let _l = DRAIN_LOCK.lock().unwrap();
16504 let mock = MockMetering::with_limits(vec![
16510 ReserveScript::Insufficient,
16511 ReserveScript::Admit { with_permit: false },
16512 ReserveScript::Blocked,
16513 ReserveScript::Admit { with_permit: true },
16514 ]);
16515 let mut st = fake_worker_state();
16516 st.metering = Some(mock.clone());
16517
16518 let metrics = get_metrics(State(st.clone()), HeaderMap::new()).await;
16520 assert_eq!(metrics.status(), StatusCode::OK);
16521 let metrics_body = axum::body::to_bytes(metrics.into_body(), usize::MAX)
16522 .await
16523 .unwrap();
16524 let metrics_body: serde_json::Value = serde_json::from_slice(&metrics_body).unwrap();
16525 assert_eq!(metrics_body["budget_source_reload_failed"], 0);
16526 assert_eq!(metrics_body["budget_source_reload_consecutive"], 0);
16527 assert_eq!(metrics_body["budget_source_available"], true);
16528
16529 let request = || {
16530 Json(
16531 serde_json::from_value::<CompletionReq>(json!({
16532 "model": "m",
16533 "prompt_ids": [1],
16534 "max_tokens": 1,
16535 }))
16536 .unwrap(),
16537 )
16538 };
16539
16540 let denied = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16541 assert_eq!(denied.status(), StatusCode::PAYMENT_REQUIRED);
16542 let denied_body = axum::body::to_bytes(denied.into_body(), usize::MAX)
16543 .await
16544 .unwrap();
16545 let denied_body: serde_json::Value = serde_json::from_slice(&denied_body).unwrap();
16546 assert_eq!(denied_body["error"]["type"], "insufficient_balance");
16547 assert_eq!(denied_body["error"]["code"], "insufficient_balance");
16548
16549 let included = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16550 assert_eq!(included.status(), StatusCode::OK);
16551
16552 let blocked = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16555 assert_eq!(blocked.status(), StatusCode::PAYMENT_REQUIRED);
16556
16557 let admitted = completions(State(st.clone()), HeaderMap::new(), None, request()).await;
16558 assert_eq!(admitted.status(), StatusCode::OK);
16559
16560 let events = mock.events();
16561 let terminal: Vec<&MeterEvent> = events
16562 .iter()
16563 .filter(|e| matches!(e, MeterEvent::Reject { .. } | MeterEvent::Complete { .. }))
16564 .collect();
16565 assert_eq!(
16566 terminal.len(),
16567 4,
16568 "four requests, four terminal settles: {events:?}"
16569 );
16570 assert!(matches!(
16571 terminal[0],
16572 MeterEvent::Reject { status: 402, .. }
16573 ));
16574 assert!(matches!(terminal[1], MeterEvent::Complete { .. }));
16575 assert!(matches!(
16576 terminal[2],
16577 MeterEvent::Reject { status: 402, .. }
16578 ));
16579 assert!(matches!(terminal[3], MeterEvent::Complete { .. }));
16580 let permits: Vec<bool> = events
16582 .iter()
16583 .filter_map(|e| match e {
16584 MeterEvent::Open { with_permit, .. } => Some(*with_permit),
16585 _ => None,
16586 })
16587 .collect();
16588 assert_eq!(
16589 permits,
16590 vec![false, false, false, true],
16591 "the permit rides the receipt exactly when reserve minted one: {events:?}"
16592 );
16593 }
16594
16595 #[tokio::test]
16599 async fn a_capped_key_answers_its_own_402_and_the_principal_crosses_the_seam() {
16600 let mock = MockMetering::with_limits(vec![ReserveScript::PrincipalCapped]);
16601 let mut st = fake_worker_state();
16602 st.metering = Some(mock.clone());
16603 let tenant = auth::TenantCtx {
16604 tenant: "acme".into(),
16605 lane_class: auth::LaneClass::Interactive,
16606 rate_limit: None,
16607 key_prefix: Some("mk-acme-testprefix00".into()),
16608 };
16609 let mut request = gate_request(1, 1);
16610 let rejection = admit_tenant_budget(&st, &tenant, &mut request)
16611 .expect_err("a capped key must be refused at admission");
16612 assert!(matches!(rejection, BudgetRejection::PrincipalCapped));
16613 let (response, outcome) = rejection.into_response();
16614 assert_eq!(outcome, "key_spend_cap_reached");
16615 assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED);
16616 let body = body_value(response).await;
16617 assert_eq!(body["error"]["code"], "key_spend_cap_reached");
16618 assert!(
16619 body["error"]["message"].as_str().unwrap().contains("cap"),
16620 "the 402 must point at the KEY's cap, not tenant credit: {body}"
16621 );
16622 let events = mock.events();
16623 assert!(
16624 events.contains(&MeterEvent::Reserve {
16625 tenant: "acme".into(),
16626 principal: Some("mk-acme-testprefix00".into()),
16627 model: "qwen/qwen3.8-27b".into(),
16628 }),
16629 "the key prefix must reach reserve: {events:?}"
16630 );
16631 }
16632
16633 #[tokio::test]
16634 async fn streaming_client_disconnect_records_partial_usage_and_cost() {
16635 let _l = DRAIN_LOCK.lock().unwrap();
16636 let mut st = fake_worker_state_with_steps(4, std::time::Duration::from_millis(100));
16637 let mock = MockMetering::admit_all();
16638 st.metering = Some(mock.clone());
16639
16640 let response = completions(
16641 State(st),
16642 HeaderMap::new(),
16643 None,
16644 Json(
16645 serde_json::from_value(json!({
16646 "model": "m",
16647 "prompt": "disconnect after one delta",
16648 "stream": true,
16649 }))
16650 .unwrap(),
16651 ),
16652 )
16653 .await;
16654 assert_eq!(response.status(), StatusCode::OK);
16655 let request_id = response.headers()["x-request-id"]
16656 .to_str()
16657 .unwrap()
16658 .to_string();
16659 let mut body = Box::pin(response.into_body().into_data_stream());
16660 let first = std::future::poll_fn(|cx| body.as_mut().poll_next(cx))
16661 .await
16662 .expect("stream ended before first delta")
16663 .expect("stream body failed");
16664 assert!(
16665 is_sse_data_frame(&first),
16666 "first frame was not SSE data: {first:?}"
16667 );
16668 drop(body);
16669
16670 let mut dropped = None;
16673 for _ in 0..500 {
16674 if let Some(event) = mock
16675 .events()
16676 .into_iter()
16677 .find(|e| matches!(e, MeterEvent::Dropped { .. }))
16678 {
16679 dropped = Some(event);
16680 break;
16681 }
16682 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
16683 }
16684 let events = mock.events();
16685 assert!(
16686 events
16687 .iter()
16688 .any(|e| matches!(e, MeterEvent::Open { request_id: id, .. } if id == &request_id)),
16689 "the receipt was opened under the caller-visible request id: {events:?}"
16690 );
16691 assert_eq!(
16692 dropped,
16693 Some(MeterEvent::Dropped {
16694 prompt: 1,
16695 cached: 0,
16696 completion: 1,
16697 }),
16698 "a client disconnect must leave the partial counts on the dropped receipt \
16699 (the implementation prices that drop): {events:?}"
16700 );
16701 }
16702
16703 #[tokio::test]
16704 async fn draining_rejects_new_requests_with_503_and_retry_after() {
16705 let _l = DRAIN_LOCK.lock().unwrap();
16706 let st = fake_worker_state();
16707 DRAINING.store(true, std::sync::atomic::Ordering::SeqCst);
16708 let resp = chat_completions(
16710 State(st.clone()),
16711 axum::http::HeaderMap::new(),
16712 None,
16713 Json(
16714 serde_json::from_value(serde_json::json!({
16715 "model": "m", "messages": [{"role": "user", "content": "t"}]
16716 }))
16717 .unwrap(),
16718 ),
16719 )
16720 .await;
16721 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16722 let ra = resp.headers()["retry-after"].to_str().unwrap().to_string();
16727 let ra_s: u64 = ra
16728 .parse()
16729 .expect("Retry-After must be integer delay-seconds");
16730 assert!(
16731 ra_s > 0 && ra_s <= 60,
16732 "Retry-After {ra_s}s is outside the honored window"
16733 );
16734 let ra_ms: u64 = resp.headers()["retry-after-ms"]
16735 .to_str()
16736 .unwrap()
16737 .parse()
16738 .unwrap();
16739 assert_eq!(ra_ms, ra_s * 1000, "the two retry headers must agree");
16740 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16741 .await
16742 .unwrap();
16743 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16744 assert!(
16745 payload["error"]["message"]
16746 .as_str()
16747 .unwrap()
16748 .contains("draining")
16749 );
16750 assert_eq!(payload["error"]["type"], "server_error");
16751 assert_eq!(payload["error"]["code"], "draining");
16752 let resp = completions(
16753 State(st.clone()),
16754 axum::http::HeaderMap::new(),
16755 None,
16756 Json(
16757 serde_json::from_value(serde_json::json!({
16758 "model": "m", "prompt": "t"
16759 }))
16760 .unwrap(),
16761 ),
16762 )
16763 .await;
16764 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16765 assert!(resp.headers().contains_key("retry-after"));
16766 assert_eq!(
16767 st.inflight[0].load(std::sync::atomic::Ordering::SeqCst),
16768 0,
16769 "rejected requests must not hold slots"
16770 );
16771 let resp = health_live(State(st.clone())).await.into_response();
16774 assert_eq!(
16775 resp.status(),
16776 StatusCode::OK,
16777 "a drain must not look like a liveness fault"
16778 );
16779 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16780 .await
16781 .unwrap();
16782 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16783 assert_eq!(payload["status"], "draining");
16784 let resp = health_ready(State(st.clone())).await.into_response();
16786 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16787 let retry_s = drain_deadline_s().clamp(1, 60);
16788 let retry_s_text = retry_s.to_string();
16789 let retry_ms_text = (retry_s * 1000).to_string();
16790 assert_eq!(retry_after(&resp).as_deref(), Some(retry_s_text.as_str()));
16791 assert_eq!(
16792 resp.headers().get("retry-after-ms").unwrap(),
16793 retry_ms_text.as_str()
16794 );
16795 assert_ne!(
16796 resp.headers()
16797 .get("x-should-retry")
16798 .and_then(|v| v.to_str().ok()),
16799 Some("false")
16800 );
16801 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16802 .await
16803 .unwrap();
16804 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16805 assert_eq!(payload["status"], "not_ready");
16806 assert!(payload["detail"].as_str().unwrap().contains("draining"));
16807 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16808 let resp = chat_completions(
16810 State(st.clone()),
16811 axum::http::HeaderMap::new(),
16812 None,
16813 Json(
16814 serde_json::from_value(serde_json::json!({
16815 "model": "m", "messages": [{"role": "user", "content": "t"}]
16816 }))
16817 .unwrap(),
16818 ),
16819 )
16820 .await;
16821 assert_eq!(resp.status(), StatusCode::OK);
16822 }
16823
16824 #[tokio::test]
16827 async fn health_is_green_only_while_the_worker_is_alive() {
16828 let _l = DRAIN_LOCK.lock().unwrap();
16831 let st = fake_worker_state();
16832 let resp = health_live(State(st.clone())).await.into_response();
16835 assert_eq!(resp.status(), StatusCode::OK);
16836 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16837 .await
16838 .unwrap();
16839 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16840 assert_eq!(payload["status"], "ok");
16841 assert_eq!(payload["worker"]["phase"], "idle");
16842 assert!(payload["worker"]["stall_threshold_ms"].as_u64().unwrap() > 0);
16843 let ready = health_ready(State(st.clone())).await.into_response();
16844 assert_eq!(ready.status(), StatusCode::OK);
16845
16846 st.health.mark_dead("worker thread panicked: test-injected");
16850 let resp = health_live(State(st.clone())).await.into_response();
16851 assert_eq!(
16852 resp.status(),
16853 StatusCode::SERVICE_UNAVAILABLE,
16854 "a dead worker MUST NOT report a healthy liveness"
16855 );
16856 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
16857 .await
16858 .unwrap();
16859 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16860 assert_eq!(payload["status"], "unhealthy");
16861 assert!(
16863 payload["detail"]
16864 .as_str()
16865 .unwrap()
16866 .contains("test-injected"),
16867 "cause not surfaced: {payload}"
16868 );
16869 let ready = health_ready(State(st.clone())).await.into_response();
16870 assert_eq!(
16871 ready.status(),
16872 StatusCode::SERVICE_UNAVAILABLE,
16873 "dead is also not ready"
16874 );
16875
16876 st.health.mark_ready();
16879 assert_eq!(
16880 health_live(State(st.clone()))
16881 .await
16882 .into_response()
16883 .status(),
16884 StatusCode::OK,
16885 "mark_ready must clear the latch (a successful respawn)"
16886 );
16887 }
16888
16889 #[tokio::test]
16890 async fn readyz_peer_probe_integrity_is_present_and_advisory() {
16891 let _l = DRAIN_LOCK.lock().unwrap();
16892 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16893 let st = fake_worker_state();
16894
16895 let ready = health_ready(State(st.clone())).await.into_response();
16896 assert_eq!(ready.status(), StatusCode::OK);
16897 let bytes = axum::body::to_bytes(ready.into_body(), usize::MAX)
16898 .await
16899 .unwrap();
16900 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16901 assert_eq!(payload["peer_probe_integrity"], "ok");
16902
16903 st.health.note_peer_probe_deferral(2, false);
16904 let deferred = health_ready(State(st.clone())).await.into_response();
16905 assert_eq!(deferred.status(), StatusCode::OK);
16906 let bytes = axum::body::to_bytes(deferred.into_body(), usize::MAX)
16907 .await
16908 .unwrap();
16909 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16910 assert_eq!(payload["peer_probe_integrity"], "deferred_2");
16911
16912 st.health.note_peer_probe_deferral(4, true);
16913 let degraded = health_ready(State(st.clone())).await.into_response();
16914 assert_eq!(
16915 degraded.status(),
16916 StatusCode::OK,
16917 "peer degradation is advisory while plain serving remains healthy"
16918 );
16919 let bytes = axum::body::to_bytes(degraded.into_body(), usize::MAX)
16920 .await
16921 .unwrap();
16922 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16923 assert_eq!(payload["peer_probe_integrity"], "degraded");
16924
16925 st.health.mark_dead("test-injected worker failure");
16926 let unready = health_ready(State(st)).await.into_response();
16927 assert_eq!(unready.status(), StatusCode::SERVICE_UNAVAILABLE);
16928 let bytes = axum::body::to_bytes(unready.into_body(), usize::MAX)
16929 .await
16930 .unwrap();
16931 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
16932 assert_eq!(
16933 payload["peer_probe_integrity"], "degraded",
16934 "the advisory field must also survive an unrelated readiness failure"
16935 );
16936 }
16937
16938 #[tokio::test]
16939 async fn liveness_failure_obeys_the_retry_contract() {
16940 let _l = DRAIN_LOCK.lock().unwrap();
16944 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16945 let st = fake_worker_state();
16946 st.health
16947 .mark_dead("worker thread panicked: retry-contract-test");
16948
16949 let resp = health_live(State(st)).await.into_response();
16950 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16951 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16952 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
16953 assert_ne!(
16954 resp.headers()
16955 .get("x-should-retry")
16956 .and_then(|v| v.to_str().ok()),
16957 Some("false")
16958 );
16959 }
16960
16961 #[tokio::test]
16962 async fn readiness_failure_obeys_the_retry_contract() {
16963 let _l = DRAIN_LOCK.lock().unwrap();
16964 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16965 let st = fake_worker_state();
16966 st.health
16967 .mark_dead("worker thread panicked: retry-contract-test");
16968
16969 let resp = health_ready(State(st)).await.into_response();
16970 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
16971 assert_eq!(retry_after(&resp).as_deref(), Some("2"));
16972 assert_eq!(resp.headers().get("retry-after-ms").unwrap(), "2000");
16973 assert_ne!(
16974 resp.headers()
16975 .get("x-should-retry")
16976 .and_then(|v| v.to_str().ok()),
16977 Some("false")
16978 );
16979 }
16980
16981 #[tokio::test]
16982 async fn a_wedged_gpu_flips_health_even_though_the_worker_thread_is_fine() {
16983 let _l = DRAIN_LOCK.lock().unwrap();
16993 DRAINING.store(false, std::sync::atomic::Ordering::SeqCst);
16994 let st = fake_worker_state();
16995 assert_eq!(
16996 health_live(State(st.clone()))
16997 .await
16998 .into_response()
16999 .status(),
17000 StatusCode::OK
17001 );
17002 st.health
17003 .mark_gpu_fault("nvidia-smi probe exceeded 10s deadline (GSP hang class)");
17004 let resp = health_live(State(st.clone())).await.into_response();
17005 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
17006 let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
17007 .await
17008 .unwrap();
17009 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
17010 assert!(
17011 payload["detail"]
17012 .as_str()
17013 .unwrap()
17014 .contains("probe exceeded")
17015 );
17016 st.health.mark_ready();
17019 assert_eq!(
17020 health_live(State(st.clone()))
17021 .await
17022 .into_response()
17023 .status(),
17024 StatusCode::SERVICE_UNAVAILABLE,
17025 "a GPU fault must not be cleared by an in-process respawn"
17026 );
17027 }
17028
17029 #[test]
17030 fn v1_models_entry_keeps_catalog_shape_with_honest_nulls() {
17031 let caps = ModelCaps {
17033 tools_branch: true,
17034 qwen_think: true,
17035 think_switch: true,
17036 chat_ok: true,
17037 context_length: 262144,
17038 tokenizer: "qwen2".into(),
17039 instruct_type: Some("chatml".into()),
17040 effort_levels: false,
17041 qwen_effort: false,
17042 gemma_think: false,
17043 dsv4: false,
17044 chat_temperature_default: None,
17045 chat_top_p_default: None,
17046 n_vocab: 151_936,
17047 };
17048 let e = model_entry_v1("main", Some(&caps), None);
17049 assert_eq!(e["id"], "main");
17050 assert_eq!(e["name"], "main");
17051 assert_eq!(e["object"], "model");
17052 assert_eq!(e["context_length"], 262144);
17053 assert!(e["pricing"]["input"].is_null());
17055 assert!(e["pricing"]["output"].is_null());
17056
17057 let meta = OpenRouterModelMetadata {
17060 pricing: OpenRouterPricing {
17061 prompt: Some("0.00000038".into()),
17062 cached_prompt: Some("0.0000002".into()),
17063 completion: Some("0.0000026".into()),
17064 ..Default::default()
17065 },
17066 input_modalities: vec!["image".into(), "video".into()],
17067 max_output_length: Some(32768),
17068 ..Default::default()
17069 };
17070 let e = model_entry_v1("main", Some(&caps), Some(&meta));
17071 assert_eq!(e["pricing"]["currency"], "USD");
17074 assert_eq!(e["pricing"]["unit"], "per_1m_tokens");
17075 assert_eq!(e["pricing"]["input"], "0.38");
17076 assert_eq!(e["pricing"]["output"], "2.60");
17077 assert_eq!(e["pricing"]["cached_input"], "0.20");
17078 assert!(e["pricing"]["cache_write"].is_null());
17079 assert_eq!(e["pricing"]["minimum_request"], "0");
17080 assert_eq!(e["owned_by"], "main");
17081 assert_eq!(e["type"], "chat");
17082 assert_eq!(e["max_output_tokens"], 32768);
17083 assert_eq!(e["endpoints"], json!(["chat/completions"]));
17084 assert_eq!(e["input_modalities"], json!(["text", "image", "video"]));
17085 assert_eq!(e["output_modalities"], json!(["text"]));
17086 assert_eq!(e["capabilities"]["streaming"], true);
17087 assert_eq!(e["capabilities"]["tools"], true);
17088 assert_eq!(e["lifecycle"]["status"], "active");
17089 assert!(e["lifecycle"]["deprecation_at"].is_null());
17090 assert_eq!(e["reliability"]["first_token_timeout_seconds"], 120);
17091 assert_eq!(e["reliability"]["capacity_scope"], "model_region");
17092 let mut keys: Vec<&str> = e.as_object().unwrap().keys().map(String::as_str).collect();
17096 keys.sort_unstable();
17097 assert_eq!(
17098 keys,
17099 [
17100 "capabilities",
17101 "context_length",
17102 "endpoints",
17103 "id",
17104 "input_modalities",
17105 "lifecycle",
17106 "max_output_tokens",
17107 "name",
17108 "object",
17109 "output_modalities",
17110 "owned_by",
17111 "pricing",
17112 "reliability",
17113 "type",
17114 ],
17115 "unexpected /v1/models entry keys"
17116 );
17117 let mut price_keys: Vec<&str> = e["pricing"]
17118 .as_object()
17119 .unwrap()
17120 .keys()
17121 .map(String::as_str)
17122 .collect();
17123 price_keys.sort_unstable();
17124 assert_eq!(
17125 price_keys,
17126 [
17127 "cache_write",
17128 "cached_input",
17129 "currency",
17130 "input",
17131 "minimum_request",
17132 "output",
17133 "unit",
17134 ],
17135 "unexpected /v1/models pricing keys"
17136 );
17137
17138 let e = model_entry_v1("m", None, None);
17140 assert!(e["context_length"].is_null());
17141 assert!(e["max_output_tokens"].is_null());
17142 let bare = ModelCaps::default(); let e = model_entry_v1("m", Some(&bare), None);
17144 assert!(e["context_length"].is_null());
17145 }
17146
17147 #[test]
17153 fn catalog_row_follows_the_declared_surface() {
17154 let caps = ModelCaps {
17155 tools_branch: true,
17156 ..Default::default()
17157 };
17158
17159 let embed = OpenRouterModelMetadata {
17160 surface: Some("embedding".into()),
17161 max_output_length: Some(1),
17162 ..Default::default()
17163 };
17164 let e = model_entry_v1("qwen", Some(&caps), Some(&embed));
17165 assert_eq!(e["type"], "embedding");
17166 assert_eq!(e["endpoints"], json!(["embeddings"]));
17167 assert_eq!(e["output_modalities"], json!(["embeddings"]));
17168 assert_eq!(e["capabilities"]["streaming"], false);
17169 assert_eq!(
17170 e["capabilities"]["tools"], false,
17171 "an embedder has no tools"
17172 );
17173 assert_eq!(e["capabilities"]["reasoning"], false);
17174 assert_eq!(e["capabilities"]["structured_output"], false);
17175 assert_eq!(e["capabilities"]["prompt_caching"], false);
17176 assert!(
17177 e["max_output_tokens"].is_null(),
17178 "a surface that emits no completion tokens must not advertise a ceiling"
17179 );
17180
17181 let rerank = OpenRouterModelMetadata {
17182 surface: Some("rerank".into()),
17183 ..Default::default()
17184 };
17185 let r = model_entry_v1("qwen", Some(&caps), Some(&rerank));
17186 assert_eq!(r["type"], "rerank");
17187 assert_eq!(r["endpoints"], json!(["rerank"]));
17188 assert_eq!(r["output_modalities"], json!(["rerank"]));
17189 assert_eq!(r["capabilities"]["tools"], false);
17190 assert_eq!(r["capabilities"]["reasoning"], false);
17191
17192 let chat = OpenRouterModelMetadata {
17195 max_output_length: Some(32768),
17196 ..Default::default()
17197 };
17198 let c = model_entry_v1("main", Some(&caps), Some(&chat));
17199 assert_eq!(c["type"], "chat");
17200 assert_eq!(c["endpoints"], json!(["chat/completions"]));
17201 assert_eq!(c["output_modalities"], json!(["text"]));
17202 assert_eq!(c["capabilities"]["tools"], true);
17203 assert_eq!(c["max_output_tokens"], 32768);
17204 }
17205
17206 #[test]
17209 fn unknown_surface_is_rejected_at_config_load() {
17210 let bad = OpenRouterModelMetadata {
17211 surface: Some("embeddings".into()), ..Default::default()
17213 };
17214 let err = validate_openrouter_metadata("qwen/qwen3-embedding-8b", &bad)
17215 .expect_err("an unknown surface must not load");
17216 assert!(err.contains("surface"), "{err}");
17217
17218 for good in ["chat", "embedding", "rerank"] {
17219 let ok = OpenRouterModelMetadata {
17220 surface: Some(good.into()),
17221 ..Default::default()
17222 };
17223 assert!(
17224 validate_openrouter_metadata("m", &ok).is_ok(),
17225 "{good} must load"
17226 );
17227 }
17228 }
17229
17230 #[test]
17231 fn per_million_price_is_exact_decimal_shift() {
17232 assert_eq!(per_million_price("0.00000038").as_deref(), Some("0.38"));
17234 assert_eq!(per_million_price("0.0000026").as_deref(), Some("2.60"));
17235 assert_eq!(per_million_price("0.0000002").as_deref(), Some("0.20"));
17236 assert_eq!(per_million_price("0").as_deref(), Some("0.00"));
17237 assert_eq!(per_million_price("1.5").as_deref(), Some("1500000.00"));
17238 assert_eq!(per_million_price("0.000000125").as_deref(), Some("0.125"));
17239 assert_eq!(per_million_price("not-a-price"), None);
17240 assert_eq!(per_million_price(""), None);
17241 }
17242
17243 #[test]
17244 fn metadata_provider_block_parses_and_validates() {
17245 let (_, provider) = OpenRouterMetadataFile::parse(
17246 r#"
17247 [provider]
17248 id = "tiyuvta"
17249 status_url = "https://status.tiyuvta.ai"
17250 support_contact = "mailto:support@tiyuvta.ai"
17251 incident_contact = "mailto:incidents@tiyuvta.ai"
17252 regions = ["eu-central"]
17253 "#,
17254 )
17255 .unwrap();
17256 let provider = provider.unwrap();
17257 assert_eq!(provider.id, "tiyuvta");
17258 assert_eq!(provider.regions, vec!["eu-central"]);
17259 let err = OpenRouterMetadataFile::parse("[provider]\nid = \"\"\n").unwrap_err();
17261 assert!(err.contains("provider.id"), "{err}");
17262 let err = OpenRouterMetadataFile::parse(
17264 "[provider]\nid = \"x\"\nsupport_contact = \"ops@example.com\"\n",
17265 )
17266 .unwrap_err();
17267 assert!(err.contains("must be a URI"), "{err}");
17268 let (_, provider) = OpenRouterMetadataFile::parse("").unwrap();
17270 assert!(provider.is_none());
17271 }
17272
17273 #[test]
17274 fn models_openai_default_body_stays_byte_identical() {
17275 let body = models_openai_body(&["main".into(), "judge".into()]);
17276 let bytes = serde_json::to_vec(&body).unwrap();
17277 assert_eq!(
17278 bytes,
17279 br#"{"object":"list","data":[{"id":"main","object":"model"},{"id":"judge","object":"model"}]}"#
17280 );
17281 }
17282
17283 #[test]
17284 fn canonical_model_id_tolerates_a_marketplace_stripping_the_vendor_prefix() {
17285 let loaded = vec![
17287 "qwen/qwen3.6-27b".to_string(),
17288 "qwen/qwen3.6-35b-a3b".to_string(),
17289 ];
17290 assert_eq!(
17291 canonical_model_id(&loaded, "qwen3.6-35b-a3b").as_deref(),
17292 Some("qwen/qwen3.6-35b-a3b"),
17293 );
17294 assert_eq!(
17295 canonical_model_id(&loaded, "qwen3.6-27b").as_deref(),
17296 Some("qwen/qwen3.6-27b"),
17297 );
17298 assert_eq!(
17300 canonical_model_id(&loaded, "qwen/qwen3.6-35b-a3b").as_deref(),
17301 Some("qwen/qwen3.6-35b-a3b"),
17302 );
17303 assert_eq!(canonical_model_id(&loaded, "gpt-4o"), None);
17305 assert_eq!(canonical_model_id(&loaded, "vendor/qwen3.6-35b-a3b"), None);
17306 assert_eq!(canonical_model_id(&loaded, ""), None);
17307 }
17308
17309 #[test]
17310 fn canonical_model_id_refuses_an_ambiguous_suffix_rather_than_guessing() {
17311 let loaded = vec!["a/shared-name".to_string(), "b/shared-name".to_string()];
17314 assert_eq!(canonical_model_id(&loaded, "shared-name"), None);
17315 assert_eq!(
17317 canonical_model_id(&loaded, "a/shared-name").as_deref(),
17318 Some("a/shared-name")
17319 );
17320 assert_eq!(
17321 canonical_model_id(&loaded, "b/shared-name").as_deref(),
17322 Some("b/shared-name")
17323 );
17324 let bare = vec!["solo".to_string()];
17326 assert_eq!(canonical_model_id(&bare, "solo").as_deref(), Some("solo"));
17327 }
17328
17329 #[test]
17330 fn openrouter_models_entry_serializes_complete_metadata() {
17331 let metadata = OpenRouterMetadataFile::from_toml(
17332 r#"
17333[models.main]
17334hugging_face_id = "Qwen/Qwen3.6-27B"
17335created = 1786032000
17336quantization = "nvfp4"
17337description = "Qwen3.6 27B served by memra."
17338max_prompt_length = 245760
17339max_output_length = 16384
17340default_output_length = 4096
17341is_ready = true
17342is_free = false
17343discount_to_user = 0.1
17344openrouter_slug = "qwen/qwen3.6-27b"
17345datacenters = [{ country_code = "US", region = "us-east-1" }]
17346zdr = true
17347hipaa = false
17348
17349[models.main.pricing]
17350prompt = "0.000000234"
17351cached_prompt = "0.0000000585"
17352cache_write = "0.000000234"
17353completion = "0.000001872"
17354internal_reasoning = "0.000001872"
17355request = "0.01"
17356
17357[models.main.capacity]
17358prompt_tpm = 1000000
17359cached_prompt_tpm = 2000000
17360completion_tpm = 500000
17361request_rpm = 1000
17362concurrency = 64
17363"#,
17364 )
17365 .unwrap();
17366 let caps = ModelCaps {
17367 tools_branch: true,
17368 qwen_think: true,
17369 think_switch: true,
17370 chat_ok: true,
17371 context_length: 262144,
17372 tokenizer: "qwen2".into(),
17373 instruct_type: Some("chatml".into()),
17374 ..Default::default()
17375 };
17376 let entry = model_entry_openrouter("main", Some(&caps), metadata.get("main"));
17377
17378 assert_eq!(entry["schema_version"], "2.4");
17379 assert_eq!(entry["id"], "main");
17380 assert_eq!(entry["name"], "main");
17381 assert_eq!(entry["hugging_face_id"], "Qwen/Qwen3.6-27B");
17382 assert_eq!(entry["created"], 1786032000u64);
17383 assert_eq!(entry["quantization"], "nvfp4");
17384 assert_eq!(entry["tokenizer"], "qwen2");
17385 assert_eq!(entry["description"], "Qwen3.6 27B served by memra.");
17386 assert!(
17387 entry.get("object").is_none(),
17388 "OpenRouter schema 2.4 rejects unknown OpenAI fields"
17389 );
17390
17391 let input = &entry["input_modalities"][0];
17392 assert_eq!(input["type"], "text");
17393 assert_eq!(
17394 input["supported_inputs"]["max_context_length"]["value"],
17395 262144
17396 );
17397 assert_eq!(
17398 input["supported_inputs"]["max_prompt_length"]["value"],
17399 245760
17400 );
17401 let input_prices = input["pricing"].as_array().unwrap();
17402 let input_price = |kind: &str| {
17403 input_prices
17404 .iter()
17405 .find(|price| price["type"] == kind)
17406 .unwrap()
17407 };
17408 assert_eq!(input_price("prompt")["cost_usd"], "0.000000234");
17409 assert_eq!(input_price("cached_prompt")["cost_usd"], "0.0000000585");
17410 assert_eq!(input_price("cache_write")["cost_usd"], "0.000000234");
17411 assert_eq!(input["capacity"][0]["value"], 1000000);
17412 assert_eq!(input["capacity"][1]["value"], 2000000);
17413
17414 let output = &entry["output_modalities"][0];
17415 assert_eq!(output["type"], "text");
17416 assert_eq!(output["max_length"]["value"], 16384);
17417 assert_eq!(output["streaming"], true);
17418 assert_eq!(output["supported_parameters"]["tools"]["type"], "boolean");
17419 assert_eq!(
17420 output["supported_parameters"]["structured_outputs"]["type"],
17421 "boolean"
17422 );
17423 assert_eq!(
17424 output["supported_parameters"]["reasoning"]["type"],
17425 "boolean"
17426 );
17427 assert_eq!(output["pricing"][0]["type"], "completion");
17428 assert_eq!(output["pricing"][0]["cost_usd"], "0.000001872");
17429 assert_eq!(output["pricing"][1]["type"], "internal_reasoning");
17430 assert_eq!(output["capacity"][0]["value"], 500000);
17431 assert_eq!(output["capacity"][1]["type"], "concurrency");
17432 assert_eq!(output["capacity"][1]["value"], 64);
17433
17434 assert_eq!(entry["pricing"][0]["type"], "request");
17435 assert_eq!(entry["pricing"][0]["cost_usd"], "0.01");
17436 assert_eq!(entry["capacity"][0]["value"], 1000);
17437 assert_eq!(entry["is_ready"], true);
17438 assert_eq!(entry["is_free"], false);
17439 assert_eq!(entry["discount_to_user"], 0.1);
17440 assert_eq!(entry["openrouter"]["slug"], "qwen/qwen3.6-27b");
17441 assert_eq!(entry["datacenters"][0]["country_code"], "US");
17442 assert_eq!(entry["compliance"]["zdr"], true);
17443 assert_eq!(entry["compliance"]["hipaa"], false);
17444 }
17445
17446 const GATEWAY_REGISTRY_FIXTURE: &str = r#"
17451[models."qwen/qwen3.6-35b-a3b"]
17452hugging_face_id = "Qwen/Qwen3.6-35B-A3B"
17453created = 1777260255
17454quantization = "int4"
17455description = "Qwen3.6 35B-A3B fixture entry."
17456max_prompt_length = 262144
17457max_output_length = 262144
17458default_output_length = 8192
17459is_ready = true
17460is_free = false
17461discount_to_user = 0.0
17462openrouter_slug = "qwen/qwen3.6-35b-a3b"
17463zdr = false
17464hipaa = false
17465
17466[[models."qwen/qwen3.6-35b-a3b".datacenters]]
17467country_code = "CA"
17468region = "Ontario"
17469
17470[models."qwen/qwen3.6-35b-a3b".pricing]
17471prompt = "0.0000000931"
17472cached_prompt = "0.0000000652"
17473completion = "0.0000009025"
17474
17475[models."qwen/qwen3.6-35b-a3b".capacity]
17476prompt_tpm = 780000
17477cached_prompt_tpm = 310000
17478completion_tpm = 9600
17479request_rpm = 160
17480concurrency = 16
17481
17482[planned_models."qwen/qwen3.8-27b"]
17483description = "Planned fixture entry; must never be emitted."
17484max_prompt_length = 262144
17485max_output_length = 262144
17486default_output_length = 8192
17487is_ready = false
17488is_free = false
17489discount_to_user = 0.0
17490openrouter_slug = "qwen/qwen3.8-27b"
17491zdr = false
17492hipaa = false
17493
17494[planned_models."qwen/qwen3.8-27b".pricing]
17495prompt = "0.0000002745"
17496cached_prompt = "0.0000001922"
17497completion = "0.0000022800"
17498
17499[planned_models."google/gemma-4-26b-a4b-it"]
17500hugging_face_id = "google/gemma-4-26B-A4B-it"
17501created = 1775227989
17502quantization = "int4"
17503description = "Planned fixture entry; must never be emitted."
17504max_prompt_length = 262144
17505max_output_length = 262144
17506default_output_length = 8192
17507is_ready = false
17508is_free = false
17509discount_to_user = 0.0
17510openrouter_slug = "google/gemma-4-26b-a4b-it"
17511zdr = false
17512hipaa = false
17513
17514[planned_models."google/gemma-4-26b-a4b-it".pricing]
17515prompt = "0.0000000665"
17516cached_prompt = "0.0000000466"
17517completion = "0.0000003230"
17518"#;
17519
17520 #[test]
17521 fn gateway_registry_generates_the_staged_active_shape() {
17522 let metadata = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17523 let caps = ModelCaps {
17524 tools_branch: true,
17525 qwen_think: true,
17526 think_switch: true,
17527 chat_ok: true,
17528 context_length: 262144,
17529 tokenizer: "qwen2".into(),
17530 instruct_type: Some("chatml".into()),
17531 ..Default::default()
17532 };
17533 let q35_entry = model_entry_openrouter(
17534 "qwen/qwen3.6-35b-a3b",
17535 Some(&caps),
17536 metadata.get("qwen/qwen3.6-35b-a3b"),
17537 );
17538 assert_eq!(q35_entry["created"], 1777260255u64);
17539 assert_eq!(q35_entry["quantization"], "int4");
17540 assert_eq!(q35_entry["is_ready"], true);
17541 assert_eq!(
17542 q35_entry["input_modalities"][0]["supported_inputs"]["max_context_length"]["value"],
17543 262144
17544 );
17545 assert_eq!(
17546 q35_entry["input_modalities"][0]["supported_inputs"]["max_prompt_length"]["value"],
17547 262144
17548 );
17549 assert_eq!(
17550 q35_entry["output_modalities"][0]["max_length"]["value"],
17551 262144
17552 );
17553 let prices = q35_entry["input_modalities"][0]["pricing"]
17554 .as_array()
17555 .unwrap();
17556 assert_eq!(prices[0]["cost_usd"], "0.0000000931");
17557 assert_eq!(prices[1]["cost_usd"], "0.0000000652");
17558 assert_eq!(
17562 q35_entry["input_modalities"][0]["capacity"][0]["value"],
17563 780000
17564 );
17565 assert_eq!(
17566 q35_entry["input_modalities"][0]["capacity"][1]["value"],
17567 310000
17568 );
17569 assert_eq!(
17570 q35_entry["output_modalities"][0]["supported_parameters"]["max_tokens"]["max"],
17571 262144
17572 );
17573 assert_eq!(
17574 q35_entry["output_modalities"][0]["capacity"][0]["value"],
17575 9600
17576 );
17577 assert_eq!(
17578 q35_entry["output_modalities"][0]["capacity"][1]["value"],
17579 16
17580 );
17581 assert_eq!(
17582 q35_entry["output_modalities"][0]["pricing"][0]["cost_usd"],
17583 "0.0000009025"
17584 );
17585 assert_eq!(q35_entry["capacity"][0]["value"], 160); assert_eq!(q35_entry["datacenters"][0]["country_code"], "CA");
17587
17588 assert_eq!(
17589 metadata.len(),
17590 1,
17591 "planned models must never enter the active map"
17592 );
17593 assert!(!metadata.contains_key("qwen/qwen3.6-27b"));
17594 assert!(!metadata.contains_key("qwen/qwen3.8-27b"));
17595 assert!(!metadata.contains_key("google/gemma-4-26b-a4b-it"));
17596
17597 let openmodels = model_entry_openmodels(
17598 "qwen/qwen3.6-35b-a3b",
17599 Some(&caps),
17600 metadata.get("qwen/qwen3.6-35b-a3b"),
17601 )
17602 .unwrap();
17603 assert_eq!(openmodels["currency"], "USD");
17604 assert_eq!(openmodels["max_output_length"], 262144);
17605 assert_eq!(openmodels["is_ready"], true);
17606 assert_eq!(openmodels["is_free"], false);
17607 assert_eq!(openmodels["discount_to_user"], 0.0);
17608 }
17609
17610 #[test]
17611 fn gateway_registry_limits_are_live_request_limits() {
17612 let metadata_file = OpenRouterMetadataFile::from_toml(GATEWAY_REGISTRY_FIXTURE).unwrap();
17613 let metadata = metadata_file.get("qwen/qwen3.6-35b-a3b").unwrap();
17614 let caps = ModelCaps {
17615 context_length: 262_144,
17616 ..Default::default()
17617 };
17618 let build = |value: serde_json::Value| {
17619 let req: CompletionReq = serde_json::from_value(value).unwrap();
17620 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
17621 build_request(&req, tx, lanes::Lane::Interactive, None)
17622 };
17623
17624 let mut omitted = build(json!({
17625 "model": "qwen/qwen3.6-35b-a3b",
17626 "prompt_ids": [1, 2, 3]
17627 }));
17628 apply_model_request_limits(&mut omitted, Some(metadata), Some(&caps)).unwrap();
17629 assert_eq!(omitted.params.max_new, 8_192);
17630 assert_eq!(omitted.max_prompt_tokens, Some(262_144));
17631
17632 let mut field_top = build(json!({
17633 "model": "qwen/qwen3.6-35b-a3b",
17634 "prompt_ids": [1],
17635 "max_tokens": 262144
17636 }));
17637 apply_model_request_limits(&mut field_top, Some(metadata), Some(&caps)).unwrap();
17638 assert_eq!(field_top.params.max_new, 262_144);
17639 assert_eq!(
17640 budget_completion_bound(&field_top, 100, Some(&caps)).unwrap(),
17641 262_044,
17642 "the field-top output request is accepted but bounded by remaining trained context",
17643 );
17644
17645 let mut too_much_output = build(json!({
17646 "model": "qwen/qwen3.6-35b-a3b",
17647 "prompt_ids": [1],
17648 "max_tokens": 262145
17649 }));
17650 let (message, param) =
17651 apply_model_request_limits(&mut too_much_output, Some(metadata), Some(&caps))
17652 .unwrap_err();
17653 assert_eq!(param, "max_tokens");
17654 assert!(message.contains("262145"));
17655
17656 let mut oversized_allocation = build(json!({
17657 "model": "qwen/qwen3.6-35b-a3b",
17658 "prompt_ids": [1],
17659 "max_tokens": 1,
17660 "max_ctx": 262145
17661 }));
17662 let (_, param) =
17663 apply_model_request_limits(&mut oversized_allocation, Some(metadata), Some(&caps))
17664 .unwrap_err();
17665 assert_eq!(param, "max_ctx");
17666 }
17667
17668 #[test]
17669 fn planned_registry_entries_are_validated_but_never_activated() {
17670 let parsed = OpenRouterMetadataFile::from_toml(
17671 r#"
17672[planned_models.future]
17673max_output_length = 262144
17674default_output_length = 8192
17675
17676[planned_models.future.pricing]
17677prompt = "0.0000001"
17678"#,
17679 )
17680 .unwrap();
17681 assert!(parsed.is_empty());
17682
17683 let error = OpenRouterMetadataFile::from_toml(
17684 r#"
17685[planned_models.future]
17686default_output_length = 8192
17687"#,
17688 )
17689 .unwrap_err();
17690 assert!(error.contains("requires max_output_length"));
17691 }
17692
17693 #[test]
17698 fn every_catalog_feed_honours_the_declared_surface() {
17699 let metadata = OpenRouterMetadataFile::from_toml(
17700 r#"
17701[models."qwen/qwen3-embedding-8b"]
17702surface = "embedding"
17703created = 1787961600
17704max_output_length = 1
17705is_ready = true
17706is_free = false
17707discount_to_user = 0.0
17708
17709[models."qwen/qwen3-embedding-8b".pricing]
17710prompt = "0.00000001"
17711cached_prompt = "0.0"
17712completion = "0.0"
17713
17714[models."main"]
17715created = 1787443200
17716max_output_length = 32768
17717is_ready = true
17718is_free = false
17719discount_to_user = 0.0
17720
17721[models."main".pricing]
17722prompt = "0.00000025"
17723cached_prompt = "0.00000009"
17724completion = "0.0000012"
17725"#,
17726 )
17727 .unwrap();
17728 let caps = ModelCaps {
17729 tools_branch: true,
17730 qwen_think: true,
17731 chat_ok: true,
17732 context_length: 32768,
17733 ..Default::default()
17734 };
17735 let embed = metadata.get("qwen/qwen3-embedding-8b");
17736 let chat = metadata.get("main");
17737
17738 let or = model_entry_openrouter("qwen/qwen3-embedding-8b", Some(&caps), embed);
17740 let out = &or["output_modalities"][0];
17741 assert_eq!(out["type"], "embeddings", "openrouter feed: {or}");
17742 assert!(
17743 out.get("streaming").is_none(),
17744 "the embeddings branch declares no streaming property (additionalProperties:false): {out}"
17745 );
17746 let params = &out["supported_parameters"];
17751 assert_eq!(
17752 params.as_object().map(|o| o.len()),
17753 Some(0),
17754 "no completion parameter belongs on an embedder row: {params}"
17755 );
17756 for field in [
17757 "tools",
17758 "tool_choice",
17759 "reasoning",
17760 "max_tokens",
17761 "json_mode",
17762 "structured_outputs",
17763 "stop",
17764 "temperature",
17765 "seed",
17766 ] {
17767 assert!(params[field].is_null(), "{field} leaked onto an embedder");
17768 }
17769 assert!(
17770 out["max_length"].is_null(),
17771 "a surface emitting no completion tokens advertises no ceiling: {out}"
17772 );
17773
17774 let om = model_entry_openmodels("qwen/qwen3-embedding-8b", Some(&caps), embed)
17776 .expect("openmodels entry builds");
17777 assert_eq!(om["output_modalities"], json!(["embeddings"]));
17778 let features = om["supported_features"].as_array().unwrap();
17779 assert!(
17780 !features
17781 .iter()
17782 .any(|f| f == "tool_calling" || f == "reasoning"),
17783 "chat-only features leaked onto an embedder: {features:?}"
17784 );
17785
17786 let v1 = model_entry_v1("qwen/qwen3-embedding-8b", Some(&caps), embed);
17788 assert_eq!(v1["type"], "embedding");
17789 assert_eq!(v1["capabilities"]["tools"], false);
17790
17791 let or_chat = model_entry_openrouter("main", Some(&caps), chat);
17793 let out_chat = &or_chat["output_modalities"][0];
17794 assert_eq!(out_chat["type"], "text");
17795 assert_eq!(out_chat["streaming"], true);
17796 assert!(!out_chat["supported_parameters"]["tools"].is_null());
17797 assert!(!out_chat["supported_parameters"]["max_tokens"].is_null());
17798 assert!(!out_chat["supported_parameters"]["structured_outputs"].is_null());
17799 assert_eq!(out_chat["max_length"]["value"], 32768u64);
17800 let om_chat = model_entry_openmodels("main", Some(&caps), chat).expect("chat entry builds");
17801 assert_eq!(om_chat["output_modalities"], json!(["text"]));
17802 assert!(
17803 om_chat["supported_features"]
17804 .as_array()
17805 .unwrap()
17806 .iter()
17807 .any(|f| f == "tool_calling")
17808 );
17809 assert_eq!(model_entry_v1("main", Some(&caps), chat)["type"], "chat");
17810 }
17811
17812 #[test]
17819 fn openrouter_output_modality_matches_the_vendored_2_4_schema() {
17820 let raw = std::fs::read_to_string(concat!(
17821 env!("CARGO_MANIFEST_DIR"),
17822 "/../../research/gateway-20260812/raw/sources/",
17823 "openrouter-provider-schema-v2.4-20260812.json"
17824 ))
17825 .expect("vendored Provider Monitor 2.4 schema is in-tree");
17826 let schema: serde_json::Value = serde_json::from_str(&raw).expect("schema parses");
17827 let branches = schema["components"]["schemas"]["OutputModality"]["oneOf"]
17828 .as_array()
17829 .expect("OutputModality is a oneOf");
17830
17831 let metadata = OpenRouterMetadataFile::from_toml(
17832 r#"
17833[models."embed"]
17834surface = "embedding"
17835created = 1787961600
17836max_output_length = 1
17837is_ready = true
17838is_free = false
17839discount_to_user = 0.0
17840
17841[models."embed".pricing]
17842prompt = "0.00000001"
17843cached_prompt = "0.0"
17844completion = "0.0"
17845
17846[models."rr"]
17847surface = "rerank"
17848created = 1787961600
17849max_output_length = 1
17850is_ready = true
17851is_free = false
17852discount_to_user = 0.0
17853
17854[models."rr".pricing]
17855prompt = "0.00000003"
17856cached_prompt = "0.0"
17857completion = "0.0"
17858
17859[models."chatty"]
17860created = 1787443200
17861max_output_length = 32768
17862is_ready = true
17863is_free = false
17864discount_to_user = 0.0
17865
17866[models."chatty".pricing]
17867prompt = "0.00000025"
17868cached_prompt = "0.00000009"
17869completion = "0.0000012"
17870"#,
17871 )
17872 .unwrap();
17873 let caps = ModelCaps {
17874 tools_branch: true,
17875 qwen_think: true,
17876 chat_ok: true,
17877 context_length: 32768,
17878 ..Default::default()
17879 };
17880
17881 for (alias, want_type) in [
17882 ("embed", "embeddings"),
17883 ("rr", "rerank"),
17884 ("chatty", "text"),
17885 ] {
17886 let row = model_entry_openrouter(alias, Some(&caps), metadata.get(alias));
17887 let modality = &row["output_modalities"][0];
17888 assert_eq!(modality["type"], want_type, "{alias}: {row}");
17889
17890 let branch = branches
17892 .iter()
17893 .find(|b| b["properties"]["type"]["enum"][0] == want_type)
17894 .unwrap_or_else(|| panic!("{want_type:?} is not an OutputModality branch"));
17895 let allowed: std::collections::BTreeSet<&str> = branch["properties"]
17896 .as_object()
17897 .expect("branch properties")
17898 .keys()
17899 .map(String::as_str)
17900 .collect();
17901 for key in modality.as_object().expect("modality object").keys() {
17902 assert!(
17903 allowed.contains(key.as_str()),
17904 "{alias}: {key:?} is not a property of the {want_type:?} branch \
17905 (additionalProperties:false); allowed = {allowed:?}"
17906 );
17907 }
17908 for req in branch["required"].as_array().into_iter().flatten() {
17909 let req = req.as_str().expect("required entry is a string");
17910 assert!(
17911 modality.get(req).is_some(),
17912 "{alias}: required property {req:?} missing from the {want_type:?} branch"
17913 );
17914 }
17915 }
17916 }
17917
17918 #[test]
17919 fn openrouter_models_entry_omits_undeclared_optional_fields() {
17920 let entry = model_entry_openrouter("minimal", None, None);
17921 let object = entry.as_object().unwrap();
17922 for field in [
17923 "hugging_face_id",
17924 "created",
17925 "quantization",
17926 "tokenizer",
17927 "description",
17928 "pricing",
17929 "capacity",
17930 "is_ready",
17931 "is_free",
17932 "discount_to_user",
17933 "openrouter",
17934 "datacenters",
17935 "compliance",
17936 ] {
17937 assert!(
17938 !object.contains_key(field),
17939 "optional field {field} must be absent, not null"
17940 );
17941 }
17942 assert_eq!(entry["schema_version"], "2.4");
17943 assert_eq!(entry["input_modalities"][0]["type"], "text");
17944 assert!(
17945 entry["input_modalities"][0]
17946 .get("supported_inputs")
17947 .is_none()
17948 );
17949 assert!(entry["input_modalities"][0].get("pricing").is_none());
17950 assert!(entry["input_modalities"][0].get("capacity").is_none());
17951 assert_eq!(entry["output_modalities"][0]["type"], "text");
17952 assert_eq!(entry["output_modalities"][0]["streaming"], true);
17953 assert!(entry["output_modalities"][0]["supported_parameters"].is_object());
17954 assert!(entry["output_modalities"][0].get("max_length").is_none());
17955 assert!(entry["output_modalities"][0].get("pricing").is_none());
17956 assert!(entry["output_modalities"][0].get("capacity").is_none());
17957 }
17958
17959 #[test]
17960 fn openmodels_entry_serializes_standard_provider_shape() {
17961 let metadata = OpenRouterMetadataFile::from_toml(
17962 r#"
17963[models."qwen/qwen3.6-27b"]
17964created = 1786032000
17965max_output_length = 16384
17966is_ready = true
17967is_free = false
17968discount_to_user = 0.05
17969
17970[models."qwen/qwen3.6-27b".pricing]
17971prompt = "0.000000291"
17972cached_prompt = "0.000000291"
17973completion = "0.000002763"
17974request = "0"
17975"#,
17976 )
17977 .unwrap();
17978 let caps = ModelCaps {
17979 tools_branch: true,
17980 qwen_think: true,
17981 chat_ok: true,
17982 context_length: 262144,
17983 ..Default::default()
17984 };
17985 let entry = model_entry_openmodels(
17986 "qwen/qwen3.6-27b",
17987 Some(&caps),
17988 metadata.get("qwen/qwen3.6-27b"),
17989 )
17990 .unwrap();
17991
17992 assert_eq!(entry["id"], "qwen/qwen3.6-27b");
17993 assert_eq!(entry["name"], "qwen/qwen3.6-27b");
17994 assert_eq!(entry["created"], 1786032000u64);
17995 assert_eq!(entry["input_modalities"], json!(["text"]));
17996 assert_eq!(entry["output_modalities"], json!(["text"]));
17997 assert_eq!(entry["context_length"], 262144u64);
17998 assert_eq!(entry["max_output_length"], 16384u64);
17999 assert_eq!(entry["currency"], "USD");
18000 assert_eq!(entry["pricing"]["prompt"], "0.000000291");
18001 assert_eq!(entry["pricing"]["completion"], "0.000002763");
18002 assert_eq!(entry["pricing"]["input_cache_read"], "0.000000291");
18003 assert_eq!(entry["pricing"]["request"], "0");
18004 assert_eq!(
18005 entry["supported_features"],
18006 json!(["tool_calling", "reasoning"])
18007 );
18008 assert_eq!(entry["is_ready"], true);
18009 assert_eq!(entry["is_free"], false);
18010 assert_eq!(entry["discount_to_user"], 0.05);
18011 assert!(entry.get("schema_version").is_none());
18012 assert!(entry.get("quantization").is_none());
18013 }
18014
18015 #[test]
18016 fn openmodels_entry_rejects_missing_operator_metadata() {
18017 let caps = ModelCaps {
18018 context_length: 262144,
18019 ..Default::default()
18020 };
18021 let error = model_entry_openmodels("qwen/qwen3.6-27b", Some(&caps), None).unwrap_err();
18022 assert_eq!(
18023 error,
18024 "OpenModels feed requires MEMRA_MODEL_METADATA for model \"qwen/qwen3.6-27b\""
18025 );
18026 }
18027
18028 #[tokio::test]
18029 async fn blocking_response_excludes_stop_text_across_token_events() {
18030 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
18031 tx.send(Event::Token {
18032 id: 1,
18033 text: "answer\nPro".into(),
18034 })
18035 .unwrap();
18036 tx.send(Event::Token {
18037 id: 2,
18038 text: "blem: leaked prompt".into(),
18039 })
18040 .unwrap();
18041 tx.send(Event::Done {
18042 stop_reason: "Callback".into(),
18043 n_tokens: 2,
18044 n_prompt: 8,
18045 n_cached: 0,
18046 elapsed_s: 0.5,
18047 spec: None,
18048 })
18049 .unwrap();
18050 drop(tx);
18051 let response = blocking_response(
18052 rx,
18053 "plain_quant".into(),
18054 false,
18055 vec!["Problem:".into()],
18056 None,
18057 Envelope::new(false),
18058 )
18059 .await;
18060 assert_eq!(response.status(), StatusCode::OK);
18061 let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
18062 .await
18063 .unwrap();
18064 let payload: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
18065 assert_eq!(payload["text"], "answer\n");
18066 assert_eq!(payload["stop_reason"], "Callback");
18067 }
18068
18069 #[test]
18073 fn step_walker_expansion_and_separator_law() {
18074 const PNG64: &str = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAY0lEQVR4nO3PQQ3AIADAQEANmlCD9IngcVnSU9DOe/b4s6UDXjWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgfeKYAYIDsx/LAAAAAElFTkSuQmCC";
18076 let uri = format!("data:image/png;base64,{PNG64}");
18077 let content = serde_json::json!([
18078 {"type": "text", "text": "look at"},
18079 {"type": "text", "text": "this:"},
18080 {"type": "image_url", "image_url": {"url": uri}},
18081 {"type": "text", "text": "what is it?"},
18082 ]);
18083 let mut pending: Vec<PendingStepImage> = Vec::new();
18084 let out = content_to_text_vision_step(&content, &mut pending).unwrap();
18085 let mut expansion = String::from("<im_start>");
18086 for _ in 0..memra_engine::vision_step::SV_MAIN_ROWS {
18087 expansion.push_str("<im_patch>");
18088 }
18089 expansion.push_str("<im_end>");
18090 assert_eq!(out, format!("look at this:{expansion}what is it?"));
18093 assert_eq!(pending.len(), 1);
18094 assert_eq!(pending[0].plan.n_tiles, 0);
18095 assert_eq!(pending[0].plan.n_prompt_tokens(), 171);
18096
18097 let vid = serde_json::json!([{ "type": "video_url", "video_url": {"url": uri} }]);
18099 assert!(content_to_text_vision_step(&vid, &mut Vec::new()).is_err());
18100 let http = serde_json::json!([
18101 {"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}
18102 ]);
18103 assert!(content_to_text_vision_step(&http, &mut Vec::new()).is_err());
18104 }
18105}