codex-helper-core 0.20.1

use serde::{Deserialize, Serialize};

use crate::provider_signals::{
    ProviderSignal, ProviderSignalConfidence, ProviderSignalKind, ProviderSignalTarget,
};
use crate::runtime_identity::ProviderEndpointKey;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum PolicyActionKind {
    Cooldown,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum PolicyActionOwner {
    CodexHelper,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "snake_case")]
pub enum PolicyActionRecoveryState {
    #[default]
    Active,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PolicyAction {
    pub id: String,
    pub kind: PolicyActionKind,
    pub owner: PolicyActionOwner,
    pub provider_endpoint_key: ProviderEndpointKey,
    pub source_signal: ProviderSignal,
    pub reason: String,
    pub confidence: ProviderSignalConfidence,
    pub created_at_ms: u64,
    pub expires_at_ms: u64,
    #[serde(default)]
    pub recovery_state: PolicyActionRecoveryState,
    #[serde(default)]
    pub generation: u64,
}

impl PolicyAction {
    pub fn cooldown_from_signal(
        signal: ProviderSignal,
        created_at_ms: u64,
        default_cooldown_secs: u64,
        generation: u64,
    ) -> Option<Self> {
        if !signal.is_high_confidence_route_facing() {
            return None;
        }
        if !matches!(
            signal.kind,
            ProviderSignalKind::Quota
                | ProviderSignalKind::RateLimit
                | ProviderSignalKind::Capacity
                | ProviderSignalKind::Transport
                | ProviderSignalKind::Balance
        ) {
            return None;
        }
        let provider_endpoint_key = match &signal.target {
            ProviderSignalTarget::ProviderEndpoint {
                provider_endpoint_key,
            } => provider_endpoint_key.clone(),
            ProviderSignalTarget::Provider { .. } | ProviderSignalTarget::Service { .. } => {
                return None;
            }
        };
        let cooldown_secs = signal.cooldown_horizon_secs().or_else(|| {
            (matches!(
                signal.kind,
                ProviderSignalKind::Capacity | ProviderSignalKind::Transport
            ) && default_cooldown_secs > 0)
                .then_some(default_cooldown_secs)
        })?;
        if cooldown_secs == 0 {
            return None;
        }

        let reason = signal
            .reason
            .clone()
            .or_else(|| signal.error_class.clone())
            .unwrap_or_else(|| format!("{:?}", signal.kind).to_ascii_lowercase());
        Some(Self {
            id: format!(
                "codex-helper:{}:{}",
                provider_endpoint_key.stable_key(),
                created_at_ms
            ),
            kind: PolicyActionKind::Cooldown,
            owner: PolicyActionOwner::CodexHelper,
            provider_endpoint_key,
            source_signal: signal.clone(),
            reason,
            confidence: signal.confidence,
            created_at_ms,
            expires_at_ms: created_at_ms.saturating_add(cooldown_secs.saturating_mul(1000)),
            recovery_state: PolicyActionRecoveryState::Active,
            generation,
        })
    }

    pub fn is_active_at(&self, now_ms: u64) -> bool {
        self.recovery_state == PolicyActionRecoveryState::Active && now_ms < self.expires_at_ms
    }

    pub fn remaining_secs_at(&self, now_ms: u64) -> Option<u64> {
        self.is_active_at(now_ms)
            .then(|| self.expires_at_ms.saturating_sub(now_ms).div_ceil(1000))
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PolicyActionProjection {
    pub provider_endpoint_key: ProviderEndpointKey,
    pub active_cooldown: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cooldown_remaining_secs: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub action_id: Option<String>,
}

impl PolicyActionProjection {
    pub fn from_action(action: &PolicyAction, now_ms: u64) -> Option<Self> {
        let cooldown_remaining_secs = action.remaining_secs_at(now_ms)?;
        Some(Self {
            provider_endpoint_key: action.provider_endpoint_key.clone(),
            active_cooldown: matches!(action.kind, PolicyActionKind::Cooldown),
            cooldown_remaining_secs: Some(cooldown_remaining_secs),
            reason: Some(action.reason.clone()),
            action_id: Some(action.id.clone()),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::provider_signals::{ProviderSignalSource, ProviderSignalTarget};

    fn quota_signal(reset_after_secs: Option<u64>) -> ProviderSignal {
        let mut signal = ProviderSignal::high_confidence_route_facing(
            ProviderSignalKind::Quota,
            ProviderSignalSource::UpstreamResponse,
            ProviderSignalTarget::ProviderEndpoint {
                provider_endpoint_key: ProviderEndpointKey::new("codex", "monthly", "default"),
            },
            100,
        );
        signal.reset_after_secs = reset_after_secs;
        signal.reason = Some("usage_limit_reached".to_string());
        signal
    }

    #[test]
    fn high_confidence_quota_signal_creates_owned_cooldown() {
        let action = PolicyAction::cooldown_from_signal(quota_signal(Some(30)), 1_000, 0, 7)
            .expect("cooldown action");

        assert_eq!(action.owner, PolicyActionOwner::CodexHelper);
        assert_eq!(action.expires_at_ms, 31_000);
        assert_eq!(action.generation, 7);
        assert!(action.is_active_at(30_999));
        assert!(!action.is_active_at(31_000));
    }

    #[test]
    fn quota_without_horizon_is_recorded_only() {
        assert!(PolicyAction::cooldown_from_signal(quota_signal(None), 1_000, 0, 1).is_none());
    }
}