kcode_k1_codex_token_usage/

lib.rs

use kcode_k1_codex_conversations::{Error, ErrorKind};
use serde_json::Value;

const METHOD: &str = "thread/tokenUsage/updated";

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TokenUsageUpdated {
    Absent,
    Malformed(Error),
    Valid(TokenUsage),
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenUsageScope {
    pub thread_id: String,
    pub turn_id: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenUsage {
    pub scope: TokenUsageScope,
    pub total: TokenUsageBreakdown,
    pub last: TokenUsageBreakdown,
    pub model_context_window: Option<i64>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenUsageBreakdown {
    pub input_tokens: i64,
    pub cached_input_tokens: i64,
    pub cache_write_input_tokens: i64,
    pub output_tokens: i64,
    pub reasoning_output_tokens: i64,
    pub total_tokens: i64,
}

pub fn decode_token_usage_updated(message: Value) -> TokenUsageUpdated {
    if message.get("method").and_then(Value::as_str) != Some(METHOD) {
        return TokenUsageUpdated::Absent;
    }
    match decode(&message) {
        Ok(usage) => TokenUsageUpdated::Valid(usage),
        Err(error) => TokenUsageUpdated::Malformed(error),
    }
}

fn decode(message: &Value) -> Result<TokenUsage, Error> {
    let params = object_field(message, "params")?;
    let scope = TokenUsageScope {
        thread_id: string_field(params, "threadId")?.to_owned(),
        turn_id: string_field(params, "turnId")?.to_owned(),
    };
    let usage = object_field(params, "tokenUsage")?;
    Ok(TokenUsage {
        scope,
        total: breakdown(object_field(usage, "total")?)?,
        last: breakdown(object_field(usage, "last")?)?,
        model_context_window: nullable_count(usage, "modelContextWindow")?,
    })
}

fn breakdown(value: &Value) -> Result<TokenUsageBreakdown, Error> {
    Ok(TokenUsageBreakdown {
        input_tokens: count(value, "inputTokens")?,
        cached_input_tokens: count(value, "cachedInputTokens")?,
        cache_write_input_tokens: optional_count(value, "cacheWriteInputTokens")?,
        output_tokens: count(value, "outputTokens")?,
        reasoning_output_tokens: count(value, "reasoningOutputTokens")?,
        total_tokens: count(value, "totalTokens")?,
    })
}

fn object_field<'a>(value: &'a Value, field: &str) -> Result<&'a Value, Error> {
    value
        .get(field)
        .filter(|field| field.is_object())
        .ok_or_else(|| protocol(format!("token usage notification has invalid {field}")))
}

fn string_field<'a>(value: &'a Value, field: &str) -> Result<&'a str, Error> {
    value
        .get(field)
        .and_then(Value::as_str)
        .ok_or_else(|| protocol(format!("token usage notification has invalid {field}")))
}

fn count(value: &Value, field: &str) -> Result<i64, Error> {
    value
        .get(field)
        .and_then(Value::as_i64)
        .filter(|count| *count >= 0)
        .ok_or_else(|| protocol(format!("token usage notification has invalid {field}")))
}

fn optional_count(value: &Value, field: &str) -> Result<i64, Error> {
    if value.get(field).is_none() {
        Ok(0)
    } else {
        count(value, field)
    }
}

fn nullable_count(value: &Value, field: &str) -> Result<Option<i64>, Error> {
    match value.get(field) {
        Some(Value::Null) => Ok(None),
        Some(_) => count(value, field).map(Some),
        None => Err(protocol(format!(
            "token usage notification has invalid {field}"
        ))),
    }
}

fn protocol(message: String) -> Error {
    Error::new(ErrorKind::Protocol, message)
}

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

    fn notification() -> Value {
        json!({
            "method": METHOD,
            "params": {
                "threadId": "thread-a",
                "turnId": "turn-b",
                "tokenUsage": {
                    "total": counts(10),
                    "last": counts(1),
                    "modelContextWindow": 128000
                }
            }
        })
    }

    fn counts(base: i64) -> Value {
        json!({
            "inputTokens": base,
            "cachedInputTokens": base + 1,
            "cacheWriteInputTokens": base + 2,
            "outputTokens": base + 3,
            "reasoningOutputTokens": base + 4,
            "totalTokens": base + 5
        })
    }

    fn malformed(value: Value) {
        assert!(matches!(
            decode_token_usage_updated(value),
            TokenUsageUpdated::Malformed(Error {
                kind: ErrorKind::Protocol,
                ..
            })
        ));
    }

    #[test]
    fn decodes_full_notification() {
        assert_eq!(
            decode_token_usage_updated(notification()),
            TokenUsageUpdated::Valid(TokenUsage {
                scope: TokenUsageScope {
                    thread_id: "thread-a".into(),
                    turn_id: "turn-b".into(),
                },
                total: TokenUsageBreakdown {
                    input_tokens: 10,
                    cached_input_tokens: 11,
                    cache_write_input_tokens: 12,
                    output_tokens: 13,
                    reasoning_output_tokens: 14,
                    total_tokens: 15,
                },
                last: TokenUsageBreakdown {
                    input_tokens: 1,
                    cached_input_tokens: 2,
                    cache_write_input_tokens: 3,
                    output_tokens: 4,
                    reasoning_output_tokens: 5,
                    total_tokens: 6,
                },
                model_context_window: Some(128000),
            })
        );
    }

    #[test]
    fn accepts_null_window_absent_cache_write_and_unknown_fields() {
        let mut value = notification();
        let usage = value.pointer_mut("/params/tokenUsage").unwrap();
        usage["modelContextWindow"] = Value::Null;
        usage["extra"] = json!(true);
        for part in ["total", "last"] {
            usage[part]
                .as_object_mut()
                .unwrap()
                .remove("cacheWriteInputTokens");
            usage[part]["extra"] = json!("ignored");
        }
        assert!(matches!(
            decode_token_usage_updated(value),
            TokenUsageUpdated::Valid(TokenUsage {
                model_context_window: None,
                total: TokenUsageBreakdown {
                    cache_write_input_tokens: 0,
                    ..
                },
                last: TokenUsageBreakdown {
                    cache_write_input_tokens: 0,
                    ..
                },
                ..
            })
        ));
    }

    #[test]
    fn rejects_wrong_method_and_nonmatching_messages_as_absent() {
        for value in [json!({"method": "thread/other"}), json!({}), Value::Null] {
            assert_eq!(decode_token_usage_updated(value), TokenUsageUpdated::Absent);
        }
    }

    #[test]
    fn rejects_every_invalid_scope_string_class() {
        for field in ["threadId", "turnId"] {
            for bad in [
                Value::Null,
                json!(1),
                json!(-1),
                json!(1.5),
                json!(u64::MAX),
            ] {
                let mut value = notification();
                value["params"].as_object_mut().unwrap().remove(field);
                malformed(value);
                let mut value = notification();
                value["params"][field] = bad;
                malformed(value);
            }
        }
    }

    #[test]
    fn rejects_every_invalid_required_count_class() {
        for part in ["total", "last"] {
            for field in [
                "inputTokens",
                "cachedInputTokens",
                "outputTokens",
                "reasoningOutputTokens",
                "totalTokens",
            ] {
                for bad in [
                    Value::Null,
                    json!("1"),
                    json!(-1),
                    json!(1.5),
                    json!(u64::MAX),
                ] {
                    let mut value = notification();
                    value["params"]["tokenUsage"][part]
                        .as_object_mut()
                        .unwrap()
                        .remove(field);
                    malformed(value);
                    let mut value = notification();
                    value["params"]["tokenUsage"][part][field] = bad;
                    malformed(value);
                }
            }
        }
    }

    #[test]
    fn rejects_every_invalid_optional_and_window_count_class() {
        for part in ["total", "last"] {
            for bad in [
                Value::Null,
                json!("1"),
                json!(-1),
                json!(1.5),
                json!(u64::MAX),
            ] {
                let mut value = notification();
                value["params"]["tokenUsage"][part]["cacheWriteInputTokens"] = bad;
                malformed(value);
            }
        }
        for bad in [
            Value::Null,
            json!("1"),
            json!(-1),
            json!(1.5),
            json!(u64::MAX),
        ] {
            let mut value = notification();
            value["params"]["tokenUsage"]["modelContextWindow"] = bad;
            if value.pointer("/params/tokenUsage/modelContextWindow") == Some(&Value::Null) {
                continue;
            }
            malformed(value);
        }
        let mut value = notification();
        value["params"]["tokenUsage"]
            .as_object_mut()
            .unwrap()
            .remove("modelContextWindow");
        malformed(value);
    }
}