prompt_cache.rs

//! OpenAI prompt-caching helpers (keys, retention, GPT-5.6 breakpoints,
//! first-party host detection, tool allowlists).
//!
//! OpenAI caches **exact contiguous prefixes**. Tools + schemas sit before
//! system/developer text. This module centralises the request-side knobs the
//! harness can set without migrating to the Responses API.

use serde_json::{json, Value};
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

/// Per-request prompt-cache controls passed into provider adapters.
#[derive(Clone, Debug, Default)]
pub struct PromptCacheRequest {
    /// Routing stickiness key (`prompt_cache_key`). Combined with the first
    /// ~256-token prefix hash server-side. Required for reliable GPT-5.6+ matching.
    pub key: Option<String>,
    /// Pre-GPT-5.6 retention policy: `"in_memory"` or `"24h"`.
    pub retention: Option<String>,
    /// `service_tier` for Flex / priority routing (`"flex"`, `"priority"`, …).
    pub service_tier: Option<String>,
    /// When true, emit GPT-5.6+ explicit breakpoints + `prompt_cache_options.mode=explicit`.
    pub explicit_breakpoints: bool,
    /// When `Some`, send the full `tools` array but restrict callable tools via
    /// `tool_choice.allowed_tools` so the tools prefix stays cache-stable.
    pub allowed_tool_names: Option<Vec<String>>,
}

impl PromptCacheRequest {
    pub fn none() -> Self {
        Self::default()
    }
}

/// True when the host is first-party OpenAI / Codex / Azure OpenAI — safe to
/// send `prompt_cache_*`, `service_tier`, and `allowed_tools`. Unknown OpenAI-
/// compatible gateways may 400 on these fields.
pub fn supports_prompt_cache_controls(base_url: &str) -> bool {
    supports_prompt_cache_controls_for(base_url, &[])
}

/// Parse the host from a base URL (scheme stripped, path dropped).
pub fn cache_control_host(base_url: &str) -> String {
    let u = base_url.to_ascii_lowercase();
    u.split("://")
        .nth(1)
        .unwrap_or(u.as_str())
        .split('/')
        .next()
        .unwrap_or("")
        .to_string()
}

/// First-party OpenAI hosts, plus any extra suffixes the operator has verified.
pub fn supports_prompt_cache_controls_for(base_url: &str, extra_hosts: &[String]) -> bool {
    let host = cache_control_host(base_url);
    if host == "api.openai.com"
        || host.ends_with(".openai.com")
        || host == "chatgpt.com"
        || host.ends_with(".chatgpt.com")
        || host.contains("openai.azure.com")
        || host.contains("cognitiveservices.azure.com")
        || host.contains("services.ai.azure.com")
    {
        return true;
    }
    extra_hosts.iter().any(|h| {
        let h = h.trim().trim_start_matches('.').to_ascii_lowercase();
        !h.is_empty() && (host == h || host.ends_with(&format!(".{h}")))
    })
}

/// GPT-5.6+ (and later major families) use breakpoint caching rather than
/// longest-unmarked-prefix matching. Pre-5.6 models reject the new fields.
pub fn model_uses_explicit_cache(model: &str) -> bool {
    let m = model.to_ascii_lowercase();
    if m.starts_with("gpt-6") || m.starts_with("gpt-7") {
        return true;
    }
    // gpt-5.6, gpt-5.6-sol, gpt-5.10, …
    if let Some(rest) = m.strip_prefix("gpt-5.") {
        let minor: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect();
        if let Ok(n) = minor.parse::<u32>() {
            return n >= 6;
        }
    }
    // Bare aliases that imply the latest 5.x line with explicit caching.
    matches!(
        m.as_str(),
        "gpt-5.6" | "gpt-5.6-pro" | "gpt-5.6-mini" | "gpt-5.6-nano"
    )
}

/// Short stable fingerprint of a tools array (names only, sorted).
pub fn tools_fingerprint(tools: &[Value]) -> String {
    let mut names: Vec<&str> = tools
        .iter()
        .filter_map(|t| {
            t.get("function")
                .and_then(|f| f.get("name"))
                .and_then(Value::as_str)
                .or_else(|| t.get("name").and_then(Value::as_str))
        })
        .collect();
    names.sort_unstable();
    names.dedup();
    let mut hasher = DefaultHasher::new();
    for n in &names {
        n.hash(&mut hasher);
    }
    format!("{:x}", hasher.finish())
}

/// Build a prompt_cache_key. Keep under ~15 RPM per key+prefix: callers that
/// fan out parallel subagents should pass distinct `scope` values (run ids).
///
/// Shape: `cc:{workspace}:{session}:{model_family}:{tools_fp}:g{gen}[:{scope}]`
pub fn build_prompt_cache_key(
    workspace_hash: &str,
    session_id: &str,
    model: &str,
    tools_fp: &str,
    generation: u64,
    scope: Option<&str>,
) -> String {
    let model_family = model_family(model);
    let session = if session_id.is_empty() {
        "nosess"
    } else {
        session_id
    };
    let ws = if workspace_hash.is_empty() {
        "ws"
    } else {
        workspace_hash
    };
    let mut key = format!("cc:{ws}:{session}:{model_family}:{tools_fp}:g{generation}");
    if let Some(s) = scope.filter(|s| !s.is_empty()) {
        key.push(':');
        key.push_str(s);
    }
    // Hard cap — OpenAI accepts reasonably long keys; keep ours tidy.
    if key.len() > 200 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        format!("cc:{:x}:g{generation}", hasher.finish())
    } else {
        key
    }
}

fn model_family(model: &str) -> String {
    let m = model.to_ascii_lowercase();
    // Keep major.minor (gpt-5.6, gpt-4o, o3, …) without date snapshots so
    // dated aliases share a key family when the user pins a snapshot mid-session
    // only intentionally via generation bump.
    let base: String = m
        .chars()
        .take_while(|c| c.is_ascii_alphanumeric() || *c == '.' || *c == '-' || *c == '_')
        .collect();
    // Drop trailing date-like -2024-08-06
    if let Some(idx) = base.find("-20") {
        let rest = &base[idx + 1..];
        if rest.len() >= 10
            && rest
                .as_bytes()
                .iter()
                .all(|b| b.is_ascii_digit() || *b == b'-')
        {
            return base[..idx].to_string();
        }
    }
    if base.is_empty() {
        "model".into()
    } else {
        base
    }
}

/// Session id derived from the session file path (stable across resumes).
pub fn session_id_from_path(path: Option<&std::path::Path>) -> String {
    path.and_then(|p| p.file_stem())
        .and_then(|s| s.to_str())
        .map(|s| {
            // Keep it short and path-safe.
            let cleaned: String = s
                .chars()
                .map(|c| {
                    if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                        c
                    } else {
                        '-'
                    }
                })
                .collect();
            if cleaned.len() > 48 {
                cleaned[..48].to_string()
            } else if cleaned.is_empty() {
                "session".into()
            } else {
                cleaned
            }
        })
        .unwrap_or_else(|| "nosess".into())
}

/// Apply OpenAI Chat Completions prompt-cache fields onto a request body.
pub fn apply_openai_chat_cache_fields(body: &mut Value, cache: &PromptCacheRequest, model: &str) {
    let Some(obj) = body.as_object_mut() else {
        return;
    };
    if let Some(ref key) = cache.key {
        if !key.is_empty() {
            obj.insert("prompt_cache_key".into(), json!(key));
        }
    }
    if let Some(ref tier) = cache.service_tier {
        if !tier.is_empty() {
            obj.insert("service_tier".into(), json!(tier));
        }
    }

    let explicit = cache.explicit_breakpoints && model_uses_explicit_cache(model);
    if explicit {
        obj.insert(
            "prompt_cache_options".into(),
            json!({
                "mode": "explicit",
                "ttl": "30m"
            }),
        );
        // Mark the first system message content block with a breakpoint.
        if let Some(messages) = obj.get_mut("messages").and_then(|m| m.as_array_mut()) {
            mark_first_system_breakpoint(messages);
        }
    } else if let Some(ref retention) = cache.retention {
        // Pre-5.6 extended / in-memory retention. Do not send on 5.6+ (deprecated).
        if !model_uses_explicit_cache(model) && !retention.is_empty() {
            obj.insert("prompt_cache_retention".into(), json!(retention));
        }
    }

    if let Some(ref allowed) = cache.allowed_tool_names {
        if obj
            .get("tools")
            .and_then(|t| t.as_array())
            .is_some_and(|a| !a.is_empty())
            && !allowed.is_empty()
        {
            // Prefer allowed_tools over bare "auto" so the full tools[] prefix
            // stays identical while per-turn gating still works. Named-function
            // force (goal_write_plan) wins when already set by the adapter.
            let already_forced = obj
                .get("tool_choice")
                .and_then(|t| t.get("type"))
                .and_then(Value::as_str)
                == Some("function");
            if !already_forced {
                let tools: Vec<Value> = allowed
                    .iter()
                    .map(|name| json!({"type": "function", "name": name}))
                    .collect();
                obj.insert(
                    "tool_choice".into(),
                    json!({
                        "type": "allowed_tools",
                        "mode": "auto",
                        "tools": tools
                    }),
                );
            }
        }
    }
}

/// Apply the same cache fields to a Codex / OpenAI Responses-shaped body.
/// (No previous_response_id migration — only cache routing/retention knobs.)
pub fn apply_responses_cache_fields(body: &mut Value, cache: &PromptCacheRequest, model: &str) {
    let Some(obj) = body.as_object_mut() else {
        return;
    };
    if let Some(ref key) = cache.key {
        if !key.is_empty() {
            obj.insert("prompt_cache_key".into(), json!(key));
        }
    }
    if let Some(ref tier) = cache.service_tier {
        if !tier.is_empty() {
            obj.insert("service_tier".into(), json!(tier));
        }
    }
    let explicit = cache.explicit_breakpoints && model_uses_explicit_cache(model);
    if explicit {
        obj.insert(
            "prompt_cache_options".into(),
            json!({
                "mode": "explicit",
                "ttl": "30m"
            }),
        );
        // Breakpoint on instructions is not a content block; mark the first
        // input message text part when present. Standing system is mapped to
        // top-level `instructions` by the Codex adapter — add a synthetic
        // leading input_text breakpoint message only when input is non-empty
        // would reshuffle history. Instead, attach breakpoint to the first
        // user/message content block if any.
        if let Some(input) = obj.get_mut("input").and_then(|v| v.as_array_mut()) {
            mark_responses_first_text_breakpoint(input);
        }
    } else if let Some(ref retention) = cache.retention {
        if !model_uses_explicit_cache(model) && !retention.is_empty() {
            obj.insert("prompt_cache_retention".into(), json!(retention));
        }
    }
    // allowed_tools on Responses uses the same tool_choice shape.
    if let Some(ref allowed) = cache.allowed_tool_names {
        if obj
            .get("tools")
            .and_then(|t| t.as_array())
            .is_some_and(|a| !a.is_empty())
            && !allowed.is_empty()
        {
            let tools: Vec<Value> = allowed
                .iter()
                .map(|name| json!({"type": "function", "name": name}))
                .collect();
            obj.insert(
                "tool_choice".into(),
                json!({
                    "type": "allowed_tools",
                    "mode": "auto",
                    "tools": tools
                }),
            );
        }
    }
}

fn mark_first_system_breakpoint(messages: &mut [Value]) {
    for msg in messages.iter_mut() {
        let role = msg.get("role").and_then(Value::as_str).unwrap_or("");
        if role != "system" && role != "developer" {
            continue;
        }
        ensure_content_array_with_breakpoint(msg);
        return;
    }
}

fn ensure_content_array_with_breakpoint(msg: &mut Value) {
    let Some(obj) = msg.as_object_mut() else {
        return;
    };
    match obj.get("content").cloned() {
        Some(Value::String(text)) => {
            obj.insert(
                "content".into(),
                json!([{
                    "type": "text",
                    "text": text,
                    "prompt_cache_breakpoint": { "mode": "explicit" }
                }]),
            );
        }
        Some(Value::Array(mut arr)) => {
            if let Some(first) = arr.iter_mut().find(|p| {
                p.get("type").and_then(Value::as_str) == Some("text") || p.get("text").is_some()
            }) {
                if let Some(block) = first.as_object_mut() {
                    block.insert(
                        "prompt_cache_breakpoint".into(),
                        json!({ "mode": "explicit" }),
                    );
                }
            } else if let Some(first) = arr.first_mut().and_then(|p| p.as_object_mut()) {
                first.insert(
                    "prompt_cache_breakpoint".into(),
                    json!({ "mode": "explicit" }),
                );
            }
            obj.insert("content".into(), Value::Array(arr));
        }
        _ => {}
    }
}

fn mark_responses_first_text_breakpoint(input: &mut [Value]) {
    for item in input.iter_mut() {
        let ty = item.get("type").and_then(Value::as_str).unwrap_or("");
        if ty != "message" {
            continue;
        }
        if let Some(content) = item.get_mut("content").and_then(|c| c.as_array_mut()) {
            for part in content.iter_mut() {
                let pty = part.get("type").and_then(Value::as_str).unwrap_or("");
                if pty == "input_text" || pty == "output_text" || pty == "text" {
                    if let Some(obj) = part.as_object_mut() {
                        obj.insert(
                            "prompt_cache_breakpoint".into(),
                            json!({ "mode": "explicit" }),
                        );
                        return;
                    }
                }
            }
        }
        // Only mark the first message item.
        return;
    }
}

/// Resolve config retention string, defaulting to 24h for long coding sessions
/// when the feature is enabled and the model is pre-5.6.
pub fn resolve_retention(configured: Option<&str>, model: &str) -> Option<String> {
    if model_uses_explicit_cache(model) {
        return None;
    }
    match configured.map(str::trim).filter(|s| !s.is_empty()) {
        Some("off") | Some("none") => None,
        Some(s) => Some(s.to_string()),
        None => Some("24h".into()),
    }
}

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

    #[test]
    fn detects_first_party_hosts() {
        assert!(supports_prompt_cache_controls("https://api.openai.com/v1"));
        assert!(supports_prompt_cache_controls(
            "https://chatgpt.com/backend-api/codex"
        ));
        assert!(supports_prompt_cache_controls(
            "https://my-res.openai.azure.com/"
        ));
        assert!(!supports_prompt_cache_controls(
            "https://ai.karutoil.site/v1"
        ));
        assert!(!supports_prompt_cache_controls("https://api.deepseek.com"));
        assert!(!supports_prompt_cache_controls_for(
            "https://openrouter.ai/api/v1",
            &[]
        ));
        assert!(supports_prompt_cache_controls_for(
            "https://openrouter.ai/api/v1",
            &["openrouter.ai".into()]
        ));
        assert!(!supports_prompt_cache_controls_for(
            "https://notopenrouter.ai/v1",
            &["openrouter.ai".into()]
        ));
        assert!(supports_prompt_cache_controls_for(
            "https://api.openai.com/v1",
            &[]
        ));
    }

    #[test]
    fn explicit_cache_model_detection() {
        assert!(model_uses_explicit_cache("gpt-5.6"));
        assert!(model_uses_explicit_cache("gpt-5.6-sol"));
        assert!(model_uses_explicit_cache("gpt-5.10"));
        assert!(!model_uses_explicit_cache("gpt-5.5"));
        assert!(!model_uses_explicit_cache("gpt-5.1"));
        assert!(!model_uses_explicit_cache("gpt-4o"));
        assert!(model_uses_explicit_cache("gpt-6"));
    }

    #[test]
    fn key_is_stable_and_scoped() {
        let a = build_prompt_cache_key("wh", "sess1", "gpt-4o", "abc", 0, None);
        let b = build_prompt_cache_key("wh", "sess1", "gpt-4o", "abc", 0, None);
        assert_eq!(a, b);
        let c = build_prompt_cache_key("wh", "sess1", "gpt-4o", "abc", 1, None);
        assert_ne!(a, c);
        let d = build_prompt_cache_key("wh", "sess1", "gpt-4o", "abc", 0, Some("sa-1"));
        assert!(d.ends_with(":sa-1"));
    }

    #[test]
    fn chat_cache_fields_explicit_breakpoint() {
        let mut body = json!({
            "model": "gpt-5.6",
            "messages": [
                {"role": "system", "content": "You are helpful."},
                {"role": "user", "content": "hi"}
            ],
            "tools": [{"type":"function","function":{"name":"bash"}}],
            "tool_choice": "auto"
        });
        let cache = PromptCacheRequest {
            key: Some("cc:test".into()),
            retention: Some("24h".into()),
            service_tier: None,
            explicit_breakpoints: true,
            allowed_tool_names: Some(vec!["bash".into()]),
        };
        apply_openai_chat_cache_fields(&mut body, &cache, "gpt-5.6");
        assert_eq!(body["prompt_cache_key"], "cc:test");
        assert_eq!(body["prompt_cache_options"]["mode"], "explicit");
        assert!(body.get("prompt_cache_retention").is_none());
        assert_eq!(
            body["messages"][0]["content"][0]["prompt_cache_breakpoint"]["mode"],
            "explicit"
        );
        assert_eq!(body["tool_choice"]["type"], "allowed_tools");
    }

    #[test]
    fn chat_cache_fields_retention_on_older_models() {
        let mut body = json!({
            "model": "gpt-4o",
            "messages": [{"role":"user","content":"hi"}]
        });
        let cache = PromptCacheRequest {
            key: Some("k".into()),
            retention: Some("24h".into()),
            service_tier: Some("flex".into()),
            explicit_breakpoints: true, // ignored — model not 5.6+
            allowed_tool_names: None,
        };
        apply_openai_chat_cache_fields(&mut body, &cache, "gpt-4o");
        assert_eq!(body["prompt_cache_retention"], "24h");
        assert_eq!(body["service_tier"], "flex");
        assert!(body.get("prompt_cache_options").is_none());
    }

    #[test]
    fn tools_fingerprint_order_independent() {
        let a = vec![
            json!({"function":{"name":"z"}}),
            json!({"function":{"name":"a"}}),
        ];
        let b = vec![
            json!({"function":{"name":"a"}}),
            json!({"function":{"name":"z"}}),
        ];
        assert_eq!(tools_fingerprint(&a), tools_fingerprint(&b));
    }

    #[test]
    fn provider_request_default_helper() {
        // Compile-time sanity: Default works.
        let _ = PromptCacheRequest::default();
        let _ = ResolvedProvider {
            name: "x".into(),
            kind: crate::config::ProviderKind::OpenAI,
            base_url: "https://api.openai.com/v1".into(),
            api_key: None,
            headers: Vec::new(),
            oauth: false,
            context_window: None,
            models_override: Vec::new(),
            models_endpoint: None,
        };
    }
}