kvca_proxy/src/cbr.rs

use serde_json::Value;

/// Canonical Byte Representation of the `messages` field.
///
/// This produces a deterministic byte sequence from the messages JSON,
/// independent of the original key ordering or whitespace in the request.
/// Semantically identical messages arrays always produce identical bytes.
pub fn canonical_messages_bytes(messages: &Value) -> Vec<u8> {
    let canonical = canonicalize_value(messages);
    // Compact (no extraneous whitespace), deterministic key order
    serde_json::to_vec(&canonical).expect("canonical JSON serialization is infallible")
}

/// Canonical cache identity for every request field that can affect prompt
/// rendering, followed by `messages` so conversation growth extends the tail.
pub fn canonical_prompt_bytes(body: &Value) -> Vec<u8> {
    let Some(object) = body.as_object() else {
        return canonical_messages_bytes(&Value::Null);
    };

    let messages = object.get("messages").unwrap_or(&Value::Null);
    let mut prompt_options = serde_json::Map::new();
    for (key, value) in object {
        if key != "messages" && !is_generation_only_field(key) {
            prompt_options.insert(key.clone(), canonicalize_value(value));
        }
    }

    let mut bytes = serde_json::to_vec(&Value::Object(prompt_options))
        .expect("canonical JSON serialization is infallible");
    bytes.push(0);
    bytes.extend(canonical_messages_bytes(messages));
    bytes
}

/// Fields known not to alter tokenized prompt construction. Unknown extension
/// fields are deliberately included so backend-specific template kwargs and
/// future prompt controls cannot silently alias an unrelated cache namespace.
fn is_generation_only_field(key: &str) -> bool {
    matches!(
        key,
        "frequency_penalty"
            | "logit_bias"
            | "logprobs"
            | "max_completion_tokens"
            | "max_tokens"
            | "metadata"
            | "n"
            | "parallel_tool_calls"
            | "presence_penalty"
            | "seed"
            | "service_tier"
            | "stop"
            | "store"
            | "stream"
            | "stream_options"
            | "temperature"
            | "top_logprobs"
            | "top_p"
            | "user"
    )
}

/// Recursively sort object keys for deterministic serialization.
fn canonicalize_value(value: &Value) -> Value {
    match value {
        Value::Object(map) => {
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort();
            let mut out = serde_json::Map::with_capacity(map.len());
            for k in keys {
                out.insert(k.clone(), canonicalize_value(&map[k]));
            }
            Value::Object(out)
        }
        Value::Array(arr) => Value::Array(arr.iter().map(canonicalize_value).collect()),
        other => other.clone(),
    }
}

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

    #[test]
    fn test_canonical_stable_key_order() {
        let a: Value =
            serde_json::from_str(r#"{"messages":[{"role":"user","content":"hi"}]}"#).unwrap();
        let b: Value =
            serde_json::from_str(r#"{"messages":[{"content":"hi","role":"user"}]}"#).unwrap();

        let messages_a = &a["messages"];
        let messages_b = &b["messages"];

        assert_eq!(
            canonical_messages_bytes(messages_a),
            canonical_messages_bytes(messages_b)
        );
    }

    #[test]
    fn test_canonical_whitespace_agnostic() {
        let a: Value =
            serde_json::from_str(r#"{"messages":[{"role":"user","content":"hello world"}]}"#)
                .unwrap();
        let b: Value = serde_json::from_str(
            r#"{
            "messages": [
                {
                    "role": "user",
                    "content": "hello world"
                }
            ]
        }"#,
        )
        .unwrap();

        assert_eq!(
            canonical_messages_bytes(&a["messages"]),
            canonical_messages_bytes(&b["messages"])
        );
    }

    #[test]
    fn test_different_content_different_bytes() {
        let a: Value =
            serde_json::from_str(r#"{"messages":[{"role":"user","content":"hi"}]}"#).unwrap();
        let b: Value =
            serde_json::from_str(r#"{"messages":[{"role":"user","content":"hello"}]}"#).unwrap();

        assert_ne!(
            canonical_messages_bytes(&a["messages"]),
            canonical_messages_bytes(&b["messages"])
        );
    }

    #[test]
    fn test_prompt_identity_includes_template_inputs() {
        let a = serde_json::json!({
            "model": "model-a",
            "messages": [{"role": "user", "content": "hi"}],
            "chat_template_kwargs": {"enable_thinking": true}
        });
        let b = serde_json::json!({
            "model": "model-a",
            "messages": [{"role": "user", "content": "hi"}],
            "chat_template_kwargs": {"enable_thinking": false}
        });
        assert_ne!(canonical_prompt_bytes(&a), canonical_prompt_bytes(&b));
    }

    #[test]
    fn test_prompt_identity_ignores_generation_controls() {
        let a = serde_json::json!({
            "model": "model-a",
            "messages": [{"role": "user", "content": "hi"}],
            "temperature": 0.1
        });
        let b = serde_json::json!({
            "model": "model-a",
            "messages": [{"role": "user", "content": "hi"}],
            "temperature": 0.9
        });
        assert_eq!(canonical_prompt_bytes(&a), canonical_prompt_bytes(&b));
    }
}