openlatch-client 0.5.2

//! Canonical I-10 Anthropic Messages interventions.
//!
//! Each actuator builds a complete JSON value and validates that only its
//! narrow field boundary changed. Callers serialize last and retain the raw
//! request bytes as the sole fallback.

use std::collections::{HashMap, HashSet};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use serde_json::{json, Value};

use crate::model_relay::billing::BillingMode;
use crate::zone_eval::types::OptimizeParams;

const CACHE_TTL: Duration = Duration::from_secs(60 * 60);
const PROTECTED_TOOLS: [&str; 3] = ["Edit", "Write", "Read"];

#[derive(Clone, Debug, PartialEq)]
pub struct AppliedIntervention {
    pub body: Value,
    pub lever: String,
    pub session_key: String,
    pub context_pair: Option<String>,
    pub effort_ceiling: Option<String>,
    pub changed: bool,
}

#[derive(Clone, Debug, Default)]
struct SessionInterventions {
    last_prefix: Option<u64>,
    last_seen: Option<Instant>,
    effort: Option<String>,
    rejected_context_pairs: HashSet<String>,
}

/// Bounded facts needed by relay actuators; prompt bytes are never retained.
#[derive(Default)]
pub struct InterventionSessions {
    sessions: Mutex<HashMap<String, SessionInterventions>>,
}

impl InterventionSessions {
    pub fn apply(
        &self,
        session_key: &str,
        body: &Value,
        params: &OptimizeParams,
        billing: BillingMode,
        prefix_churn: bool,
        now: Instant,
    ) -> Result<AppliedIntervention, &'static str> {
        let mut sessions = self.sessions.lock().map_err(|_| "guarded")?;
        if sessions.len() >= 1_024 && !sessions.contains_key(session_key) {
            return Err("guarded");
        }
        let state = sessions.entry(session_key.to_string()).or_default();
        match params {
            OptimizeParams::EffortClamp(params) => {
                effort_clamp(session_key, body, params, state, now)
            }
            OptimizeParams::ContextEdit(params) => context_edit(session_key, body, params, state),
            OptimizeParams::PrefixGuard(params) => {
                prefix_guard(session_key, body, params, billing, prefix_churn)
            }
            _ => Err("incapable"),
        }
    }

    pub fn accepted(&self, applied: &AppliedIntervention, now: Instant) {
        let Ok(mut sessions) = self.sessions.lock() else {
            return;
        };
        let state = sessions.entry(applied.session_key.clone()).or_default();
        state.last_seen = Some(now);
        state.last_prefix = Some(prefix_hash(&applied.body));
        if applied.lever == "effort_clamp" {
            state.effort = applied.effort_ceiling.clone();
        }
    }

    pub fn rejected(&self, applied: &AppliedIntervention) {
        let Some(pair) = applied.context_pair.as_ref() else {
            return;
        };
        let Ok(mut sessions) = self.sessions.lock() else {
            return;
        };
        sessions
            .entry(applied.session_key.clone())
            .or_default()
            .rejected_context_pairs
            .insert(pair.clone());
    }
}

fn effort_clamp(
    session_key: &str,
    body: &Value,
    params: &serde_json::Map<String, Value>,
    state: &SessionInterventions,
    now: Instant,
) -> Result<AppliedIntervention, &'static str> {
    let current = body
        .pointer("/output_config/effort")
        .and_then(Value::as_str)
        .ok_or("incapable")?;
    let requested = params
        .get("effort")
        .and_then(Value::as_str)
        .ok_or("guarded")?;
    let floor = params
        .get("quality_floor")
        .and_then(Value::as_str)
        .ok_or("guarded")?;
    let rank = |v: &str| match v {
        "low" => Some(0),
        "medium" => Some(1),
        "high" => Some(2),
        "max" => Some(3),
        _ => None,
    };
    let (Some(current_rank), Some(requested_rank), Some(floor_rank)) =
        (rank(current), rank(requested), rank(floor))
    else {
        return Err("guarded");
    };
    let cold = state.last_seen.is_none()
        || state
            .last_seen
            .is_some_and(|last| now.saturating_duration_since(last) >= CACHE_TTL)
        || is_compaction(body)
        || state
            .last_prefix
            .is_some_and(|hash| hash != prefix_hash(body));
    let selected_rank = requested_rank.max(floor_rank);
    let ceiling_rank = if cold {
        selected_rank
    } else {
        state.effort.as_deref().and_then(rank).ok_or("guarded")?
    };
    let target_rank = current_rank.min(ceiling_rank);
    let target = ["low", "medium", "high", "max"][target_rank as usize];
    let mut edited = body.clone();
    edited["output_config"]["effort"] = Value::String(target.to_string());
    Ok(AppliedIntervention {
        changed: target != current,
        body: edited,
        lever: "effort_clamp".into(),
        session_key: session_key.into(),
        context_pair: None,
        effort_ceiling: Some(["low", "medium", "high", "max"][ceiling_rank as usize].into()),
    })
}

fn context_edit(
    session_key: &str,
    body: &Value,
    params: &serde_json::Map<String, Value>,
    state: &SessionInterventions,
) -> Result<AppliedIntervention, &'static str> {
    if is_compaction(body) {
        return Err("guarded");
    }
    let cm = body
        .get("context_management")
        .and_then(Value::as_object)
        .ok_or("incapable")?;
    let edits = cm
        .get("edits")
        .and_then(Value::as_array)
        .ok_or("incapable")?;
    let edit_index = edits
        .iter()
        .position(|edit| {
            edit.get("type").and_then(Value::as_str) == Some("clear_tool_uses_20250919")
        })
        .ok_or("incapable")?;
    let edit = edits[edit_index].as_object().ok_or("guarded")?;
    let mut edited_cm = cm.clone();
    let incoming_trigger = number(params, "trigger")?;
    let incoming_keep = number(params, "keep")?;
    let incoming_clear = number(params, "clear_at_least")?;
    if !context_edit_net_positive(body, incoming_keep, incoming_clear, params)? {
        return Err("guarded");
    }
    let current_trigger = typed_control(edit.get("trigger"), "input_tokens")?;
    let current_keep = typed_control(edit.get("keep"), "tool_uses")?;
    let current_clear = typed_control(edit.get("clear_at_least"), "input_tokens")?;
    let mut edited_edit = edit.clone();
    edited_edit.insert(
        "trigger".into(),
        json!({"type":"input_tokens","value":incoming_trigger.min(current_trigger)}),
    );
    edited_edit.insert(
        "keep".into(),
        json!({"type":"tool_uses","value":incoming_keep.min(current_keep)}),
    );
    edited_edit.insert(
        "clear_at_least".into(),
        json!({"type":"input_tokens","value":incoming_clear.max(current_clear)}),
    );
    let mut excluded: Vec<String> = edit
        .get("exclude_tools")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .filter_map(Value::as_str)
        .map(str::to_string)
        .collect();
    excluded.extend(
        params
            .get("exclude_tools")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
            .filter_map(Value::as_str)
            .map(str::to_string),
    );
    excluded.extend(PROTECTED_TOOLS.map(str::to_string));
    excluded.sort();
    excluded.dedup();
    edited_edit.insert("exclude_tools".into(), json!(excluded));
    let mut edited_edits = edits.clone();
    edited_edits[edit_index] = Value::Object(edited_edit);
    edited_cm.insert("edits".into(), Value::Array(edited_edits));
    let pair = Value::Object(edited_cm.clone()).to_string();
    if state.rejected_context_pairs.contains(&pair) {
        return Err("guarded");
    }
    let mut edited = body.clone();
    edited["context_management"] = Value::Object(edited_cm);
    if edited == *body {
        return Err("guarded");
    }
    if !context_edit_is_structurally_valid(body, &edited, edit_index) {
        return Err("guarded");
    }
    Ok(AppliedIntervention {
        changed: true,
        body: edited,
        lever: "context_edit".into(),
        session_key: session_key.into(),
        context_pair: Some(pair),
        effort_ceiling: None,
    })
}

fn prefix_guard(
    session_key: &str,
    body: &Value,
    params: &serde_json::Map<String, Value>,
    billing: BillingMode,
    prefix_churn: bool,
) -> Result<AppliedIntervention, &'static str> {
    if billing != BillingMode::ApiKey {
        return Err("guarded");
    }
    if !prefix_churn {
        return Err("guarded");
    }
    let ttl = params
        .get("pin_ttl")
        .and_then(Value::as_str)
        .unwrap_or("1h");
    if ttl != "1h" {
        return Err("guarded");
    }
    let blocks = body
        .get("system")
        .and_then(Value::as_array)
        .ok_or("incapable")?;
    let index = blocks
        .iter()
        .rposition(|block| block.get("cache_control").is_some())
        .ok_or("guarded")?;
    let mut edited = body.clone();
    edited["system"][index]["cache_control"] = json!({"type":"ephemeral", "ttl":"1h"});
    if edited == *body {
        return Err("guarded");
    }
    Ok(AppliedIntervention {
        changed: true,
        body: edited,
        lever: "prefix_guard".into(),
        session_key: session_key.into(),
        context_pair: None,
        effort_ceiling: None,
    })
}

fn number(params: &serde_json::Map<String, Value>, key: &str) -> Result<u64, &'static str> {
    params.get(key).and_then(Value::as_u64).ok_or("guarded")
}

fn typed_control(value: Option<&Value>, expected_type: &str) -> Result<u64, &'static str> {
    let value = value.and_then(Value::as_object).ok_or("guarded")?;
    if value.get("type").and_then(Value::as_str) != Some(expected_type) {
        return Err("guarded");
    }
    value.get("value").and_then(Value::as_u64).ok_or("guarded")
}

/**
 * Measure only tool-result blocks the configured strategy can actually clear,
 * then feed those gross/retained token estimates into the shared net equation.
 */
fn context_edit_net_positive(
    body: &Value,
    keep: u64,
    clear_at_least: u64,
    params: &serde_json::Map<String, Value>,
) -> Result<bool, &'static str> {
    let messages = body
        .get("messages")
        .and_then(Value::as_array)
        .ok_or("guarded")?;
    let excluded: HashSet<&str> = params
        .get("exclude_tools")
        .and_then(Value::as_array)
        .ok_or("guarded")?
        .iter()
        .filter_map(Value::as_str)
        .chain(PROTECTED_TOOLS)
        .collect();
    let mut tool_names: HashMap<&str, &str> = HashMap::new();
    for message in messages {
        for block in message
            .get("content")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
        {
            if block.get("type").and_then(Value::as_str) == Some("tool_use") {
                if let (Some(id), Some(name)) = (
                    block.get("id").and_then(Value::as_str),
                    block.get("name").and_then(Value::as_str),
                ) {
                    tool_names.insert(id, name);
                }
            }
        }
    }
    let mut eligible: Vec<(usize, u64)> = Vec::new();
    for (message_index, message) in messages.iter().enumerate() {
        for block in message
            .get("content")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
        {
            if block.get("type").and_then(Value::as_str) != Some("tool_result") {
                continue;
            }
            let id = block
                .get("tool_use_id")
                .and_then(Value::as_str)
                .ok_or("guarded")?;
            let name = tool_names.get(id).copied().ok_or("guarded")?;
            if !excluded.contains(name) {
                let bytes = serde_json::to_vec(block).map_err(|_| "guarded")?.len() as u64;
                eligible.push((message_index, bytes.div_ceil(4)));
            }
        }
    }
    let removable = eligible.len().saturating_sub(keep as usize);
    if removable == 0 {
        return Ok(false);
    }
    let gross: u64 = eligible[..removable]
        .iter()
        .map(|(_, tokens)| *tokens)
        .sum();
    if gross < clear_at_least {
        return Ok(false);
    }
    let first_removed = eligible[0].0;
    let retained_bytes = serde_json::to_vec(&messages[first_removed + 1..])
        .map_err(|_| "guarded")?
        .len() as u64;
    Ok(super::net::tokens_net_hundredths(
        gross,
        retained_bytes.div_ceil(4),
        super::net::write_multiplier_for(body),
    ) > 0)
}

fn context_edit_is_structurally_valid(original: &Value, edited: &Value, index: usize) -> bool {
    let (Some(before), Some(after)) = (original.as_object(), edited.as_object()) else {
        return false;
    };
    if before.len() != after.len()
        || before
            .iter()
            .any(|(key, value)| key != "context_management" && after.get(key) != Some(value))
    {
        return false;
    }
    let (Some(before_cm), Some(after_cm)) = (
        before.get("context_management").and_then(Value::as_object),
        after.get("context_management").and_then(Value::as_object),
    ) else {
        return false;
    };
    if before_cm.len() != after_cm.len()
        || before_cm
            .iter()
            .any(|(key, value)| key != "edits" && after_cm.get(key) != Some(value))
    {
        return false;
    }
    let (Some(before_edits), Some(after_edits)) = (
        before_cm.get("edits").and_then(Value::as_array),
        after_cm.get("edits").and_then(Value::as_array),
    ) else {
        return false;
    };
    if before_edits.len() != after_edits.len()
        || before_edits
            .iter()
            .enumerate()
            .any(|(i, value)| i != index && after_edits.get(i) != Some(value))
    {
        return false;
    }
    let (Some(before_edit), Some(after_edit)) = (
        before_edits.get(index).and_then(Value::as_object),
        after_edits.get(index).and_then(Value::as_object),
    ) else {
        return false;
    };
    before_edit.len() == after_edit.len()
        && before_edit.iter().all(|(key, value)| {
            matches!(
                key.as_str(),
                "trigger" | "keep" | "clear_at_least" | "exclude_tools"
            ) || after_edit.get(key) == Some(value)
        })
        && after_edit.get("type").and_then(Value::as_str) == Some("clear_tool_uses_20250919")
        && typed_control(after_edit.get("trigger"), "input_tokens").is_ok()
        && typed_control(after_edit.get("keep"), "tool_uses").is_ok()
        && typed_control(after_edit.get("clear_at_least"), "input_tokens").is_ok()
}

fn is_compaction(body: &Value) -> bool {
    body.pointer("/context_management/compaction")
        .is_some_and(|v| v != &Value::Null && v != &Value::Bool(false))
        || body.get("is_compaction").and_then(Value::as_bool) == Some(true)
        || body
            .pointer("/context_management/edits")
            .and_then(Value::as_array)
            .is_some_and(|edits| {
                edits.iter().any(|edit| {
                    edit.get("type").and_then(Value::as_str) == Some("compact_20260112")
                })
            })
}

fn prefix_hash(body: &Value) -> u64 {
    use std::hash::{Hash, Hasher};
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    body.get("model").hash(&mut hasher);
    body.get("tools").hash(&mut hasher);
    body.get("system").hash(&mut hasher);
    hasher.finish()
}

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

    #[test]
    fn effort_clamp_never_raises_and_only_changes_a_cold_session() {
        let sessions = InterventionSessions::default();
        let body = json!({"model":"claude-opus-4-8","output_config":{"effort":"high"},"system":[],"messages":[]});
        let params = OptimizeParams::EffortClamp(
            json!({"effort":"low","quality_floor":"medium"})
                .as_object()
                .unwrap()
                .clone(),
        );
        let applied = sessions
            .apply(
                "s",
                &body,
                &params,
                BillingMode::ApiKey,
                false,
                Instant::now(),
            )
            .unwrap();
        assert_eq!(applied.body["output_config"]["effort"], "medium");
        sessions.accepted(&applied, Instant::now());
        let warm = sessions
            .apply(
                "s",
                &body,
                &params,
                BillingMode::ApiKey,
                false,
                Instant::now(),
            )
            .unwrap();
        assert_eq!(warm.body["output_config"]["effort"], "medium");

        let lower = json!({"model":"claude-opus-4-8","output_config":{"effort":"low"},"system":[],"messages":[]});
        let warm_lower = sessions
            .apply(
                "s",
                &lower,
                &params,
                BillingMode::ApiKey,
                false,
                Instant::now(),
            )
            .unwrap();
        assert_eq!(warm_lower.body["output_config"]["effort"], "low");
        assert!(!warm_lower.changed, "a lower caller effort is never raised");
    }

    #[test]
    fn prefix_guard_only_rewrites_the_last_existing_breakpoint() {
        let sessions = InterventionSessions::default();
        let body = json!({"system":[
            {"type":"text","text":"a","cache_control":{"type":"ephemeral"}},
            {"type":"text","text":"b"},
            {"type":"text","text":"c","cache_control":{"type":"ephemeral"}}
        ],"messages":[]});
        let params =
            OptimizeParams::PrefixGuard(json!({"pin_ttl":"1h"}).as_object().unwrap().clone());
        let applied = sessions
            .apply(
                "s",
                &body,
                &params,
                BillingMode::ApiKey,
                true,
                Instant::now(),
            )
            .unwrap();
        assert_eq!(applied.body["system"][0], body["system"][0]);
        assert_eq!(applied.body["system"][1], body["system"][1]);
        assert_eq!(
            applied.body["system"][2]["cache_control"],
            json!({"type":"ephemeral","ttl":"1h"})
        );
    }

    #[test]
    fn context_edit_is_monotone_and_protects_file_and_thinking_tools() {
        let sessions = InterventionSessions::default();
        let body = json!({
            "context_management":{"edits":[
                {"type":"clear_thinking_20251015","keep":{"type":"thinking_turns","value":2}},
                {"type":"clear_tool_uses_20250919","trigger":{"type":"input_tokens","value":100},"keep":{"type":"tool_uses","value":1},"clear_at_least":{"type":"input_tokens","value":20},"exclude_tools":["memory"]}
            ]},
            "messages":[
                {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Search","input":{}}]},
                {"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":"x".repeat(4000)}]},
                {"role":"user","content":"tail"}
            ]
        });
        let params = OptimizeParams::ContextEdit(
            json!({"trigger":80,"keep":0,"clear_at_least":20,"exclude_tools":[]})
                .as_object()
                .unwrap()
                .clone(),
        );
        let applied = sessions
            .apply(
                "s",
                &body,
                &params,
                BillingMode::ApiKey,
                false,
                Instant::now(),
            )
            .unwrap();
        let edit = &applied.body["context_management"]["edits"][1];
        assert_eq!(edit["trigger"], json!({"type":"input_tokens","value":80}));
        assert_eq!(edit["keep"], json!({"type":"tool_uses","value":0}));
        assert_eq!(
            edit["clear_at_least"],
            json!({"type":"input_tokens","value":20})
        );
        let excluded = edit["exclude_tools"].as_array().unwrap();
        for protected in ["Edit", "Write", "Read", "memory"] {
            assert!(excluded.contains(&json!(protected)));
        }
        assert_eq!(
            applied.body["context_management"]["edits"][0],
            body["context_management"]["edits"][0]
        );
    }
}