openlatch-client 0.1.18

//! Capture from the terminal usage chunk (D-10, D-15) + pricing-input capture
//! (F-31) + the `unknown_model` gap (D-12).
//!
//! The usage lives in the terminal `message_delta` SSE event (streaming) or the
//! response body (non-streaming). It is read **in passing** while the stream
//! forwards — never held. The scan is chunk-by-chunk and never buffers more than
//! the current chunk; a usage line split across a chunk boundary is simply
//! missed, which correctly degrades to the `tokenizer_estimated` path rather than
//! stalling the stream (REJECTED: collecting the whole SSE stream then parsing —
//! it buffers and kills TTFT).

use serde_json::Value;

/// Frozen enum `cost_basis = provider_reported | tokenizer_estimated | interpolated`.
///
/// A property of **capture**, not of pricing. `interpolated` is produced
/// platform-side (F-36); the client emits only the first two.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CostBasis {
    /// The terminal usage chunk arrived cleanly (2xx).
    ProviderReported,
    /// The stream was interrupted/unparseable — token counts are a local estimate.
    TokenizerEstimated,
    /// Tokens are provider-reported but no pricebook row matched (platform-set).
    Interpolated,
}

impl CostBasis {
    pub fn as_str(&self) -> &'static str {
        match self {
            CostBasis::ProviderReported => "provider_reported",
            CostBasis::TokenizerEstimated => "tokenizer_estimated",
            CostBasis::Interpolated => "interpolated",
        }
    }
}

/// Frozen enum `capture_gap = unknown_wire_format | unknown_model | provider_error | stream_interrupted`.
/// Nullable on the wire — set only when capture was incomplete.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CaptureGap {
    /// Body opaque / over the 32 MB ceiling (set by the forwarder, plan 01).
    UnknownWireFormat,
    /// The request model is not in the known (D-21) set.
    UnknownModel,
    /// The provider returned a non-2xx (F-36) — event emitted, tokens zero.
    ProviderError,
    /// The response stream ended before a usable terminal usage chunk.
    StreamInterrupted,
}

impl CaptureGap {
    pub fn as_str(&self) -> &'static str {
        match self {
            CaptureGap::UnknownWireFormat => "unknown_wire_format",
            CaptureGap::UnknownModel => "unknown_model",
            CaptureGap::ProviderError => "provider_error",
            CaptureGap::StreamInterrupted => "stream_interrupted",
        }
    }
}

/// The five raw token counts the client emits (C-3). **`input_tokens` is
/// post-last-breakpoint only — never the total.** Total input is
/// `input_tokens + cache_creation + cache_read` and is computed platform-side.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Usage {
    /// `gen_ai.usage.input_tokens` — post-last-breakpoint only.
    pub input_tokens: u64,
    /// `gen_ai.usage.cache_read.input_tokens`.
    pub cache_read: u64,
    /// `gen_ai.usage.cache_creation.input_tokens` — the sum of the two buckets.
    pub cache_write: u64,
    /// `ai.openlatch.cache.ephemeral_5m_input_tokens` (priced 1.25×).
    pub eph_5m: u64,
    /// `ai.openlatch.cache.ephemeral_1h_input_tokens` (priced 2×).
    pub eph_1h: u64,
    /// `gen_ai.usage.output_tokens`.
    pub output_tokens: u64,
}

impl Usage {
    /// Field-wise **max** merge. Anthropic splits usage across `message_start`
    /// (final input/cache, preliminary `output_tokens = 1`) and the terminal
    /// `message_delta` (final cumulative output), so each field takes the larger
    /// of the two — input/cache land once, output grows to its final value.
    fn merged_max(self, other: Usage) -> Usage {
        Usage {
            input_tokens: self.input_tokens.max(other.input_tokens),
            cache_read: self.cache_read.max(other.cache_read),
            cache_write: self.cache_write.max(other.cache_write),
            eph_5m: self.eph_5m.max(other.eph_5m),
            eph_1h: self.eph_1h.max(other.eph_1h),
            output_tokens: self.output_tokens.max(other.output_tokens),
        }
    }
}

/// Accumulates usage across streamed chunks. Anthropic splits usage across the
/// `message_start` event (input + cache fields, `output_tokens = 1`) and the
/// terminal `message_delta` event (final cumulative `output_tokens`), so fields
/// are merged by **max** — input/cache appear once (message_start), output grows
/// to its final value in message_delta. Non-streaming responses carry a single
/// top-level `usage` object, handled by the same merge.
#[derive(Clone, Copy, Debug, Default)]
pub struct UsageAccumulator {
    usage: Usage,
    /// True once any usage object has been observed (message_start, message_delta,
    /// or a non-streaming body) — used only for the "did this stream carry any
    /// usage at all" diagnostic, NOT for the cost-basis decision.
    seen: bool,
    /// True once a **terminal** usage object has been observed: a streaming
    /// `message_delta` (final cumulative output) or a complete non-streaming
    /// response body. `message_start` — which carries FINAL input/cache but a
    /// PRELIMINARY `output_tokens = 1` — deliberately does NOT set this. This flag
    /// (not `seen`) is what separates `provider_reported` from
    /// `tokenizer_estimated`: a stream that ends before `message_delta` is only
    /// partially measured and must fall back to a local estimate.
    terminal: bool,
}

impl UsageAccumulator {
    /// Scan one forwarded chunk for usage and merge whatever is found. Returns
    /// `true` if this chunk contributed usage. **Read-only over the chunk** — the
    /// bytes are never mutated and never retained.
    pub fn scan_chunk(&mut self, chunk: &[u8]) -> bool {
        match scan_usage(chunk) {
            Some(found) => {
                self.merge(found.usage);
                self.seen = true;
                if found.terminal {
                    self.terminal = true;
                }
                true
            }
            None => false,
        }
    }

    fn merge(&mut self, u: Usage) {
        self.usage = self.usage.merged_max(u);
    }

    /// True once a usage object has been observed at least once.
    pub fn has_usage(&self) -> bool {
        self.seen
    }

    /// True once a **terminal** usage object has been observed — a streaming
    /// `message_delta` (final cumulative output) or a complete non-streaming
    /// response body. `message_start` (final input/cache but a preliminary
    /// `output_tokens = 1`) does NOT set this, so a stream interrupted before the
    /// terminal chunk correctly reports "not fully measured" and degrades to a
    /// local estimate rather than emitting the preliminary output as final.
    pub fn is_terminal(&self) -> bool {
        self.terminal
    }

    /// The accumulated usage.
    pub fn usage(&self) -> Usage {
        self.usage
    }
}

/// The outcome of scanning one forwarded chunk: the merged usage found and
/// whether any of it came from a **terminal** usage object (a `message_delta` or
/// a non-streaming response body) rather than the preliminary `message_start`.
struct ScanResult {
    usage: Usage,
    terminal: bool,
}

/// Extract a `Usage` from one chunk, if it carries a usage object, and classify
/// whether the chunk carried terminal usage.
///
/// Handles both SSE (`data: {…}` lines, usage under `.usage` or `.message.usage`)
/// and a raw non-streaming JSON body (top-level `.usage`).
fn scan_usage(chunk: &[u8]) -> Option<ScanResult> {
    let text = std::str::from_utf8(chunk).ok()?;
    let mut best: Option<Usage> = None;
    let mut terminal = false;

    // SSE data lines first.
    for line in text.lines() {
        let line = line.trim_start();
        let payload = line.strip_prefix("data:").map(str::trim).unwrap_or(line);
        if !payload.starts_with('{') {
            continue;
        }
        // Cheap pre-filter before the serde parse: `usage` only appears in
        // `message_start` / `message_delta`, so skip the bulk `content_block_delta`
        // lines entirely rather than parse-and-throw-away. A line that does contain
        // the literal "usage" still parses exactly as before — zero behavior change.
        if !payload.contains("usage") {
            continue;
        }
        if let Ok(v) = serde_json::from_str::<Value>(payload) {
            if let Some((u, term)) = usage_and_terminal(&v) {
                best = Some(merge_pick(best, u));
                terminal |= term;
            }
        }
    }

    // Non-streaming: the whole chunk may be one JSON object with `.usage`. Guard
    // the parse on a leading `{` so a non-JSON chunk is never fed to serde (a bare
    // number/array/string could parse yet never carry `.usage`, so this is a pure
    // cost cut — zero behavior change).
    if best.is_none() && text.trim_start().starts_with('{') {
        if let Ok(v) = serde_json::from_str::<Value>(text.trim()) {
            if let Some((u, term)) = usage_and_terminal(&v) {
                best = Some(u);
                terminal |= term;
            }
        }
    }

    best.map(|usage| ScanResult { usage, terminal })
}

/// Prefer the usage object carrying the most signal (larger output/input),
/// merging field-wise by max so message_start + message_delta both contribute.
fn merge_pick(prev: Option<Usage>, cur: Usage) -> Usage {
    match prev {
        None => cur,
        Some(p) => p.merged_max(cur),
    }
}

/// Pull a `Usage` out of a parsed SSE/response value and classify whether it is
/// **terminal**.
///
/// - `message_start` carries usage under `.message.usage` with FINAL input/cache
///   but a PRELIMINARY `output_tokens = 1` → **not terminal**. The stream is not
///   fully measured until the terminal chunk arrives.
/// - A streaming `message_delta` (`.usage`, final cumulative output) and a
///   non-streaming response body (top-level `.usage`, all-final) are **terminal**.
///
/// The `type` discriminator is what distinguishes the two: only `message_start`
/// is treated as preliminary; every other value carrying a top-level `.usage`
/// (message_delta and the non-streaming body, which has no `message_start` type)
/// is a complete measurement.
fn usage_and_terminal(v: &Value) -> Option<(Usage, bool)> {
    if v.get("type").and_then(Value::as_str) == Some("message_start") {
        let u = v.get("message").and_then(|m| m.get("usage"))?;
        return Some((usage_fields(u), false));
    }
    let u = v.get("usage")?;
    Some((usage_fields(u), true))
}

/// Read the six raw token fields out of a `usage` object.
fn usage_fields(u: &Value) -> Usage {
    let cache_creation = u.get("cache_creation");
    let eph_5m = cache_creation
        .and_then(|c| c.get("ephemeral_5m_input_tokens"))
        .and_then(Value::as_u64)
        .unwrap_or(0);
    let eph_1h = cache_creation
        .and_then(|c| c.get("ephemeral_1h_input_tokens"))
        .and_then(Value::as_u64)
        .unwrap_or(0);

    Usage {
        input_tokens: u.get("input_tokens").and_then(Value::as_u64).unwrap_or(0),
        cache_read: u
            .get("cache_read_input_tokens")
            .and_then(Value::as_u64)
            .unwrap_or(0),
        cache_write: u
            .get("cache_creation_input_tokens")
            .and_then(Value::as_u64)
            .unwrap_or(0),
        eph_5m,
        eph_1h,
        output_tokens: u.get("output_tokens").and_then(Value::as_u64).unwrap_or(0),
    }
}

/// The pricing-input modifiers derived from the request (F-31). `batch` and
/// `fast_mode` are NOT-NULL wire booleans; `inference_geo` is nullable.
#[derive(Clone, Debug, Default)]
pub struct PricingInputs {
    pub batch: bool,
    pub fast_mode: bool,
    pub inference_geo: Option<String>,
}

/// Derive the pricing inputs from the request body + headers.
///
/// ⚠️ Conservative by design. `/v1/messages` (the captured path) is not a batch
/// endpoint, so `batch` is essentially always false (the PRD flags whether batch
/// traffic transits the listener at all as unverified). `fast_mode` and
/// `inference_geo` have **no confirmed wire source**; they default false/None and
/// are only set when an explicit, unambiguous signal is present.
pub fn derive_pricing_inputs(body: &Value, headers: &axum::http::HeaderMap) -> PricingInputs {
    // batch: only true on an explicit request-body flag (defensive — normally
    // false on /v1/messages).
    let batch = body.get("batch").and_then(Value::as_bool).unwrap_or(false);

    // fast_mode: Anthropic exposes no confirmed "fast" flag on /v1/messages.
    // Recognise only an explicit body boolean; default false otherwise.
    let fast_mode = body
        .get("fast_mode")
        .and_then(Value::as_bool)
        .unwrap_or(false);

    // inference_geo: no confirmed source. Read an explicit header if a deployment
    // sets one, else null.
    let inference_geo = headers
        .get("x-openlatch-inference-geo")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.trim().to_ascii_lowercase())
        .filter(|s| !s.is_empty());

    PricingInputs {
        batch,
        fast_mode,
        inference_geo,
    }
}

/// Extract the `model` string from the request body.
pub fn model_of(body: &Value) -> Option<String> {
    body.get("model")
        .and_then(Value::as_str)
        .map(str::to_string)
}

/// Whether the request body carries at least one `cache_control` breakpoint.
/// Used as context for the (weak) `cache.preserved` signal (D-15).
pub fn has_cache_breakpoint(raw_body: &[u8]) -> bool {
    // A substring scan is sufficient and avoids re-parsing the whole body; the
    // key only appears as a JSON object key on a real breakpoint.
    memmem(raw_body, b"\"cache_control\"")
}

/// Infer `cache.preserved` (D-15) — a **weak/open** signal.
///
/// ⚠️ Cold start, 5-minute TTL expiry, and genuine customer churn all produce
/// `cache_read = 0` legitimately, so a `false` here does not prove the breakpoint
/// was lost. Recorded as an open question (I-1 OQ2); this is a plan-03 release
/// gate, not a settled fact. Inferred `true` only when we actually read cache.
pub fn infer_cache_preserved(usage: &Usage) -> bool {
    usage.cache_read > 0
}

/// Tiny substring search (no `memchr` dependency needed for this hot-but-small path).
fn memmem(haystack: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() || haystack.len() < needle.len() {
        return false;
    }
    haystack.windows(needle.len()).any(|w| w == needle)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn c3_input_is_not_the_total() {
        // C-3: input_tokens=50, cache_read=100000 → total input MUST be 100050.
        // A test that fails if anyone treats input_tokens as the total.
        let chunk = br#"data: {"type":"message_start","message":{"usage":{"input_tokens":50,"cache_read_input_tokens":100000,"cache_creation_input_tokens":0,"output_tokens":1}}}"#;
        let mut acc = UsageAccumulator::default();
        assert!(acc.scan_chunk(chunk));
        let u = acc.usage();
        assert_eq!(u.input_tokens, 50);
        assert_eq!(u.cache_read, 100_000);
        let total_input = u.input_tokens + u.cache_write + u.cache_read;
        assert_eq!(
            total_input, 100_050,
            "total input must be input + cache_creation + cache_read (C-3)"
        );
    }

    #[test]
    fn merges_message_start_and_message_delta() {
        // message_start carries input+cache, output=1; message_delta carries the
        // final output. Merge-by-max yields the complete usage.
        let start = br#"data: {"type":"message_start","message":{"usage":{"input_tokens":10,"cache_read_input_tokens":5,"cache_creation_input_tokens":8,"cache_creation":{"ephemeral_5m_input_tokens":6,"ephemeral_1h_input_tokens":2},"output_tokens":1}}}"#;
        let delta = br#"data: {"type":"message_delta","usage":{"output_tokens":321}}"#;
        let mut acc = UsageAccumulator::default();
        acc.scan_chunk(start);
        acc.scan_chunk(delta);
        let u = acc.usage();
        assert_eq!(u.input_tokens, 10);
        assert_eq!(u.cache_read, 5);
        assert_eq!(u.cache_write, 8);
        assert_eq!(u.eph_5m, 6);
        assert_eq!(u.eph_1h, 2);
        assert_eq!(u.output_tokens, 321);
        assert!(acc.has_usage());
    }

    #[test]
    fn message_start_is_not_terminal_until_message_delta() {
        // FIX 1: message_start carries FINAL input/cache but a PRELIMINARY
        // output_tokens=1, so it must NOT count as terminal. A stream that ends
        // here is only partially measured (→ tokenizer_estimated in finalize).
        let start = br#"data: {"type":"message_start","message":{"usage":{"input_tokens":10,"cache_read_input_tokens":5,"output_tokens":1}}}"#;
        let mut acc = UsageAccumulator::default();
        assert!(
            acc.scan_chunk(start),
            "message_start contributes input/cache usage"
        );
        assert!(acc.has_usage(), "usage WAS observed");
        assert!(
            !acc.is_terminal(),
            "but message_start is NOT terminal — output is preliminary (=1)"
        );

        // The terminal message_delta flips the flag and carries the final output.
        let delta = br#"data: {"type":"message_delta","usage":{"output_tokens":321}}"#;
        acc.scan_chunk(delta);
        assert!(acc.is_terminal(), "message_delta IS terminal");
        assert_eq!(
            acc.usage().output_tokens,
            321,
            "the terminal output overrides the preliminary 1"
        );
    }

    #[test]
    fn non_streaming_body_is_terminal() {
        // A complete non-streaming response body (top-level .usage, no
        // message_start type) is a full measurement → terminal.
        let body =
            br#"{"id":"msg_1","type":"message","usage":{"input_tokens":42,"output_tokens":7}}"#;
        let mut acc = UsageAccumulator::default();
        assert!(acc.scan_chunk(body));
        assert!(acc.is_terminal());
    }

    #[test]
    fn ephemeral_5m_1h_split_captured() {
        let chunk = br#"data: {"usage":{"input_tokens":0,"cache_creation_input_tokens":100,"cache_creation":{"ephemeral_5m_input_tokens":80,"ephemeral_1h_input_tokens":20},"output_tokens":0}}"#;
        let mut acc = UsageAccumulator::default();
        acc.scan_chunk(chunk);
        let u = acc.usage();
        assert_eq!(u.eph_5m, 80);
        assert_eq!(u.eph_1h, 20);
        assert_eq!(u.eph_5m + u.eph_1h, u.cache_write);
    }

    #[test]
    fn non_streaming_body_usage() {
        let body = br#"{"id":"msg_1","usage":{"input_tokens":42,"output_tokens":7}}"#;
        let mut acc = UsageAccumulator::default();
        assert!(acc.scan_chunk(body));
        assert_eq!(acc.usage().input_tokens, 42);
        assert_eq!(acc.usage().output_tokens, 7);
    }

    #[test]
    fn non_usage_chunk_is_ignored() {
        let mut acc = UsageAccumulator::default();
        assert!(!acc.scan_chunk(b"data: {\"type\":\"content_block_delta\"}\n\n"));
        assert!(!acc.has_usage());
    }

    #[test]
    fn cache_preserved_is_read_gated() {
        assert!(infer_cache_preserved(&Usage {
            cache_read: 1,
            ..Default::default()
        }));
        assert!(!infer_cache_preserved(&Usage::default()));
    }

    #[test]
    fn pricing_inputs_default_conservative() {
        let body = serde_json::json!({"model":"claude-opus-4-8","messages":[]});
        let p = derive_pricing_inputs(&body, &axum::http::HeaderMap::new());
        assert!(!p.batch);
        assert!(!p.fast_mode);
        assert!(p.inference_geo.is_none());
    }

    #[test]
    fn breakpoint_detection() {
        assert!(has_cache_breakpoint(
            br#"{"system":[{"type":"text","cache_control":{"type":"ephemeral"}}]}"#
        ));
        assert!(!has_cache_breakpoint(br#"{"messages":[]}"#));
    }
}