o402 0.1.5

//! OpenAI JSON `usage`: cached ⊂ prompt, reasoning ⊂ completion.

use serde_json::json;

/// Token counts after reading a 2xx JSON body.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(crate) struct Usage {
    /// `prompt_tokens` / `input_tokens`.
    pub(crate) prompt_tokens: u128,
    /// `completion_tokens` / `output_tokens`. Absent on embeddings.
    pub(crate) completion_tokens: u128,
    /// `prompt_tokens_details.cached_tokens` (subset of prompt).
    pub(crate) cached_tokens: u128,
    /// `completion_tokens_details.reasoning_tokens` (subset of completion).
    pub(crate) reasoning_tokens: u128,
}

impl Usage {
    /// Cached tokens, clamped to prompt.
    #[must_use]
    pub(crate) fn cached(self) -> u128 {
        self.cached_tokens.min(self.prompt_tokens)
    }

    /// Reasoning tokens, clamped to completion.
    #[must_use]
    pub(crate) fn reasoning(self) -> u128 {
        self.reasoning_tokens.min(self.completion_tokens)
    }

    /// Uncached prompt tokens.
    #[must_use]
    pub(crate) fn uncached(self) -> u128 {
        self.prompt_tokens.saturating_sub(self.cached())
    }

    /// Non-reasoning completion tokens.
    #[must_use]
    pub(crate) fn output(self) -> u128 {
        self.completion_tokens.saturating_sub(self.reasoning())
    }
}

/// Parses `usage` from an OpenAI JSON body. `None` if the object is missing.
#[must_use]
pub(crate) fn from_json(body: &[u8]) -> Option<Usage> {
    let value: serde_json::Value = serde_json::from_slice(body).ok()?;
    from_value(&value)
}

/// Incremental SSE scanner. Last `usage` object wins.
#[derive(Debug, Default)]
pub(crate) struct SseUsageParser {
    buf: Vec<u8>,
    usage: Option<Usage>,
}

impl SseUsageParser {
    /// Empty buffer, no usage yet.
    #[must_use]
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Consumes a body chunk. Incomplete trailing lines stay buffered.
    pub(crate) fn push(&mut self, chunk: &[u8]) {
        self.buf.extend_from_slice(chunk);
        self.drain_lines();
    }

    /// Parses a leftover line that never saw a trailing newline.
    pub(crate) fn finish(&mut self) {
        if self.buf.is_empty() {
            return;
        }
        let mut line = std::mem::take(&mut self.buf);
        trim_line_ending(&mut line);
        self.ingest_line(&line);
    }

    /// Last parsed usage, if any chunk contained one.
    #[must_use]
    pub(crate) const fn usage(&self) -> Option<Usage> {
        self.usage
    }

    fn drain_lines(&mut self) {
        loop {
            let Some(idx) = self.buf.iter().position(|&b| b == b'\n') else {
                return;
            };
            let rest = self.buf.split_off(idx.saturating_add(1));
            let mut line = std::mem::replace(&mut self.buf, rest);
            trim_line_ending(&mut line);
            self.ingest_line(&line);
        }
    }

    fn ingest_line(&mut self, line: &[u8]) {
        let Ok(text) = std::str::from_utf8(line) else {
            return;
        };
        let Some(data) = text.strip_prefix("data:") else {
            return;
        };
        let data = data.trim();
        if data.is_empty() || data == "[DONE]" {
            return;
        }
        if let Some(usage) = from_sse_data(data) {
            self.usage = Some(usage);
        }
    }
}

fn from_sse_data(data: &str) -> Option<Usage> {
    let value: serde_json::Value = serde_json::from_str(data).ok()?;
    from_value(&value)
}

/// Sets `stream_options.include_usage` to true. `None` if `body` is not a JSON object.
#[must_use]
pub(crate) fn force_include_usage(body: &[u8]) -> Option<Vec<u8>> {
    let mut value: serde_json::Value = serde_json::from_slice(body).ok()?;
    let obj = value.as_object_mut()?;
    match obj.get_mut("stream_options") {
        Some(serde_json::Value::Object(opts)) => {
            opts.insert("include_usage".to_owned(), serde_json::Value::Bool(true));
        }
        _ => {
            obj.insert("stream_options".to_owned(), json!({"include_usage": true}));
        }
    }
    serde_json::to_vec(&value).ok()
}

fn trim_line_ending(line: &mut Vec<u8>) {
    if line.last() == Some(&b'\n') {
        line.pop();
    }
    if line.last() == Some(&b'\r') {
        line.pop();
    }
}

fn from_value(value: &serde_json::Value) -> Option<Usage> {
    let usage = value.get("usage")?;
    if !usage.is_object() {
        return None;
    }
    let prompt =
        json_tokens(usage, "prompt_tokens").or_else(|| json_tokens(usage, "input_tokens"))?;
    let completion = json_tokens(usage, "completion_tokens")
        .or_else(|| json_tokens(usage, "output_tokens"))
        .unwrap_or(0);
    let cached = usage
        .get("prompt_tokens_details")
        .and_then(|details| json_tokens(details, "cached_tokens"))
        .or_else(|| {
            usage
                .get("input_tokens_details")
                .and_then(|details| json_tokens(details, "cached_tokens"))
        })
        .unwrap_or(0);
    let reasoning = usage
        .get("completion_tokens_details")
        .and_then(|details| json_tokens(details, "reasoning_tokens"))
        .or_else(|| {
            usage
                .get("output_tokens_details")
                .and_then(|details| json_tokens(details, "reasoning_tokens"))
        })
        .unwrap_or(0);
    Some(Usage {
        prompt_tokens: prompt,
        completion_tokens: completion,
        cached_tokens: cached,
        reasoning_tokens: reasoning,
    })
}

fn json_tokens(value: &serde_json::Value, key: &str) -> Option<u128> {
    value
        .get(key)
        .and_then(serde_json::Value::as_u64)
        .map(u128::from)
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::{SseUsageParser, Usage, force_include_usage, from_json};

    #[test]
    fn cached_and_reasoning_are_clamped_subsets() {
        let usage = from_json(
            br#"{
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 8,
                    "prompt_tokens_details": {"cached_tokens": 99},
                    "completion_tokens_details": {"reasoning_tokens": 99}
                }
            }"#,
        )
        .expect("usage");
        assert_eq!(usage.cached(), 10, "cached clamp");
        assert_eq!(usage.uncached(), 0, "uncached");
        assert_eq!(usage.reasoning(), 8, "reasoning clamp");
        assert_eq!(usage.output(), 0, "output");
    }

    #[test]
    fn subtracts_cached_and_reasoning_before_billable_buckets() {
        let usage = from_json(
            br#"{
                "usage": {
                    "prompt_tokens": 1000,
                    "completion_tokens": 100,
                    "prompt_tokens_details": {"cached_tokens": 200},
                    "completion_tokens_details": {"reasoning_tokens": 40}
                }
            }"#,
        )
        .expect("usage");
        assert_eq!(usage.uncached(), 800, "uncached");
        assert_eq!(usage.cached(), 200, "cached");
        assert_eq!(usage.output(), 60, "output");
        assert_eq!(usage.reasoning(), 40, "reasoning");
    }

    #[test]
    fn responses_cached_tokens_are_not_billed_as_uncached() {
        let usage = from_json(
            br#"{
                "usage": {
                    "input_tokens": 1000,
                    "output_tokens": 100,
                    "input_tokens_details": {"cached_tokens": 200},
                    "output_tokens_details": {"reasoning_tokens": 40}
                }
            }"#,
        )
        .expect("usage");
        assert_eq!(usage.prompt_tokens, 1000, "prompt");
        assert_eq!(usage.uncached(), 800, "uncached");
        assert_eq!(usage.cached(), 200, "cached");
        assert_eq!(usage.output(), 60, "output");
        assert_eq!(usage.reasoning(), 40, "reasoning");
    }

    #[test]
    fn embeddings_have_prompt_only() {
        let usage =
            from_json(br#"{"usage":{"prompt_tokens":8191,"total_tokens":8191}}"#).expect("usage");
        assert_eq!(
            usage,
            Usage {
                prompt_tokens: 8191,
                completion_tokens: 0,
                cached_tokens: 0,
                reasoning_tokens: 0,
            },
            "embeddings"
        );
        assert_eq!(usage.output(), 0, "no output");
        assert_eq!(usage.reasoning(), 0, "no reasoning");
    }

    #[test]
    fn missing_usage_is_none() {
        assert!(from_json(br#"{"id":"ok"}"#).is_none(), "missing");
        assert!(from_json(br#"{"usage":null}"#).is_none(), "null");
        assert!(from_json(br#"{"usage":{}}"#).is_none(), "empty");
        assert!(from_json(b"not-json").is_none(), "invalid");
    }

    #[test]
    fn sse_parser_takes_last_usage_chunk() {
        let mut parser = SseUsageParser::new();
        parser.push(b"data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n");
        parser.push(
            b"data: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":4,\"prompt_tokens_details\":{\"cached_tokens\":3},\"completion_tokens_details\":{\"reasoning_tokens\":1}}}\n\ndata: [DONE]\n\n",
        );
        parser.finish();
        let usage = parser.usage().expect("usage");
        assert_eq!(usage.prompt_tokens, 10, "prompt");
        assert_eq!(usage.completion_tokens, 4, "completion");
        assert_eq!(usage.cached(), 3, "cached");
        assert_eq!(usage.reasoning(), 1, "reasoning");
    }

    #[test]
    fn sse_parser_joins_split_lines() {
        let mut parser = SseUsageParser::new();
        parser.push(b"data: {\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":2}");
        assert!(parser.usage().is_none(), "incomplete");
        parser.push(b"}\n");
        parser.finish();
        let usage = parser.usage().expect("usage");
        assert_eq!(usage.prompt_tokens, 8, "prompt");
        assert_eq!(usage.completion_tokens, 2, "completion");
    }

    #[test]
    fn sse_parser_finish_without_newline() {
        let mut parser = SseUsageParser::new();
        parser.push(b"data: {\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}");
        parser.finish();
        assert_eq!(
            parser.usage().map(|usage| usage.prompt_tokens),
            Some(1),
            "finish"
        );
    }

    #[test]
    fn force_include_usage_merges_object_and_replaces_otherwise() {
        let merged = force_include_usage(
            br#"{"model":"m","stream":true,"stream_options":{"include_usage":false,"foo":1}}"#,
        )
        .expect("merge");
        let merged_json: serde_json::Value = serde_json::from_slice(&merged).expect("json");
        assert_eq!(
            merged_json.pointer("/stream_options/include_usage"),
            Some(&serde_json::Value::Bool(true)),
            "forced true"
        );
        assert_eq!(
            merged_json.pointer("/stream_options/foo"),
            Some(&serde_json::Value::Number(1.into())),
            "kept extra"
        );

        let replaced =
            force_include_usage(br#"{"model":"m","stream_options":"nope"}"#).expect("replace");
        let replaced_json: serde_json::Value = serde_json::from_slice(&replaced).expect("json");
        assert_eq!(
            replaced_json.get("stream_options"),
            Some(&json!({"include_usage": true})),
            "replaced non-object"
        );

        let inserted = force_include_usage(br#"{"model":"m","stream":true}"#).expect("insert");
        let inserted_json: serde_json::Value = serde_json::from_slice(&inserted).expect("json");
        assert_eq!(
            inserted_json.get("stream_options"),
            Some(&json!({"include_usage": true})),
            "inserted"
        );

        assert!(force_include_usage(b"not-json").is_none(), "invalid");
    }
}