klieo-runlog 3.2.0

//! Live [`CaptureSink`] that accumulates replayable LLM I/O (ADR-048).
//!
//! Wire one into `RunOptions::with_capture_sink`, run the agent, then
//! `Capture::from_runlog(run_log, sink.drain())` — the single production
//! populator of the previously-empty `Capture.llm_io` sidecar.

use std::sync::Mutex;

use klieo_core::llm::{ChatRequest, ChatResponse, Message, Role};
use klieo_core::runtime::CaptureSink;

use crate::fingerprint::request_fingerprint;
use crate::types::LlmIo;

/// Buffers one [`LlmIo`] per LLM call during a run; [`Self::drain`] hands the
/// recording to [`Capture::from_runlog`](crate::Capture::from_runlog).
#[derive(Default)]
pub struct RunLogCaptureSink {
    llm_io: Mutex<Vec<LlmIo>>,
    provider: Option<String>,
}

impl RunLogCaptureSink {
    /// Empty sink, ready to record; wire into `RunOptions::with_capture_sink`.
    pub fn new() -> Self {
        Self::default()
    }

    /// Record `provider` on every captured call so the projector can price a
    /// live run. [`crate::projector::project_with_price_table`] only attaches a
    /// USD cost when the [`LlmIo`] carries provider + model + token counts;
    /// without a provider here, live captures stay unpriced.
    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
        self.provider = Some(provider.into());
        self
    }

    /// Drains the recording in call order and re-empties the sink, so a second
    /// call after a run returns nothing rather than re-yielding stale I/O.
    pub fn drain(&self) -> Vec<LlmIo> {
        std::mem::take(
            &mut *self
                .llm_io
                .lock()
                .expect("RunLogCaptureSink mutex poisoned"),
        )
    }
}

/// The recorded prompt is the most recent `User` turn in the request (empty if
/// none) — a deliberate single-string reduction of the multi-message request.
fn last_user_content(messages: &[Message]) -> String {
    messages
        .iter()
        .rev()
        .find(|m| m.role == Role::User)
        .map(|m| m.content.clone())
        .unwrap_or_default()
}

impl CaptureSink for RunLogCaptureSink {
    fn record_llm_call(&self, request: &ChatRequest, response: &ChatResponse) {
        let mut io = LlmIo::new(
            last_user_content(&request.messages),
            response.message.content.clone(),
        )
        .with_response_shape(response.finish_reason, response.message.tool_calls.clone())
        .with_request_fingerprint(request_fingerprint(request))
        .with_model(response.model.clone());
        if let Some(provider) = &self.provider {
            io = io
                .with_cost(
                    provider.clone(),
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                )
                .with_cache_tokens(
                    response.usage.cache_read_tokens,
                    response.usage.cache_creation_tokens,
                );
        }
        self.llm_io
            .lock()
            .expect("RunLogCaptureSink mutex poisoned")
            .push(io);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::llm::{FinishReason, Message, Role, Usage};

    fn msg(role: Role, content: &str) -> Message {
        Message {
            role,
            content: content.into(),
            tool_calls: vec![],
            tool_call_id: None,
        }
    }

    fn response(content: &str, model: &str) -> ChatResponse {
        ChatResponse::new(
            msg(Role::Assistant, content),
            Usage::default(),
            FinishReason::Stop,
            model,
        )
    }

    #[test]
    fn last_user_content_picks_the_most_recent_user_turn() {
        let messages = [
            msg(Role::User, "first"),
            msg(Role::Assistant, "reply"),
            msg(Role::User, "second"),
        ];
        assert_eq!(last_user_content(&messages), "second");
    }

    #[test]
    fn last_user_content_is_empty_when_no_user_message() {
        let messages = [msg(Role::System, "sys"), msg(Role::Assistant, "a")];
        assert_eq!(last_user_content(&messages), "");
    }

    fn response_with_usage(content: &str, model: &str, usage: Usage) -> ChatResponse {
        ChatResponse::new(
            msg(Role::Assistant, content),
            usage,
            FinishReason::Stop,
            model,
        )
    }

    #[test]
    fn with_provider_records_cost_sidecar_for_pricing() {
        let sink = RunLogCaptureSink::new().with_provider("ollama");
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "hi")]),
            &response_with_usage("ok", "ollama:qwen2.5:14b", Usage::new(100, 50)),
        );
        let io = sink.drain();
        assert_eq!(io[0].provider.as_deref(), Some("ollama"));
        assert_eq!(io[0].prompt_tokens, Some(100));
        assert_eq!(io[0].completion_tokens, Some(50));
    }

    #[test]
    fn with_provider_records_cache_tokens_for_pricing() {
        let sink = RunLogCaptureSink::new().with_provider("anthropic");
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "hi")]),
            &response_with_usage(
                "ok",
                "claude-sonnet-4-6",
                Usage::new(10, 50).with_cache_tokens(100, 8),
            ),
        );
        let io = sink.drain();
        assert_eq!(io[0].cache_read_tokens, Some(100));
        assert_eq!(io[0].cache_creation_tokens, Some(8));
    }

    #[test]
    fn without_provider_leaves_cost_sidecar_unset() {
        let sink = RunLogCaptureSink::new();
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "hi")]),
            &response_with_usage("ok", "ollama:qwen2.5:14b", Usage::new(100, 50)),
        );
        let io = sink.drain();
        assert_eq!(io[0].provider, None);
        assert_eq!(io[0].prompt_tokens, None);
    }

    #[test]
    fn record_then_drain_yields_the_call_then_empties() {
        let sink = RunLogCaptureSink::new();
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "hi")]),
            &response("ok", "ollama:qwen2.5:14b"),
        );
        let first = sink.drain();
        assert_eq!(first.len(), 1);
        assert_eq!(first[0].prompt, "hi");
        assert_eq!(first[0].completion, "ok");
        assert_eq!(first[0].finish_reason, Some(FinishReason::Stop));
        assert!(first[0].request_fingerprint.is_some());
        assert!(sink.drain().is_empty(), "second drain yields nothing");
    }

    #[test]
    fn record_threads_model_into_llm_io() {
        let sink = RunLogCaptureSink::new();
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "hi")]),
            &response("ok", "ollama:qwen2.5:14b"),
        );
        let io = sink.drain();
        assert_eq!(
            io[0].model.as_deref(),
            Some("ollama:qwen2.5:14b"),
            "model must flow from ChatResponse into LlmIo"
        );
    }
}