dynamo-llm 1.3.0

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use serde::{Deserialize, Serialize};

use crate::protocols::common::extensions::AgentContext;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestTraceRecord {
    pub schema: RequestTraceSchema,
    pub event_type: RequestTraceEventType,
    pub event_time_unix_ms: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub event_source: Option<RequestTraceEventSource>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_context: Option<AgentContext>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub request: Option<RequestTraceMetrics>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool: Option<RequestTraceToolEvent>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RequestTraceSchema {
    #[serde(rename = "dynamo.request.trace.v1")]
    V1,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RequestTraceEventType {
    #[serde(rename = "request_end")]
    RequestEnd,
    #[serde(rename = "tool_start")]
    ToolStart,
    #[serde(rename = "tool_end")]
    ToolEnd,
    #[serde(rename = "tool_error")]
    ToolError,
}

impl RequestTraceEventType {
    pub fn is_tool_event(self) -> bool {
        matches!(self, Self::ToolStart | Self::ToolEnd | Self::ToolError)
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RequestTraceEventSource {
    #[serde(rename = "dynamo")]
    Dynamo,
    #[serde(rename = "harness")]
    Harness,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestTraceMetrics {
    pub request_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub x_request_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub input_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cached_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub request_received_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prefill_wait_time_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prefill_time_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ttft_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total_time_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub avg_itl_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub kv_hit_rate: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub kv_transfer_estimated_latency_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub queue_depth: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub worker: Option<RequestTraceWorkerInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub replay: Option<RequestReplayMetrics>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason_metadata: Option<FinishReasonMetadata>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct RequestReplayMetrics {
    pub trace_block_size: usize,
    pub input_length: usize,
    pub input_sequence_hashes: Vec<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RequestTraceWorkerInfo {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prefill_worker_id: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prefill_dp_rank: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub decode_worker_id: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub decode_dp_rank: Option<u32>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct FinishReasonMetadata {
    /// Final OpenAI-compatible finish reason after parser/post-processing rewrites.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<dynamo_protocols::types::FinishReason>,

    /// Raw backend finish reason before parser/post-processing rewrites.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub backend_finish_reason: Option<String>,

    /// Backend-provided stop condition that caused `finish_reason=stop`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_reason: Option<dynamo_protocols::types::StopReason>,

    /// Complete tool calls emitted by the model. Arguments are intentionally
    /// omitted so request traces remain metadata, not payload logs.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCallMetadata>,

    /// Per-choice finish metadata for multi-choice requests. The top-level
    /// fields remain as a compact summary for the common single-choice case.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub choices: Vec<ChoiceFinishReasonMetadata>,
}

impl FinishReasonMetadata {
    pub fn is_empty(&self) -> bool {
        self.finish_reason.is_none()
            && self.backend_finish_reason.is_none()
            && self.stop_reason.is_none()
            && self.tool_calls.is_empty()
            && self.choices.is_empty()
    }

    pub fn record_choice_finish_reason(
        &mut self,
        choice_index: u32,
        finish_reason: dynamo_protocols::types::FinishReason,
    ) {
        self.choice_metadata_mut(choice_index).finish_reason = Some(finish_reason);
    }

    pub fn record_choice_backend_finish_reason(
        &mut self,
        choice_index: u32,
        backend_finish_reason: Option<String>,
        stop_reason: Option<dynamo_protocols::types::StopReason>,
    ) {
        if backend_finish_reason.is_none() && stop_reason.is_none() {
            return;
        }
        let choice = self.choice_metadata_mut(choice_index);
        if let Some(backend_finish_reason) = backend_finish_reason {
            choice.backend_finish_reason = Some(backend_finish_reason);
        }
        if let Some(stop_reason) = stop_reason {
            choice.stop_reason = Some(stop_reason);
        }
    }

    fn choice_metadata_mut(&mut self, choice_index: u32) -> &mut ChoiceFinishReasonMetadata {
        if let Some(position) = self
            .choices
            .iter()
            .position(|choice| choice.choice_index == choice_index)
        {
            return &mut self.choices[position];
        }

        let position = self.choices.len();
        self.choices.push(ChoiceFinishReasonMetadata {
            choice_index,
            finish_reason: None,
            backend_finish_reason: None,
            stop_reason: None,
        });
        &mut self.choices[position]
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ToolCallMetadata {
    pub choice_index: u32,
    pub tool_call_index: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ChoiceFinishReasonMetadata {
    pub choice_index: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<dynamo_protocols::types::FinishReason>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub backend_finish_reason: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_reason: Option<dynamo_protocols::types::StopReason>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestTraceToolEvent {
    pub tool_call_id: String,
    pub tool_class: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub started_at_unix_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ended_at_unix_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub status: Option<RequestTraceToolStatus>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_bytes: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_name_hash: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error_type: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestTraceToolEventIngress {
    pub schema: RequestTraceSchema,
    pub event_type: RequestTraceEventType,
    pub event_time_unix_ms: u64,
    pub session_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parent_session_id: Option<String>,
    pub tool: RequestTraceToolEvent,
}

impl From<RequestTraceToolEventIngress> for RequestTraceRecord {
    fn from(ingress: RequestTraceToolEventIngress) -> Self {
        Self {
            schema: ingress.schema,
            event_type: ingress.event_type,
            event_time_unix_ms: ingress.event_time_unix_ms,
            event_source: Some(RequestTraceEventSource::Harness),
            agent_context: Some(AgentContext {
                session_id: ingress.session_id,
                parent_session_id: ingress.parent_session_id,
                session_final: None,
                kv_hints: None,
            }),
            request: None,
            tool: Some(ingress.tool),
        }
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RequestTraceToolStatus {
    #[serde(rename = "running")]
    Running,
    #[serde(rename = "succeeded", alias = "ok", alias = "success")]
    Succeeded,
    #[serde(rename = "error", alias = "failed")]
    Error,
    #[serde(rename = "cancelled", alias = "canceled", alias = "timeout")]
    Cancelled,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn request_trace_has_only_replay_fields() {
        let record = RequestTraceRecord {
            schema: RequestTraceSchema::V1,
            event_type: RequestTraceEventType::RequestEnd,
            event_time_unix_ms: 1_100,
            event_source: None,
            agent_context: None,
            request: Some(RequestTraceMetrics {
                request_id: "req-1".to_string(),
                x_request_id: None,
                model: None,
                input_tokens: None,
                output_tokens: Some(4),
                cached_tokens: None,
                request_received_ms: Some(1_000),
                prefill_wait_time_ms: None,
                prefill_time_ms: None,
                ttft_ms: None,
                total_time_ms: None,
                avg_itl_ms: None,
                kv_hit_rate: None,
                kv_transfer_estimated_latency_ms: None,
                queue_depth: None,
                worker: None,
                replay: Some(RequestReplayMetrics {
                    trace_block_size: 2,
                    input_length: 4,
                    input_sequence_hashes: vec![11, 22],
                }),
                finish_reason_metadata: None,
            }),
            tool: None,
        };

        let value = serde_json::to_value(record).unwrap();
        assert_eq!(value["schema"], "dynamo.request.trace.v1");
        assert_eq!(value["event_type"], "request_end");
        assert!(value.get("event_source").is_none());
        assert!(value.get("agent_context").is_none());
        assert!(value.get("tool").is_none());
        assert!(value["request"].get("model").is_none());
        assert!(value["request"].get("finish_reason_metadata").is_none());
        assert!(value["request"].get("payload").is_none());
    }

    #[test]
    fn request_trace_tool_status_accepts_documented_aliases() {
        for (input, expected) in [
            ("running", RequestTraceToolStatus::Running),
            ("succeeded", RequestTraceToolStatus::Succeeded),
            ("ok", RequestTraceToolStatus::Succeeded),
            ("success", RequestTraceToolStatus::Succeeded),
            ("error", RequestTraceToolStatus::Error),
            ("failed", RequestTraceToolStatus::Error),
            ("cancelled", RequestTraceToolStatus::Cancelled),
            ("canceled", RequestTraceToolStatus::Cancelled),
            ("timeout", RequestTraceToolStatus::Cancelled),
        ] {
            let parsed: RequestTraceToolStatus =
                serde_json::from_str(&format!("\"{input}\"")).unwrap();
            assert_eq!(parsed, expected);
        }
    }

    #[test]
    fn empty_choice_backend_finish_reason_is_ignored() {
        let mut metadata = FinishReasonMetadata::default();

        metadata.record_choice_backend_finish_reason(0, None, None);

        assert!(metadata.is_empty());
        assert!(metadata.choices.is_empty());
    }
}