Fallback to HTTP on UPGRADE_REQUIRED (#10824)

Allow the server to trigger a connection downgrade in case the protocol
changes in incompatible ways.
This commit is contained in:
pakrym-oai
2026-02-07 21:06:33 -08:00
committed by GitHub
parent d68e9c0f19
commit 6d08298f4e
3 changed files with 114 additions and 38 deletions

View File

@@ -225,6 +225,11 @@ pub struct ModelClientSession {
turn_state: Arc<OnceLock<String>>,
}
enum WebsocketStreamOutcome {
Stream(ResponseStream),
FallbackToHttp,
}
impl ModelClient {
#[allow(clippy::too_many_arguments)]
/// Creates a new session-scoped `ModelClient`.
@@ -926,7 +931,7 @@ impl ModelClientSession {
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
turn_metadata_header: Option<&str>,
) -> Result<ResponseStream> {
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.auth_manager.clone();
let api_prompt = Self::build_responses_request(prompt)?;
@@ -957,6 +962,11 @@ impl ModelClientSession {
.await
{
Ok(_) => {}
Err(ApiError::Transport(TransportError::Http { status, .. }))
if status == StatusCode::UPGRADE_REQUIRED =>
{
return Ok(WebsocketStreamOutcome::FallbackToHttp);
}
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
@@ -992,7 +1002,10 @@ impl ModelClientSession {
}
});
return Ok(map_response_stream(stream_result, otel_manager.clone()));
return Ok(WebsocketStreamOutcome::Stream(map_response_stream(
stream_result,
otel_manager.clone(),
)));
}
}
@@ -1036,26 +1049,33 @@ impl ModelClientSession {
self.client.responses_websocket_enabled() && !self.client.disable_websockets();
if websocket_enabled {
self.stream_responses_websocket(
prompt,
model_info,
otel_manager,
effort,
summary,
turn_metadata_header,
)
.await
} else {
self.stream_responses_api(
prompt,
model_info,
otel_manager,
effort,
summary,
turn_metadata_header,
)
.await
match self
.stream_responses_websocket(
prompt,
model_info,
otel_manager,
effort,
summary,
turn_metadata_header,
)
.await?
{
WebsocketStreamOutcome::Stream(stream) => return Ok(stream),
WebsocketStreamOutcome::FallbackToHttp => {
self.try_switch_fallback_transport(otel_manager);
}
}
}
self.stream_responses_api(
prompt,
model_info,
otel_manager,
effort,
summary,
turn_metadata_header,
)
.await
}
}
}

View File

@@ -558,9 +558,8 @@ impl TurnContext {
}
async fn build_turn_metadata_header(&self) -> Option<String> {
let cwd = self.cwd.clone();
self.turn_metadata_header
.get_or_init(|| async { build_turn_metadata_header(cwd).await })
.get_or_init(|| async { build_turn_metadata_header(self.cwd.clone()).await })
.await
.clone()
}
@@ -1098,14 +1097,14 @@ impl Session {
// Warm a websocket in the background so the first turn can reuse it.
// This performs only connection setup; user input is still sent later via response.create
// when submit_turn() runs.
sess.services.model_client.pre_establish_connection(
sess.services.otel_manager.clone(),
resolve_turn_metadata_header_with_timeout(
build_turn_metadata_header(session_configuration.cwd.clone()),
None,
)
.boxed(),
);
let turn_metadata_header = resolve_turn_metadata_header_with_timeout(
build_turn_metadata_header(session_configuration.cwd.clone()),
None,
)
.boxed();
sess.services
.model_client
.pre_establish_connection(sess.services.otel_manager.clone(), turn_metadata_header);
// Dispatch the SessionConfiguredEvent first and then report any errors.
// If resuming, include converted initial messages in the payload so UIs can render them immediately.