Files
codex/codex-rs/core/src/client_common.rs
T
pakrym-oai 12de0e395d Preserve WebSocket continuations when steering a turn (#48508)
## Why

Steering an active WebSocket response previously dropped the connection and
resent the full history. Draining the response preserves the connection and
allows the follow-up request to continue with `previous_response_id`.

## What changed

- Drain WebSocket responses when steering. For models using
  `use_responses_lite`, first send `response.interrupt` with
  `mode: "discard_partial_items"` once the response ID is available.
- Treat `response.incomplete` with reason `interrupted` as completion with
  `end_turn: false`, preserving token usage and allowing the turn to continue.
  Other incomplete reasons remain errors.

## Testing

Update the steering integration test to cover completed and discarded reasoning
items, asserting connection reuse, incremental follow-up input, and token usage
from the interrupted response.

GitOrigin-RevId: bc714b713e797cf1a67e3b26ffbf7664cb92eb33
2026-09-26 18:36:55 +00:00

144 lines
4.7 KiB
Rust

pub use codex_api::ResponseEvent;
use codex_protocol::error::Result;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ImageDetail;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_tools::ToolSpec;
use futures::Stream;
use serde_json::Value;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
/// API request payload for a single model turn
#[derive(Debug, Clone)]
pub struct Prompt {
/// Conversation context input items.
pub input: Vec<ResponseItem>,
/// Tools available to the model, including additional tools sourced from
/// external MCP servers.
pub(crate) tools: Arc<[ToolSpec]>,
/// Whether parallel tool calls are permitted for this prompt.
pub(crate) parallel_tool_calls: bool,
pub base_instructions: BaseInstructions,
/// Optional the output schema for the model's response.
pub output_schema: Option<Value>,
/// Whether the Responses API should strictly validate `output_schema`.
pub output_schema_strict: bool,
pub(crate) cyber_access_program: Option<codex_protocol::turn_input::CyberAccessProgram>,
}
impl Default for Prompt {
fn default() -> Self {
Self {
input: Vec::new(),
tools: Arc::default(),
parallel_tool_calls: false,
base_instructions: BaseInstructions::default(),
output_schema: None,
output_schema_strict: true,
cyber_access_program: None,
}
}
}
impl Prompt {
pub(crate) fn get_formatted_input_for_request(
&self,
model_info: &ModelInfo,
) -> Vec<ResponseItem> {
let mut input = self.input.clone();
normalize_image_details(&mut input, model_info);
input
}
}
fn normalize_image_details(items: &mut [ResponseItem], model_info: &ModelInfo) {
for item in items {
match item {
ResponseItem::Message { content, .. } => {
for content_item in content {
if let ContentItem::InputImage { detail, .. } = content_item {
normalize_image_detail(detail, model_info);
}
}
}
ResponseItem::FunctionCallOutput { output, .. }
| ResponseItem::CustomToolCallOutput { output, .. } => {
if let Some(content) = output.content_items_mut() {
for content_item in content {
if let FunctionCallOutputContentItem::InputImage { detail, .. } =
content_item
{
normalize_image_detail(detail, model_info);
}
}
}
}
ResponseItem::AdditionalTools { .. }
| ResponseItem::Reasoning { .. }
| ResponseItem::AgentMessage { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::ToolSearchCall { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::ToolSearchOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::ImageGenerationCall { .. }
| ResponseItem::Compaction { .. }
| ResponseItem::ConfigurationUpdate { .. }
| ResponseItem::CompactionTrigger { .. }
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other => {}
}
}
}
fn normalize_image_detail(detail: &mut Option<ImageDetail>, model_info: &ModelInfo) {
if model_info.use_responses_lite {
*detail = None;
} else if *detail == Some(ImageDetail::Original) && !model_info.supports_image_detail_original {
*detail = Some(DEFAULT_IMAGE_DETAIL);
}
}
pub struct ResponseStream {
pub(crate) rx_event: mpsc::Receiver<Result<ResponseEvent>>,
pub(crate) interrupt: Option<oneshot::Sender<()>>,
/// Signals the mapper task that the consumer stopped polling before the
/// provider stream reached its own terminal event.
pub(crate) consumer_dropped: CancellationToken,
}
impl Stream for ResponseStream {
type Item = Result<ResponseEvent>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.rx_event.poll_recv(cx)
}
}
impl Drop for ResponseStream {
fn drop(&mut self) {
self.consumer_dropped.cancel();
}
}
#[cfg(test)]
#[path = "client_common_tests.rs"]
mod tests;