klieo-runlog 3.3.1

//! 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::{Arc, Mutex};

use klieo_core::llm::{ChatRequest, ChatResponse, Message, Role};
use klieo_core::redact::AuditRedactor;
use klieo_core::runtime::CaptureSink;
use serde_json::Value;

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).
///
/// When a redactor is wired via [`Self::with_redactor`], the recorded prompt
/// and completion are scrubbed of PII *before* they enter the buffer, so a
/// persisted or exported [`Capture`](crate::Capture) never carries raw user
/// text. Without a redactor the text is stored verbatim — required for
/// byte-exact replay, but not safe to persist under a compliance regime that
/// forbids raw-PII sinks (EU AI Act Art. 10 data governance).
#[derive(Default)]
pub struct RunLogCaptureSink {
    llm_io: Mutex<Vec<LlmIo>>,
    provider: Option<String>,
    redactor: Option<Arc<dyn AuditRedactor>>,
}

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
    }

    /// Redact PII from every recorded prompt and completion before it enters
    /// the buffer. Wire this for any deployment that persists or exports
    /// captures under a data-governance obligation; leave it unset only for
    /// throwaway replay where byte-exact prompt text is required.
    pub fn with_redactor(mut self, redactor: Arc<dyn AuditRedactor>) -> Self {
        self.redactor = Some(redactor);
        self
    }

    /// Scrub PII from `text` when a redactor is wired, else return it verbatim.
    /// A redactor that returns a non-string projection (e.g. the fail-closed
    /// digest fallback) is serialized to its JSON form — still PII-free.
    fn redact_text(&self, text: String) -> String {
        let Some(redactor) = &self.redactor else {
            return text;
        };
        match redactor.redact(&Value::String(text)) {
            Value::String(redacted) => redacted,
            other => other.to_string(),
        }
    }

    /// 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) {
        // ponytail: redact prompt + completion only — the named PII sink. The
        // assistant tool-call arguments are left intact because the re-drive
        // double dispatches them verbatim on replay; redacting them would
        // corrupt reproduction. Scrub those upstream if they carry PII.
        let mut io = LlmIo::new(
            self.redact_text(last_user_content(&request.messages)),
            self.redact_text(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 redactor_scrubs_pii_from_prompt_and_completion() {
        use klieo_ops::redactor::DefaultRedactor;
        let sink = RunLogCaptureSink::new().with_redactor(Arc::new(DefaultRedactor::new()));
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "mail me at alice@example.com")]),
            &response("reply to bob@example.com then", "ollama:qwen2.5:14b"),
        );
        let io = sink.drain();
        assert_eq!(io[0].prompt, "mail me at [REDACTED:EMAIL]");
        assert_eq!(io[0].completion, "reply to [REDACTED:EMAIL] then");
    }

    #[test]
    fn redact_text_serializes_non_string_redactor_output() {
        use klieo_core::redact::AuditRedactor;
        use serde_json::{json, Value};

        struct DigestRedactor;
        impl AuditRedactor for DigestRedactor {
            fn redact(&self, _value: &Value) -> Value {
                json!({"redacted": "sha256", "digest": "deadbeef"})
            }
        }

        let sink = RunLogCaptureSink::new().with_redactor(Arc::new(DigestRedactor));
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "alice@example.com")]),
            &response("bob@example.com", "m"),
        );
        let io = sink.drain();
        assert!(!io[0].prompt.contains("alice@example.com"), "PII gone");
        assert!(
            io[0].prompt.contains("redacted"),
            "non-string output serialized to JSON"
        );
        assert!(!io[0].completion.contains("bob@example.com"), "PII gone");
    }

    #[test]
    fn tool_call_arguments_are_not_redacted_even_with_redactor() {
        use klieo_core::llm::{FinishReason, ToolCall};
        use klieo_ops::redactor::DefaultRedactor;

        let sink = RunLogCaptureSink::new().with_redactor(Arc::new(DefaultRedactor::new()));
        let args = serde_json::json!({ "to": "alice@example.com" });
        let mut assistant = msg(Role::Assistant, "sent");
        assistant.tool_calls = vec![ToolCall::new("1", "send_email", args.clone())];
        let resp = ChatResponse::new(assistant, Usage::default(), FinishReason::ToolCalls, "m");
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "email alice")]),
            &resp,
        );

        let io = sink.drain();
        assert_eq!(io[0].tool_calls.len(), 1);
        assert_eq!(
            io[0].tool_calls[0].args, args,
            "tool-call args are recorded verbatim for replay even with a redactor wired"
        );
    }

    #[test]
    fn without_redactor_keeps_prompt_and_completion_verbatim() {
        let sink = RunLogCaptureSink::new();
        sink.record_llm_call(
            &ChatRequest::new(vec![msg(Role::User, "mail me at alice@example.com")]),
            &response("reply to bob@example.com", "ollama:qwen2.5:14b"),
        );
        let io = sink.drain();
        assert_eq!(io[0].prompt, "mail me at alice@example.com");
        assert_eq!(io[0].completion, "reply to bob@example.com");
    }

    #[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"
        );
    }
}