mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Fall back to http when websockets fail (#10139)
I expect not all proxies work with websockets, fall back to http if websockets fail.
This commit is contained in:
committed by
GitHub
Unverified
parent
798c4b3260
commit
3b1cddf001
@@ -65,6 +65,7 @@ use crate::model_provider_info::ModelProviderInfo;
|
||||
use crate::model_provider_info::WireApi;
|
||||
use crate::tools::spec::create_tools_json_for_chat_completions_api;
|
||||
use crate::tools::spec::create_tools_json_for_responses_api;
|
||||
use crate::transport_manager::TransportManager;
|
||||
|
||||
pub const WEB_SEARCH_ELIGIBLE_HEADER: &str = "x-oai-web-search-eligible";
|
||||
pub const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
|
||||
@@ -80,6 +81,7 @@ struct ModelClientState {
|
||||
effort: Option<ReasoningEffortConfig>,
|
||||
summary: ReasoningSummaryConfig,
|
||||
session_source: SessionSource,
|
||||
transport_manager: TransportManager,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -91,6 +93,7 @@ pub struct ModelClientSession {
|
||||
state: Arc<ModelClientState>,
|
||||
connection: Option<ApiWebSocketConnection>,
|
||||
websocket_last_items: Vec<ResponseItem>,
|
||||
transport_manager: TransportManager,
|
||||
/// Turn state for sticky routing.
|
||||
///
|
||||
/// This is an `OnceLock` that stores the turn state value received from the server
|
||||
@@ -116,6 +119,7 @@ impl ModelClient {
|
||||
summary: ReasoningSummaryConfig,
|
||||
conversation_id: ThreadId,
|
||||
session_source: SessionSource,
|
||||
transport_manager: TransportManager,
|
||||
) -> Self {
|
||||
Self {
|
||||
state: Arc::new(ModelClientState {
|
||||
@@ -128,6 +132,7 @@ impl ModelClient {
|
||||
effort,
|
||||
summary,
|
||||
session_source,
|
||||
transport_manager,
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -137,6 +142,7 @@ impl ModelClient {
|
||||
state: Arc::clone(&self.state),
|
||||
connection: None,
|
||||
websocket_last_items: Vec::new(),
|
||||
transport_manager: self.state.transport_manager.clone(),
|
||||
turn_state: Arc::new(OnceLock::new()),
|
||||
}
|
||||
}
|
||||
@@ -171,6 +177,10 @@ impl ModelClient {
|
||||
self.state.session_source.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn transport_manager(&self) -> TransportManager {
|
||||
self.state.transport_manager.clone()
|
||||
}
|
||||
|
||||
/// Returns the currently configured model slug.
|
||||
pub fn get_model(&self) -> String {
|
||||
self.state.model_info.slug.clone()
|
||||
@@ -250,7 +260,10 @@ impl ModelClientSession {
|
||||
/// For Chat providers, the underlying stream is optionally aggregated
|
||||
/// based on the `show_raw_agent_reasoning` flag in the config.
|
||||
pub async fn stream(&mut self, prompt: &Prompt) -> Result<ResponseStream> {
|
||||
match self.state.provider.wire_api {
|
||||
let wire_api = self
|
||||
.transport_manager
|
||||
.effective_wire_api(self.state.provider.wire_api);
|
||||
match wire_api {
|
||||
WireApi::Responses => self.stream_responses_api(prompt).await,
|
||||
WireApi::ResponsesWebsocket => self.stream_responses_websocket(prompt).await,
|
||||
WireApi::Chat => {
|
||||
@@ -271,6 +284,24 @@ impl ModelClientSession {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn try_switch_fallback_transport(&mut self) -> bool {
|
||||
let activated = self
|
||||
.transport_manager
|
||||
.activate_http_fallback(self.state.provider.wire_api);
|
||||
if activated {
|
||||
warn!("falling back to HTTP");
|
||||
self.state.otel_manager.counter(
|
||||
"codex.transport.fallback_to_http",
|
||||
1,
|
||||
&[("from_wire_api", "responses_websocket")],
|
||||
);
|
||||
|
||||
self.connection = None;
|
||||
self.websocket_last_items.clear();
|
||||
}
|
||||
activated
|
||||
}
|
||||
|
||||
fn build_responses_request(&self, prompt: &Prompt) -> Result<ApiPrompt> {
|
||||
let instructions = prompt.base_instructions.text.clone();
|
||||
let tools_json: Vec<Value> = create_tools_json_for_responses_api(&prompt.tools)?;
|
||||
|
||||
@@ -30,6 +30,7 @@ use crate::stream_events_utils::HandleOutputCtx;
|
||||
use crate::stream_events_utils::handle_non_tool_response_item;
|
||||
use crate::stream_events_utils::handle_output_item_done;
|
||||
use crate::terminal;
|
||||
use crate::transport_manager::TransportManager;
|
||||
use crate::truncate::TruncationPolicy;
|
||||
use crate::user_notification::UserNotifier;
|
||||
use crate::util::error_or_panic;
|
||||
@@ -624,6 +625,7 @@ impl Session {
|
||||
model_info: ModelInfo,
|
||||
conversation_id: ThreadId,
|
||||
sub_id: String,
|
||||
transport_manager: TransportManager,
|
||||
) -> TurnContext {
|
||||
let otel_manager = otel_manager.clone().with_model(
|
||||
session_configuration.collaboration_mode.model(),
|
||||
@@ -640,6 +642,7 @@ impl Session {
|
||||
session_configuration.model_reasoning_summary,
|
||||
conversation_id,
|
||||
session_configuration.session_source.clone(),
|
||||
transport_manager,
|
||||
);
|
||||
|
||||
let tools_config = ToolsConfig::new(&ToolsConfigParams {
|
||||
@@ -869,6 +872,7 @@ impl Session {
|
||||
skills_manager,
|
||||
agent_control,
|
||||
state_db: state_db_ctx.clone(),
|
||||
transport_manager: TransportManager::new(),
|
||||
};
|
||||
|
||||
let sess = Arc::new(Session {
|
||||
@@ -1188,6 +1192,7 @@ impl Session {
|
||||
model_info,
|
||||
self.conversation_id,
|
||||
sub_id,
|
||||
self.services.transport_manager.clone(),
|
||||
);
|
||||
if let Some(final_schema) = final_output_json_schema {
|
||||
turn_context.final_output_json_schema = final_schema;
|
||||
@@ -3042,6 +3047,7 @@ async fn spawn_review_thread(
|
||||
per_turn_config.model_reasoning_summary,
|
||||
sess.conversation_id,
|
||||
parent_turn_context.client.get_session_source(),
|
||||
parent_turn_context.client.transport_manager(),
|
||||
);
|
||||
|
||||
let review_turn_context = TurnContext {
|
||||
@@ -3537,7 +3543,9 @@ async fn run_sampling_request(
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(output) => return Ok(output),
|
||||
Ok(output) => {
|
||||
return Ok(output);
|
||||
}
|
||||
Err(CodexErr::ContextWindowExceeded) => {
|
||||
sess.set_total_tokens_full(&turn_context).await;
|
||||
return Err(CodexErr::ContextWindowExceeded);
|
||||
@@ -3558,6 +3566,17 @@ async fn run_sampling_request(
|
||||
|
||||
// Use the configured provider-specific stream retry budget.
|
||||
let max_retries = turn_context.client.get_provider().stream_max_retries();
|
||||
if retries >= max_retries && client_session.try_switch_fallback_transport() {
|
||||
sess.send_event(
|
||||
&turn_context,
|
||||
EventMsg::Warning(WarningEvent {
|
||||
message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
retries = 0;
|
||||
continue;
|
||||
}
|
||||
if retries < max_retries {
|
||||
retries += 1;
|
||||
let delay = match &err {
|
||||
@@ -4753,6 +4772,7 @@ mod tests {
|
||||
skills_manager,
|
||||
agent_control,
|
||||
state_db: None,
|
||||
transport_manager: TransportManager::new(),
|
||||
};
|
||||
|
||||
let turn_context = Session::make_turn_context(
|
||||
@@ -4764,6 +4784,7 @@ mod tests {
|
||||
model_info,
|
||||
conversation_id,
|
||||
"turn_id".to_string(),
|
||||
services.transport_manager.clone(),
|
||||
);
|
||||
|
||||
let session = Session {
|
||||
@@ -4865,6 +4886,7 @@ mod tests {
|
||||
skills_manager,
|
||||
agent_control,
|
||||
state_db: None,
|
||||
transport_manager: TransportManager::new(),
|
||||
};
|
||||
|
||||
let turn_context = Arc::new(Session::make_turn_context(
|
||||
@@ -4876,6 +4898,7 @@ mod tests {
|
||||
model_info,
|
||||
conversation_id,
|
||||
"turn_id".to_string(),
|
||||
services.transport_manager.clone(),
|
||||
));
|
||||
|
||||
let session = Arc::new(Session {
|
||||
|
||||
@@ -38,6 +38,7 @@ pub mod landlock;
|
||||
pub mod mcp;
|
||||
mod mcp_connection_manager;
|
||||
pub mod models_manager;
|
||||
mod transport_manager;
|
||||
pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY;
|
||||
pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD;
|
||||
pub use mcp_connection_manager::SandboxState;
|
||||
@@ -113,6 +114,7 @@ pub use rollout::list::parse_cursor;
|
||||
pub use rollout::list::read_head_for_summary;
|
||||
pub use rollout::list::read_session_meta_line;
|
||||
pub use rollout::rollout_date_parts;
|
||||
pub use transport_manager::TransportManager;
|
||||
mod function_tool;
|
||||
mod state;
|
||||
mod tasks;
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::models_manager::manager::ModelsManager;
|
||||
use crate::skills::SkillsManager;
|
||||
use crate::state_db::StateDbHandle;
|
||||
use crate::tools::sandboxing::ApprovalStore;
|
||||
use crate::transport_manager::TransportManager;
|
||||
use crate::unified_exec::UnifiedExecProcessManager;
|
||||
use crate::user_notification::UserNotifier;
|
||||
use codex_otel::OtelManager;
|
||||
@@ -32,4 +33,5 @@ pub(crate) struct SessionServices {
|
||||
pub(crate) skills_manager: Arc<SkillsManager>,
|
||||
pub(crate) agent_control: AgentControl,
|
||||
pub(crate) state_db: Option<StateDbHandle>,
|
||||
pub(crate) transport_manager: TransportManager,
|
||||
}
|
||||
|
||||
@@ -781,6 +781,7 @@ mod tests {
|
||||
turn.client.get_reasoning_summary(),
|
||||
session.conversation_id,
|
||||
session_source,
|
||||
session.services.transport_manager.clone(),
|
||||
);
|
||||
|
||||
let invocation = invocation(
|
||||
@@ -1221,6 +1222,7 @@ mod tests {
|
||||
let mut base_config = (*turn.client.config()).clone();
|
||||
base_config.user_instructions = Some("base-user".to_string());
|
||||
turn.user_instructions = Some("resolved-user".to_string());
|
||||
let transport_manager = turn.client.transport_manager();
|
||||
turn.client = ModelClient::new(
|
||||
Arc::new(base_config.clone()),
|
||||
Some(session.services.auth_manager.clone()),
|
||||
@@ -1231,6 +1233,7 @@ mod tests {
|
||||
turn.client.get_reasoning_summary(),
|
||||
session.conversation_id,
|
||||
session_source,
|
||||
transport_manager,
|
||||
);
|
||||
let base_instructions = BaseInstructions {
|
||||
text: "base".to_string(),
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use crate::model_provider_info::WireApi;
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct TransportManager {
|
||||
fallback_to_http: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl TransportManager {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn effective_wire_api(&self, provider_wire_api: WireApi) -> WireApi {
|
||||
if self.fallback_to_http.load(Ordering::Relaxed)
|
||||
&& provider_wire_api == WireApi::ResponsesWebsocket
|
||||
{
|
||||
WireApi::Responses
|
||||
} else {
|
||||
provider_wire_api
|
||||
}
|
||||
}
|
||||
|
||||
pub fn activate_http_fallback(&self, provider_wire_api: WireApi) -> bool {
|
||||
provider_wire_api == WireApi::ResponsesWebsocket
|
||||
&& !self.fallback_to_http.swap(true, Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user