Skip to main content

codex_helper_core/provider_signals/
model.rs

1use serde::{Deserialize, Serialize};
2
3use crate::runtime_identity::ProviderEndpointKey;
4
5#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)]
6#[serde(rename_all = "snake_case")]
7pub enum ProviderSignalConfidence {
8    Low,
9    #[default]
10    Medium,
11    High,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
15#[serde(rename_all = "snake_case")]
16pub enum ProviderSignalKind {
17    Quota,
18    RateLimit,
19    Capacity,
20    Transport,
21    Balance,
22    ServiceStatus,
23    Capability,
24    LocalConcurrency,
25    Unknown,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
29#[serde(rename_all = "snake_case")]
30pub enum ProviderSignalSource {
31    UpstreamResponse,
32    ResponseHeaders,
33    BalanceSnapshot,
34    ServiceStatus,
35    CapabilityProbe,
36    LocalScheduler,
37    RouteAttempt,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
41#[serde(rename_all = "snake_case")]
42pub enum ProviderSignalTarget {
43    ProviderEndpoint {
44        provider_endpoint_key: ProviderEndpointKey,
45    },
46    Provider {
47        service: String,
48        provider_id: String,
49    },
50    Service {
51        service: String,
52    },
53}
54
55impl ProviderSignalTarget {
56    pub fn provider_endpoint_key(&self) -> Option<&ProviderEndpointKey> {
57        match self {
58            Self::ProviderEndpoint {
59                provider_endpoint_key,
60            } => Some(provider_endpoint_key),
61            Self::Provider { .. } | Self::Service { .. } => None,
62        }
63    }
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
67pub struct ProviderSignalTrace {
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub trace_id: Option<String>,
70    #[serde(skip_serializing_if = "Option::is_none")]
71    pub cf_ray: Option<String>,
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub upstream_request_id: Option<String>,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
77pub struct ProviderSignal {
78    pub kind: ProviderSignalKind,
79    pub source: ProviderSignalSource,
80    pub target: ProviderSignalTarget,
81    pub confidence: ProviderSignalConfidence,
82    pub observed_at_ms: u64,
83    #[serde(default, skip_serializing_if = "bool_is_false")]
84    pub route_facing: bool,
85    #[serde(skip_serializing_if = "Option::is_none")]
86    pub retry_after_secs: Option<u64>,
87    #[serde(skip_serializing_if = "Option::is_none")]
88    pub reset_after_secs: Option<u64>,
89    #[serde(skip_serializing_if = "Option::is_none")]
90    pub reason: Option<String>,
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub error_class: Option<String>,
93    #[serde(default, skip_serializing_if = "ProviderSignalTrace::is_empty")]
94    pub trace: ProviderSignalTrace,
95}
96
97impl ProviderSignal {
98    pub fn high_confidence_route_facing(
99        kind: ProviderSignalKind,
100        source: ProviderSignalSource,
101        target: ProviderSignalTarget,
102        observed_at_ms: u64,
103    ) -> Self {
104        Self {
105            kind,
106            source,
107            target,
108            confidence: ProviderSignalConfidence::High,
109            observed_at_ms,
110            route_facing: true,
111            retry_after_secs: None,
112            reset_after_secs: None,
113            reason: None,
114            error_class: None,
115            trace: ProviderSignalTrace::default(),
116        }
117    }
118
119    pub fn cooldown_horizon_secs(&self) -> Option<u64> {
120        self.reset_after_secs.or(self.retry_after_secs)
121    }
122
123    pub fn is_high_confidence_route_facing(&self) -> bool {
124        self.route_facing && self.confidence >= ProviderSignalConfidence::High
125    }
126}
127
128impl ProviderSignalTrace {
129    pub fn is_empty(&self) -> bool {
130        self.trace_id.is_none() && self.cf_ray.is_none() && self.upstream_request_id.is_none()
131    }
132}
133
134fn bool_is_false(value: &bool) -> bool {
135    !*value
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn provider_endpoint_target_exposes_stable_key() {
144        let key = ProviderEndpointKey::new("codex", "monthly", "default");
145        let target = ProviderSignalTarget::ProviderEndpoint {
146            provider_endpoint_key: key.clone(),
147        };
148
149        assert_eq!(target.provider_endpoint_key(), Some(&key));
150    }
151
152    #[test]
153    fn cooldown_horizon_prefers_reset_then_retry_after() {
154        let mut signal = ProviderSignal::high_confidence_route_facing(
155            ProviderSignalKind::Quota,
156            ProviderSignalSource::UpstreamResponse,
157            ProviderSignalTarget::ProviderEndpoint {
158                provider_endpoint_key: ProviderEndpointKey::new("codex", "monthly", "default"),
159            },
160            100,
161        );
162        signal.retry_after_secs = Some(30);
163        signal.reset_after_secs = Some(60);
164
165        assert_eq!(signal.cooldown_horizon_secs(), Some(60));
166        assert!(signal.is_high_confidence_route_facing());
167    }
168}