klieo-runlog 3.3.1

//! `Episode` stream → `RunLog` aggregate projection.
//!
//! Pure function. No I/O. Walks a `&[Episode]` and assembles the spec §8.2
//! aggregate view. Intentionally tolerant — missing `Started`/`Completed`
//! episodes do not panic; they leave the corresponding fields with sensible
//! defaults (`status = Running`, `finished_at = None`).
//!
//! ## Live-API carryovers
//!
//! `klieo_core::Episode::LlmCall` carries only `tokens` + `latency_ms`; it
//! does not expose `prompt`/`completion`/`model_name`, nor a per-event
//! timestamp. Until those fields land on the `klieo-core` trait surface
//! (Wave 7+ carryover; Plan #11 freeze in effect), the projector records
//! `input` / `output` as `Value::Null`, accumulates all reported tokens into
//! `Usage::completion_tokens`, and falls back to `Utc::now()` for
//! `started_at` / `finished_at`.
//!
//! ## Replay-friendly projection via [`project_with_llm_io`]
//!
//! Callers that retain the original prompts + completions can pass a
//! parallel `&[LlmIo]` slice to [`project_with_llm_io`]; the i-th [`LlmIo`]
//! populates the i-th `Episode::LlmCall`'s `input` (`{"prompt": ...}`),
//! `output` (the completion text), and `name` (the model). This makes the
//! projected `RunLog` round-trip cleanly through [`crate::replay::replay`].

use crate::pricing::PriceTable;
use crate::types::{Cost, LlmIo, RunLog, RunStatus, Step, StepKind, Usage};
use chrono::Utc;
use klieo_core::ids::RunId;
use klieo_core::memory::{Episode, ToolResult};
use std::time::Duration;

/// Project a `Vec<Episode>` into a `RunLog` aggregate.
///
/// `agent` is the agent name; pass the same string the agent was registered with.
/// The function never fails — projection of incomplete or malformed episode
/// streams produces a best-effort `RunLog` (with `status = Running` and
/// `finished_at = None` when the stream is open-ended).
///
/// LLM steps end up with `input` / `output` set to `Value::Null` because
/// `Episode::LlmCall` does not carry the prompt or completion text. Use
/// [`project_with_llm_io`] when those values are available alongside.
pub fn project(run_id: RunId, agent: impl Into<String>, episodes: &[Episode]) -> RunLog {
    project_with_llm_io(run_id, agent, episodes, &[])
}

/// Project a `Vec<Episode>` into a `RunLog`, populating LLM `input` / `output`
/// from a parallel `&[LlmIo]` sidecar.
///
/// The i-th `Episode::LlmCall` in `episodes` is paired with `llm_io[i]`. If
/// `llm_io` has fewer entries than there are `Episode::LlmCall` events, the
/// remaining LLM steps fall back to `Value::Null` (matching [`project`]).
/// Extra `LlmIo` entries beyond the last `Episode::LlmCall` are ignored.
///
/// Leaves `cost_estimate = None`. Use [`project_with_price_table`] to compute
/// USD costs from a [`PriceTable`].
pub fn project_with_llm_io(
    run_id: RunId,
    agent: impl Into<String>,
    episodes: &[Episode],
    llm_io: &[LlmIo],
) -> RunLog {
    project_with_price_table(run_id, agent, episodes, llm_io, &PriceTable::new())
}

/// Project a `Vec<Episode>` into a `RunLog`, populating LLM `input` / `output`
/// from a parallel `&[LlmIo]` sidecar **and** computing a per-run USD
/// [`Cost`] estimate from a caller-supplied [`PriceTable`].
///
/// Pairing rules match [`project_with_llm_io`]: the i-th `Episode::LlmCall`
/// is paired with `llm_io[i]`.
///
/// For each pair where the [`LlmIo`] carries `provider`, `model`,
/// `prompt_tokens`, **and** `completion_tokens`, `prices.lookup(provider,
/// model)` is consulted. On a hit, prompt + completion USD are accumulated
/// into the run-level totals. If at least one LLM step contributed a
/// successful lookup (zero-rate hits count), `log.cost_estimate` is set to
/// `Some(Cost { … })`; otherwise it stays `None`.
///
/// Steps with missing provider/model/token data are silently skipped for
/// cost accounting. They still appear as `Step`s in the projected `RunLog`.
pub fn project_with_price_table(
    run_id: RunId,
    agent: impl Into<String>,
    episodes: &[Episode],
    llm_io: &[LlmIo],
    prices: &PriceTable,
) -> RunLog {
    let now = Utc::now();
    let mut log = RunLog {
        run_id,
        agent: agent.into(),
        started_at: now,
        finished_at: None,
        status: RunStatus::Running,
        steps: Vec::new(),
        tokens: Usage::default(),
        cost_estimate: None,
    };

    let mut idx: u32 = 0;
    let mut llm_idx: usize = 0;
    let mut prompt_usd_acc: f64 = 0.0;
    let mut completion_usd_acc: f64 = 0.0;
    let mut cache_usd_acc: f64 = 0.0;
    let mut any_priced_step: bool = false;
    for ep in episodes {
        match ep {
            Episode::Started { .. } => {
                // Episode::Started doesn't carry a timestamp on the trait surface
                // today (see Wave 7+ carryover). Use Utc::now() as a fallback.
                log.started_at = now;
            }
            Episode::LlmCall {
                tokens,
                latency_ms,
                provider,
                model,
                prompt_tokens,
                completion_tokens,
            } => {
                let io = llm_io.get(llm_idx);
                let cost = io.and_then(|io| step_cost(io, prices)).or_else(|| {
                    episode_cost(
                        provider.as_deref(),
                        model.as_deref(),
                        *prompt_tokens,
                        *completion_tokens,
                        prices,
                    )
                });
                let (input, output, name) = match io {
                    Some(io) => (
                        serde_json::json!({ "prompt": io.prompt }),
                        serde_json::Value::String(io.completion.clone()),
                        // Prefer LlmIo sidecar's model name when present
                        // (it's typically the run-time-bound model); fall
                        // back to the value Episode carries.
                        io.model.clone().or_else(|| model.clone()),
                    ),
                    None => (
                        serde_json::Value::Null,
                        serde_json::Value::Null,
                        model.clone(),
                    ),
                };
                log.steps.push(Step {
                    idx,
                    kind: StepKind::LlmCall,
                    name,
                    prompt_tokens: *prompt_tokens,
                    completion_tokens: *completion_tokens,
                    cost_usd: cost.as_ref().map(|c| c.total_usd),
                    input,
                    output,
                    error: None,
                    latency: Duration::from_millis(u64::from(*latency_ms)),
                    span_id: None,
                });
                // Prefer the prompt/completion split: from the Episode when it
                // carries one, else from the paired LlmIo sidecar. Only when no
                // split is available anywhere does the scalar total fall back to
                // completion_tokens (it would otherwise misattribute prompt
                // tokens as completion).
                let split = match (*prompt_tokens, *completion_tokens) {
                    (Some(pt), Some(ct)) => Some((pt, ct)),
                    _ => io.and_then(|io| io.prompt_tokens.zip(io.completion_tokens)),
                };
                match split {
                    Some((pt, ct)) => {
                        log.tokens.prompt_tokens = log.tokens.prompt_tokens.saturating_add(pt);
                        log.tokens.completion_tokens =
                            log.tokens.completion_tokens.saturating_add(ct);
                    }
                    None => {
                        log.tokens.completion_tokens =
                            log.tokens.completion_tokens.saturating_add(*tokens);
                    }
                }

                // Cost accounting — only when the step was priced above.
                if let Some(c) = &cost {
                    prompt_usd_acc += c.prompt_usd;
                    completion_usd_acc += c.completion_usd;
                    cache_usd_acc += c.cache_usd;
                    any_priced_step = true;
                }

                idx += 1;
                llm_idx += 1;
            }
            Episode::ToolCall { name, args, result } => {
                let (output, error) = match result {
                    ToolResult::Ok { value } => (value.clone(), None),
                    ToolResult::Err { message } => (serde_json::Value::Null, Some(message.clone())),
                };
                log.steps.push(Step {
                    idx,
                    kind: StepKind::ToolCall,
                    name: Some(name.clone()),
                    prompt_tokens: None,
                    completion_tokens: None,
                    cost_usd: None,
                    input: args.clone(),
                    output,
                    error,
                    latency: Duration::ZERO, // Episode::ToolCall doesn't carry latency yet.
                    span_id: None,
                });
                idx += 1;
            }
            Episode::BusPublish { .. } | Episode::BusReceive { .. } => {
                // Bus episodes are not Steps in the spec §8.2 model.
            }
            Episode::Completed => {
                log.status = RunStatus::Completed;
                log.finished_at = Some(Utc::now());
            }
            Episode::Failed { .. } => {
                log.status = RunStatus::Failed;
                log.finished_at = Some(Utc::now());
            }
            Episode::SummaryCheckpoint {
                input_message_count,
                summary_chars,
                latency_ms,
                tokens,
            } => {
                // Project as a dedicated `StepKind::SummaryCheckpoint`
                // so cost / latency dashboards can isolate summarizer
                // overhead from substantive reasoning. `input` /
                // `output` carry the lightweight metrics the variant
                // already exposes — the actual transcript / summary
                // text is intentionally not stored here.
                log.steps.push(Step {
                    idx,
                    kind: StepKind::SummaryCheckpoint,
                    name: Some("summary".into()),
                    prompt_tokens: None,
                    completion_tokens: None,
                    cost_usd: None,
                    input: serde_json::json!({
                        "input_message_count": input_message_count,
                    }),
                    output: serde_json::json!({
                        "summary_chars": summary_chars,
                    }),
                    error: None,
                    latency: Duration::from_millis(u64::from(*latency_ms)),
                    span_id: None,
                });
                // Summarizer tokens accumulate separately into
                // `completion_tokens` for now (same constraint as
                // LlmCall — single-bucket usage on the trait surface).
                log.tokens.completion_tokens = log.tokens.completion_tokens.saturating_add(*tokens);
                idx += 1;
            }
            Episode::Ops(payload) => {
                // Project as `StepKind::OpsEvent`. The `kind` field
                // from the OpsEvent JSON tag surfaces as `Step.name`
                // for dashboards; the full payload is stored in
                // `Step.input` for downstream consumption (e.g. the
                // evidence-verifier and klieo-ops-evidence-verify).
                let kind_name = payload
                    .get("kind")
                    .and_then(|k| k.as_str())
                    .map(str::to_owned);
                log.steps.push(Step {
                    idx,
                    kind: StepKind::OpsEvent,
                    name: kind_name,
                    prompt_tokens: None,
                    completion_tokens: None,
                    cost_usd: None,
                    input: payload.clone(),
                    output: serde_json::Value::Null,
                    error: None,
                    latency: Duration::ZERO,
                    span_id: None,
                });
                idx += 1;
            }
            Episode::MemoryRecall {
                query,
                k,
                returned_fact_ids,
            } => {
                // Project as `StepKind::MemoryRecall` so the run view can
                // surface graphRAG recalls alongside LLM/tool steps.
                // `query` arrives already redacted + length-bounded by the
                // recall-recording wrapper (see `Episode::MemoryRecall`
                // doc), so it is safe to store verbatim in `Step.input`.
                log.steps.push(Step {
                    idx,
                    kind: StepKind::MemoryRecall,
                    name: Some("recall".to_string()),
                    prompt_tokens: None,
                    completion_tokens: None,
                    cost_usd: None,
                    input: serde_json::json!({ "query": query, "k": k }),
                    output: serde_json::json!({ "returned_fact_ids": returned_fact_ids }),
                    error: None,
                    latency: Duration::ZERO,
                    span_id: None,
                });
                idx += 1;
            }
            // The Episode enum is `#[non_exhaustive]`; future additive
            // variants project as no-ops until explicitly handled.
            _ => {}
        }
    }

    if any_priced_step {
        log.cost_estimate = Some(Cost::new(prompt_usd_acc, completion_usd_acc, cache_usd_acc));
    }

    log
}

/// Per-step cost lookup for [`project_with_price_table`].
///
/// Returns `Some(Cost)` only when all four cost-accounting fields on `io`
/// (`provider`, `model`, `prompt_tokens`, `completion_tokens`) are populated
/// AND `prices.lookup` hits. Any missing field or table miss yields `None` —
/// the caller treats that as "skip this step for cost accounting".
fn step_cost(io: &LlmIo, prices: &PriceTable) -> Option<Cost> {
    let provider = io.provider.as_deref()?;
    let model = io.model.as_deref()?;
    let pt = io.prompt_tokens?;
    let ct = io.completion_tokens?;
    let rates = prices.lookup(provider, model)?;
    // Cache tokens are optional (only providers with prompt caching report
    // them); absent counts price as 0 at the cache tiers.
    let cache_read = io.cache_read_tokens.unwrap_or(0);
    let cache_creation = io.cache_creation_tokens.unwrap_or(0);
    Some(rates.cost_for_cached(pt, ct, cache_read, cache_creation))
}

/// Per-step cost from the episode's own `provider`/`model`/token split — the
/// fallback when no priced `LlmIo` sidecar exists (a live run with no
/// `CaptureSink`). Same all-or-nothing contract as [`step_cost`]; the
/// `LlmIo`-derived cost takes precedence, so this never double-counts.
fn episode_cost(
    provider: Option<&str>,
    model: Option<&str>,
    prompt_tokens: Option<u32>,
    completion_tokens: Option<u32>,
    prices: &PriceTable,
) -> Option<Cost> {
    let rates = prices.lookup(provider?, model?)?;
    Some(rates.cost_for(prompt_tokens?, completion_tokens?))
}