fix(core/models): sync upstream remote model parity

This commit is contained in:
James Peter
2026-04-05 16:07:57 +10:00
parent 18351ae86a
commit eff1ef4f6d
4 changed files with 189 additions and 10 deletions
+38 -3
View File
@@ -107,6 +107,7 @@ fn preferred_ws_version_from_env() -> ResponsesWebsocketVersion {
// be replayed on every subsequent request within the same turn (retries,
// continuations, websocket reconnects).
const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
const X_CODEX_WINDOW_ID_HEADER: &str = "x-codex-window-id";
const MODEL_CAP_MODEL_HEADER: &str = "x-codex-model-cap-model";
const MODEL_CAP_RESET_AFTER_HEADER: &str = "x-codex-model-cap-reset-after-seconds";
@@ -406,6 +407,10 @@ impl ModelClient {
&self.config.code_home
}
fn current_window_id(&self, session_id: Uuid) -> String {
format!("{session_id}:0")
}
pub(crate) fn config(&self) -> &crate::config::Config {
&self.config
}
@@ -829,6 +834,9 @@ impl ModelClient {
req_builder = req_builder
.header("conversation_id", session_id_str.clone())
.header("session_id", session_id_str.clone());
if let Ok(window_id) = HeaderValue::from_str(&self.current_window_id(session_id)) {
req_builder = req_builder.header(X_CODEX_WINDOW_ID_HEADER, window_id);
}
if let Some(auth) = auth.as_ref()
&& auth.mode.is_chatgpt()
@@ -889,9 +897,13 @@ impl ModelClient {
}
let ws_payload_text = serde_json::to_string(&serde_json::Value::Object(ws_payload))?;
let connect = tokio_tungstenite::connect_async(ws_request).await;
let connect = timeout(
self.provider.websocket_connect_timeout(),
tokio_tungstenite::connect_async(ws_request),
)
.await;
match connect {
Ok((mut ws_stream, response)) => {
Ok(Ok((mut ws_stream, response))) => {
let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(1600);
if let Some(value) = response
@@ -1041,7 +1053,7 @@ impl ModelClient {
return Ok(ResponseStream { rx_event });
}
Err(err) => {
Ok(Err(err)) => {
if websocket_connect_is_upgrade_required(&err) {
self.websockets_disabled.store(true, Ordering::Relaxed);
warn!("responses websocket upgrade required; falling back to HTTP responses transport");
@@ -1060,6 +1072,22 @@ impl ModelClient {
self.websockets_disabled.store(true, Ordering::Relaxed);
return Err(err);
}
Err(_) => {
let err = CodexErr::Stream(
format!(
"[ws] timed out connecting after {} ms",
self.provider.websocket_connect_timeout().as_millis()
),
None,
Some(request_id.clone()),
);
if (attempt as u64) < max_retries {
tokio::time::sleep(backoff(attempt as u64)).await;
continue;
}
self.websockets_disabled.store(true, Ordering::Relaxed);
return Err(err);
}
}
}
}
@@ -1272,6 +1300,9 @@ impl ModelClient {
.header("session_id", session_id_str.clone())
.header(reqwest::header::ACCEPT, "text/event-stream")
.json(&payload_json);
if let Ok(window_id) = HeaderValue::from_str(&self.current_window_id(session_id)) {
req_builder = req_builder.header(X_CODEX_WINDOW_ID_HEADER, window_id);
}
if let Some(auth) = auth.as_ref()
&& auth.mode.is_chatgpt()
@@ -1844,6 +1875,7 @@ impl ModelClient {
.clone()
.or_else(|| find_family_for_model(model_slug))
.unwrap_or_else(|| self.config.model_family.clone());
let session_id = prompt.session_id_override.unwrap_or(self.session_id);
let instructions = prompt.get_full_instructions(&family).into_owned();
let payload = CompactHistoryRequest {
model: model_slug,
@@ -1884,6 +1916,9 @@ impl ModelClient {
request = attach_openai_subagent_header(request);
request = attach_codex_beta_features_header(request, &self.config);
if let Ok(window_id) = HeaderValue::from_str(&self.current_window_id(session_id)) {
request = request.header(X_CODEX_WINDOW_ID_HEADER, window_id);
}
if let Some(auth) = auth.as_ref()
&& auth.mode.is_chatgpt()
+36 -2
View File
@@ -40,11 +40,25 @@ const MAX_OUTPUT_DEFAULT: u64 = 128_000;
const IMAGE_GENERATION_TOOL: &str = "image_generation";
static UPSTREAM_MODELS: Lazy<Vec<ModelInfo>> = Lazy::new(|| {
serde_json::from_str::<ModelsResponse>(include_str!("../../../codex-rs/core/models.json"))
serde_json::from_str::<ModelsResponse>(include_str!("../../../codex-rs/models-manager/models.json"))
.map(|response| response.models)
.unwrap_or_else(|err| panic!("failed to parse upstream models.json: {err}"))
});
fn namespaced_model_suffix(model: &str) -> Option<&str> {
let (namespace, suffix) = model.split_once('/')?;
if suffix.contains('/') {
return None;
}
if !namespace
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
{
return None;
}
Some(suffix)
}
pub const STANDARD_CONTEXT_WINDOW_272K: u64 = CONTEXT_WINDOW_272K;
pub const EXTENDED_CONTEXT_WINDOW_1M: u64 = CONTEXT_WINDOW_1M;
@@ -102,6 +116,10 @@ pub struct ModelFamily {
/// a tool call instead of just a bash command
pub apply_patch_tool_type: Option<ApplyPatchToolType>,
/// This should be set when the model expects a `shell_command` tool that
/// accepts a shell script string instead of argv-style arguments.
pub uses_shell_command_tool: bool,
/// Whether web_search should request text-only or multimodal results.
pub web_search_tool_type: WebSearchToolType,
@@ -160,6 +178,7 @@ macro_rules! model_family {
prefer_websockets: false,
uses_local_shell_tool: false,
apply_patch_tool_type: None,
uses_shell_command_tool: false,
web_search_tool_type: WebSearchToolType::Text,
supports_image_detail_original: false,
supports_image_generation: false,
@@ -174,7 +193,12 @@ macro_rules! model_family {
}
fn apply_upstream_model_overrides(mut family: ModelFamily) -> ModelFamily {
let Some(model_info) = UPSTREAM_MODELS.iter().find(|model| model.slug == family.slug) else {
let model_slug = family
.slug
.strip_prefix("openai/")
.or_else(|| namespaced_model_suffix(&family.slug))
.unwrap_or(&family.slug);
let Some(model_info) = UPSTREAM_MODELS.iter().find(|model| model.slug == model_slug) else {
return family;
};
@@ -198,6 +222,8 @@ fn apply_upstream_model_overrides(mut family: ModelFamily) -> ModelFamily {
.iter()
.any(|tool| tool == IMAGE_GENERATION_TOOL);
family.uses_local_shell_tool = matches!(model_info.shell_type, ConfigShellToolType::Local);
family.uses_shell_command_tool =
matches!(model_info.shell_type, ConfigShellToolType::ShellCommand);
family.auto_compact_token_limit = model_info.auto_compact_token_limit();
family.truncation_policy = match model_info.truncation_policy.mode {
TruncationMode::Bytes => TruncationPolicy::Bytes(
@@ -214,6 +240,13 @@ fn apply_upstream_model_overrides(mut family: ModelFamily) -> ModelFamily {
/// Returns a `ModelFamily` for the given model slug, or `None` if the slug
/// does not match any known model family.
pub fn find_family_for_model(slug: &str) -> Option<ModelFamily> {
if let Some(suffix) = namespaced_model_suffix(slug)
&& let Some(mut family) = find_family_for_model(suffix)
{
family.slug = slug.to_string();
return Some(family);
}
if slug.starts_with("o3") {
model_family!(
slug, "o3",
@@ -430,6 +463,7 @@ pub fn derive_default_model_family(model: &str) -> ModelFamily {
prefer_websockets: false,
uses_local_shell_tool: false,
apply_patch_tool_type: None,
uses_shell_command_tool: false,
web_search_tool_type: WebSearchToolType::Text,
supports_image_detail_original: false,
supports_image_generation: false,
+18
View File
@@ -25,6 +25,7 @@ use std::time::Instant;
use tokio::process::Command;
use crate::error::EnvVarError;
const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000;
pub(crate) const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000;
const DEFAULT_STREAM_MAX_RETRIES: u64 = 5;
const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4;
/// Hard cap for user-configured `stream_max_retries`.
@@ -131,6 +132,9 @@ pub struct ModelProviderInfo {
/// the connection as lost.
pub stream_idle_timeout_ms: Option<u64>,
/// Timeout (in milliseconds) when establishing a websocket transport connection.
pub websocket_connect_timeout_ms: Option<u64>,
/// Whether this provider requires some form of standard authentication (API key, ChatGPT token).
#[serde(default)]
pub requires_openai_auth: bool,
@@ -573,6 +577,12 @@ impl ModelProviderInfo {
.unwrap_or(Duration::from_millis(DEFAULT_STREAM_IDLE_TIMEOUT_MS))
}
pub fn websocket_connect_timeout(&self) -> Duration {
self.websocket_connect_timeout_ms
.map(Duration::from_millis)
.unwrap_or(Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS))
}
pub fn base_url_for_probe(&self) -> String {
self.base_url
.clone()
@@ -744,6 +754,7 @@ pub fn built_in_model_providers(
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: true,
openrouter: None,
},
@@ -791,6 +802,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str) -> ModelProviderInfo {
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
}
@@ -837,6 +849,7 @@ base_url = "http://localhost:11434/v1"
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
@@ -869,6 +882,7 @@ query_params = { api-version = "2025-04-01-preview" }
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
@@ -904,6 +918,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
@@ -929,6 +944,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
}
@@ -964,6 +980,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
@@ -1066,6 +1083,7 @@ refresh_interval_ms = 0
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
+97 -5
View File
@@ -267,11 +267,7 @@ impl RemoteModelsManager {
let info = {
let state = self.state.read().await;
state
.models
.iter()
.find(|info| info.slug.eq_ignore_ascii_case(model))
.cloned()
find_remote_model_info(&state.models, model)
};
let Some(info) = info else {
return family;
@@ -358,6 +354,39 @@ impl RemoteModelsManager {
}
}
fn namespaced_model_suffix(model: &str) -> Option<&str> {
let (namespace, suffix) = model.split_once('/')?;
if suffix.contains('/') {
return None;
}
if !namespace
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
{
return None;
}
Some(suffix)
}
fn find_remote_model_info(models: &[ModelInfo], model: &str) -> Option<ModelInfo> {
models
.iter()
.filter(|info| model.eq_ignore_ascii_case(&info.slug) || model.starts_with(&info.slug))
.cloned()
.max_by_key(|info| info.slug.len())
.or_else(|| {
namespaced_model_suffix(model).and_then(|suffix| {
models
.iter()
.filter(|info| {
suffix.eq_ignore_ascii_case(&info.slug) || suffix.starts_with(&info.slug)
})
.cloned()
.max_by_key(|info| info.slug.len())
})
})
}
pub fn apply_model_info_overrides(info: &ModelInfo, family: ModelFamily) -> ModelFamily {
apply_model_info_overrides_with_personality(info, family, None)
}
@@ -412,6 +441,69 @@ fn apply_model_info_overrides_with_personality(
family
}
#[cfg(test)]
mod tests {
use super::find_remote_model_info;
use code_protocol::config_types::ReasoningSummary;
use code_protocol::openai_models::ConfigShellToolType;
use code_protocol::openai_models::ModelInfo;
use code_protocol::openai_models::ModelVisibility;
use code_protocol::openai_models::TruncationPolicyConfig;
use code_protocol::openai_models::WebSearchToolType;
use code_protocol::openai_models::default_input_modalities;
fn model(slug: &str) -> ModelInfo {
ModelInfo {
slug: slug.to_string(),
display_name: slug.to_string(),
description: None,
default_reasoning_level: None,
supported_reasoning_levels: Vec::new(),
shell_type: ConfigShellToolType::Default,
visibility: ModelVisibility::None,
supported_in_api: true,
priority: 0,
availability_nux: None,
upgrade: None,
base_instructions: String::new(),
model_messages: None,
supports_reasoning_summaries: false,
default_reasoning_summary: ReasoningSummary::Auto,
support_verbosity: false,
default_verbosity: None,
apply_patch_tool_type: None,
web_search_tool_type: WebSearchToolType::Text,
truncation_policy: TruncationPolicyConfig::bytes(10_000),
supports_parallel_tool_calls: false,
supports_image_detail_original: false,
context_window: None,
auto_compact_token_limit: None,
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: default_input_modalities(),
used_fallback_model_metadata: false,
supports_search_tool: false,
}
}
#[test]
fn find_remote_model_info_matches_namespaced_suffix() {
let models = vec![model("gpt-5.3-codex")];
let found = find_remote_model_info(&models, "custom/gpt-5.3-codex")
.expect("namespaced slug should resolve");
assert_eq!(found.slug, "gpt-5.3-codex");
}
#[test]
fn find_remote_model_info_rejects_multi_segment_namespace() {
let models = vec![model("gpt-5.3-codex")];
assert!(find_remote_model_info(&models, "foo/bar/gpt-5.3-codex").is_none());
}
}
fn map_web_search_tool_type(tool_type: WebSearchToolType) -> WebSearchToolType {
match tool_type {
WebSearchToolType::Text => WebSearchToolType::Text,