ferrum-types 0.12.2

//! Request and response types for inference

mod native_projection;
mod xml_parameter_schema;
mod xml_tool_calls;

pub use native_projection::{NativeChatOutputProjection, NativeChatOutputProjector};

use crate::{
    ids::*, models::TokenUsage, FinishReason, Priority, ResponseCompletionEnvelope, SamplingParams,
    TokenId,
};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

pub const PROMPT_TOKENS_METADATA_KEY: &str = "ferrum_prompt_tokens";
pub const DEFAULT_MAX_TOKENS_METADATA_KEY: &str = "ferrum_default_max_tokens";
/// Actual rendered generation-suffix state, supplied by the prompt renderer.
pub const PROMPT_OPENED_REASONING_METADATA_KEY: &str = "ferrum_prompt_opened_reasoning";

/// Explicit request for execution evidence that is expensive or sensitive to retain.
///
/// The default keeps the inference hot path unchanged. Product entrypoints opt in
/// only when the user enables a diagnostic artifact such as a request replay dump.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct InferenceEvidenceRequest {
    #[serde(default)]
    pub capture_prompt_token_ids: bool,
    /// Capture engine token-commit timestamps from the request's monotonic
    /// clock. Disabled by default so ordinary inference does not allocate or
    /// read an additional clock in the token hot path.
    #[serde(default)]
    pub capture_engine_token_timing: bool,
}

/// Engine-boundary timing for one inference request.
///
/// Every token timestamp is an offset from the same `std::time::Instant`
/// captured at request admission. The wall anchor exists only to correlate
/// this monotonic domain with profile events; durations and ITL must be
/// computed from `token_commit_nanos_since_request_start`, never by
/// subtracting wall-clock timestamps.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum EngineDecodeStage {
    /// Host scheduling work that selected this request for a decode wave.
    DecodeScheduling,
    /// The model-executor call for one decode wave, including preparation,
    /// device submission/completion, readback, and executor-side retirement.
    DecodeExecution,
    /// Host-side validation, sampling, and state commit after device execution.
    DecodePostprocess,
}

/// A measured request-local engine decode stage in the request's monotonic
/// clock domain. These are explicit producer boundaries, not intervals inferred
/// by filling gaps between other profile events.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct EngineDecodeStageInterval {
    pub stage: EngineDecodeStage,
    pub start_nanos_since_request_start: u64,
    pub end_nanos_since_request_start: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct EngineTokenTimingEvidence {
    pub clock_source: String,
    pub wall_anchor_unix_nanos: i64,
    pub wall_anchor_max_error_nanos: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decode_ready_nanos_since_request_start: Option<u64>,
    pub token_commit_nanos_since_request_start: Vec<u64>,
    /// Opt-in decode-stage evidence captured only with engine token timing.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub decode_stage_intervals: Vec<EngineDecodeStageInterval>,
}

impl EngineTokenTimingEvidence {
    pub fn validate(&self, output_tokens: usize) -> Result<(), String> {
        if self.clock_source != "rust_std_instant" {
            return Err("engine token timing clock_source must be rust_std_instant".to_string());
        }
        if self.wall_anchor_unix_nanos <= 0 {
            return Err("engine token timing wall anchor must be positive".to_string());
        }
        if self.token_commit_nanos_since_request_start.len() != output_tokens {
            return Err(format!(
                "engine token timing has {} commits for {output_tokens} output tokens",
                self.token_commit_nanos_since_request_start.len()
            ));
        }
        if self
            .token_commit_nanos_since_request_start
            .windows(2)
            .any(|window| window[1] < window[0])
        {
            return Err("engine token commit timestamps must be monotonic".to_string());
        }
        if self.decode_stage_intervals.iter().any(|interval| {
            interval.end_nanos_since_request_start < interval.start_nanos_since_request_start
        }) {
            return Err("engine decode stage interval end precedes start".to_string());
        }
        if self.decode_stage_intervals.windows(2).any(|window| {
            window[1].start_nanos_since_request_start < window[0].start_nanos_since_request_start
        }) {
            return Err("engine decode stage intervals must be ordered by start".to_string());
        }
        Ok(())
    }

    pub fn ttft_nanos(&self) -> Option<u64> {
        self.token_commit_nanos_since_request_start.first().copied()
    }

    pub fn inter_token_nanos(&self) -> Vec<u64> {
        self.token_commit_nanos_since_request_start
            .windows(2)
            .map(|window| window[1].saturating_sub(window[0]))
            .collect()
    }

    pub fn decode_wall_nanos(&self) -> Option<u64> {
        let start = self.decode_ready_nanos_since_request_start?;
        let end = self
            .token_commit_nanos_since_request_start
            .last()
            .copied()?;
        (end >= start).then_some(end - start)
    }
}

/// Evidence captured at the engine execution boundary.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InferenceExecutionEvidence {
    #[serde(default)]
    pub prompt_token_ids: Vec<TokenId>,
    /// Complete generated-token history at the engine boundary. Streaming
    /// chunks cannot reconstruct this reliably because special tokens and
    /// deferred multi-byte pieces may not produce a visible text delta.
    #[serde(default)]
    pub output_token_ids: Vec<TokenId>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub engine_token_timing: Option<EngineTokenTimingEvidence>,
}

/// Inference request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceRequest {
    /// Unique request identifier
    pub id: RequestId,
    /// Input prompt text
    pub prompt: String,
    /// Model to use for inference
    pub model_id: ModelId,
    /// Sampling parameters
    pub sampling_params: SamplingParams,
    /// Whether to stream response
    pub stream: bool,
    /// Request priority
    pub priority: Priority,
    /// Client identifier
    pub client_id: Option<ClientId>,
    /// Session identifier for stateful interactions
    pub session_id: Option<SessionId>,
    /// Request creation timestamp
    pub created_at: DateTime<Utc>,
    /// Structured product/API request context. `prompt` remains the rendered
    /// model input for current engines; this carries the original semantic
    /// request boundary for API features such as tools and response formats.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub api_request: Option<ApiRequest>,
    /// Explicitly requested execution evidence. Disabled by default so normal
    /// inference does not retain or copy prompt token IDs after completion.
    #[serde(default)]
    pub evidence_request: InferenceEvidenceRequest,
    /// Additional metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ApiRequest {
    Chat(ApiChatRequest),
    Completion(ApiCompletionRequest),
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ApiResponse {
    Chat(ApiChatResponse),
    Completion(ApiCompletionResponse),
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiChatRequest {
    pub messages: Vec<ApiChatMessage>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tools: Vec<ApiTool>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ApiToolChoice>,
    /// Wire protocol emitted by the model's chat template for tool calls.
    /// Native protocols retain their complete envelope for forced/named calls;
    /// `Json` preserves the legacy bare-argument fallback.
    #[serde(default)]
    pub tool_call_protocol: ApiToolCallProtocol,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub legacy_functions: Vec<ApiFunction>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub legacy_function_call: Option<ApiFunctionCallChoice>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ApiResponseFormat>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stream_options: Option<ApiStreamOptions>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiCompletionRequest {
    pub prompt: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ApiResponseFormat>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiChatResponse {
    pub message: ApiChatMessage,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiCompletionResponse {
    pub text: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiChatMessage {
    pub role: ApiMessageRole,
    pub content: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ApiToolCall>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub function_call: Option<ApiFunctionCall>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ApiMessageRole {
    System,
    User,
    Assistant,
    Function,
    Tool,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiTool {
    #[serde(rename = "type")]
    pub tool_type: String,
    pub function: ApiFunction,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiFunction {
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub parameters: Option<serde_json::Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ApiToolChoice {
    Mode(String),
    Function {
        #[serde(rename = "type")]
        tool_type: String,
        function: ApiToolChoiceFunction,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiToolChoiceFunction {
    pub name: String,
}

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ApiToolCallProtocol {
    #[default]
    Json,
    /// A tools-aware template emits a named JSON call, not bare arguments.
    NativeJson,
    FunctionParameterXml,
}

impl ApiToolCallProtocol {
    /// Atomic tokenizer entries that this wire protocol may need to emit.
    ///
    /// The sampling layer resolves these texts against the loaded tokenizer
    /// and allows only exact matches. Other added-vocabulary tokens remain
    /// masked.
    pub const fn generated_control_token_texts(self) -> &'static [&'static str] {
        match self {
            Self::Json | Self::NativeJson => &[],
            Self::FunctionParameterXml => &["<tool_call>", "</tool_call>"],
        }
    }

    /// Lexical envelope that may satisfy a pending response boundary.
    pub fn generated_response_envelope(self) -> Option<ResponseCompletionEnvelope> {
        match self {
            Self::Json | Self::NativeJson => None,
            Self::FunctionParameterXml => Some(ResponseCompletionEnvelope {
                open_token_text: "<tool_call>".to_string(),
                close_token_text: "</tool_call>".to_string(),
                max_envelopes: MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE,
            }),
        }
    }
}

impl ApiChatRequest {
    /// A forced native call must retain its declared envelope and tool name.
    /// A bare argument-object grammar would mask the model's native wire format.
    pub fn requires_native_tool_call(&self) -> bool {
        matches!(
            self.tool_call_protocol,
            ApiToolCallProtocol::NativeJson | ApiToolCallProtocol::FunctionParameterXml
        ) && match self.tool_choice.as_ref() {
            Some(ApiToolChoice::Mode(mode)) => mode.eq_ignore_ascii_case("required"),
            Some(ApiToolChoice::Function { tool_type, .. }) => tool_type == "function",
            None => false,
        }
    }

    pub fn allows_tool_name(&self, name: &str) -> bool {
        api_tool_name_allowed(self, name)
    }

    /// Automatic tool calls and hard final JSON are alternative output branches.
    /// A final object must never acquire tool-call semantics merely because its
    /// fields happen to match a function's argument schema.
    pub fn automatic_tools_with_hard_response_format(&self) -> bool {
        !self.tools.is_empty()
            && self.tool_choice.as_ref().is_none_or(|choice| {
                matches!(choice, ApiToolChoice::Mode(mode) if mode.eq_ignore_ascii_case("auto"))
            })
            && self.response_format.as_ref().is_some_and(|format| {
                format.format_type == "json_object"
                    || (format.format_type == "json_schema"
                        && format
                            .json_schema
                            .as_ref()
                            .is_some_and(|schema| schema.strict == Some(true)))
            })
    }

    /// Protocol controls are generatable only when this request can emit a
    /// modern tool call. `tool_choice: none` must not widen token sampling.
    pub fn generated_control_token_texts(&self) -> &'static [&'static str] {
        if self.tools.is_empty() || api_tool_choice_is_none(self) {
            return &[];
        }
        if self.automatic_tools_with_hard_response_format() || self.requires_native_tool_call() {
            return &["<tool_call>", "</tool_call>"];
        }
        self.tool_call_protocol.generated_control_token_texts()
    }

    /// Alternate complete response envelope available to this request.
    pub fn generated_response_envelope(&self) -> Option<ResponseCompletionEnvelope> {
        if self.tools.is_empty() || api_tool_choice_is_none(self) {
            return None;
        }
        if self.automatic_tools_with_hard_response_format() || self.requires_native_tool_call() {
            return Some(ResponseCompletionEnvelope {
                open_token_text: "<tool_call>".to_string(),
                close_token_text: "</tool_call>".to_string(),
                max_envelopes: MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE,
            });
        }
        self.tool_call_protocol.generated_response_envelope()
    }
}

impl ApiRequest {
    pub fn generated_control_token_texts(&self) -> &'static [&'static str] {
        match self {
            Self::Chat(request) => request.generated_control_token_texts(),
            Self::Completion(_) => &[],
        }
    }

    pub fn generated_response_envelope(&self) -> Option<ResponseCompletionEnvelope> {
        match self {
            Self::Chat(request) => request.generated_response_envelope(),
            Self::Completion(_) => None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ApiFunctionCallChoice {
    Mode(String),
    Function { name: String },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiToolCall {
    pub id: String,
    #[serde(rename = "type")]
    pub tool_type: String,
    pub function: ApiFunctionCall,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiFunctionCall {
    pub name: String,
    pub arguments: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiResponseFormat {
    #[serde(rename = "type")]
    pub format_type: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub json_schema: Option<ApiJsonSchema>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ApiJsonSchema {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    pub schema: serde_json::Value,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ApiStreamOptions {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub include_usage: Option<bool>,
}

const MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE: usize = 32;

pub fn api_response_from_generated_text(
    request: &InferenceRequest,
    text: &str,
    finish_reason: FinishReason,
) -> Option<ApiResponse> {
    if let Some(mut projector) = NativeChatOutputProjector::for_request(request) {
        projector.push(text);
        return projector
            .finish(finish_reason)
            .api_response
            .map(ApiResponse::Chat);
    }
    let ApiRequest::Chat(chat_request) = request.api_request.as_ref()? else {
        return None;
    };
    chat_api_response_from_generated_text(chat_request, text, finish_reason).map(ApiResponse::Chat)
}

/// Branch established by the structured-output grammar for a complete result.
/// When both the final and tool languages accept the same bytes, the grammar
/// owner must select `Final` before invoking the classified response helper.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StructuredOutputBranch {
    Final,
    ToolCall,
}

/// Preserve an authoritative grammar decision through product response shaping.
///
/// `text` is the complete result payload, excluding reasoning and its headers.
/// This function does not validate the final schema or infer branch ownership
/// from JSON fields. Only a caller that established full grammar acceptance
/// and final precedence may supply `branch`.
///
/// A final answer deliberately returns `Some`, even when its fields resemble a
/// call, so downstream text fallbacks cannot reinterpret it. Classified calls
/// accept complete registered envelopes or a named JSON object with explicit
/// arguments; they never infer the sole tool from an unnamed argument object.
pub fn api_response_from_classified_generated_text(
    request: &InferenceRequest,
    text: &str,
    finish_reason: FinishReason,
    branch: StructuredOutputBranch,
) -> crate::Result<Option<ApiResponse>> {
    let Some(ApiRequest::Chat(chat_request)) = request.api_request.as_ref() else {
        return Ok(None);
    };
    if !matches!(finish_reason, FinishReason::Stop | FinishReason::EOS) {
        return Ok(None);
    }
    if !chat_request.automatic_tools_with_hard_response_format()
        && !chat_request.requires_native_tool_call()
    {
        return Err(crate::FerrumError::invalid_request(
            "classified structured output requires a composed tool output contract",
        ));
    }
    if chat_request.requires_native_tool_call() && branch == StructuredOutputBranch::Final {
        return Err(crate::FerrumError::invalid_request(
            "a required native tool call cannot be classified as a final answer",
        ));
    }

    let (content, tool_calls, wire_finish_reason) = match branch {
        StructuredOutputBranch::Final => (text.to_string(), Vec::new(), "stop"),
        StructuredOutputBranch::ToolCall => {
            let calls = parse_explicit_tool_call_envelopes(text, chat_request)
                .or_else(|| parse_classified_named_json_call(text, chat_request))
                .filter(|calls| {
                    calls.iter().all(|call| {
                        serde_json::from_str::<serde_json::Value>(&call.function.arguments)
                            .is_ok_and(|arguments| arguments.is_object())
                    })
                })
                .ok_or_else(|| {
                    crate::FerrumError::invalid_format(
                        "classified tool output is not a complete declared function call",
                    )
                })?;
            (String::new(), calls, "tool_calls")
        }
    };
    Ok(Some(ApiResponse::Chat(ApiChatResponse {
        message: ApiChatMessage {
            role: ApiMessageRole::Assistant,
            content,
            name: None,
            tool_calls,
            tool_call_id: None,
            function_call: None,
        },
        finish_reason: Some(wire_finish_reason.to_string()),
    })))
}

fn parse_classified_named_json_call(
    text: &str,
    chat_request: &ApiChatRequest,
) -> Option<Vec<ApiToolCall>> {
    if !matches!(
        chat_request.tool_call_protocol,
        ApiToolCallProtocol::Json | ApiToolCallProtocol::NativeJson
    ) {
        return None;
    }
    // Consume the whole JSON value, without searching surrounding prose,
    // fences, native control markers, or a reasoning section for a call.
    let value: serde_json::Value = serde_json::from_str(text).ok()?;
    let name = value.get("name")?.as_str()?;
    if !api_tool_name_allowed(chat_request, name) {
        return None;
    }
    let arguments = match (value.get("arguments"), value.get("parameters")) {
        (Some(arguments), None) | (None, Some(arguments)) => arguments,
        _ => return None,
    };
    if !arguments.is_object() {
        return None;
    }
    let call = ApiToolCall {
        id: "call_0".to_string(),
        tool_type: "function".to_string(),
        function: ApiFunctionCall {
            name: name.to_string(),
            arguments: serde_json::to_string(arguments).ok()?,
        },
    };
    validate_parsed_tool_calls(vec![call])
}

pub fn chat_api_may_emit_tool_or_function_call(chat_request: &ApiChatRequest) -> bool {
    (!chat_request.tools.is_empty() && !api_tool_choice_is_none(chat_request))
        || (!chat_request.legacy_functions.is_empty()
            && !api_function_call_choice_is_none(chat_request))
}

pub fn chat_api_response_from_generated_text(
    chat_request: &ApiChatRequest,
    text: &str,
    finish_reason: FinishReason,
) -> Option<ApiChatResponse> {
    if !matches!(finish_reason, FinishReason::Stop | FinishReason::EOS) {
        return None;
    }

    if !chat_request.tools.is_empty() && !api_tool_choice_is_none(chat_request) {
        if let Some((content, tool_calls)) =
            parse_tool_calls_from_generated_text(text, chat_request)
        {
            return Some(ApiChatResponse {
                message: ApiChatMessage {
                    role: ApiMessageRole::Assistant,
                    content,
                    name: None,
                    tool_calls,
                    tool_call_id: None,
                    function_call: None,
                },
                finish_reason: Some("tool_calls".to_string()),
            });
        }
    }

    if !chat_request.legacy_functions.is_empty() && !api_function_call_choice_is_none(chat_request)
    {
        if let Some(function_call) =
            parse_legacy_function_call_from_generated_text(text, chat_request)
        {
            return Some(ApiChatResponse {
                message: ApiChatMessage {
                    role: ApiMessageRole::Assistant,
                    content: String::new(),
                    name: None,
                    tool_calls: Vec::new(),
                    tool_call_id: None,
                    function_call: Some(function_call),
                },
                finish_reason: Some("function_call".to_string()),
            });
        }
    }

    None
}

fn api_tool_choice_is_none(chat_request: &ApiChatRequest) -> bool {
    matches!(
        chat_request.tool_choice.as_ref(),
        Some(ApiToolChoice::Mode(mode)) if mode.eq_ignore_ascii_case("none")
    )
}

fn api_function_call_choice_is_none(chat_request: &ApiChatRequest) -> bool {
    matches!(
        chat_request.legacy_function_call.as_ref(),
        Some(ApiFunctionCallChoice::Mode(mode)) if mode.eq_ignore_ascii_case("none")
    )
}

fn parse_tool_calls_from_generated_text(
    text: &str,
    chat_request: &ApiChatRequest,
) -> Option<(String, Vec<ApiToolCall>)> {
    if chat_request.automatic_tools_with_hard_response_format()
        || chat_request.requires_native_tool_call()
    {
        return parse_explicit_tool_call_envelopes(text, chat_request)
            .or_else(|| {
                // Only a forced native JSON contract makes a complete named
                // object unambiguously a call. Automatic hard final JSON still
                // requires the sampler's branch classification.
                chat_request
                    .requires_native_tool_call()
                    .then(|| parse_classified_named_json_call(text, chat_request))
                    .flatten()
            })
            .map(|calls| (String::new(), calls));
    }
    if chat_request.tool_call_protocol == ApiToolCallProtocol::FunctionParameterXml {
        // Only native framing establishes tool intent for this protocol.
        // Ordinary JSON or JSON inside malformed framing is not a fallback.
        return xml_tool_calls::parse_with_content(text, chat_request, false)
            .map(|parsed| (parsed.content, parsed.calls));
    }

    let value = parse_json_value_from_generated_text(text)?;
    parse_json_tool_call_value(&value, chat_request, 0, true).map(|calls| (String::new(), calls))
}

fn parse_json_tool_call_value(
    value: &serde_json::Value,
    chat_request: &ApiChatRequest,
    index_offset: usize,
    allow_unwrapped_arguments: bool,
) -> Option<Vec<ApiToolCall>> {
    if let Some(calls) = value.get("tool_calls").and_then(|value| value.as_array()) {
        if calls.len() > MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE {
            return None;
        }
        let parsed = calls
            .iter()
            .enumerate()
            .map(|(index, value)| parse_tool_call_value(value, index_offset + index, chat_request))
            .collect::<Option<Vec<_>>>()?;
        return validate_parsed_tool_calls(parsed);
    }
    if let Some(tool_call) = value.get("tool_call") {
        return parse_tool_call_value(tool_call, index_offset, chat_request)
            .and_then(|call| validate_parsed_tool_calls(vec![call]));
    }
    if let Some(tool_call) = parse_wrapped_tool_call_value(value, index_offset, chat_request) {
        return validate_parsed_tool_calls(vec![tool_call]);
    }
    parse_tool_call_value(value, index_offset, chat_request)
        .or_else(|| {
            allow_unwrapped_arguments
                .then(|| parse_forced_tool_arguments_value(value, index_offset, chat_request))
                .flatten()
        })
        .and_then(|call| validate_parsed_tool_calls(vec![call]))
}

/// Only explicit, fully consumed protocol envelopes establish a tool branch.
/// Do not search arbitrary text for tags: they may be literal final JSON data.
fn parse_explicit_tool_call_envelopes(
    text: &str,
    chat_request: &ApiChatRequest,
) -> Option<Vec<ApiToolCall>> {
    if chat_request.tool_call_protocol == ApiToolCallProtocol::FunctionParameterXml {
        return xml_tool_calls::parse(text, chat_request, true);
    }
    const OPEN: &str = "<tool_call>";
    const CLOSE: &str = "</tool_call>";
    let mut remaining = text.trim();
    let mut calls = Vec::new();
    while !remaining.is_empty() {
        if calls.len() >= MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE {
            return None;
        }
        let payload = remaining.strip_prefix(OPEN)?.trim_start();
        // The closing marker can legitimately occur inside a JSON
        // argument string; consume JSON before interpreting framing.
        let mut values =
            serde_json::Deserializer::from_str(payload).into_iter::<serde_json::Value>();
        let value = values.next()?.ok()?;
        remaining = payload[values.byte_offset()..]
            .trim_start()
            .strip_prefix(CLOSE)?
            .trim_start();
        let parsed = parse_json_tool_call_value(&value, chat_request, calls.len(), false)?;
        if calls.len() + parsed.len() > MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE {
            return None;
        }
        for call in parsed {
            if calls
                .iter()
                .any(|previous: &ApiToolCall| previous.id == call.id)
            {
                return None;
            }
            calls.push(call);
        }
    }
    validate_parsed_tool_calls(calls)
}

fn validate_parsed_tool_calls(calls: Vec<ApiToolCall>) -> Option<Vec<ApiToolCall>> {
    if calls.is_empty() || calls.len() > MAX_PARALLEL_TOOL_CALLS_PER_RESPONSE {
        return None;
    }
    for (index, call) in calls.iter().enumerate() {
        if calls[..index].iter().any(|previous| {
            previous.function.name == call.function.name
                && previous.function.arguments == call.function.arguments
        }) {
            return None;
        }
    }
    Some(calls)
}

fn parse_wrapped_tool_call_value(
    value: &serde_json::Value,
    index: usize,
    chat_request: &ApiChatRequest,
) -> Option<ApiToolCall> {
    for key in ["auto", "tool", "tool_call", "auto_tool_response"] {
        if let Some(wrapped) = value.get(key) {
            if let Some(call) = parse_tool_call_value(wrapped, index, chat_request) {
                return Some(call);
            }
        }
    }
    None
}

fn parse_tool_call_value(
    value: &serde_json::Value,
    index: usize,
    chat_request: &ApiChatRequest,
) -> Option<ApiToolCall> {
    let tool_type = value
        .get("type")
        .and_then(|value| value.as_str())
        .unwrap_or("function");
    if tool_type != "function" {
        return None;
    }
    let function = value.get("function").unwrap_or(value);
    let name = function
        .as_str()
        .or_else(|| function.get("name").and_then(|value| value.as_str()))
        .or_else(|| function.get("tool").and_then(|value| value.as_str()))
        .or_else(|| value.get("name").and_then(|value| value.as_str()))?;
    if !api_tool_name_allowed(chat_request, name) {
        return None;
    }
    let arguments = api_arguments_to_string(
        function
            .get("arguments")
            .or_else(|| function.get("parameters"))
            .or_else(|| value.get("arguments"))
            .or_else(|| value.get("parameters")),
    );
    let id = value
        .get("id")
        .and_then(|value| value.as_str())
        .map(str::to_string)
        .unwrap_or_else(|| format!("call_{index}"));

    Some(ApiToolCall {
        id,
        tool_type: "function".to_string(),
        function: ApiFunctionCall {
            name: name.to_string(),
            arguments,
        },
    })
}

fn parse_forced_tool_arguments_value(
    value: &serde_json::Value,
    index: usize,
    chat_request: &ApiChatRequest,
) -> Option<ApiToolCall> {
    let tool = unwrapped_tool_arguments_target(chat_request, value)?;
    if value.get("tool_calls").is_some()
        || value.get("tool_call").is_some()
        || value.get("function").is_some()
        || value.get("name").is_some()
    {
        return None;
    }

    Some(ApiToolCall {
        id: format!("call_{index}"),
        tool_type: "function".to_string(),
        function: ApiFunctionCall {
            name: tool.function.name.clone(),
            arguments: serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string()),
        },
    })
}

fn unwrapped_tool_arguments_target<'a>(
    chat_request: &'a ApiChatRequest,
    value: &serde_json::Value,
) -> Option<&'a ApiTool> {
    if let Some(name) = forced_tool_choice_name(chat_request) {
        return chat_request
            .tools
            .iter()
            .find(|tool| tool.tool_type == "function" && tool.function.name == name);
    }

    if matches!(
        chat_request.tool_choice.as_ref(),
        Some(ApiToolChoice::Mode(mode)) if !mode.eq_ignore_ascii_case("auto")
    ) {
        return None;
    }

    let tool = single_function_tool(chat_request)?;
    if !value_looks_like_tool_arguments(value, tool) {
        return None;
    }
    Some(tool)
}

fn value_looks_like_tool_arguments(value: &serde_json::Value, tool: &ApiTool) -> bool {
    let Some(arguments) = value.as_object() else {
        return false;
    };
    if arguments.is_empty() {
        return false;
    }
    let Some(properties) = tool
        .function
        .parameters
        .as_ref()
        .and_then(|parameters| parameters.get("properties"))
        .and_then(|properties| properties.as_object())
    else {
        return false;
    };
    arguments.keys().all(|key| properties.contains_key(key))
}

fn forced_tool_choice_name(chat_request: &ApiChatRequest) -> Option<&str> {
    match chat_request.tool_choice.as_ref() {
        Some(ApiToolChoice::Function {
            tool_type,
            function,
        }) if tool_type == "function" && api_tool_name_allowed(chat_request, &function.name) => {
            Some(function.name.as_str())
        }
        Some(ApiToolChoice::Mode(mode)) if mode.eq_ignore_ascii_case("required") => {
            single_function_tool(chat_request).map(|tool| tool.function.name.as_str())
        }
        _ => None,
    }
}

fn single_function_tool(chat_request: &ApiChatRequest) -> Option<&ApiTool> {
    let mut tools = chat_request
        .tools
        .iter()
        .filter(|tool| tool.tool_type == "function");
    let tool = tools.next()?;
    tools.next().is_none().then_some(tool)
}

fn parse_legacy_function_call_from_generated_text(
    text: &str,
    chat_request: &ApiChatRequest,
) -> Option<ApiFunctionCall> {
    let value = parse_json_value_from_generated_text(text)?;
    let function = value.get("function_call").unwrap_or(&value);
    let name = function.get("name").and_then(|value| value.as_str())?;
    if !api_function_name_allowed(chat_request, name) {
        return None;
    }
    Some(ApiFunctionCall {
        name: name.to_string(),
        arguments: api_arguments_to_string(function.get("arguments")),
    })
}

fn api_tool_name_allowed(chat_request: &ApiChatRequest, name: &str) -> bool {
    match chat_request.tool_choice.as_ref() {
        Some(ApiToolChoice::Mode(mode)) if mode.eq_ignore_ascii_case("none") => false,
        Some(ApiToolChoice::Function {
            tool_type,
            function,
        }) => {
            tool_type == "function"
                && function.name == name
                && chat_request
                    .tools
                    .iter()
                    .any(|tool| tool.function.name == name)
        }
        _ => chat_request
            .tools
            .iter()
            .any(|tool| tool.function.name == name),
    }
}

fn api_function_name_allowed(chat_request: &ApiChatRequest, name: &str) -> bool {
    match chat_request.legacy_function_call.as_ref() {
        Some(ApiFunctionCallChoice::Mode(mode)) if mode.eq_ignore_ascii_case("none") => false,
        Some(ApiFunctionCallChoice::Function { name: selected }) => {
            selected == name
                && chat_request
                    .legacy_functions
                    .iter()
                    .any(|function| function.name == name)
        }
        _ => chat_request
            .legacy_functions
            .iter()
            .any(|function| function.name == name),
    }
}

fn parse_json_value_from_generated_text(text: &str) -> Option<serde_json::Value> {
    let trimmed = strip_single_json_fence(text.trim());
    serde_json::from_str(trimmed).ok().or_else(|| {
        let start = trimmed.find('{')?;
        let end = trimmed.rfind('}')?;
        (start <= end)
            .then(|| serde_json::from_str(&trimmed[start..=end]).ok())
            .flatten()
    })
}

fn strip_single_json_fence(text: &str) -> &str {
    let Some(rest) = text.strip_prefix("```") else {
        return text;
    };
    let rest = rest.strip_prefix("json").unwrap_or(rest).trim_start();
    rest.strip_suffix("```").map(str::trim).unwrap_or(text)
}

fn api_arguments_to_string(arguments: Option<&serde_json::Value>) -> String {
    match arguments {
        Some(serde_json::Value::String(raw)) => raw.clone(),
        Some(value) => serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string()),
        None => "{}".to_string(),
    }
}

impl InferenceRequest {
    /// Whether sampling must compile a grammar, including native tool-only output.
    pub fn requires_structured_output(&self) -> bool {
        !matches!(
            self.sampling_params.response_format,
            crate::ResponseFormat::Text
        ) || matches!(self.api_request.as_ref(), Some(ApiRequest::Chat(chat)) if chat.requires_native_tool_call())
    }

    /// Create a new inference request
    pub fn new(prompt: impl Into<String>, model_id: impl Into<ModelId>) -> Self {
        Self {
            id: RequestId::new(),
            prompt: prompt.into(),
            model_id: model_id.into(),
            sampling_params: SamplingParams::default(),
            stream: false,
            priority: Priority::default(),
            client_id: None,
            session_id: None,
            created_at: Utc::now(),
            api_request: None,
            evidence_request: InferenceEvidenceRequest::default(),
            metadata: HashMap::new(),
        }
    }

    /// Set sampling parameters
    pub fn with_sampling_params(mut self, params: SamplingParams) -> Self {
        self.sampling_params = params;
        self
    }

    /// Enable streaming
    pub fn with_stream(mut self, stream: bool) -> Self {
        self.stream = stream;
        self
    }

    /// Set priority
    pub fn with_priority(mut self, priority: Priority) -> Self {
        self.priority = priority;
        self
    }

    /// Set client ID
    pub fn with_client_id(mut self, client_id: impl Into<ClientId>) -> Self {
        self.client_id = Some(client_id.into());
        self
    }

    /// Set session ID
    pub fn with_session_id(mut self, session_id: SessionId) -> Self {
        self.session_id = Some(session_id);
        self
    }

    /// Set structured product/API request context.
    pub fn with_api_request(mut self, api_request: ApiRequest) -> Self {
        self.api_request = Some(api_request);
        self
    }

    /// Request prompt-token evidence from the execution boundary.
    pub fn with_prompt_token_evidence(mut self) -> Self {
        self.evidence_request.capture_prompt_token_ids = true;
        self
    }

    /// Request engine token-commit timing evidence.
    pub fn with_engine_token_timing_evidence(mut self) -> Self {
        self.evidence_request.capture_engine_token_timing = true;
        self
    }

    /// Add metadata
    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }
}

/// Inference response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceResponse {
    /// Request ID this response corresponds to
    pub request_id: RequestId,
    /// Generated text
    pub text: String,
    /// Generated token IDs
    pub tokens: Vec<TokenId>,
    /// Reason for completion
    pub finish_reason: FinishReason,
    /// Token usage statistics
    pub usage: TokenUsage,
    /// Total latency in milliseconds
    pub latency_ms: u64,
    /// Response creation timestamp
    pub created_at: DateTime<Utc>,
    /// Additional response metadata
    pub metadata: HashMap<String, serde_json::Value>,
    /// Structured product/API response context. Engines that can produce
    /// product-native outputs, such as assistant tool calls, can populate
    /// this without overloading plain text or ad hoc metadata.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub api_response: Option<ApiResponse>,
    /// Optional engine-boundary evidence requested by the caller.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub execution_evidence: Option<InferenceExecutionEvidence>,
}

/// Streaming response chunk
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamChunk {
    /// Request ID this chunk corresponds to
    pub request_id: RequestId,
    /// Text delta for this chunk
    pub text: String,
    /// Token ID for this chunk (if available)
    pub token: Option<TokenId>,
    /// Finish reason if this is the final chunk
    pub finish_reason: Option<FinishReason>,
    /// Token usage (typically only in final chunk)
    pub usage: Option<TokenUsage>,
    /// Chunk creation timestamp
    pub created_at: DateTime<Utc>,
    /// Chunk metadata
    pub metadata: HashMap<String, serde_json::Value>,
    /// Structured product/API response context for final streaming chunks.
    /// This mirrors `InferenceResponse::api_response` so streaming endpoints
    /// can return native tool/function-call payloads without reparsing text.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub api_response: Option<ApiResponse>,
    /// Optional engine-boundary evidence, emitted on the final chunk.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub execution_evidence: Option<InferenceExecutionEvidence>,
}

/// Batch request for processing multiple requests together
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchRequest {
    /// Batch identifier
    pub batch_id: BatchId,
    /// Requests in this batch
    pub requests: Vec<InferenceRequest>,
    /// Maximum sequence length for this batch
    pub max_sequence_length: usize,
    /// Batch creation timestamp
    pub created_at: DateTime<Utc>,
}

impl BatchRequest {
    /// Create a new batch request
    pub fn new(requests: Vec<InferenceRequest>) -> Self {
        let max_sequence_length = requests
            .iter()
            .map(|r| r.sampling_params.max_tokens)
            .max()
            .unwrap_or(512);

        Self {
            batch_id: BatchId::new(),
            requests,
            max_sequence_length,
            created_at: Utc::now(),
        }
    }

    /// Get the number of requests in this batch
    pub fn size(&self) -> usize {
        self.requests.len()
    }

    /// Check if batch is empty
    pub fn is_empty(&self) -> bool {
        self.requests.is_empty()
    }
}

/// Request state in the scheduler
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RequestState {
    /// Request is waiting in queue
    Waiting,
    /// Request is being processed
    Running,
    /// Request was preempted and is waiting to resume
    Preempted,
    /// Request completed successfully
    Completed,
    /// Request failed with error
    Failed,
    /// Request was cancelled
    Cancelled,
}

/// Scheduled request with additional state information
#[derive(Debug, Clone)]
pub struct ScheduledRequest {
    /// The original request
    pub request: InferenceRequest,
    /// Current state in scheduler
    pub state: RequestState,
    /// Allocated cache blocks
    pub allocated_blocks: Vec<crate::BlockId>,
    /// Number of tokens processed so far
    pub tokens_processed: usize,
    /// Estimated completion time
    pub estimated_completion: Option<DateTime<Utc>>,
}

impl ScheduledRequest {
    /// Create a new scheduled request
    pub fn new(request: InferenceRequest) -> Self {
        Self {
            request,
            state: RequestState::Waiting,
            allocated_blocks: Vec::new(),
            tokens_processed: 0,
            estimated_completion: None,
        }
    }

    /// Update request state
    pub fn set_state(&mut self, state: RequestState) {
        self.state = state;
    }

    /// Add allocated cache blocks
    pub fn add_blocks(&mut self, blocks: Vec<crate::BlockId>) {
        self.allocated_blocks.extend(blocks);
    }

    /// Update tokens processed
    pub fn update_progress(&mut self, tokens_processed: usize) {
        self.tokens_processed = tokens_processed;
    }
}