Skip to main content

codex_helper_core/proxy/
codex_relay_capabilities.rs

1use axum::http::StatusCode;
2use serde::Deserialize;
3use serde::Serialize;
4use serde_json::Value;
5
6use crate::codex_capability_profile::{
7    CodexCapabilityDecision, CodexCapabilityProfile, CodexCapabilityProfileInput,
8    CodexCapabilitySupport, CodexModelCatalogProfile, CodexPatchModeRecommendation,
9    CodexPatchModeRecommendationInput,
10};
11use crate::codex_integration::CodexPatchMode;
12
13use super::codex_relay_probe::CodexRelayProbeObservation;
14use super::codex_relay_probe::codex_relay_probe_cases;
15use super::codex_relay_target::{CodexRelayTargetSelection, select_codex_relay_target};
16use super::models_compat::maybe_decode_models_response_body;
17use super::{
18    CodexRelayProbeClient, CodexRelayProbeKind, CodexRelayProbeResult, ProxyControlError,
19    ProxyService,
20};
21
22#[derive(Debug, Clone, Default, Serialize, Deserialize)]
23pub struct CodexRelayCapabilitiesRequest {
24    #[serde(default)]
25    pub station_name: Option<String>,
26    #[serde(default)]
27    pub provider_id: Option<String>,
28    #[serde(default)]
29    pub endpoint_id: Option<String>,
30    #[serde(default)]
31    pub upstream_index: Option<usize>,
32    #[serde(default)]
33    pub model: Option<String>,
34    #[serde(default, alias = "patch_preset")]
35    pub patch_mode: Option<CodexPatchMode>,
36    #[serde(default)]
37    pub responses_websocket: Option<bool>,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct CodexRelayCapabilitiesResponse {
42    pub api_version: u32,
43    pub service_name: String,
44    pub station_name: String,
45    pub upstream_index: usize,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub provider_id: Option<String>,
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub endpoint_id: Option<String>,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub provider_endpoint_key: Option<String>,
52    pub upstream_base_url: String,
53    pub patch_mode: CodexPatchMode,
54    pub responses_websocket: bool,
55    pub model: Option<String>,
56    pub expected: CodexCapabilityProfile,
57    pub observed: CodexRelayCapabilitiesObserved,
58    pub recommendation: CodexPatchModeRecommendation,
59    pub mismatches: Vec<CodexRelayCapabilityMismatch>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct CodexRelayCapabilitiesObserved {
64    pub models: CodexRelayProbeResult,
65    pub responses: CodexRelayProbeResult,
66    pub responses_compact: CodexRelayProbeResult,
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct CodexRelayCapabilityMismatch {
71    pub capability: String,
72    pub expected: String,
73    pub observed: String,
74    pub reason: String,
75}
76
77pub(super) async fn codex_relay_capabilities_for_proxy(
78    proxy: &ProxyService,
79    payload: CodexRelayCapabilitiesRequest,
80) -> Result<CodexRelayCapabilitiesResponse, ProxyControlError> {
81    if proxy.service_name != "codex" {
82        return Err(ProxyControlError::new(
83            StatusCode::BAD_REQUEST,
84            "Codex relay capabilities are only available for the codex service",
85        ));
86    }
87
88    let cfg = proxy.config.snapshot().await;
89    let mgr = proxy.service_manager(cfg.as_ref());
90    let target = select_codex_relay_target(
91        mgr,
92        CodexRelayTargetSelection {
93            station_name: payload.station_name.as_deref(),
94            upstream_index: payload.upstream_index,
95            provider_id: payload.provider_id.as_deref(),
96            endpoint_id: payload.endpoint_id.as_deref(),
97        },
98    )?;
99    let patch_mode = payload
100        .patch_mode
101        .or_else(current_codex_switch_patch_mode)
102        .or_else(|| {
103            crate::config::codex_client_patch_config_from_config_file()
104                .ok()
105                .map(|cfg| cfg.preset)
106        })
107        .unwrap_or_default();
108    let responses_websocket = payload
109        .responses_websocket
110        .or_else(current_codex_switch_responses_websocket)
111        .or_else(|| {
112            crate::config::codex_client_patch_config_from_config_file()
113                .ok()
114                .map(|cfg| cfg.options.responses_websocket)
115        })
116        .unwrap_or(false);
117    let model = payload
118        .model
119        .as_deref()
120        .map(str::trim)
121        .filter(|value| !value.is_empty())
122        .map(ToOwned::to_owned);
123
124    let probe_client = CodexRelayProbeClient::new(proxy.client.clone());
125    let observations = run_capability_probe_cases(&probe_client, &target.upstream).await;
126    let models_observation = observation_for_kind(&observations, CodexRelayProbeKind::Models);
127
128    let expected = build_expected_profile(
129        patch_mode,
130        responses_websocket,
131        model.as_deref(),
132        models_observation,
133    );
134    let observed = build_observed_from_probe_observations(&observations);
135    let recommendation = build_recommendation(patch_mode, &expected, &observed);
136    let mismatches = build_mismatches(&expected, &observed);
137
138    let response = CodexRelayCapabilitiesResponse {
139        api_version: 1,
140        service_name: proxy.service_name.to_string(),
141        station_name: target.station_name,
142        upstream_index: target.upstream_index,
143        provider_id: target.provider_id,
144        endpoint_id: target.endpoint_id,
145        provider_endpoint_key: target.provider_endpoint_key,
146        upstream_base_url: target.upstream.base_url,
147        patch_mode,
148        responses_websocket,
149        model,
150        expected,
151        observed,
152        recommendation,
153        mismatches,
154    };
155    if let Err(error) = super::codex_relay_evidence::append_codex_relay_capabilities_evidence(
156        &response,
157        "proxy_service",
158    ) {
159        tracing::warn!("failed to write Codex relay capability evidence: {}", error);
160    }
161    Ok(response)
162}
163
164async fn run_capability_probe_cases(
165    probe_client: &CodexRelayProbeClient,
166    upstream: &crate::config::UpstreamConfig,
167) -> Vec<CodexRelayProbeObservation> {
168    let mut observations = Vec::with_capacity(codex_relay_probe_cases().len());
169    for case in codex_relay_probe_cases() {
170        observations.push(
171            probe_client
172                .probe_upstream_observation(upstream, &case.spec())
173                .await,
174        );
175    }
176    observations
177}
178
179fn build_observed_from_probe_observations(
180    observations: &[CodexRelayProbeObservation],
181) -> CodexRelayCapabilitiesObserved {
182    CodexRelayCapabilitiesObserved {
183        models: observation_for_kind(observations, CodexRelayProbeKind::Models)
184            .result
185            .clone(),
186        responses: observation_for_kind(observations, CodexRelayProbeKind::Responses)
187            .result
188            .clone(),
189        responses_compact: observation_for_kind(
190            observations,
191            CodexRelayProbeKind::ResponsesCompact,
192        )
193        .result
194        .clone(),
195    }
196}
197
198fn observation_for_kind(
199    observations: &[CodexRelayProbeObservation],
200    kind: CodexRelayProbeKind,
201) -> &CodexRelayProbeObservation {
202    observations
203        .iter()
204        .find(|observation| observation.result.kind == kind)
205        .expect("Codex relay probe registry must include all observed response fields")
206}
207
208fn current_codex_switch_patch_mode() -> Option<CodexPatchMode> {
209    crate::codex_integration::codex_switch_status()
210        .ok()
211        .and_then(|status| status.patch_mode)
212}
213
214fn current_codex_switch_responses_websocket() -> Option<bool> {
215    crate::codex_integration::codex_switch_status()
216        .ok()
217        .and_then(|status| status.supports_websockets)
218}
219
220fn build_expected_profile(
221    patch_mode: CodexPatchMode,
222    responses_websocket: bool,
223    model: Option<&str>,
224    models_observation: &CodexRelayProbeObservation,
225) -> CodexCapabilityProfile {
226    let model_catalog = translated_models_catalog(models_observation, model).unwrap_or_else(|| {
227        CodexModelCatalogProfile::unknown(models_observation.result.reason.clone())
228    });
229    CodexCapabilityProfile::for_input(CodexCapabilityProfileInput::from_patch_config(
230        patch_mode,
231        crate::codex_integration::CodexSwitchOptions {
232            responses_websocket,
233        },
234        model_catalog,
235    ))
236}
237
238fn translated_models_catalog(
239    models_observation: &CodexRelayProbeObservation,
240    model: Option<&str>,
241) -> Option<CodexModelCatalogProfile> {
242    let status = models_observation.status?;
243    if !status.is_success() {
244        return None;
245    }
246    let body = maybe_decode_models_response_body(
247        "codex",
248        "/models",
249        &models_observation.headers,
250        models_observation.body.clone(),
251    );
252    let value = serde_json::from_slice::<Value>(body.as_ref()).ok()?;
253    Some(CodexModelCatalogProfile::from_models_response_json(
254        &value, model,
255    ))
256}
257
258fn build_mismatches(
259    expected: &CodexCapabilityProfile,
260    observed: &CodexRelayCapabilitiesObserved,
261) -> Vec<CodexRelayCapabilityMismatch> {
262    let mut out = Vec::new();
263    push_endpoint_mismatch(
264        &mut out,
265        "responses",
266        &CodexCapabilityDecision::supported("Codex model requests require a /responses endpoint"),
267        &observed.responses,
268    );
269    push_endpoint_mismatch(
270        &mut out,
271        "remote_compaction_v1",
272        &expected.remote_compaction_v1,
273        &observed.responses_compact,
274    );
275    if observed.models.translation_required {
276        out.push(CodexRelayCapabilityMismatch {
277            capability: "model_catalog".to_string(),
278            expected: "codex_models".to_string(),
279            observed: "openai_data_list".to_string(),
280            reason: "relay returned an OpenAI models list; helper translation is required before Codex can see model metadata".to_string(),
281        });
282    }
283    out
284}
285
286fn build_recommendation(
287    current_patch_mode: CodexPatchMode,
288    expected: &CodexCapabilityProfile,
289    observed: &CodexRelayCapabilitiesObserved,
290) -> CodexPatchModeRecommendation {
291    CodexPatchModeRecommendation::for_input(CodexPatchModeRecommendationInput {
292        current_patch_mode,
293        model_catalog: expected.model_catalog.clone(),
294        responses: observed_support_to_capability_support(observed.responses.support),
295        responses_compact: observed_support_to_capability_support(
296            observed.responses_compact.support,
297        ),
298    })
299}
300
301fn observed_support_to_capability_support(
302    support: super::CodexRelayProbeSupport,
303) -> CodexCapabilitySupport {
304    match support {
305        super::CodexRelayProbeSupport::Supported => CodexCapabilitySupport::Supported,
306        super::CodexRelayProbeSupport::Unsupported => CodexCapabilitySupport::Unsupported,
307        super::CodexRelayProbeSupport::Unknown => CodexCapabilitySupport::Unknown,
308    }
309}
310
311fn push_endpoint_mismatch(
312    out: &mut Vec<CodexRelayCapabilityMismatch>,
313    capability: &str,
314    expected: &CodexCapabilityDecision,
315    observed: &CodexRelayProbeResult,
316) {
317    let expected_label = support_label(expected.support);
318    let observed_label = format!(
319        "{} via {}",
320        probe_support_label(observed.support),
321        probe_confidence_label(observed.confidence)
322    );
323    if expected.support == crate::codex_capability_profile::CodexCapabilitySupport::Supported
324        && observed.support != super::CodexRelayProbeSupport::Supported
325    {
326        out.push(CodexRelayCapabilityMismatch {
327            capability: capability.to_string(),
328            expected: expected_label.to_string(),
329            observed: observed_label,
330            reason: observed.reason.clone(),
331        });
332    }
333}
334
335fn support_label(support: crate::codex_capability_profile::CodexCapabilitySupport) -> &'static str {
336    match support {
337        crate::codex_capability_profile::CodexCapabilitySupport::Unknown => "unknown",
338        crate::codex_capability_profile::CodexCapabilitySupport::Supported => "supported",
339        crate::codex_capability_profile::CodexCapabilitySupport::Unsupported => "unsupported",
340    }
341}
342
343fn probe_support_label(support: super::CodexRelayProbeSupport) -> &'static str {
344    match support {
345        super::CodexRelayProbeSupport::Supported => "supported",
346        super::CodexRelayProbeSupport::Unsupported => "unsupported",
347        super::CodexRelayProbeSupport::Unknown => "unknown",
348    }
349}
350
351fn probe_confidence_label(confidence: super::CodexRelayProbeConfidence) -> &'static str {
352    match confidence {
353        super::CodexRelayProbeConfidence::SuccessStatus => "success_status",
354        super::CodexRelayProbeConfidence::EndpointValidation => "endpoint_validation",
355        super::CodexRelayProbeConfidence::ErrorClassification => "error_classification",
356        super::CodexRelayProbeConfidence::Transport => "transport",
357        super::CodexRelayProbeConfidence::Malformed => "malformed",
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use std::collections::{BTreeMap, HashMap};
364    use std::sync::{Arc, Mutex};
365
366    use reqwest::Client;
367
368    use super::*;
369    use crate::config::{
370        ProviderConfigV4, ProxyConfigV4, RoutingConfigV4, ServiceViewV4, UpstreamAuth,
371    };
372    use crate::lb::LbState;
373
374    fn probe_result(
375        kind: CodexRelayProbeKind,
376        support: super::super::CodexRelayProbeSupport,
377    ) -> CodexRelayProbeResult {
378        CodexRelayProbeResult {
379            kind,
380            support,
381            confidence: super::super::CodexRelayProbeConfidence::SuccessStatus,
382            status_code: Some(200),
383            response_shape: Some("ok".to_string()),
384            translation_required: false,
385            error_class: None,
386            reason: "ok".to_string(),
387        }
388    }
389
390    fn observation(kind: CodexRelayProbeKind) -> CodexRelayProbeObservation {
391        CodexRelayProbeObservation {
392            result: probe_result(kind, super::super::CodexRelayProbeSupport::Supported),
393            status: Some(StatusCode::OK),
394            headers: axum::http::HeaderMap::new(),
395            body: axum::body::Bytes::new(),
396        }
397    }
398
399    #[test]
400    fn codex_relay_capabilities_observed_shape_is_built_from_probe_registry() {
401        let observations = codex_relay_probe_cases()
402            .iter()
403            .map(|case| observation(case.kind))
404            .collect::<Vec<_>>();
405
406        let observed = build_observed_from_probe_observations(&observations);
407
408        assert_eq!(observed.models.kind, CodexRelayProbeKind::Models);
409        assert_eq!(observed.responses.kind, CodexRelayProbeKind::Responses);
410        assert_eq!(
411            observed.responses_compact.kind,
412            CodexRelayProbeKind::ResponsesCompact
413        );
414    }
415
416    #[tokio::test]
417    async fn codex_relay_capabilities_targets_route_graph_provider_id() {
418        let v4 = ProxyConfigV4 {
419            codex: ServiceViewV4 {
420                providers: BTreeMap::from([
421                    (
422                        "input8".to_string(),
423                        ProviderConfigV4 {
424                            base_url: Some("http://127.0.0.1:9/v1".to_string()),
425                            inline_auth: UpstreamAuth::default(),
426                            ..ProviderConfigV4::default()
427                        },
428                    ),
429                    (
430                        "ciii".to_string(),
431                        ProviderConfigV4 {
432                            base_url: Some("http://127.0.0.1:10/v1".to_string()),
433                            inline_auth: UpstreamAuth::default(),
434                            ..ProviderConfigV4::default()
435                        },
436                    ),
437                ]),
438                routing: Some(RoutingConfigV4::ordered_failover(vec![
439                    "input8".to_string(),
440                    "ciii".to_string(),
441                ])),
442                ..ServiceViewV4::default()
443            },
444            ..ProxyConfigV4::default()
445        };
446        let runtime = crate::config::compile_v4_to_runtime(&v4).expect("compile v4 runtime");
447        let proxy = ProxyService::new_with_v4_source(
448            Client::new(),
449            Arc::new(runtime),
450            Some(Arc::new(v4)),
451            "codex",
452            Arc::new(Mutex::new(HashMap::<String, LbState>::new())),
453        );
454
455        let response = codex_relay_capabilities_for_proxy(
456            &proxy,
457            CodexRelayCapabilitiesRequest {
458                provider_id: Some("ciii".to_string()),
459                model: Some("gpt-5.5".to_string()),
460                ..Default::default()
461            },
462        )
463        .await
464        .expect("capabilities response");
465
466        assert_eq!(response.station_name, "routing");
467        assert_eq!(response.upstream_index, 1);
468        assert_eq!(response.provider_id.as_deref(), Some("ciii"));
469        assert_eq!(response.endpoint_id.as_deref(), Some("default"));
470        assert_eq!(
471            response.provider_endpoint_key.as_deref(),
472            Some("codex/ciii/default")
473        );
474        assert_eq!(response.upstream_base_url, "http://127.0.0.1:10/v1");
475    }
476}