dynamo-llm 1.3.0-dev.1

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

use serde::{Deserialize, Serialize};

use crate::agents::context::AgentContext;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentTraceRecord {
    pub schema: TraceSchema,
    pub event_type: TraceEventType,
    pub event_time_unix_ms: u64,
    pub event_source: TraceEventSource,
    pub agent_context: AgentContext,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub request: Option<AgentRequestMetrics>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool: Option<AgentToolEvent>,
}

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

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

impl TraceEventType {
    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 TraceEventSource {
    #[serde(rename = "dynamo")]
    Dynamo,
    #[serde(rename = "harness", alias = "ms_agent")]
    Harness,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRequestMetrics {
    pub request_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub x_request_id: Option<String>,
    pub model: 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<WorkerInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub replay: Option<AgentReplayMetrics>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason_metadata: Option<FinishReasonMetadata>,
}

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

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct WorkerInfo {
    #[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 agent 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>,
    ) {
        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 AgentToolEvent {
    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<AgentToolStatus>,
    #[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, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum AgentToolStatus {
    #[serde(rename = "running")]
    Running,
    #[serde(rename = "succeeded", alias = "ok", alias = "success")]
    Succeeded,
    #[serde(rename = "error", alias = "failed")]
    Error,
    #[serde(rename = "cancelled", alias = "timeout", alias = "canceled")]
    Cancelled,
}