mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
pakrym-oai
12de0e395d
## 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
106 lines
3.4 KiB
Rust
106 lines
3.4 KiB
Rust
//! Interceptors can transform incremental events and retain upstream cancellation.
|
|
use super::*;
|
|
use crate::client_common::ResponseStream;
|
|
use codex_extension_api::ResponseEvent;
|
|
use futures::StreamExt;
|
|
use pretty_assertions::assert_eq;
|
|
use std::time::Duration;
|
|
use tokio::sync::mpsc;
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
struct Append(&'static str);
|
|
impl ModelResponseInterceptor for Append {
|
|
fn intercept(self: Box<Self>, stream: ModelResponseStream) -> ModelResponseStream {
|
|
Box::pin(stream.map(move |event| {
|
|
event.map(|event| match event {
|
|
ResponseEvent::OutputTextDelta(text) => {
|
|
ResponseEvent::OutputTextDelta(format!("{text}{}", self.0))
|
|
}
|
|
other => other,
|
|
})
|
|
}))
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn interceptors_forward_before_completion_in_order_and_cancel_upstream() {
|
|
let (tx_event, rx_event) = mpsc::channel(1);
|
|
let consumer_dropped = CancellationToken::new();
|
|
let upstream_cancelled = consumer_dropped.clone();
|
|
let mut stream = intercept_stream(
|
|
Box::pin(
|
|
ResponseStream {
|
|
rx_event,
|
|
interrupt: None,
|
|
consumer_dropped,
|
|
}
|
|
.map(|event| {
|
|
event.map_err(|error| {
|
|
codex_extension_api::ModelResponseError::Stream(error.to_string())
|
|
})
|
|
}),
|
|
),
|
|
vec![Box::new(Append("a")), Box::new(Append("b"))],
|
|
);
|
|
tx_event
|
|
.send(Ok(ResponseEvent::OutputTextDelta("text".into())))
|
|
.await
|
|
.unwrap();
|
|
// Upstream remains open and has not sent Completed.
|
|
let event = tokio::time::timeout(Duration::from_secs(1), stream.next())
|
|
.await
|
|
.unwrap()
|
|
.unwrap()
|
|
.unwrap();
|
|
let ResponseEvent::OutputTextDelta(text) = event else {
|
|
panic!("expected text delta");
|
|
};
|
|
assert_eq!(text, "textab");
|
|
drop(stream);
|
|
tokio::time::timeout(Duration::from_secs(1), upstream_cancelled.cancelled())
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct MetadataContributor;
|
|
impl ModelRequestContributor for MetadataContributor {
|
|
fn request(&self, input: ModelRequestInput<'_>) -> Option<Box<dyn ModelResponseInterceptor>> {
|
|
*input.client_metadata = Some(HashMap::from([
|
|
("thread_id".into(), "replaced".into()),
|
|
("guardian_credits_requested".into(), "true".into()),
|
|
("parent_response_id".into(), "forged".into()),
|
|
("existing".into(), "replaced".into()),
|
|
("x-request-id".into(), "request".into()),
|
|
]));
|
|
None
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn contributor_metadata_preserves_core_fields_and_existing_values() {
|
|
let mut metadata = Some(HashMap::from([
|
|
("thread_id".into(), "thread".into()),
|
|
("existing".into(), "original".into()),
|
|
]));
|
|
let contributors: Vec<Arc<dyn ModelRequestContributor>> = vec![Arc::new(MetadataContributor)];
|
|
assert!(
|
|
prepare(
|
|
&contributors,
|
|
"thread",
|
|
"model",
|
|
ModelRequestKind::Generation,
|
|
&mut metadata
|
|
)
|
|
.is_empty()
|
|
);
|
|
assert_eq!(
|
|
metadata,
|
|
Some(HashMap::from([
|
|
("thread_id".into(), "thread".into()),
|
|
("existing".into(), "original".into()),
|
|
("x-request-id".into(), "request".into()),
|
|
]))
|
|
);
|
|
}
|