klieo-runlog 0.41.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 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,
                model,
                prompt_tokens,
                completion_tokens,
                ..
            } => {
                let io = llm_io.get(llm_idx);
                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,
                    input,
                    output,
                    error: None,
                    latency: Duration::from_millis(u64::from(*latency_ms)),
                    span_id: None,
                });
                // Use the prompt/completion split when Episode carries it;
                // otherwise accumulate the total into completion_tokens
                // (legacy behaviour preserved for callers that haven't
                // populated the new fields yet).
                match (*prompt_tokens, *completion_tokens) {
                    (Some(pt), Some(ct)) => {
                        log.tokens.prompt_tokens = log.tokens.prompt_tokens.saturating_add(pt);
                        log.tokens.completion_tokens =
                            log.tokens.completion_tokens.saturating_add(ct);
                    }
                    _ => {
                        log.tokens.completion_tokens =
                            log.tokens.completion_tokens.saturating_add(*tokens);
                    }
                }

                // Cost accounting — only when caller provides full info.
                if let Some(io) = io {
                    if let Some(c) = step_cost(io, prices) {
                        prompt_usd_acc += c.prompt_usd;
                        completion_usd_acc += c.completion_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()),
                    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()),
                    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,
                    input: payload.clone(),
                    output: serde_json::Value::Null,
                    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 {
            prompt_usd: prompt_usd_acc,
            completion_usd: completion_usd_acc,
            total_usd: prompt_usd_acc + completion_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)?;
    Some(rates.cost_for(pt, ct))
}