caretta 0.17.0

//! Optional Caretta learning-engine integration.
//!
//! This module is deliberately fail-open. The engine can rank candidates and
//! provide evidence, but every response is checked locally and the engine is
//! never given a path to execute commands or widen Caretta's safety envelope.

use crate::agent::config_store::{LEARNING_ENGINE_TOKEN_SLOT, load_secret};
use crate::agent::process::emit_event;
use crate::agent::types::{
    AgentEvent, Config, DecisionRequest, DecisionResponse, LEARNING_PROTOCOL_VERSION,
    LearningCandidate, LearningEvent, LearningKey, LearningKeySet, LearningMode, ObservationRecord,
    PolicyManifest, RankedCandidate, SignedPolicyManifest,
};
use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};
use std::collections::HashSet;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

const DEFAULT_DECISION_POINT: &str = "agent.run";
const MAX_TIMEOUT_MS: u64 = 30_000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LearningError(pub String);

impl std::fmt::Display for LearningError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

pub trait LearningAdapter {
    fn recommend(&self, request: &DecisionRequest) -> Result<DecisionResponse, LearningError>;
    fn observe(&self, record: &ObservationRecord) -> Result<(), LearningError>;
}

#[derive(Clone, Debug)]
pub struct AcceptedRecommendation {
    pub ranked_candidates: Vec<RankedCandidate>,
    pub policy: Option<PolicyManifest>,
}

/// The host-side decision gate. Keep this separate from the HTTP client so
/// malformed, stale, over-permissive, or mis-signed responses are testable
/// without a running service.
pub struct DecisionGate;

impl DecisionGate {
    pub fn validate(
        request: &DecisionRequest,
        response: &DecisionResponse,
        key: Option<&LearningKey>,
        configured_key_uuid: Option<&str>,
        max_age: Duration,
        now_ms: u64,
    ) -> Result<AcceptedRecommendation, LearningError> {
        if response.protocol_version != LEARNING_PROTOCOL_VERSION {
            return Err(LearningError("learning response protocol mismatch".into()));
        }
        if response.run_id != request.run_id
            || response.decision_id != request.decision_id
            || response.decision_point != request.decision_point
        {
            return Err(LearningError(
                "learning response correlation mismatch".into(),
            ));
        }
        if response.generated_at_ms > now_ms
            || now_ms.saturating_sub(response.generated_at_ms) > max_age.as_millis() as u64
        {
            return Err(LearningError("learning response is stale".into()));
        }
        if response.expires_at_ms <= now_ms {
            return Err(LearningError("learning response has expired".into()));
        }
        if !response
            .expected_value
            .is_none_or(|value| value.is_finite())
            || !response
                .uncertainty
                .is_none_or(|value| value.is_finite() && (0.0..=1.0).contains(&value))
        {
            return Err(LearningError(
                "learning response contains unbounded values".into(),
            ));
        }

        let candidate_ids: HashSet<&str> =
            request.candidates.iter().map(|c| c.id.as_str()).collect();
        let mut seen = HashSet::new();
        for candidate in &response.ranked_candidates {
            if !candidate_ids.contains(candidate.candidate_id.as_str()) {
                return Err(LearningError(format!(
                    "learning response named unknown candidate {}",
                    candidate.candidate_id
                )));
            }
            if !seen.insert(candidate.candidate_id.as_str()) {
                return Err(LearningError(
                    "learning response contains duplicate candidates".into(),
                ));
            }
            if !candidate.expected_value.is_finite()
                || !candidate.uncertainty.is_finite()
                || !(0.0..=1.0).contains(&candidate.uncertainty)
                || !candidate.parameter_patch.is_object()
            {
                return Err(LearningError(
                    "learning candidate has invalid values or patch".into(),
                ));
            }
            let Some(baseline) = request
                .candidates
                .iter()
                .find(|request_candidate| request_candidate.id == candidate.candidate_id)
                .map(|request_candidate| &request_candidate.parameters)
            else {
                return Err(LearningError(
                    "learning candidate baseline is missing".into(),
                ));
            };
            let Some(patch) = candidate.parameter_patch.as_object() else {
                return Err(LearningError(
                    "learning candidate patch is not an object".into(),
                ));
            };
            if patch.len() > baseline.as_object().map_or(0, serde_json::Map::len)
                || patch.iter().any(|(name, value)| {
                    baseline.get(name).is_none_or(|original| {
                        std::mem::discriminant(original) != std::mem::discriminant(value)
                    })
                })
            {
                return Err(LearningError(
                    "learning candidate patch violates its parameter schema".into(),
                ));
            }
        }

        let Some(signed) = response.policy_manifest.as_ref() else {
            if response.ranked_candidates.is_empty() {
                return Ok(AcceptedRecommendation {
                    ranked_candidates: vec![],
                    policy: None,
                });
            }
            return Err(LearningError(
                "ranked learning response has no policy manifest".into(),
            ));
        };
        validate_manifest(request, response, signed, key, configured_key_uuid, now_ms)?;

        Ok(AcceptedRecommendation {
            ranked_candidates: response.ranked_candidates.clone(),
            policy: Some(signed.manifest.clone()),
        })
    }
}

fn validate_manifest(
    request: &DecisionRequest,
    response: &DecisionResponse,
    signed: &SignedPolicyManifest,
    key: Option<&LearningKey>,
    configured_key_uuid: Option<&str>,
    now_ms: u64,
) -> Result<(), LearningError> {
    let manifest = &signed.manifest;
    if manifest.protocol_version != LEARNING_PROTOCOL_VERSION
        || manifest.run_id != request.run_id
        || manifest.decision_id != request.decision_id
        || manifest.decision_point != request.decision_point
        || manifest.local_safety_envelope_digest != request.local_safety_envelope_digest
        || manifest.expires_at_ms != response.expires_at_ms
        || manifest.issued_at_ms > now_ms
        || manifest.expires_at_ms <= now_ms
    {
        return Err(LearningError(
            "learning policy manifest integrity check failed".into(),
        ));
    }
    if configured_key_uuid.is_some_and(|configured| configured != manifest.key_uuid) {
        return Err(LearningError("learning policy key UUID mismatch".into()));
    }
    let key = key.ok_or_else(|| LearningError("learning policy signing key unavailable".into()))?;
    if key.key_uuid != manifest.key_uuid || key.algorithm != "EdDSA" {
        return Err(LearningError("learning policy signing key mismatch".into()));
    }
    let header = decode_header(&signed.signature)
        .map_err(|_| LearningError("learning policy signature header is invalid".into()))?;
    if header.alg != Algorithm::EdDSA || header.kid.as_deref() != Some(manifest.key_uuid.as_str()) {
        return Err(LearningError(
            "learning policy signature algorithm or key mismatch".into(),
        ));
    }
    let decoding_key = DecodingKey::from_ed_pem(key.public_key_pem.as_bytes())
        .map_err(|_| LearningError("learning policy public key is invalid".into()))?;
    let validation = Validation::new(Algorithm::EdDSA);
    let token = decode::<PolicyManifest>(&signed.signature, &decoding_key, &validation)
        .map_err(|_| LearningError("learning policy signature verification failed".into()))?;
    if token.claims != *manifest {
        return Err(LearningError(
            "signed learning policy payload differs from manifest".into(),
        ));
    }
    if !mode_is_within_cap(manifest.allowed_mode, request.mode_cap) {
        return Err(LearningError(
            "learning policy exceeds local mode cap".into(),
        ));
    }
    let allowed: HashSet<&str> = manifest
        .allowed_candidate_ids
        .iter()
        .map(String::as_str)
        .collect();
    if response
        .ranked_candidates
        .iter()
        .any(|candidate| !allowed.contains(candidate.candidate_id.as_str()))
    {
        return Err(LearningError(
            "learning policy does not cover ranked candidates".into(),
        ));
    }
    if !manifest.max_risk.is_finite()
        || !(0.0..=1.0).contains(&manifest.max_risk)
        || !manifest.max_budget_usd.is_finite()
        || manifest.max_budget_usd < 0.0
    {
        return Err(LearningError(
            "learning policy has invalid local bounds".into(),
        ));
    }
    Ok(())
}

fn mode_is_within_cap(mode: LearningMode, cap: LearningMode) -> bool {
    fn level(mode: LearningMode) -> u8 {
        match mode {
            LearningMode::Abstain => 0,
            LearningMode::Shadow => 1,
            LearningMode::Advisory => 2,
            LearningMode::ApprovalRequired => 3,
            LearningMode::Experimental => 4,
            LearningMode::Autonomous => 5,
        }
    }
    level(mode) <= level(cap)
}

#[cfg(not(target_arch = "wasm32"))]
#[derive(Clone)]
pub struct HttpLearningAdapter {
    endpoint: String,
    token: String,
    key_uuid: Option<String>,
    timeout: Duration,
    max_response_age: Duration,
}

#[cfg(not(target_arch = "wasm32"))]
impl HttpLearningAdapter {
    pub fn from_config(config: &Config) -> Option<Self> {
        let endpoint = config
            .learning
            .endpoint
            .as_deref()?
            .trim()
            .trim_end_matches('/');
        if endpoint.is_empty()
            || !(endpoint.starts_with("https://") || endpoint.starts_with("http://"))
        {
            return None;
        }
        let token = std::env::var("CARETTA_LEARNING_ENGINE_TOKEN")
            .ok()
            .map(|value| value.trim().to_string())
            .filter(|value| !value.is_empty())
            .or_else(|| load_secret(&config.root, LEARNING_ENGINE_TOKEN_SLOT))?;
        Some(Self {
            endpoint: endpoint.to_string(),
            token,
            key_uuid: config.learning.key_uuid.clone(),
            timeout: Duration::from_millis(config.learning.timeout_ms.clamp(1, MAX_TIMEOUT_MS)),
            max_response_age: Duration::from_millis(config.learning.max_response_age_ms.max(1)),
        })
    }

    fn agent(&self) -> ureq::Agent {
        ureq::AgentBuilder::new().timeout(self.timeout).build()
    }

    fn request_json<T: serde::Serialize, R: DeserializeOwned>(
        &self,
        method: &str,
        path: &str,
        body: Option<&T>,
        idempotency_key: Option<&str>,
    ) -> Result<R, LearningError> {
        let url = format!("{}{}", self.endpoint, path);
        let mut request = match method {
            "GET" => self.agent().get(&url),
            "POST" => self.agent().post(&url),
            _ => return Err(LearningError("unsupported learning HTTP method".into())),
        }
        .set("authorization", &format!("Bearer {}", self.token))
        .set("accept", "application/json")
        .set("content-type", "application/json");
        if let Some(idempotency_key) = idempotency_key {
            request = request.set("idempotency-key", idempotency_key);
        }
        let response = if let Some(body) = body {
            request.send_json(serde_json::to_value(body).map_err(|e| LearningError(e.to_string()))?)
        } else {
            request.call()
        }
        .map_err(|error| LearningError(safe_http_error(error)))?;
        response
            .into_json::<R>()
            .map_err(|error| LearningError(format!("invalid learning JSON response: {error}")))
    }

    fn keys(&self) -> Result<LearningKeySet, LearningError> {
        self.request_json::<(), LearningKeySet>(
            "GET",
            "/.well-known/caretta-learning-keys.json",
            None,
            None,
        )
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl LearningAdapter for HttpLearningAdapter {
    fn recommend(&self, request: &DecisionRequest) -> Result<DecisionResponse, LearningError> {
        let idempotency_key = format!("caretta-recommendation-{}", request.decision_id);
        let response: DecisionResponse = self.request_json(
            "POST",
            "/v1/recommendation",
            Some(request),
            Some(&idempotency_key),
        )?;
        let keys = self.keys()?;
        let key = self
            .key_uuid
            .as_deref()
            .and_then(|uuid| keys.keys.iter().find(|key| key.key_uuid == uuid))
            .or_else(|| keys.keys.first());
        DecisionGate::validate(
            request,
            &response,
            key,
            self.key_uuid.as_deref(),
            self.max_response_age,
            now_ms(),
        )?;
        Ok(response)
    }

    fn observe(&self, record: &ObservationRecord) -> Result<(), LearningError> {
        let key = format!("caretta-learning-{}", record.decision_id);
        let url = format!("{}{}", self.endpoint, "/v1/observation");
        self.agent()
            .post(&url)
            .set("authorization", &format!("Bearer {}", self.token))
            .set("content-type", "application/json")
            .set("idempotency-key", &key)
            .send_json(serde_json::to_value(record).map_err(|e| LearningError(e.to_string()))?)
            .map_err(|error| LearningError(safe_http_error(error)))?;
        Ok(())
    }
}

#[cfg(not(target_arch = "wasm32"))]
fn safe_http_error(error: ureq::Error) -> String {
    match error {
        ureq::Error::Status(status, _) => format!("learning engine returned HTTP {status}"),
        ureq::Error::Transport(error) => format!("learning engine unavailable: {error}"),
    }
}

pub fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

pub fn safety_envelope_digest(config: &Config) -> String {
    let mut digest = Sha256::new();
    digest.update(
        serde_json::to_vec(&(
            LEARNING_PROTOCOL_VERSION,
            config.auto_mode,
            config.dry_run,
            config.project_name.as_str(),
        ))
        .unwrap_or_default(),
    );
    digest
        .finalize()
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect()
}

pub fn new_run_id() -> String {
    format!("run-{}", now_ms())
}

pub fn new_decision_id(run_id: &str) -> String {
    format!("{run_id}-decision")
}

pub fn build_request(
    config: &Config,
    prompt: &str,
    run_id: &str,
    decision_id: &str,
) -> DecisionRequest {
    let decision_point = DEFAULT_DECISION_POINT.to_string();
    DecisionRequest {
        protocol_version: LEARNING_PROTOCOL_VERSION.to_string(),
        run_id: run_id.to_string(),
        decision_id: decision_id.to_string(),
        project_id: config.project_name.clone(),
        decision_point: decision_point.clone(),
        candidates: vec![LearningCandidate {
            id: "default".to_string(),
            parameters: serde_json::json!({
                "agent": config.agent.to_string(),
                "model": config.model,
                "prompt_present": !prompt.trim().is_empty(),
            }),
        }],
        experiment_context: serde_json::json!({ "project": config.project_name }),
        local_safety_envelope_digest: safety_envelope_digest(config),
        mode_cap: config.learning.mode_for(&decision_point),
        requested_at_ms: now_ms(),
    }
}

pub fn emit_recommendation(
    request: &DecisionRequest,
    response: Option<&DecisionResponse>,
    status: &str,
) {
    let (candidate_id, expected_value, uncertainty, evidence_refs) = response
        .and_then(|response| response.ranked_candidates.first())
        .map(|candidate| {
            (
                Some(candidate.candidate_id.clone()),
                Some(candidate.expected_value),
                Some(candidate.uncertainty),
                candidate.evidence_refs.clone(),
            )
        })
        .unwrap_or((None, None, None, Vec::new()));
    emit_event(AgentEvent::Learning(LearningEvent::Recommendation {
        run_id: request.run_id.clone(),
        decision_id: request.decision_id.clone(),
        decision_point: request.decision_point.clone(),
        mode: request.mode_cap,
        candidate_id,
        expected_value,
        uncertainty,
        evidence_refs,
        status: status.to_string(),
    }));
}

pub fn observation_for_run(
    request: &DecisionRequest,
    selected_candidate_id: Option<String>,
    ok: bool,
    elapsed_ms: u128,
    approval_path: &str,
) -> ObservationRecord {
    ObservationRecord {
        protocol_version: LEARNING_PROTOCOL_VERSION.to_string(),
        run_id: request.run_id.clone(),
        decision_id: request.decision_id.clone(),
        project_id: request.project_id.clone(),
        decision_point: request.decision_point.clone(),
        decisions: serde_json::json!({ "mode": request.mode_cap }),
        selected_candidate_id,
        rejected_candidate_ids: Vec::new(),
        actions: serde_json::json!({ "kind": "agent_run" }),
        costs: serde_json::json!({ "elapsed_ms": elapsed_ms }),
        outcome: serde_json::json!({ "success": ok }),
        approval_path: Some(approval_path.to_string()),
        rollback_status: Some("not_applicable".to_string()),
        policy_version: "local-envelope-v1".to_string(),
        experiment_version: "none".to_string(),
        provenance: serde_json::json!({ "source": "caretta-host" }),
        recorded_at_ms: now_ms(),
    }
}

/// Apply only the narrow host-owned patch schema. The engine cannot change the
/// provider, command, permissions, working directory, or safety flags.
pub fn apply_safe_candidate_patch(
    config: &Config,
    mode: LearningMode,
    auto_approved: bool,
    candidate: Option<&RankedCandidate>,
) -> Config {
    let mut effective = config.clone();
    let can_apply = matches!(mode, LearningMode::Experimental | LearningMode::Autonomous)
        || (mode == LearningMode::ApprovalRequired && auto_approved);
    if !can_apply || config.dry_run {
        return effective;
    }
    if let Some(candidate) = candidate
        && candidate.candidate_id == "default"
        && let Some(model) = candidate
            .parameter_patch
            .get("model")
            .and_then(|value| value.as_str())
        && !model.trim().is_empty()
        && model.len() <= 256
    {
        effective.model = model.to_string();
    }
    effective
}

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

    fn request() -> DecisionRequest {
        DecisionRequest {
            protocol_version: LEARNING_PROTOCOL_VERSION.into(),
            run_id: "run-1".into(),
            decision_id: "decision-1".into(),
            project_id: "project".into(),
            decision_point: DEFAULT_DECISION_POINT.into(),
            candidates: vec![LearningCandidate {
                id: "default".into(),
                parameters: serde_json::json!({}),
            }],
            experiment_context: serde_json::json!({}),
            local_safety_envelope_digest: "digest".into(),
            mode_cap: LearningMode::Shadow,
            requested_at_ms: 100,
        }
    }

    #[test]
    fn disabled_config_has_no_endpoint_and_shadow_mode() {
        let config = crate::agent::types::LearningConfig::default();
        assert!(config.endpoint.is_none());
        assert_eq!(
            config.mode_for(DEFAULT_DECISION_POINT),
            LearningMode::Shadow
        );
    }

    #[test]
    fn invalid_candidate_is_rejected_before_any_execution() {
        let request = request();
        let response = DecisionResponse {
            protocol_version: LEARNING_PROTOCOL_VERSION.into(),
            run_id: request.run_id.clone(),
            decision_id: request.decision_id.clone(),
            decision_point: request.decision_point.clone(),
            ranked_candidates: vec![RankedCandidate {
                candidate_id: "unknown".into(),
                parameter_patch: serde_json::json!({}),
                expected_value: 0.5,
                uncertainty: 0.1,
                evidence_refs: vec![],
            }],
            expected_value: Some(0.5),
            uncertainty: Some(0.1),
            evidence_refs: vec![],
            adapter_name: "test".into(),
            adapter_version: "1".into(),
            policy_manifest: None,
            expires_at_ms: 200,
            abstention_reason: None,
            generated_at_ms: 100,
        };
        assert!(
            DecisionGate::validate(&request, &response, None, None, Duration::from_secs(1), 100)
                .is_err()
        );
    }
}