aprender-serve 0.70.2

/// #169: SafeTensors CUDA backend — format parity with GGUF GPU path.
#[cfg(feature = "cuda")]
fn try_safetensors_cuda_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    start: Instant,
    cancel: &CancelToken,
) -> Option<Response> {
    // Residency FIRST: this backend is tried before the CPU quantized one, so a
    // refusal ahead of the `?` refused `ignore_eos` for every model in a cuda
    // build, not just SafeTensors CUDA ones (aprender#3956).
    let model_lock = state.safetensors_cuda_model()?;
    // PERF-039: fail closed rather than silently dropping `ignore_eos`.
    if let Some(r) = super::openai_handlers::reject_unsupported_ignore_eos(
        state,
        request,
        "SafeTensors CUDA",
    ) {
        return Some(r);
    }
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };

    // #4007: the loaded model's architecture, not the client's `model` string.
    // #3723: the request's thinking mode; an ON the template cannot express is refused by name.
    let prompt = match crate::api::realize_handlers::format_chat_messages_for_state_thinking_tools(
        state,
        &request.messages,
        Some(&request.model),
        request.thinking(),
        request.tools.as_deref(),
    ) {
        Ok(p) => p,
        Err(e) => return Some(fail_response(state, StatusCode::BAD_REQUEST, e.to_string())),
    };
    let input_ids = tokenizer.encode(&prompt);
    let max_tokens = request.max_tokens.unwrap_or(256).min(4096) as usize;

    // Qwen2 EOS: 151645 (<|endoftext|>)
    let eos_id = 151645u32;

    let mut model = match model_lock.lock() {
        Ok(m) => m,
        Err(e) => {
            return Some(
                (
                    axum::http::StatusCode::INTERNAL_SERVER_ERROR,
                    axum::Json(serde_json::json!({"error": format!("Model lock failed: {e}")})),
                )
                    .into_response(),
            );
        }
    };

    let output_ids = match model.generate(&input_ids, max_tokens, eos_id) {
        Ok(ids) => ids,
        Err(e) => {
            let msg = format!("SafeTensors CUDA generation failed: {e}");
            return Some(
                (
                    crate::api::generation_error_status(&e),
                    axum::Json(serde_json::json!({"error": msg})),
                )
                    .into_response(),
            );
        }
    };

    let output_text = tokenizer.decode(&output_ids).unwrap_or_else(|_| String::from("[decode error]"));
    let completion_tokens = output_ids.len();
    let prompt_tokens = input_ids.len();

    // PMAT-756: this backend builds the chat JSON inline (not via build_chat_response) and
    // previously ignored request.stop entirely — the returned message kept the stop string
    // and ran to max_tokens. Reuse the shared finalize_chat_text helper so it matches every
    // other non-streaming chat backend (stop truncation + "stop" finish_reason precedence).
    let (output_text, finish_reason) =
        finalize_chat_text(output_text, request.stop.as_deref(), completion_tokens, max_tokens);

    let body = format!(
        r#"{{"id":"{}","object":"chat.completion","model":"{}","choices":[{{"index":0,"message":{{"role":"assistant","content":{}}},"finish_reason":"{}"}}],"usage":{{"prompt_tokens":{},"completion_tokens":{},"total_tokens":{}}}}}"#,
        request_id,
        request.model,
        serde_json::to_string(&output_text).unwrap_or_default(),
        finish_reason,
        prompt_tokens,
        completion_tokens,
        prompt_tokens + completion_tokens,
    );

    Some((
        [(axum::http::header::CONTENT_TYPE, "application/json")],
        body,
    ).into_response())
}

/// CUDA-optimized backend (true streaming support).
#[cfg(feature = "cuda")]
async fn try_cuda_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    trace_level: Option<&str>,
    start: Instant,
    cancel: &CancelToken,
) -> Option<Response> {
    // PMAT-821: config now built by chat_quantized_config (in openai_handlers.rs).
    let ttft_trace = std::env::var("TTFT_TRACE").is_ok();
    let t0 = if ttft_trace { Some(std::time::Instant::now()) } else { None };

    let cuda_model_lock = state.cuda_model()?;
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };
    // GH-319: Use actual model architecture for chat template detection
    let arch_hint = state.model_architecture();
    // D5: refused before any path (streaming, batch scheduler, direct) takes the
    // model, so the client gets a 400 rather than a 500 or a mid-stream error.
    let tokenized =
        tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), request.tools.as_deref(), state);
    let prompt_ids =
        match fit_serving_context(state, tokenized) {
            Ok(ids) => ids,
            Err(r) => return Some(r),
        };
    if let Some(t) = t0 {
        eprintln!("[TTFT] {:>20}: {:>7.2}ms ({}tok)", "tokenize", t.elapsed().as_secs_f64() * 1000.0, prompt_ids.len());
    }
    let prompt_tokens = prompt_ids.len();
    // PMAT-821: build the config via chat_quantized_config so ALL request sampling
    // params (top_p/repeat_penalty/repeat_last_n/seed) reach the sampler — the prior
    // inline builder dropped them, leaving the chat endpoint on neutral defaults.
    let q_config = chat_quantized_config(
        request,
        &tokenizer,
        state.model_eos_token_id(),
        state.should_trace(trace_level),
        cancel,
    );
    let max_tokens = q_config.max_tokens;

    if request.stream {
        let (tx, rx) = tokio::sync::mpsc::channel::<Result<u32, String>>(16);
        // §3: the engine's return path for the server-measured phase split.
        // The handler returns the SSE response before generation finishes, so
        // the measurement cannot be a return value; the terminal chunk awaits
        // this oneshot after the token channel closes.
        let (timing_tx, timing_rx) =
            tokio::sync::oneshot::channel::<crate::api::PhaseTimings>();

        if let Err(r) =
            dispatch_cuda_stream(state, cuda_model_lock, prompt_ids, q_config, tx, timing_tx)
        {
            return Some(r);
        }

        return Some(true_streaming_sse_response(
            rx,
            tokenizer,
            request_id.to_string(),
            request.model.clone(),
            state.metrics.clone(),
            start,
            max_tokens,
            prompt_tokens,
            Some(timing_rx),
            request.stop.as_deref(),
        ));
    }

    // Non-streaming CUDA — route through batch scheduler when available (realizr#211)
    let (timing_tx, timing_rx) = tokio::sync::oneshot::channel::<crate::api::PhaseTimings>();
    let generated = if let Some(batch_tx) = state.cuda_batch_tx() {
        cuda_batch_collect(state, batch_tx, prompt_ids, q_config, timing_tx, &tokenizer).await
    } else {
        cuda_direct_collect(state, cuda_model_lock, &prompt_ids, &q_config, timing_tx, &tokenizer, prompt_tokens)
    };
    let (token_ids, completion_tokens, response_text) = match generated {
        Ok(g) => g,
        Err(r) => return Some(r),
    };

    let latency = start.elapsed();
    state.metrics.record_success(completion_tokens, latency);
    // §3: absent unless the engine measured BOTH phases for THIS request.
    let timings = timing_rx
        .await
        .ok()
        .and_then(|phases| phases.to_timings(prompt_tokens, completion_tokens));
    Some(build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        response_text,
        prompt_tokens,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        trace_level,
        latency,
        request.tools.as_deref(),
        request_tool_choice(request),
        timings,
        None,
    ))
}

/// `try_cuda_backend`'s streaming dispatch: the batch scheduler when one is
/// running, else the direct RwLock path on a blocking task. `Err` is the 503
/// for a full batch queue.
#[cfg(feature = "cuda")]
#[allow(clippy::result_large_err)]
fn dispatch_cuda_stream(
    state: &AppState,
    cuda_model_lock: &Arc<std::sync::RwLock<crate::gguf::OwnedQuantizedModelCuda>>,
    prompt_ids: Vec<u32>,
    q_config: crate::gguf::QuantizedGenerateConfig,
    tx: tokio::sync::mpsc::Sender<Result<u32, String>>,
    timing_tx: tokio::sync::oneshot::Sender<crate::api::PhaseTimings>,
) -> Result<(), Response> {
    // PMAT-044: Use batch scheduler if available (continuous batching)
    if let Some(batch_tx) = state.cuda_batch_tx() {
        let batch_req = super::cuda_batch_scheduler::CudaBatchRequest {
            prompt_ids,
            config: q_config,
            token_tx: tx,
            non_streaming: false,
            enqueue_time: std::time::Instant::now(),
            timing_tx: Some(timing_tx),
        };
        if let Err(e) = batch_tx.try_send(batch_req) {
            // §5.2: this 503 is the one admission REFUSAL this server has
            // (the policy is otherwise `queue`), so it is counted where it
            // is returned — `kv.admission_rejected` reads this counter.
            state.record_admission_rejected();
            return Err(fail_response(
                state,
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Batch queue full: {e}"),
            ));
        }
    } else {
        // Fallback: direct RwLock path (serialized)
        let cuda_model_clone = cuda_model_lock.clone();
        let sink_metrics = state.metrics.clone();

        tokio::task::spawn_blocking(move || {
            let mut cuda_model = cuda_model_clone.write().expect("operation failed");
            let generate_start = std::time::Instant::now();
            // Stops when the client goes away — see `streaming_token_sink`.
            let sink =
                crate::api::openai_handlers::streaming_token_sink(tx.clone(), sink_metrics);
            let result = dense_cuda_turn(&mut cuda_model, &prompt_ids, &q_config, sink);
            // Taken under the SAME write lock the request ran under, so the
            // split belongs to this request and to no other.
            let _ = timing_tx.send(phase_split(&mut cuda_model, generate_start));
            if let Err(e) = result {
                let _ = tx.blocking_send(Err(e.to_string()));
            }
        });
    }
    Ok(())
}

/// `try_cuda_backend`'s non-streaming turn through the batch scheduler
/// (realizr#211): submit, then collect every token. Returns
/// `(tokens, completion_tokens, cleaned text)`.
#[cfg(feature = "cuda")]
async fn cuda_batch_collect(
    state: &AppState,
    batch_tx: &tokio::sync::mpsc::Sender<super::cuda_batch_scheduler::CudaBatchRequest>,
    prompt_ids: Vec<u32>,
    q_config: crate::gguf::QuantizedGenerateConfig,
    timing_tx: tokio::sync::oneshot::Sender<crate::api::PhaseTimings>,
    tokenizer: &BPETokenizer,
) -> Result<(Vec<u32>, usize, String), Response> {
    // Use batch scheduler: submit request and collect all tokens
    // realizr#212: capacity 512 for bulk-send after non-streaming generation
    let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<u32, String>>(512);
    let batch_req = super::cuda_batch_scheduler::CudaBatchRequest {
        prompt_ids,
        config: q_config,
        token_tx: tx,
        non_streaming: true, // realizr#212: scheduler accumulates + bulk-sends
        enqueue_time: std::time::Instant::now(),
        timing_tx: Some(timing_tx),
    };
    if let Err(e) = batch_tx.try_send(batch_req) {
        // §5.2: counted at the refusal, as in `dispatch_cuda_stream`.
        state.record_admission_rejected();
        return Err(fail_response(
            state,
            StatusCode::SERVICE_UNAVAILABLE,
            format!("Batch queue full: {e}"),
        ));
    }
    // Collect all tokens via async receive (realizr#211)
    let mut tokens = Vec::new();
    while let Some(result) = rx.recv().await {
        match result {
            Ok(token_id) => tokens.push(token_id),
            // D5: `fit_serving_context` already refused an over-context
            // prompt with a 400; the scheduler's errors are strings, and what
            // reaches here is a server fault.
            Err(e) => return Err(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)),
        }
    }
    let n = tokens.len();
    let text = tokenizer.decode(&tokens).unwrap_or_else(|_| String::new());
    Ok((tokens, n, clean_chat_output(&text)))
}

/// `try_cuda_backend`'s non-streaming turn when no batch scheduler runs: the
/// direct RwLock path (serialized). Same return shape as `cuda_batch_collect`.
#[cfg(feature = "cuda")]
#[allow(clippy::result_large_err)]
fn cuda_direct_collect(
    state: &AppState,
    cuda_model_lock: &Arc<std::sync::RwLock<crate::gguf::OwnedQuantizedModelCuda>>,
    prompt_ids: &[u32],
    q_config: &crate::gguf::QuantizedGenerateConfig,
    timing_tx: tokio::sync::oneshot::Sender<crate::api::PhaseTimings>,
    tokenizer: &BPETokenizer,
    prompt_tokens: usize,
) -> Result<(Vec<u32>, usize, String), Response> {
    let mut cuda_model = cuda_model_lock.write().expect("operation failed");
    let generate_start = std::time::Instant::now();
    let generated = match dense_cuda_turn(&mut cuda_model, prompt_ids, q_config, |_| true) {
        Ok(g) => g,
        Err(e) => return Err(fail_response(state, crate::api::generation_error_status(&e), e)),
    };
    let _ = timing_tx.send(phase_split(&mut cuda_model, generate_start));
    let tokens: Vec<u32> = generated.iter().skip(prompt_tokens).copied().collect();
    let n = tokens.len();
    let text = tokenizer.decode(&tokens).unwrap_or_else(|_| String::new());
    Ok((tokens, n, clean_chat_output(&text)))
}

/// §3: pair the engine's prefill measurement with the decode remainder.
///
/// `take_phase_timings` clears as it reads, so a request that took a
/// prefix-cache hit and never ran prefill gets `None` rather than the previous
/// request's number. `decode_ms` is left `None` in that case too — one measured
/// phase is not a phase split, and `PhaseTimings::to_timings` refuses to build
/// a wire block from it.
/// #4268: one serve request's dense CUDA turn, run on the one engine over the
/// model borrowed under the request's write lock. A GPU failure is an error,
/// as it always was here: serve has no CPU copy to fall back to. `--trace`
/// keeps the instrumented loop. Stale phase timings are cleared first so
/// [`phase_split`] reads this request's.
#[cfg(feature = "cuda")]
fn dense_cuda_turn(
    cuda_model: &mut crate::gguf::OwnedQuantizedModelCuda,
    prompt: &[u32],
    config: &crate::gguf::QuantizedGenerateConfig,
    mut on_token: impl FnMut(u32) -> bool,
) -> crate::error::Result<Vec<u32>> {
    if config.trace {
        return cuda_model.generate_gpu_resident_streaming(prompt, config, on_token);
    }
    let _ = cuda_model.take_phase_timings();
    let mut session = crate::session::Session::new(
        crate::gguf::dense_session_borrowed::BorrowedCudaForward::new(cuda_model),
    );
    crate::gguf::dense_session::dense_stream(&mut session, prompt, config, &mut on_token)
        .map(|(tokens, _)| tokens)
}

#[cfg(feature = "cuda")]
fn phase_split(
    cuda_model: &mut crate::gguf::OwnedQuantizedModelCuda,
    generate_start: std::time::Instant,
) -> crate::api::PhaseTimings {
    let mut phases = cuda_model.take_phase_timings();
    if let Some(prefill_ms) = phases.prefill_ms {
        let total_ms = generate_start.elapsed().as_secs_f64() * 1000.0;
        phases.decode_ms = Some((total_ms - prefill_ms).max(0.0));
    }
    phases
}

/// #4268: one request's dense CPU session over the shared model. Its KV cache
/// lives for the request, as `generate_with_cache`'s did.
fn dense_cpu_session(
    model: &Arc<crate::gguf::OwnedQuantizedModel>,
) -> crate::gguf::dense_session::DenseSession {
    crate::gguf::dense_session::DenseSession::new(
        crate::gguf::dense_session::DenseForward::cpu(Arc::clone(model)),
    )
}

/// Quantized model (GGUF serve mode) backend with true streaming.
fn try_quantized_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    trace_level: Option<&str>,
    start: Instant,
    cancel: &CancelToken,
) -> Option<Response> {
    // PMAT-821: config now built by chat_quantized_config (in openai_handlers.rs).
    let quantized_model = state.quantized_model()?;
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };
    // GH-319: Use actual model architecture for chat template detection
    let arch_hint = state.model_architecture();
    let prompt_ids =
        match tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), request.tools.as_deref(), state) {
            Ok(ids) => ids,
            Err(r) => return Some(r),
        };
    let prompt_tokens = prompt_ids.len();
    // PMAT-821: build the config via chat_quantized_config so ALL request sampling
    // params (top_p/repeat_penalty/repeat_last_n/seed) reach the sampler — the prior
    // inline builder dropped them, leaving the chat endpoint on neutral defaults.
    let q_config = chat_quantized_config(
        request,
        &tokenizer,
        state.model_eos_token_id(),
        state.should_trace(trace_level),
        cancel,
    );
    // #3718: the dense CPU loop spends `effective_max_tokens` (the request clamped
    // to the context room), so the finish is judged against THAT budget. Judged
    // against the request's number, every context-clamped cut read as "stop". A
    // prompt that cannot fit at all is the request's fault: refused here with 400,
    // not a 500 after the fact or a streamed error that still ends "stop" at 200.
    let max_tokens = match quantized_model.effective_max_tokens(prompt_tokens, q_config.max_tokens)
    {
        Ok(budget) => budget,
        Err(e) => return Some(fail_response(state, crate::api::generation_error_status(&e), e)),
    };

    if request.stream {
        let (tx, rx) = tokio::sync::mpsc::channel::<Result<u32, String>>(16);
        let quantized_model_clone = quantized_model.clone();
        let prompt_ids_clone = prompt_ids.clone();
        let q_config_clone = q_config.clone();
        let sink_metrics = state.metrics.clone();

        tokio::task::spawn_blocking(move || {
            // Stops when the client goes away — see `streaming_token_sink`.
            let mut sink =
                crate::api::openai_handlers::streaming_token_sink(tx.clone(), sink_metrics);
            let result = if q_config_clone.trace {
                quantized_model_clone
                    .generate_with_cache_streaming(&prompt_ids_clone, &q_config_clone, sink)
                    .map(drop)
            } else {
                // #4268: the dense CPU turn runs on the one engine.
                crate::gguf::dense_session::dense_stream(
                    &mut dense_cpu_session(&quantized_model_clone),
                    &prompt_ids_clone,
                    &q_config_clone,
                    &mut sink,
                )
                .map(drop)
            };
            if let Err(e) = result {
                let _ = tx.blocking_send(Err(e.to_string()));
            }
        });

        return Some(true_streaming_sse_response(
            rx,
            tokenizer,
            request_id.to_string(),
            request.model.clone(),
            state.metrics.clone(),
            start,
            max_tokens,
            prompt_tokens,
            // The CPU quantized decode loop does not separate prefill from
            // decode, so §3 timings are absent rather than zero.
            None,
            request.stop.as_deref(),
        ));
    }

    // Non-streaming quantized
    // #4268: the dense CPU turn runs on the one engine; `--trace` keeps the
    // instrumented loop.
    let generated = if q_config.trace {
        quantized_model.generate_with_cache(&prompt_ids, &q_config)
    } else {
        crate::gguf::dense_session::dense_turn(
            &mut dense_cpu_session(quantized_model),
            &prompt_ids,
            &q_config,
        )
        .map(|(tokens, _)| tokens)
    };
    let generated = match generated {
        Ok(g) => g,
        Err(e) => return Some(fail_response(state, crate::api::generation_error_status(&e), e)),
    };

    let token_ids: Vec<u32> = generated.iter().skip(prompt_tokens).copied().collect();
    let completion_tokens = token_ids.len();
    let text = match tokenizer.decode(&token_ids) {
        Ok(t) => clean_chat_output(&t),
        Err(e) => return Some(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)),
    };

    let latency = start.elapsed();
    state.metrics.record_success(completion_tokens, latency);
    Some(build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        text,
        prompt_tokens,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        trace_level,
        latency,
        request.tools.as_deref(),
        request_tool_choice(request),
        None,
        None,
    ))
}

/// `AprTransformer` (f32 APR / SafeTensors CPU) backend for `/v1/chat/completions`
/// and `/v1/chat/completions/stream`.
///
/// aprender#2609: the chat backend chain ran CUDA → cached → quantized →
/// `registry_fallback`, and `registry_fallback` resolves the dense f32
/// [`Model`](crate::layers::Model), which is `None` on an `AprTransformer` server.
/// So both chat routes — one of them named in the server's own startup banner —
/// were dead there while `/generate` and `/batch/generate` on the same process
/// answered 200 with real text, because only those two had grown
/// `try_apr_generate` / `try_apr_batch_generate`.
///
/// Streaming reuses [`pregenerated_sse_response`], exactly as `registry_fallback`
/// does, so `stream: true` gets the same wire format as every other backend.
///
/// Returns `None` when no `AprTransformer` is resident, leaving the dense
/// fallback unchanged.
fn try_apr_transformer_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    trace_level: Option<&str>,
    start: Instant,
    cancel: &CancelToken,
) -> Option<Response> {
    use crate::apr_transformer::GenerateConfig;

    // Residency first, as in `try_safetensors_cuda_backend` (aprender#3956).
    let apr_transformer = state.apr_transformer()?;
    // PERF-039: fail closed rather than silently dropping `ignore_eos`.
    if let Some(r) = super::openai_handlers::reject_unsupported_ignore_eos(
        state,
        request,
        "APR transformer (f32)",
    ) {
        return Some(r);
    }
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };
    let arch_hint = state.model_architecture();
    let prompt_ids =
        match tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), request.tools.as_deref(), state) {
            Ok(ids) => ids,
            Err(r) => return Some(r),
        };
    let prompt_tokens = prompt_ids.len();
    let max_tokens = request.max_tokens.unwrap_or(256);

    let gen_config = GenerateConfig {
        max_tokens,
        temperature: request.temperature.unwrap_or(0.7),
        cancel: cancel.clone(),
        ..Default::default()
    };

    let generated = match apr_transformer.generate_with_cache(&prompt_ids, &gen_config) {
        Ok(g) => g,
        Err(e) => {
            return Some(fail_response(
                state,
                super::generation_error_status(&e),
                format!("APR generation failed: {e}"),
            ))
        },
    };

    let token_ids: Vec<u32> = generated.iter().skip(prompt_tokens).copied().collect();
    let completion_tokens = token_ids.len();

    if request.stream {
        state
            .metrics
            .record_success(completion_tokens, start.elapsed());
        return Some(pregenerated_sse_response(
            token_ids,
            tokenizer,
            request_id.to_string(),
            request.model.clone(),
            request.stop.as_deref(),
            max_tokens,
            prompt_tokens,
        ));
    }

    let text = match tokenizer.decode(&token_ids) {
        Ok(t) => clean_chat_output(&t),
        Err(e) => return Some(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)),
    };

    let latency = start.elapsed();
    state.metrics.record_success(completion_tokens, latency);
    Some(build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        text,
        prompt_tokens,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        trace_level,
        latency,
        request.tools.as_deref(),
        request_tool_choice(request),
        None,
        None,
    ))
}

/// Convert usize token IDs to u32, returning error string on overflow
fn convert_token_ids(ids: &[usize]) -> Result<Vec<u32>, String> {
    ids.iter()
        .map(|&id| u32::try_from(id).map_err(|_| format!("Token ID {id} exceeds u32 range")))
        .collect()
}

/// Build generation config from request parameters.
///
/// #2375: `temperature: 0` — the OpenAI-canonical deterministic request —
/// reached `apply_temperature` unchanged here and made this backend answer
/// HTTP 500 ("Temperature must be a positive finite number") for every dense
/// model. The resolution now lives in ONE place, shared with `/v1/completions`.
fn build_gen_config(request: &ChatCompletionRequest) -> GenerationConfig {
    crate::api::realize_handlers::resolve_dense_generation_config(
        request.temperature.unwrap_or(0.7),
        request.top_p,
        request.max_tokens.unwrap_or(256),
    )
}

/// The registry chat prompt, rendered with the request's thinking mode and tokenized; an ON the
/// template cannot express, or an empty prompt, is the client's error (#3723).
/// Extracted from `registry_fallback` (complexity ratchet, #4046); behaviour unchanged.
#[allow(clippy::result_large_err)]
fn registry_prompt_ids(
    state: &AppState,
    request: &ChatCompletionRequest,
    tokenizer: &BPETokenizer,
) -> Result<Vec<u32>, Response> {
    let prompt_text = match crate::api::realize_handlers::format_chat_messages_for_state_thinking_tools(
        state,
        &request.messages,
        Some(&request.model),
        request.thinking(),
        request.tools.as_deref(),
    ) {
        Ok(p) => p,
        Err(e) => return Err(fail_response(state, StatusCode::BAD_REQUEST, e.to_string())),
    };
    let prompt_ids = tokenizer.encode(&prompt_text);
    if prompt_ids.is_empty() {
        return Err(fail_response(state, StatusCode::BAD_REQUEST, "Messages cannot be empty"));
    }
    Ok(prompt_ids)
}

/// A registry model's generated ids as u32, or the failure response.
/// Extracted from `registry_fallback` (complexity ratchet, #4046); behaviour unchanged.
#[allow(clippy::result_large_err)]
fn registry_token_ids<E: std::fmt::Display>(
    state: &AppState,
    generated: Result<Vec<usize>, E>,
) -> Result<Vec<u32>, Response> {
    let generated = match generated {
        Ok(g) => g,
        Err(e) => return Err(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)),
    };
    convert_token_ids(&generated).map_err(|e| fail_response(state, StatusCode::BAD_REQUEST, e))
}

/// Registry-based model fallback (no specialized backend).
fn registry_fallback(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    start: Instant,
    cancel: &CancelToken,
) -> Response {
    // PERF-039: fail closed rather than silently dropping `ignore_eos`.
    if let Some(r) = super::openai_handlers::reject_unsupported_ignore_eos(
        state,
        request,
        "dense registry",
    ) {
        return r;
    }
    let model_id = if request.model == "default" || request.model.is_empty() {
        None
    } else {
        Some(request.model.as_str())
    };

    let (model, tokenizer) = match state.get_model(model_id) {
        Ok((m, t)) => (m, t),
        // #2375(4): a hardcoded 404 told clients this mounted route did not
        // exist whenever the server simply had no model resident. That is a
        // server-side condition — 503, the status `/v1/predict` and
        // `/v1/gpu/warmup` already answer for it. An unknown model NAME stays
        // 404. `model_resolution_status` is the single rule (aprender#2376(5)).
        // aprender#2609 reached the identical conclusion independently: the
        // same `RegistryError("No model available")` answered 503 on
        // /generate, /stream/generate and /batch/generate but 404 here.
        Err(e) => return fail_response(state, super::model_resolution_status(&e), e),
    };

    let prompt_ids = match registry_prompt_ids(state, request, &tokenizer) {
        Ok(ids) => ids,
        Err(r) => return r,
    };

    let prompt_tokens = prompt_ids.len();
    let prompt: Vec<usize> = prompt_ids.iter().map(|&id| id as usize).collect();
    let config = build_gen_config(request).with_cancel(cancel.clone());

    let token_ids: Vec<u32> = match registry_token_ids(state, model.generate(&prompt, &config)) {
        Ok(ids) => ids,
        Err(r) => return r,
    };

    let generated_ids: Vec<u32> = token_ids[prompt.len()..].to_vec();
    let completion_tokens = generated_ids.len();

    if request.stream {
        state
            .metrics
            .record_success(completion_tokens, start.elapsed());
        return pregenerated_sse_response(
            generated_ids,
            tokenizer,
            request_id.to_string(),
            request.model.clone(),
            request.stop.as_deref(),
            request.max_tokens.unwrap_or(256),
            prompt_tokens,
        );
    }

    let response_text = match tokenizer.decode(&generated_ids) {
        Ok(t) => t,
        Err(e) => return fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e),
    };

    let duration = start.elapsed();
    state.metrics.record_success(completion_tokens, duration);

    let max_tokens = request.max_tokens.unwrap_or(256);
    build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        response_text,
        prompt_tokens,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        None,
        duration,
        request.tools.as_deref(),
        request_tool_choice(request),
        None,
        None,
    )
}

// ============================================================================
// Handlers
// ============================================================================

/// Model-load timestamp (Unix seconds).
///
/// CRUX-C-33 §created_timestamp_domain: `created` must represent model load
/// time — it MUST be stable across requests (not `SystemTime::now()` at each
/// call). Discharges FALSIFY-CRUX-C-33-004.
///
/// PP-30: it now reads the ONE process clock, latched at the first `AppState`
/// construction (i.e. when the model finished loading) rather than at the first
/// CHAT REQUEST, which is what the old private `OnceLock<i64>` latched on while
/// its own doc called itself "model load time". Same stability guarantee, and
/// the same instant `/health`'s `uptime_sec` and `/v1/effective-config`'s
/// `started_utc` report.
fn model_loaded_at_unix_secs(state: &AppState) -> i64 {
    state.clock().started_unix_secs()
}

/// OpenAI-compatible models listing handler
///
/// Returns available models in OpenAI API format (GET /v1/models).
/// Contract: `contracts/crux-C-33-v1.yaml` — envelope `{object:"list",data:[...]}`;
/// per-model `{id, object:"model", created>0, owned_by}`; `created` stable
/// across requests (model-load time, not request time).
pub async fn openai_models_handler(State(state): State<AppState>) -> Json<OpenAIModelsResponse> {
    let created = model_loaded_at_unix_secs(&state);
    let models = if let Some(registry) = &state.registry {
        registry
            .list()
            .into_iter()
            .map(|m| OpenAIModel {
                id: m.id,
                object: "model".to_string(),
                created,
                owned_by: "realizar".to_string(),
            })
            .collect()
    } else {
        // Single model mode
        vec![OpenAIModel {
            id: "default".to_string(),
            object: "model".to_string(),
            created,
            owned_by: "realizar".to_string(),
        }]
    };

    Json(OpenAIModelsResponse {
        object: "list".to_string(),
        data: models,
    })
}

/// ALB-110: APR Q4K GPU chat backend via dedicated inference thread.
///
/// Mirrors `try_apr_q4k_completions` from gpu_completions_handler.rs but for
/// chat format. Tokenizes chat messages, sends to Q4K scheduler, returns
/// chat-formatted response.
#[cfg(feature = "cuda")]
async fn try_apr_q4k_chat_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    trace_level: Option<&str>,
    start: Instant,
    cancel: &crate::generate::CancelToken,
) -> Option<Response> {
    // aprender#2465(1): this backend hands the work to the Q4K scheduler THREAD,
    // and handing a loop to another thread does not stop it. The dropped response
    // future cannot reach that thread, and the scheduler returns one accumulated
    // AprQ4kResponse so there is no per-token send left to fail either — the
    // request carries the token and the decode loop polls it. The previous comment
    // here claimed the hand-off was itself the answer; it was not.
    use crate::api::apr_q4k_scheduler::AprQ4kRequest;

    let q4k_tx = state.apr_q4k_tx()?;
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };
    let arch_hint = state.model_architecture();
    let prompt_ids =
        match tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), request.tools.as_deref(), state) {
            Ok(ids) => ids,
            Err(r) => return Some(r),
        };
    let prompt_tokens = prompt_ids.len();
    let (max_tokens, temperature, _eos_single) =
        chat_gen_params(request, &tokenizer, state.model_eos_token_id());
    // PERF-039: ignore_eos empties the stop set for this backend too.
    let eos_ids = if request.ignore_eos.unwrap_or(false) {
        Vec::new()
    } else {
        state.model_eos_ids()
    };

    let (response_tx, response_rx) = tokio::sync::oneshot::channel();

    if q4k_tx
        .send(AprQ4kRequest {
            prompt_ids,
            max_tokens,
            temperature,
            // #3786: the request seed reaches the APR Q4K sampler.
            seed: request.seed.unwrap_or(crate::sampling::DEFAULT_SEED),
            eos_ids,
            cancel: cancel.clone(),
            response_tx,
        })
        .await
        .is_err()
    {
        return Some(fail_response(
            state,
            StatusCode::INTERNAL_SERVER_ERROR,
            "Q4K thread unavailable",
        ));
    }

    let result = match response_rx.await {
        Ok(r) => r,
        Err(_) => {
            return Some(fail_response(
                state,
                StatusCode::INTERNAL_SERVER_ERROR,
                "Q4K thread dropped response",
            ))
        }
    };

    let resp = match result {
        Ok(r) => r,
        Err(e) => {
            return Some(fail_response(
                state,
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Q4K generation failed: {e}"),
            ))
        }
    };

    let text = match tokenizer.decode(&resp.output_tokens) {
        Ok(t) => clean_chat_output(&t),
        Err(e) => {
            return Some(fail_response(
                state,
                StatusCode::INTERNAL_SERVER_ERROR,
                e,
            ))
        }
    };
    let completion_tokens = resp.tokens_generated;
    state
        .metrics
        .record_success(completion_tokens, start.elapsed());

    Some(build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        text,
        prompt_tokens,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        trace_level,
        start.elapsed(),
        request.tools.as_deref(),
        request_tool_choice(request),
        None,
        None,
    ))
}

/// OpenAI-compatible /v1/chat/completions endpoint (supports streaming)
pub async fn openai_chat_completions_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
    Extension(cancel): Extension<CancelToken>,
    Json(request): Json<ChatCompletionRequest>,
) -> Response {
    let start = Instant::now();
    // GH-152: Verbose request logging
    if state.is_verbose() {
        let msg_count = request.messages.len();
        let last_msg = request
            .messages
            .last()
            .map(|m| m.content.chars().take(50).collect::<String>())
            .unwrap_or_default();
        eprintln!(
            "[VERBOSE] POST /v1/chat/completions model={} messages={} last={:?}",
            request.model, msg_count, last_msg
        );
    }

    let trace_level = headers
        .get("X-Trace-Level")
        .and_then(|v| v.to_str().ok())
        .map(str::to_lowercase);

    let request_id = format!(
        "chatcmpl-q4k-{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis()
    );

    // #3723: two spellings of the thinking toggle that disagree are refused, never picked between.
    if let Some(reason) = request.thinking_conflict() {
        return fail_response(&state, StatusCode::BAD_REQUEST, reason);
    }

    // #3571: a Qwen3.5 hybrid is answered from its resident session or not at all.
    if let Some(r) = try_qwen35_backend(&state, &request, &request_id, start, &cancel).await {
        return r;
    }

    if let Some(r) = try_qwen3_moe_backend(&state, &request, &request_id, start, &cancel) {
        return r;
    }

    #[cfg(feature = "gpu")]
    if let Some(r) = try_gpu_backend(
        &state,
        &request,
        &request_id,
        trace_level.as_deref(),
        start,
        &cancel,
    ) {
        return r;
    }

    #[cfg(feature = "gpu")]
    if let Some(r) = try_cached_backend(
        &state,
        &request,
        &request_id,
        trace_level.as_deref(),
        start,
        &cancel,
    ) {
        return r;
    }

    #[cfg(feature = "cuda")]
    if let Some(r) =
        try_cuda_backend(&state, &request, &request_id, trace_level.as_deref(), start, &cancel).await
    {
        return r;
    }

    // ALB-110: APR Q4K GPU backend via dedicated inference thread
    #[cfg(feature = "cuda")]
    if let Some(r) = try_apr_q4k_chat_backend(&state, &request, &request_id, trace_level.as_deref(), start, &cancel).await {
        return r;
    }

    // #169: SafeTensors CUDA backend (format parity)
    #[cfg(feature = "cuda")]
    if let Some(r) = try_safetensors_cuda_backend(&state, &request, &request_id, start, &cancel) {
        return r;
    }

    cpu_chat_backends(
        &state,
        &request,
        &request_id,
        trace_level.as_deref(),
        start,
        &cancel,
    )
}

/// The CPU tail of the chat backend chain: quantized, then the f32
/// `AprTransformer`, then the dense registry fallback.
///
/// Lifted out of [`openai_chat_completions_handler`] when aprender#2609 added the
/// APR arm — the handler is a dispatch table and every backend added to it costs
/// the same two branches, so the always-compiled tail lives here instead. The
/// order mirrors what `/generate` walks (`try_quantized_generate`,
/// `try_apr_generate`, `registry_generate`), which is the point: the two routes
/// must be alive on exactly the same set of resident models.
fn cpu_chat_backends(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    trace_level: Option<&str>,
    start: Instant,
    cancel: &CancelToken,
) -> Response {
    if let Some(r) =
        try_quantized_backend(state, request, request_id, trace_level, start, cancel)
    {
        return r;
    }
    if let Some(r) =
        try_apr_transformer_backend(state, request, request_id, trace_level, start, cancel)
    {
        return r;
    }
    registry_fallback(state, request, request_id, start, cancel)
}

/// PERF-039: `ignore_eos: true` empties the stop set, which every decode loop
/// reads as "never stop on a token". `max_tokens` still bounds the loop.
///
/// Extracted rather than written inline so that merging PERF-039 does not move
/// `try_qwen3_moe_backend`'s cognitive complexity (26 on main, already over the
/// pre-commit hook's threshold of 25) even by one.
fn stop_tokens_unless_ignore_eos(
    request: &ChatCompletionRequest,
    eos: impl IntoIterator<Item = u32>,
) -> Vec<u32> {
    if request.ignore_eos.unwrap_or(false) {
        Vec::new()
    } else {
        eos.into_iter().collect()
    }
}

/// aprender#1789 Option B: qwen3_moe MoE-aware dispatch for /v1/chat/completions.
///
/// Detects qwen3_moe architecture + dispatches inference through
/// `run_qwen3_moe_generate` (the same path used by the `apr run` CLI),
/// which correctly indexes per-expert FFN tensors from the mmap.
///
/// For non-qwen3_moe archs returns `None` — handler falls through to the
/// dense backend chain (CUDA / cached / quantized / registry-fallback).
///
/// For qwen3_moe archs where AppState was constructed WITHOUT
/// `with_mapped_gguf_model` (no retained mmap), returns NOT_IMPLEMENTED
/// with the same actionable error class Option A surfaced. The
/// defensive guard from aprender#1790's `validate_matmul_weight_shape`
/// will NOT fire because we never reach the dense FFN matmul.
///
/// Discharges FALSIFY-QWEN3_MOE_SERVE_DISPATCH_V1_001 + V1_003 in
/// `contracts/qwen3-moe-serve-dispatch-v1.yaml`.
/// The retained mapped GGUF and the quantized model a qwen3_moe chat needs, or the refusal (#1789).
/// Extracted from `try_qwen3_moe_backend` (complexity ratchet, #4046); behaviour unchanged.
#[allow(clippy::result_large_err)]
fn moe_models(
    state: &AppState,
    raw_arch: &str,
) -> Result<(Arc<crate::gguf::MappedGGUFModel>, Arc<crate::gguf::OwnedQuantizedModel>), Response> {
    let mapped = match state.mapped_gguf_model() {
        Some(m) => m,
        None => {
            eprintln!(
                "[WARN] aprender#1789: qwen3_moe arch detected at \
                 /v1/chat/completions (raw_arch={raw_arch}, canonical=qwen3_moe) \
                 but AppState has no retained MappedGGUFModel. This means the \
                 CLI server-command load path didn't call \
                 .with_mapped_gguf_model(). Returning NOT_IMPLEMENTED. \
                 See contracts/qwen3-moe-serve-dispatch-v1.yaml + \
                 https://github.com/paiml/aprender/issues/1789"
            );
            return Err(fail_response(
                state,
                StatusCode::NOT_IMPLEMENTED,
                "qwen3_moe arch detected but mapped GGUF not retained in AppState. \
                 See aprender#1789 + contracts/qwen3-moe-serve-dispatch-v1.yaml.",
            ));
        }
    };
    let quantized = match state.quantized_model() {
        Some(q) => q.clone(),
        None => {
            return Err(fail_response(
                state,
                StatusCode::NOT_IMPLEMENTED,
                "qwen3_moe arch detected but no OwnedQuantizedModel in AppState. \
                 See aprender#1789.",
            ));
        }
    };
    Ok((mapped, quantized))
}

/// The qwen3_moe generation config from the request (3-knob toolkit + EOS fallback chain).
/// Extracted from `try_qwen3_moe_backend` (complexity ratchet, #4046); behaviour unchanged.
fn moe_gen_config(
    state: &AppState,
    request: &ChatCompletionRequest,
    tokenizer: &BPETokenizer,
    max_tokens: usize,
    cancel: &CancelToken,
) -> crate::gguf::QuantizedGenerateConfig {
    use crate::gguf::QuantizedGenerateConfig;
    let defaults = QuantizedGenerateConfig::default();
    let eos_id = state.model_eos_token_id().or_else(|| {
        tokenizer
            .get_token_id("<|im_end|>")
            .or_else(|| tokenizer.get_token_id("<|endoftext|>"))
    });
    let stop_tokens: Vec<u32> = stop_tokens_unless_ignore_eos(request, eos_id);
    QuantizedGenerateConfig {
        max_tokens,
        temperature: request.temperature.unwrap_or(defaults.temperature),
        top_k: request.top_k.unwrap_or(defaults.top_k),
        top_p: request.top_p.unwrap_or(defaults.top_p),
        repeat_penalty: request.repeat_penalty.unwrap_or(defaults.repeat_penalty),
        repeat_last_n: request.repeat_last_n.unwrap_or(defaults.repeat_last_n),
        seed: request.seed.unwrap_or(defaults.seed),
        stop_tokens,
        cancel: cancel.clone(),
        ..defaults
    }
}

/// The CPU (`moe_no_gpu`) per-token SSE stream for a qwen3_moe chat.
/// Extracted from `try_qwen3_moe_backend` (complexity ratchet, #4046); behaviour unchanged.
#[allow(clippy::too_many_arguments)]
fn moe_stream_cpu(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    start: Instant,
    mapped: &Arc<crate::gguf::MappedGGUFModel>,
    quantized: &Arc<crate::gguf::OwnedQuantizedModel>,
    input_ids: &[u32],
    gen_config: &crate::gguf::QuantizedGenerateConfig,
    tokenizer: Arc<BPETokenizer>,
    max_tokens: usize,
    prompt_token_count: usize,
) -> Response {
    let (tx, rx) = tokio::sync::mpsc::channel::<Result<u32, String>>(64);
    let mapped_clone = mapped.clone();
    let quantized_clone = quantized.clone();
    let input_ids_clone = input_ids.to_vec();
    let gen_config_clone = gen_config.clone();
    let sink_metrics = state.metrics.clone();

    tokio::task::spawn_blocking(move || {
        let result = crate::infer::qwen3_moe_generate::run_qwen3_moe_generate_streaming(
            &mapped_clone,
            &quantized_clone,
            &input_ids_clone,
            &gen_config_clone,
            // Stops when the client goes away — see `streaming_token_sink`.
            crate::api::openai_handlers::streaming_token_sink(tx.clone(), sink_metrics),
        );
        if let Err(e) = result {
            let _ = tx.blocking_send(Err(e.to_string()));
        }
    });

    crate::api::openai_handlers::true_streaming_sse_response(
        rx,
        tokenizer,
        request_id.to_string(),
        request.model.clone(),
        state.metrics.clone(),
        start,
        max_tokens,
        prompt_token_count,
        // The MoE generator reports no phase split; §3 timings are absent.
        None,
        request.stop.as_deref(),
    )
}

fn try_qwen3_moe_backend(
    state: &AppState,
    request: &ChatCompletionRequest,
    request_id: &str,
    start: Instant,
    cancel: &CancelToken,
) -> Option<Response> {
    use crate::gguf::QuantizedGenerateConfig;

    let raw_arch = state.model_architecture()?;
    if !is_qwen3_moe_arch(&raw_arch) {
        return None;
    }

    let (mapped, quantized) = match moe_models(state, &raw_arch) {
        Ok(m) => m,
        Err(r) => return Some(r),
    };
    let tokenizer = match require_tokenizer(state) {
        Ok(t) => t,
        Err(r) => return Some(r),
    };

    let input_ids = match tokenize_chat_prompt(
        &tokenizer,
        &request.messages,
        Some(&request.model),
        request.thinking(),
        request.tools.as_deref(),
        state,
    ) {
        Ok(ids) => ids,
        Err(r) => return Some(r),
    };
    let prompt_token_count = input_ids.len();

    let max_tokens = request.max_tokens.unwrap_or(256).min(4096) as usize;
    // 3-knob toolkit (qwen3-moe-sampling-v1 + qwen3-moe-repetition-penalty-v1):
    // thread top_k/top_p/repeat_penalty/repeat_last_n/seed from the HTTP
    // request through to QuantizedGenerateConfig. Defaults match the dense
    // path's chat-completion behavior (greedy when unspecified).
    //
    // EOS stop-token: mirror the dense path's chat_gen_params fallback chain.
    // Generation halts on natural turn-end (model EOS or ChatML boundary).
    // Without this, qwen3_moe burns the full max_tokens budget per turn,
    // allowing self-prompted "Human:" runaway text — the root cause of
    // paiml/claude-code-parity-apr M287's verbosity pattern.
    //
    // Fallback order (matches chat_gen_params at openai_handlers.rs:97):
    //   1. state.model_eos_token_id() — from GGUF metadata
    //   2. tokenizer "<|im_end|>" — ChatML standard (Qwen, OpenHermes, Yi)
    //   3. tokenizer "<|endoftext|>" — GPT-style alternative
    //   4. None → empty stop_tokens (no behavior change from pre-fix)
    let gen_config = moe_gen_config(state, request, &tokenizer, max_tokens, cancel);

    // qwen3-moe-streaming-sse-v1: per-token SSE when stream=true.
    // Dispatches to the callback variant + builds an SSE response from
    // a tokio mpsc channel. Non-streaming path falls through below.
    //
    // #3987: that callback variant is the CPU forward. A CUDA server
    // (`with_moe_gpu`) must stream what the GPU generated, so it goes through
    // the ONE dispatch below and replays the result — the same
    // `stream_mode: "replayed"` every other CUDA chat stream declares.
    if request.stream && state.moe_no_gpu() {
        return Some(moe_stream_cpu(
            state, request, request_id, start, &mapped, &quantized, &input_ids, &gen_config, tokenizer,
            max_tokens, prompt_token_count,
        ));
    }

    // #3987: the ONE dispatch `apr run` uses (#3714), not the CPU-only generator this
    // used to call directly. On a CUDA server (`with_moe_gpu`) it serves on the GPU; a GPU
    // that cannot serve prints its reason and the CPU chain runs, and `used_gpu` says so.
    let (tokens, used_gpu) = match crate::infer::qwen3_moe_dispatch::run_qwen3_moe_generate_dispatch(
        &mapped,
        &quantized,
        &input_ids,
        &gen_config,
        state.moe_no_gpu(),
    ) {
        Ok(t) => t,
        Err(e) => {
            state.metrics.record_failure();
            return Some(fail_response(
                state,
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("qwen3_moe generation failed: {e}"),
            ));
        }
    };

    let generated_ids: Vec<u32> = tokens[input_ids.len()..].to_vec();
    let completion_tokens = generated_ids.len();

    if request.stream {
        state.metrics.record_success(completion_tokens, start.elapsed());
        return Some(pregenerated_sse_response(
            generated_ids,
            tokenizer,
            request_id.to_string(),
            request.model.clone(),
            request.stop.as_deref(),
            max_tokens,
            prompt_token_count,
        ));
    }

    // Apply clean_chat_output to strip self-emitted "Human:" / "User:" /
    // "<|im_end|>" / etc. prefixes from response text. Mirrors the dense
    // path at line 295 (PMAT-088). Without this, the M287 "Human: I need..."
    // runaway leaks into the chat response even after EOS detection.
    let response_text = match tokenizer.decode(&generated_ids) {
        Ok(t) => clean_chat_output(&t),
        Err(e) => {
            state.metrics.record_failure();
            return Some(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e));
        }
    };

    let duration = start.elapsed();
    state.metrics.record_success(completion_tokens, duration);

    Some(build_chat_response(
        request_id.to_string(),
        request.model.clone(),
        response_text,
        prompt_token_count,
        completion_tokens,
        max_tokens,
        request.stop.as_deref(),
        None,
        duration,
        request.tools.as_deref(),
        request_tool_choice(request),
        None,
        Some(used_gpu),
    ))
}

/// Predicate: does this raw architecture string canonicalize to qwen3_moe?
///
/// Extracted for unit testing the dispatch classification independently of
/// the full handler. See contracts/qwen3-moe-serve-dispatch-v1.yaml.
fn is_qwen3_moe_arch(raw_arch: &str) -> bool {
    crate::tensor_names::normalize_architecture(raw_arch) == "qwen3_moe"
}

#[cfg(test)]
mod qwen3_moe_dispatch_guard_tests {
    use super::is_qwen3_moe_arch;

    #[test]
    fn canonical_qwen3_moe_matches() {
        assert!(is_qwen3_moe_arch("qwen3_moe"));
    }

    #[test]
    fn huggingface_class_names_canonicalize() {
        assert!(is_qwen3_moe_arch("Qwen3MoeForCausalLM"));
        assert!(is_qwen3_moe_arch("Qwen3MoEForCausalLM"));
        assert!(is_qwen3_moe_arch("Qwen3CoderForCausalLM"));
        assert!(is_qwen3_moe_arch("Qwen3_5MoeForCausalLM"));
        assert!(is_qwen3_moe_arch("Qwen3_5MoeForConditionalGeneration"));
    }

    #[test]
    fn lowercase_underscore_variants_match() {
        assert!(is_qwen3_moe_arch("qwen3moe"));
    }

    #[test]
    fn dense_archs_do_not_match() {
        // FALSIFY-QWEN3_MOE_SERVE_DISPATCH_V1_002 negative cases: the guard
        // MUST NOT fire for dense architectures, otherwise it regresses
        // every existing chat-completions request.
        assert!(!is_qwen3_moe_arch("qwen2"));
        assert!(!is_qwen3_moe_arch("qwen3"));
        assert!(!is_qwen3_moe_arch("llama"));
        assert!(!is_qwen3_moe_arch("mistral"));
        assert!(!is_qwen3_moe_arch("phi"));
        assert!(!is_qwen3_moe_arch("gemma"));
    }

    #[test]
    fn unknown_arch_does_not_match() {
        // normalize_architecture defaults unknowns to "llama", which is not
        // qwen3_moe — so the guard should NOT fire for unknown archs.
        assert!(!is_qwen3_moe_arch("some-future-arch-3000"));
        assert!(!is_qwen3_moe_arch(""));
    }
}

#[cfg(all(test, feature = "gpu"))]
#[path = "tests/dense_session_4268.rs"]
mod dense_session_4268;