mirror of
https://github.com/openai/codex.git
synced 2026-10-30 19:18:49 +00:00
jif
e742db9f7b
## Why Shared routing cookies can change after a parent connection is established. Guardian reviews need the routing affinity captured with their parent response so unrelated traffic does not replace that hint. ## What changed - Capture `__oailb` and `__cflb` cookies from successful ChatGPT WebSocket handshakes and pair them with the parent response ID. - Retain the snapshot throughout a review and its retries, and apply it only to Guardian requests with the matching parent response. Reconnect cached reviewer sockets when the parent affinity changes. - Preserve cookie expiry and restrict reuse to the original secure origin and endpoint path. Replace routing cookies while retaining unrelated cookies, and keep the captured affinity out of prompts, history, and serialized metadata. ## Testing Add coverage for cookie scope, expiry, duplicate rejection, and merging; replaying a captured pin after shared cookies change or are deleted; and retaining the matching parent snapshot through compaction without exposing cookies in review prompts or metadata. GitOrigin-RevId: d180679e6265b612a3617f20ca1f5dec89e73aa8
107 lines
3.4 KiB
Rust
107 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 {
|
|
routing_affinity: None,
|
|
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()),
|
|
]))
|
|
);
|
|
}
|