klieo-runlog 3.4.0

//! Aggregate types for the RunLog projection.

use chrono::{DateTime, Utc};
use klieo_core::ids::RunId;
use klieo_core::llm::{FinishReason, ToolCall};
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Run lifecycle status.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum RunStatus {
    /// Run is still in progress.
    Running,
    /// Run finished successfully.
    Completed,
    /// Run failed.
    Failed,
    /// Run was cancelled by the operator or scheduler.
    Cancelled,
}

/// Discriminator for the kind of side effect a `Step` represents.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum StepKind {
    /// LLM completion or chat call.
    LlmCall,
    /// Tool invocation through `ToolInvoker`.
    ToolCall,
    /// Summarizer checkpoint — distinct from a top-level `LlmCall` so
    /// summarizer overhead does not pollute substantive-reasoning cost
    /// and latency attribution.
    SummaryCheckpoint,
    /// Operational-layer event emitted by klieo-ops (supervisor state
    /// changes, gate decisions, governor denials, etc.). The `input`
    /// field of the corresponding [`Step`] carries the raw
    /// `serde_json::Value` payload from `Episode::Ops`.
    OpsEvent,
    /// A graphRAG recall performed during the run (from
    /// `Episode::MemoryRecall`). `Step::input` carries `{ "query", "k" }`;
    /// `Step::output` carries `{ "returned_fact_ids": [...] }`.
    MemoryRecall,
}

/// One unit of agent work — either an LLM call or a tool call.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Step {
    /// Zero-based index of the step within the run.
    pub idx: u32,
    /// Discriminator.
    pub kind: StepKind,
    /// Tool name (for `ToolCall`) or model name (for `LlmCall`); `None` if unknown.
    pub name: Option<String>,
    /// Prompt-side token count for an `LlmCall` step; `None` for other kinds or
    /// when the provider did not report a split.
    #[serde(default)]
    pub prompt_tokens: Option<u32>,
    /// Completion-side token count for an `LlmCall` step; `None` for other kinds
    /// or when the provider did not report a split.
    #[serde(default)]
    pub completion_tokens: Option<u32>,
    /// Estimated USD cost for an `LlmCall` step, computed by
    /// [`project_with_price_table`](crate::project_with_price_table) from the
    /// step's `LlmIo` and a `PriceTable`. `None` for other kinds, when the
    /// provider/model/token split is unknown, or when no price matched.
    #[serde(default)]
    pub cost_usd: Option<f64>,
    /// JSON inputs (tool args, prompt structure).
    pub input: serde_json::Value,
    /// JSON outputs (tool result, completion text).
    pub output: serde_json::Value,
    /// Error message if this step failed (`None` on success).
    pub error: Option<String>,
    /// Wall-clock latency serialised as milliseconds.
    #[serde(with = "duration_millis")]
    #[cfg_attr(feature = "schemars", schemars(with = "u64"))]
    pub latency: Duration,
    /// Tracing span id (stringly-typed for now; richer linkage is a Wave 7 carryover).
    pub span_id: Option<String>,
}

/// Sidecar I/O record for a single `Episode::LlmCall`.
///
/// `klieo_core::Episode::LlmCall` carries only `tokens` + `latency_ms` (the
/// trait surface is frozen at 0.1.0 — see `docs/SEMVER.md`). To make
/// projected `RunLog`s replayable, callers that have access to the original
/// prompt + completion can pass a parallel `&[LlmIo]` slice to
/// [`crate::projector::project_with_llm_io`]: the i-th `LlmIo` is paired with
/// the i-th `Episode::LlmCall` in the episode stream.
///
/// Without this sidecar, projected `Step.input` / `Step.output` for LLM
/// steps remain `Value::Null`, and [`crate::replay::replay`] cannot reproduce
/// the recorded prompt.
///
/// ## Cost-accounting fields
///
/// `provider`, `prompt_tokens`, and `completion_tokens` are optional sidecar
/// fields used by [`crate::projector::project_with_price_table`] to compute a
/// per-run [`Cost`]. `Episode::LlmCall::tokens` is a single `u32` total with
/// no provider info, so callers populate these fields here when they want
/// USD totals on the projected `RunLog`. Each new field carries
/// `#[serde(default)]` so older serialised records still deserialise.
///
/// ## Trust model
///
/// `provider`, `prompt_tokens`, and `completion_tokens` are **caller-supplied**
/// — `LlmIo` is a passive carrier and the projector trusts the sidecar
/// verbatim.
///
/// In compliance-grade audit-trail use cases (signed evidence, regulator
/// review, billing), these fields MUST originate directly from the
/// `LlmClient`'s response (response headers, parsed usage block, or the
/// transport layer). They MUST NOT be inferred or asserted by the agent
/// layer: a compromised tool or agent that constructs `LlmIo` itself can
/// spoof `provider = "ollama"` (zero rate) to under-report cost, or attribute
/// usage to a cheaper model tier.
///
/// The projector also does not cross-check these fields against
/// `Episode::LlmCall.tokens`. Adding that consistency check (e.g. asserting
/// `prompt_tokens + completion_tokens == LlmCall.tokens`) is a deferred
/// follow-up; for now, callers wanting that guarantee should assert it
/// themselves before passing the sidecar to
/// [`crate::projector::project_with_price_table`].
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct LlmIo {
    /// Recorded prompt text (the `User` message content sent to the LLM).
    pub prompt: String,
    /// Recorded completion text (the assistant response).
    pub completion: String,
    /// Optional model name; populates `Step.name` when present.
    pub model: Option<String>,
    /// Optional provider name (e.g. `"anthropic"`, `"openai"`). Required for
    /// price-table lookup in [`crate::projector::project_with_price_table`].
    #[serde(default)]
    pub provider: Option<String>,
    /// Optional prompt-token count. Required for price-table cost accounting.
    #[serde(default)]
    pub prompt_tokens: Option<u32>,
    /// Optional completion-token count. Required for price-table cost
    /// accounting.
    #[serde(default)]
    pub completion_tokens: Option<u32>,
    /// Optional cache-read token count (prompt tokens served from the
    /// provider's cache; billed at a reduced tier). `None` for providers
    /// without caching or pre-cache-aware captures.
    #[serde(default)]
    pub cache_read_tokens: Option<u32>,
    /// Optional cache-creation token count (prompt tokens written to the
    /// provider's cache; billed at a premium tier). `None` as above.
    #[serde(default)]
    pub cache_creation_tokens: Option<u32>,
    /// Recorded `finish_reason` of the response — lets a re-drive double
    /// (ADR-048) reconstruct whether the call ended in `Stop` or `ToolCalls`,
    /// so the agent loop takes the same branch. `None` for pre-2.0 captures.
    #[serde(default)]
    #[cfg_attr(feature = "schemars", schemars(with = "Option<String>"))]
    pub finish_reason: Option<FinishReason>,
    /// Recorded assistant tool calls — replayed so the loop dispatches the same
    /// tools. Empty for `Stop` responses and pre-2.0 captures.
    #[serde(default)]
    #[cfg_attr(feature = "schemars", schemars(with = "Vec<serde_json::Value>"))]
    pub tool_calls: Vec<ToolCall>,
    /// Fingerprint of the request that produced this call
    /// ([`crate::request_fingerprint`]); the re-drive double keys responses on
    /// it. `None` for pre-2.0 captures (double falls back to order).
    #[serde(default)]
    pub request_fingerprint: Option<String>,
}

impl LlmIo {
    /// A recorded call with the minimum replayable fields; enrich with the
    /// setters below.
    pub fn new(prompt: impl Into<String>, completion: impl Into<String>) -> Self {
        Self {
            prompt: prompt.into(),
            completion: completion.into(),
            ..Self::default()
        }
    }

    /// Set the model name (populates `Step.name` on projection).
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// Set the cost-accounting sidecar (provider + token split) used by
    /// [`crate::projector::project_with_price_table`].
    pub fn with_cost(
        mut self,
        provider: impl Into<String>,
        prompt_tokens: u32,
        completion_tokens: u32,
    ) -> Self {
        self.provider = Some(provider.into());
        self.prompt_tokens = Some(prompt_tokens);
        self.completion_tokens = Some(completion_tokens);
        self
    }

    /// Set the prompt-cache token split (read = served from cache, creation =
    /// written to cache) so [`crate::projector::project_with_price_table`] can
    /// price the cache tiers distinctly from the base prompt rate.
    pub fn with_cache_tokens(mut self, cache_read_tokens: u32, cache_creation_tokens: u32) -> Self {
        self.cache_read_tokens = Some(cache_read_tokens);
        self.cache_creation_tokens = Some(cache_creation_tokens);
        self
    }

    /// Set the recorded response shape so a re-drive can reconstruct the
    /// assistant decision (ADR-048).
    pub fn with_response_shape(
        mut self,
        finish_reason: FinishReason,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        self.finish_reason = Some(finish_reason);
        self.tool_calls = tool_calls;
        self
    }

    /// Set the request fingerprint the re-drive double keys on (ADR-048).
    pub fn with_request_fingerprint(mut self, fingerprint: impl Into<String>) -> Self {
        self.request_fingerprint = Some(fingerprint.into());
        self
    }
}

/// Aggregate token usage across all `LlmCall` steps in a run.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Usage {
    /// Total prompt tokens.
    pub prompt_tokens: u32,
    /// Total completion tokens.
    pub completion_tokens: u32,
}

/// Optional cost estimate for the run, in USD.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct Cost {
    /// Base (uncached) prompt cost in USD.
    pub prompt_usd: f64,
    /// Completion cost in USD.
    pub completion_usd: f64,
    /// Prompt-cache cost in USD (cache reads + cache creation, billed at their
    /// own tiers). 0 when no cache tokens were used or priced. `#[serde(default)]`
    /// keeps pre-cache `Cost` JSON deserializable (the field reads back as 0).
    #[serde(default)]
    pub cache_usd: f64,
    /// Total cost in USD: `prompt_usd + completion_usd + cache_usd`.
    pub total_usd: f64,
}

impl Cost {
    /// Build a cost from its tier components; `total_usd` is their sum. Prefer
    /// this over a struct literal — `Cost` is `#[non_exhaustive]`, so external
    /// crates cannot construct it directly.
    pub fn new(prompt_usd: f64, completion_usd: f64, cache_usd: f64) -> Self {
        Self {
            prompt_usd,
            completion_usd,
            cache_usd,
            total_usd: prompt_usd + completion_usd + cache_usd,
        }
    }
}

/// The aggregate RunLog view per spec §8.2.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RunLog {
    /// Run identifier.
    #[cfg_attr(feature = "schemars", schemars(with = "String"))]
    pub run_id: RunId,
    /// Agent name.
    pub agent: String,
    /// First-event timestamp.
    pub started_at: DateTime<Utc>,
    /// Last-event timestamp, populated when the run is no longer `Running`.
    pub finished_at: Option<DateTime<Utc>>,
    /// Lifecycle status.
    pub status: RunStatus,
    /// Ordered list of steps.
    pub steps: Vec<Step>,
    /// Aggregate usage.
    pub tokens: Usage,
    /// Optional cost estimate; `None` until a pricing table is wired (Wave 7+ carryover).
    pub cost_estimate: Option<Cost>,
}

mod duration_millis {
    use serde::{Deserialize, Deserializer, Serializer};
    use std::time::Duration;

    pub fn serialize<S: Serializer>(d: &Duration, s: S) -> Result<S::Ok, S::Error> {
        s.serialize_u64(d.as_millis() as u64)
    }
    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
        let ms = u64::deserialize(d)?;
        Ok(Duration::from_millis(ms))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use klieo_core::ids::RunId;
    use std::time::Duration;

    #[test]
    fn llm_io_roundtrips_structured_fields() {
        let io = LlmIo::new("p", "c")
            .with_model("gpt-4o")
            .with_response_shape(
                FinishReason::ToolCalls,
                vec![ToolCall::new(
                    "call_1",
                    "search",
                    serde_json::json!({"q": "x"}),
                )],
            )
            .with_request_fingerprint("abc123");
        let json = serde_json::to_string(&io).unwrap();
        let back: LlmIo = serde_json::from_str(&json).unwrap();
        assert_eq!(back, io);
        assert_eq!(back.finish_reason, Some(FinishReason::ToolCalls));
        assert_eq!(back.tool_calls.len(), 1);
        assert_eq!(back.request_fingerprint.as_deref(), Some("abc123"));
    }

    #[test]
    fn with_cost_sets_provider_and_token_split() {
        let io = LlmIo::new("p", "c").with_cost("openai", 100, 50);
        assert_eq!(io.provider.as_deref(), Some("openai"));
        assert_eq!(io.prompt_tokens, Some(100));
        assert_eq!(io.completion_tokens, Some(50));
    }

    #[test]
    fn cache_tokens_default_to_none_and_set_via_builder() {
        let bare = LlmIo::new("p", "c");
        assert_eq!(bare.cache_read_tokens, None);
        assert_eq!(bare.cache_creation_tokens, None);

        let cached = LlmIo::new("p", "c").with_cache_tokens(100, 8);
        assert_eq!(cached.cache_read_tokens, Some(100));
        assert_eq!(cached.cache_creation_tokens, Some(8));
    }

    #[test]
    fn old_record_without_new_fields_deserializes() {
        // A pre-2.0 serialized LlmIo lacks finish_reason / tool_calls /
        // request_fingerprint; #[serde(default)] must fill them.
        let legacy = r#"{"prompt":"p","completion":"c","model":null}"#;
        let io: LlmIo = serde_json::from_str(legacy).unwrap();
        assert_eq!(io.prompt, "p");
        assert_eq!(io.finish_reason, None);
        assert!(io.tool_calls.is_empty());
        assert_eq!(io.request_fingerprint, None);
    }

    #[test]
    fn run_status_serialises_snake_case() {
        let json = serde_json::to_string(&RunStatus::Running).unwrap();
        assert_eq!(json, "\"running\"");
    }

    #[test]
    fn step_kind_round_trip() {
        for k in [StepKind::LlmCall, StepKind::ToolCall] {
            let json = serde_json::to_string(&k).unwrap();
            let back: StepKind = serde_json::from_str(&json).unwrap();
            assert_eq!(k, back);
        }
    }

    #[test]
    fn step_round_trip() {
        let step = Step {
            idx: 0,
            kind: StepKind::ToolCall,
            name: Some("calc".into()),
            prompt_tokens: Some(3),
            completion_tokens: Some(9),
            cost_usd: None,
            input: serde_json::json!({"a": 1}),
            output: serde_json::json!({"b": 2}),
            error: None,
            latency: Duration::from_millis(42),
            span_id: Some("span-xyz".into()),
        };
        let json = serde_json::to_string(&step).unwrap();
        let back: Step = serde_json::from_str(&json).unwrap();
        assert_eq!(step.idx, back.idx);
        assert_eq!(step.latency, back.latency);
        assert_eq!(step.name, back.name);
    }

    #[test]
    fn run_log_round_trip_empty_steps() {
        let log = RunLog {
            run_id: RunId::new(),
            agent: "writer".into(),
            started_at: Utc::now(),
            finished_at: None,
            status: RunStatus::Running,
            steps: vec![],
            tokens: Usage::default(),
            cost_estimate: None,
        };
        let json = serde_json::to_string(&log).unwrap();
        assert!(json.contains("\"steps\":[]"));
        let _back: RunLog = serde_json::from_str(&json).unwrap();
    }

    #[test]
    fn usage_default_is_zero() {
        let u = Usage::default();
        assert_eq!(u.prompt_tokens, 0);
        assert_eq!(u.completion_tokens, 0);
    }

    #[test]
    fn cost_round_trip() {
        let c = Cost {
            prompt_usd: 0.001,
            completion_usd: 0.002,
            cache_usd: 0.0004,
            total_usd: 0.0034,
        };
        let json = serde_json::to_string(&c).unwrap();
        let back: Cost = serde_json::from_str(&json).unwrap();
        assert!((back.total_usd - 0.0034).abs() < 1e-9);
        assert!((back.cache_usd - 0.0004).abs() < 1e-9);
    }

    #[test]
    fn cost_deserializes_legacy_json_without_cache_field() {
        // Pre-cache `Cost` JSON lacks `cache_usd`; `#[serde(default)]` must keep
        // it deserializable (reads back as 0).
        let legacy = r#"{"prompt_usd":0.001,"completion_usd":0.002,"total_usd":0.003}"#;
        let c: Cost = serde_json::from_str(legacy).unwrap();
        assert!(c.cache_usd.abs() < 1e-12);
        assert!((c.total_usd - 0.003).abs() < 1e-9);
    }
}