codex_helper_core/policy_actions/
model.rs

1use serde::{Deserialize, Serialize};
2
3use crate::provider_signals::{
4    ProviderSignal, ProviderSignalConfidence, ProviderSignalKind, ProviderSignalTarget,
5};
6use crate::runtime_identity::ProviderEndpointKey;
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "snake_case")]
10pub enum PolicyActionKind {
11    Cooldown,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
15#[serde(rename_all = "snake_case")]
16pub enum PolicyActionOwner {
17    CodexHelper,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
21#[serde(rename_all = "snake_case")]
22pub enum PolicyActionRecoveryState {
23    #[default]
24    Active,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
28pub struct PolicyAction {
29    pub id: String,
30    pub kind: PolicyActionKind,
31    pub owner: PolicyActionOwner,
32    pub provider_endpoint_key: ProviderEndpointKey,
33    pub source_signal: ProviderSignal,
34    pub reason: String,
35    pub confidence: ProviderSignalConfidence,
36    pub created_at_ms: u64,
37    pub expires_at_ms: u64,
38    #[serde(default)]
39    pub recovery_state: PolicyActionRecoveryState,
40    #[serde(default)]
41    pub generation: u64,
42}
43
44impl PolicyAction {
45    pub fn cooldown_from_signal(
46        signal: ProviderSignal,
47        created_at_ms: u64,
48        default_cooldown_secs: u64,
49        generation: u64,
50    ) -> Option<Self> {
51        if !signal.is_high_confidence_route_facing() {
52            return None;
53        }
54        if !matches!(
55            signal.kind,
56            ProviderSignalKind::Quota
57                | ProviderSignalKind::RateLimit
58                | ProviderSignalKind::Capacity
59                | ProviderSignalKind::Transport
60                | ProviderSignalKind::Balance
61        ) {
62            return None;
63        }
64        let provider_endpoint_key = match &signal.target {
65            ProviderSignalTarget::ProviderEndpoint {
66                provider_endpoint_key,
67            } => provider_endpoint_key.clone(),
68            ProviderSignalTarget::Provider { .. } | ProviderSignalTarget::Service { .. } => {
69                return None;
70            }
71        };
72        let cooldown_secs = signal.cooldown_horizon_secs().or_else(|| {
73            (matches!(
74                signal.kind,
75                ProviderSignalKind::Capacity | ProviderSignalKind::Transport
76            ) && default_cooldown_secs > 0)
77                .then_some(default_cooldown_secs)
78        })?;
79        if cooldown_secs == 0 {
80            return None;
81        }
82
83        let reason = signal
84            .reason
85            .clone()
86            .or_else(|| signal.error_class.clone())
87            .unwrap_or_else(|| format!("{:?}", signal.kind).to_ascii_lowercase());
88        Some(Self {
89            id: format!(
90                "codex-helper:{}:{}",
91                provider_endpoint_key.stable_key(),
92                created_at_ms
93            ),
94            kind: PolicyActionKind::Cooldown,
95            owner: PolicyActionOwner::CodexHelper,
96            provider_endpoint_key,
97            source_signal: signal.clone(),
98            reason,
99            confidence: signal.confidence,
100            created_at_ms,
101            expires_at_ms: created_at_ms.saturating_add(cooldown_secs.saturating_mul(1000)),
102            recovery_state: PolicyActionRecoveryState::Active,
103            generation,
104        })
105    }
106
107    pub fn is_active_at(&self, now_ms: u64) -> bool {
108        self.recovery_state == PolicyActionRecoveryState::Active && now_ms < self.expires_at_ms
109    }
110
111    pub fn remaining_secs_at(&self, now_ms: u64) -> Option<u64> {
112        self.is_active_at(now_ms)
113            .then(|| self.expires_at_ms.saturating_sub(now_ms).div_ceil(1000))
114    }
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
118pub struct PolicyActionProjection {
119    pub provider_endpoint_key: ProviderEndpointKey,
120    pub active_cooldown: bool,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub cooldown_remaining_secs: Option<u64>,
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub reason: Option<String>,
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub action_id: Option<String>,
127}
128
129impl PolicyActionProjection {
130    pub fn from_action(action: &PolicyAction, now_ms: u64) -> Option<Self> {
131        let cooldown_remaining_secs = action.remaining_secs_at(now_ms)?;
132        Some(Self {
133            provider_endpoint_key: action.provider_endpoint_key.clone(),
134            active_cooldown: matches!(action.kind, PolicyActionKind::Cooldown),
135            cooldown_remaining_secs: Some(cooldown_remaining_secs),
136            reason: Some(action.reason.clone()),
137            action_id: Some(action.id.clone()),
138        })
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::provider_signals::{ProviderSignalSource, ProviderSignalTarget};
146
147    fn quota_signal(reset_after_secs: Option<u64>) -> ProviderSignal {
148        let mut signal = ProviderSignal::high_confidence_route_facing(
149            ProviderSignalKind::Quota,
150            ProviderSignalSource::UpstreamResponse,
151            ProviderSignalTarget::ProviderEndpoint {
152                provider_endpoint_key: ProviderEndpointKey::new("codex", "monthly", "default"),
153            },
154            100,
155        );
156        signal.reset_after_secs = reset_after_secs;
157        signal.reason = Some("usage_limit_reached".to_string());
158        signal
159    }
160
161    #[test]
162    fn high_confidence_quota_signal_creates_owned_cooldown() {
163        let action = PolicyAction::cooldown_from_signal(quota_signal(Some(30)), 1_000, 0, 7)
164            .expect("cooldown action");
165
166        assert_eq!(action.owner, PolicyActionOwner::CodexHelper);
167        assert_eq!(action.expires_at_ms, 31_000);
168        assert_eq!(action.generation, 7);
169        assert!(action.is_active_at(30_999));
170        assert!(!action.is_active_at(31_000));
171    }
172
173    #[test]
174    fn quota_without_horizon_is_recorded_only() {
175        assert!(PolicyAction::cooldown_from_signal(quota_signal(None), 1_000, 0, 1).is_none());
176    }
177}
codex_helper_core/policy_actions/model.rs

codex_helper_core/policy_actions/
model.rs
