mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(openai): isolate account context and scope websocket interrupts (#298)
Filter account-owned routing and identity context at API and Provider boundaries. Bind WebSocket interrupts to the active response on its original account and socket, preserving valid continuations. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -584,6 +584,8 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
|
||||
| "chatgpt-project-id"
|
||||
| "openai-organization"
|
||||
| "openai-project"
|
||||
| "x-openai-account-routing-override"
|
||||
| "x-openai-fedramp"
|
||||
// 上游指纹必须由运行时画像统一生成,客户端 originator/User-Agent/version
|
||||
// 不能作为不透明头透传覆盖,避免不同下游客户端暴露不一致的设备指纹。
|
||||
| "originator"
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界。
|
||||
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
fmt,
|
||||
sync::{
|
||||
Arc,
|
||||
@@ -14,7 +15,7 @@ use futures::{Sink, SinkExt, Stream, StreamExt};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use thiserror::Error;
|
||||
use tokio::{
|
||||
sync::{mpsc, oneshot},
|
||||
sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
|
||||
task::JoinHandle,
|
||||
time::{Instant, timeout},
|
||||
};
|
||||
@@ -48,6 +49,21 @@ pub enum ConnectionEvent {
|
||||
Exited(PumpExitReason),
|
||||
}
|
||||
|
||||
/// 接收队列与活动期间暂存的请求共用容量;移动帧不释放它占用的名额。
|
||||
pub(super) struct PendingConnectionEvent {
|
||||
pub(super) event: ConnectionEvent,
|
||||
_permit: Option<OwnedSemaphorePermit>,
|
||||
}
|
||||
|
||||
impl PendingConnectionEvent {
|
||||
fn exited(reason: PumpExitReason) -> Self {
|
||||
Self {
|
||||
event: ConnectionEvent::Exited(reason),
|
||||
_permit: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 下游写入阶段;名称刻意使用 write,而不是暗示客户端已消费的 delivery。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FramePhase {
|
||||
@@ -213,7 +229,8 @@ pub struct ResponsesWebSocketConnection {
|
||||
opened_at: Instant,
|
||||
expired: Arc<AtomicBool>,
|
||||
commands: Option<mpsc::Sender<ConnectionCommand>>,
|
||||
incoming: mpsc::Receiver<ConnectionEvent>,
|
||||
incoming: mpsc::Receiver<PendingConnectionEvent>,
|
||||
deferred: VecDeque<PendingConnectionEvent>,
|
||||
exited: oneshot::Receiver<PumpExitReason>,
|
||||
pump_task: Option<JoinHandle<()>>,
|
||||
stats: Arc<ConnectionStats>,
|
||||
@@ -251,17 +268,43 @@ impl ResponsesWebSocketConnection {
|
||||
|
||||
/// 等待下一个需要业务层处理的客户端事件。
|
||||
pub async fn next_event(&mut self) -> Option<ConnectionEvent> {
|
||||
if self.exit_reason.is_none() {
|
||||
match self.exited.try_recv() {
|
||||
Ok(reason) => {
|
||||
self.exit_reason = Some(reason);
|
||||
}
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
|
||||
self.exit_reason = Some(PumpExitReason::PumpStopped);
|
||||
}
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {}
|
||||
}
|
||||
}
|
||||
if let Some(reason) = self.exit_reason {
|
||||
return Some(ConnectionEvent::Exited(reason));
|
||||
}
|
||||
if let Some(event) = self.deferred.pop_front() {
|
||||
return Some(event.event);
|
||||
}
|
||||
self.next_active_event().await.map(|event| event.event)
|
||||
}
|
||||
|
||||
/// 活动响应期间只消费新到的帧,已排队的下一轮请求继续保持串行。
|
||||
pub(super) async fn next_active_event(&mut self) -> Option<PendingConnectionEvent> {
|
||||
if let Some(reason) = self.exit_reason {
|
||||
return Some(PendingConnectionEvent::exited(reason));
|
||||
}
|
||||
let event = tokio::select! {
|
||||
biased;
|
||||
reason = &mut self.exited => {
|
||||
Some(ConnectionEvent::Exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
|
||||
Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
|
||||
}
|
||||
event = self.incoming.recv() => event,
|
||||
};
|
||||
if let Some(ConnectionEvent::Exited(reason)) = event.as_ref() {
|
||||
if let Some(PendingConnectionEvent {
|
||||
event: ConnectionEvent::Exited(reason),
|
||||
..
|
||||
}) = event.as_ref()
|
||||
{
|
||||
self.exit_reason.get_or_insert(*reason);
|
||||
} else if event.is_none() {
|
||||
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
|
||||
@@ -269,6 +312,10 @@ impl ResponsesWebSocketConnection {
|
||||
event
|
||||
}
|
||||
|
||||
pub(super) fn defer(&mut self, event: PendingConnectionEvent) {
|
||||
self.deferred.push_back(event);
|
||||
}
|
||||
|
||||
/// 等待连接退出,不消费留给后续串行请求的业务帧。
|
||||
pub async fn wait_for_exit(&mut self) -> PumpExitReason {
|
||||
if let Some(reason) = self.exit_reason {
|
||||
@@ -494,6 +541,7 @@ where
|
||||
expired,
|
||||
commands: Some(command_tx),
|
||||
incoming: incoming_rx,
|
||||
deferred: VecDeque::new(),
|
||||
exited: exit_rx,
|
||||
pump_task: Some(pump_task),
|
||||
stats,
|
||||
@@ -506,7 +554,7 @@ where
|
||||
async fn run_pump<S, E>(
|
||||
mut socket: S,
|
||||
mut commands: mpsc::Receiver<ConnectionCommand>,
|
||||
incoming: mpsc::Sender<ConnectionEvent>,
|
||||
incoming: mpsc::Sender<PendingConnectionEvent>,
|
||||
exited: oneshot::Sender<PumpExitReason>,
|
||||
connection_id: Arc<str>,
|
||||
cancellation: CancellationToken,
|
||||
@@ -518,6 +566,7 @@ async fn run_pump<S, E>(
|
||||
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin,
|
||||
E: fmt::Display,
|
||||
{
|
||||
let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER));
|
||||
let deadline = tokio::time::sleep_until(opened_at + config.max_age);
|
||||
tokio::pin!(deadline);
|
||||
let mut deadline_elapsed = false;
|
||||
@@ -536,11 +585,9 @@ async fn run_pump<S, E>(
|
||||
expired.store(true, Ordering::Release);
|
||||
// 满队列本身已能唤醒空闲协调层;到期标志仍阻止启动下一轮。
|
||||
// 不能因无法追加 Expired 通知而中断正在收尾的响应。
|
||||
match incoming.try_send(ConnectionEvent::Expired) {
|
||||
Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
|
||||
Err(mpsc::error::TrySendError::Closed(_)) => {
|
||||
break PumpExitReason::CoordinatorDropped;
|
||||
}
|
||||
match emit_incoming(&incoming, &event_slots, ConnectionEvent::Expired) {
|
||||
Ok(()) | Err(PumpExitReason::InboundOverload) => {}
|
||||
Err(reason) => break reason,
|
||||
}
|
||||
}
|
||||
command = commands.recv() => {
|
||||
@@ -608,12 +655,12 @@ async fn run_pump<S, E>(
|
||||
}
|
||||
match message {
|
||||
Some(Ok(Message::Text(payload))) => {
|
||||
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Text(payload.to_string())) {
|
||||
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Text(payload.to_string())) {
|
||||
break reason;
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Binary(_))) => {
|
||||
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Binary) {
|
||||
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Binary) {
|
||||
break reason;
|
||||
}
|
||||
}
|
||||
@@ -662,11 +709,20 @@ async fn run_pump<S, E>(
|
||||
}
|
||||
|
||||
fn emit_incoming(
|
||||
incoming: &mpsc::Sender<ConnectionEvent>,
|
||||
incoming: &mpsc::Sender<PendingConnectionEvent>,
|
||||
slots: &Arc<Semaphore>,
|
||||
event: ConnectionEvent,
|
||||
) -> Result<(), PumpExitReason> {
|
||||
incoming.try_send(event).map_err(|error| match error {
|
||||
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
|
||||
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
|
||||
})
|
||||
let permit = Arc::clone(slots)
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| PumpExitReason::InboundOverload)?;
|
||||
incoming
|
||||
.try_send(PendingConnectionEvent {
|
||||
event,
|
||||
_permit: Some(permit),
|
||||
})
|
||||
.map_err(|error| match error {
|
||||
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
|
||||
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ use gateway_core::engine::middleware::{
|
||||
MiddlewareBody, MiddlewareError, MiddlewareFrame, MiddlewareFraming, MiddlewareHeader,
|
||||
MiddlewareResponse,
|
||||
};
|
||||
use gateway_core::engine::response_control::ResponseControl;
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::event::ProviderResponseHeader;
|
||||
use gateway_core::operation::ProviderSessionState;
|
||||
@@ -22,8 +23,10 @@ use crate::openai::responses::validation::{ResponseValidationFacts, ResponsesDel
|
||||
|
||||
use super::{
|
||||
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody},
|
||||
connection::{FramePhase, ResponsesWebSocketConnection, WriteContext},
|
||||
protocol::{error_event, initial_engine_error_event, response_metadata_event},
|
||||
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
|
||||
protocol::{
|
||||
decode_response_interrupt, error_event, initial_engine_error_event, response_metadata_event,
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
@@ -144,6 +147,7 @@ pub(super) async fn forward_response(
|
||||
replay: &mut ConnectionReplaySnapshot,
|
||||
capture: ReplayCaptureHandle,
|
||||
validation: ResponseValidationFacts,
|
||||
response_control: ResponseControl,
|
||||
) -> ForwardOutcome {
|
||||
let (protocol, status, headers, mut body, _) = response.into_parts();
|
||||
if protocol != "openai" || StatusCode::from_u16(status).is_err() {
|
||||
@@ -156,7 +160,9 @@ pub(super) async fn forward_response(
|
||||
ProviderResponseHeader::new(name, value)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let first = match next_body_input(connection, body.as_mut()).await {
|
||||
let first = match next_body_input(connection, body.as_mut(), &response_control, &request_id)
|
||||
.await
|
||||
{
|
||||
BodyInput::Frame(Ok(Some(frame))) => frame,
|
||||
BodyInput::Frame(Ok(None)) => {
|
||||
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
|
||||
@@ -200,10 +206,33 @@ pub(super) async fn forward_response(
|
||||
loop {
|
||||
let frame = match current.take() {
|
||||
Some(frame) => frame,
|
||||
None => match next_body_input(connection, body.as_mut()).await {
|
||||
BodyInput::Frame(Ok(Some(frame))) => frame,
|
||||
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
|
||||
if delivery_validator.finish_websocket_delivery().is_err() {
|
||||
None => {
|
||||
match next_body_input(connection, body.as_mut(), &response_control, &request_id)
|
||||
.await
|
||||
{
|
||||
BodyInput::Frame(Ok(Some(frame))) => frame,
|
||||
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
|
||||
if delivery_validator.finish_websocket_delivery().is_err() {
|
||||
return fail_body(
|
||||
connection,
|
||||
body,
|
||||
MiddlewareError::InvalidState,
|
||||
&request_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
{
|
||||
let captured = capture
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(response_id) = captured.response_id.clone() {
|
||||
replay.commit(response_id, captured.provider_state.clone());
|
||||
}
|
||||
}
|
||||
body.close().await;
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
BodyInput::Frame(Ok(None)) => {
|
||||
return fail_body(
|
||||
connection,
|
||||
body,
|
||||
@@ -212,29 +241,15 @@ pub(super) async fn forward_response(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
{
|
||||
let captured = capture
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(response_id) = captured.response_id.clone() {
|
||||
replay.commit(response_id, captured.provider_state.clone());
|
||||
}
|
||||
BodyInput::Frame(Err(error)) => {
|
||||
return fail_body(connection, body, error, &request_id).await;
|
||||
}
|
||||
BodyInput::Disconnect => {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
body.close().await;
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
BodyInput::Frame(Ok(None)) => {
|
||||
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id)
|
||||
.await;
|
||||
}
|
||||
BodyInput::Frame(Err(error)) => {
|
||||
return fail_body(connection, body, error, &request_id).await;
|
||||
}
|
||||
BodyInput::Disconnect => {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
},
|
||||
}
|
||||
};
|
||||
if terminal_seen || validate_frame(&frame).is_err() {
|
||||
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
|
||||
@@ -293,11 +308,45 @@ enum BodyInput {
|
||||
async fn next_body_input(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
body: &mut dyn MiddlewareBody,
|
||||
response_control: &ResponseControl,
|
||||
request_id: &Arc<str>,
|
||||
) -> BodyInput {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = connection.wait_for_exit() => BodyInput::Disconnect,
|
||||
frame = body.next_frame() => BodyInput::Frame(frame),
|
||||
// 处理控制帧时继续持有同一个读取 future,不能取消正在加工正文的中间件。
|
||||
let frame = body.next_frame();
|
||||
tokio::pin!(frame);
|
||||
loop {
|
||||
tokio::select! {
|
||||
event = connection.next_active_event() => {
|
||||
let Some(event) = event else { return BodyInput::Disconnect; };
|
||||
match &event.event {
|
||||
ConnectionEvent::Text(payload) => {
|
||||
let error = match decode_response_interrupt(payload) {
|
||||
Ok(Some(response_id)) => response_control.interrupt(&response_id).err().map(|_| {
|
||||
super::super::RequestDecodeError::InvalidValue { field: "response_id".to_owned() }.protocol_body()
|
||||
}),
|
||||
Ok(None) => {
|
||||
connection.defer(event);
|
||||
continue;
|
||||
}
|
||||
Err(error) => Some(error.protocol_body()),
|
||||
};
|
||||
if let Some(error) = error
|
||||
&& send_protocol_error(connection, StatusCode::BAD_REQUEST, error, request_id).await == ForwardOutcome::Disconnect
|
||||
{
|
||||
return BodyInput::Disconnect;
|
||||
}
|
||||
}
|
||||
ConnectionEvent::Expired => {
|
||||
// 允许当前响应收尾,外层仍会在下一轮准入前关闭过期连接。
|
||||
}
|
||||
ConnectionEvent::Binary => {
|
||||
connection.defer(event);
|
||||
}
|
||||
ConnectionEvent::Exited(_) => return BodyInput::Disconnect,
|
||||
}
|
||||
}
|
||||
frame = &mut frame => return BodyInput::Frame(frame),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -254,6 +254,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
|
||||
}
|
||||
};
|
||||
let request_id = Arc::<str>::from(prepared.request_id().to_string());
|
||||
let response_control = prepared.response_control();
|
||||
let capture = new_replay_capture();
|
||||
let validation = ResponseValidationFacts::default();
|
||||
let input = HttpMiddlewareInput {
|
||||
@@ -343,6 +344,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
|
||||
&mut replay,
|
||||
capture,
|
||||
validation,
|
||||
response_control,
|
||||
)
|
||||
.await
|
||||
== ForwardOutcome::Disconnect
|
||||
|
||||
@@ -31,6 +31,44 @@ pub fn decode_response_create_with_context(
|
||||
decode_response_create_inner(payload, request_headers)
|
||||
}
|
||||
|
||||
/// 中断只引用活动响应;普通创建帧仍交给原有解码与串行准入路径。
|
||||
pub(super) fn decode_response_interrupt(
|
||||
payload: &str,
|
||||
) -> Result<Option<String>, ResponseCreateFrameError> {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct FrameType {
|
||||
#[serde(rename = "type")]
|
||||
message_type: Option<String>,
|
||||
}
|
||||
// 非控制帧仍在原来的串行请求边界校验,不能抢先拒绝排队的普通请求。
|
||||
// 只投影类型,避免为排队的大型 response.create 再构造完整 JSON 树。
|
||||
let Ok(frame) = serde_json::from_str::<FrameType>(payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if frame.message_type.as_deref() != Some("response.interrupt") {
|
||||
return Ok(None);
|
||||
}
|
||||
let value: Value =
|
||||
serde_json::from_str(payload).map_err(|_| ResponseCreateFrameError::InvalidJson)?;
|
||||
if value.get("mode").and_then(Value::as_str) != Some("discard_partial_items") {
|
||||
return Err(ResponseCreateFrameError::Request(
|
||||
RequestDecodeError::InvalidValue {
|
||||
field: "mode".to_owned(),
|
||||
},
|
||||
));
|
||||
}
|
||||
let id = value
|
||||
.get("response_id")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|id| !id.is_empty())
|
||||
.ok_or_else(|| {
|
||||
ResponseCreateFrameError::Request(RequestDecodeError::InvalidValue {
|
||||
field: "response_id".to_owned(),
|
||||
})
|
||||
})?;
|
||||
Ok(Some(id.to_owned()))
|
||||
}
|
||||
|
||||
fn decode_response_create_inner(
|
||||
payload: &str,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
|
||||
@@ -143,6 +143,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/openai/responses/request.rs",
|
||||
"tests/openai/responses/websocket/connection.rs",
|
||||
"tests/openai/responses/websocket/forward.rs",
|
||||
"tests/openai/responses/websocket/interrupt.rs",
|
||||
"tests/openai/responses/websocket/mod.rs",
|
||||
"tests/openai/responses/websocket/protocol.rs",
|
||||
"tests/openai/router.rs",
|
||||
|
||||
@@ -313,6 +313,10 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
HeaderValue::from_static("client-attestation"),
|
||||
);
|
||||
headers.insert("x-oai-is", HeaderValue::from_static("client-is"));
|
||||
for name in ["x-openai-account-routing-override", "x-openai-fedramp"] {
|
||||
headers.append(name, HeaderValue::from_static("first"));
|
||||
headers.append(name, HeaderValue::from_static("second"));
|
||||
}
|
||||
headers.insert(
|
||||
"x-oai-is-update",
|
||||
HeaderValue::from_static("client-is-update"),
|
||||
@@ -381,6 +385,8 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
"cookie",
|
||||
"chatgpt-account-id",
|
||||
"chatgpt-project-id",
|
||||
"x-openai-account-routing-override",
|
||||
"x-openai-fedramp",
|
||||
// 上游指纹由运行时画像统一生成,客户端不得覆盖。
|
||||
"user-agent",
|
||||
"originator",
|
||||
|
||||
@@ -916,6 +916,33 @@ async fn queued_binary_frame_should_be_rejected_after_the_active_response_finish
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn requests_deferred_during_active_response_share_the_inbound_queue_limit() {
|
||||
let (trace, mut socket, server) = start_active_response().await;
|
||||
for _ in 0..33 {
|
||||
if socket
|
||||
.send(ClientMessage::Text(
|
||||
json!({"type":"response.create","model":"model-a","input":"queued"})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
trace.cancellation_observed.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("queue overload cancels the active execution");
|
||||
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
|
||||
assert!(trace.cancelled.load(Ordering::Acquire));
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_close_should_cancel_active_execution_without_starting_queued_requests() {
|
||||
let (trace, mut socket, server) = start_active_response().await;
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
use super::*;
|
||||
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
|
||||
|
||||
#[derive(Default)]
|
||||
struct InterruptProvider {
|
||||
calls: AtomicUsize,
|
||||
controls: Mutex<Vec<ResponseControl>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for InterruptProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"openai"
|
||||
}
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::default()
|
||||
}
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
self: Arc<Self>,
|
||||
request: ProviderRequest,
|
||||
context: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
let number = self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
let id = format!("resp_interrupt_{number}");
|
||||
let control = context
|
||||
.response_control()
|
||||
.expect("root control propagated through Core")
|
||||
.clone();
|
||||
let active = control.activate(id.clone()).unwrap();
|
||||
self.controls.lock().unwrap().push(control);
|
||||
let candidate = request.candidate();
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
candidate.provider().clone(),
|
||||
candidate.upstream_model().unwrap().clone(),
|
||||
ProviderAccountId::new("acct_api_test").unwrap(),
|
||||
UpstreamTransport::new("websocket").unwrap(),
|
||||
);
|
||||
let body = futures::stream::unfold(
|
||||
(0, Some(active), id),
|
||||
|(phase, mut active, id)| async move {
|
||||
let (event_type, status, event) = match phase {
|
||||
0 => (
|
||||
"response.created",
|
||||
"in_progress",
|
||||
GatewayEvent::Started(ResponseMeta::new(&id, "model-a")),
|
||||
),
|
||||
1 => {
|
||||
active.as_ref().unwrap().requested().await;
|
||||
drop(active.take());
|
||||
(
|
||||
"response.incomplete",
|
||||
"incomplete",
|
||||
GatewayEvent::Completed(ResponseMeta::new(&id, "model-a")),
|
||||
)
|
||||
}
|
||||
_ => return None,
|
||||
};
|
||||
let mut response = json!({"id":id,"model":"model-a","status":status,"output":[]});
|
||||
if phase == 1 {
|
||||
response["incomplete_details"] = json!({"reason":"interrupted"});
|
||||
}
|
||||
let event = ProviderEvent::canonical_with_wire(
|
||||
vec![event],
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some(event_type.to_owned()),
|
||||
json!({
|
||||
"type":event_type, "response":response
|
||||
}),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
Some((Ok(event), (phase + 1, active, id)))
|
||||
},
|
||||
);
|
||||
Ok(ProviderStream::new(metadata, body, ()))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct InterruptAdmissions {
|
||||
active: AtomicUsize,
|
||||
}
|
||||
|
||||
impl ClientAdmissionPort for InterruptAdmissions {
|
||||
fn abandon(&self, _: &ClientApiKeyId, _: &ModelRequestId) {
|
||||
self.active.fetch_sub(1, Ordering::SeqCst);
|
||||
}
|
||||
fn admit(
|
||||
&self,
|
||||
_: ClientAdmissionRequest,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
|
||||
Box::pin(async {
|
||||
self.active.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(ClientAdmissionDecision::Granted)
|
||||
})
|
||||
}
|
||||
fn release<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
|
||||
Box::pin(async {
|
||||
self.active.fetch_sub(1, Ordering::SeqCst);
|
||||
Ok(true)
|
||||
})
|
||||
}
|
||||
fn restore(
|
||||
&self,
|
||||
_: ClientAdmissionRecovery,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
|
||||
Box::pin(async { unreachable!("no recovery in live connection test") })
|
||||
}
|
||||
}
|
||||
|
||||
type TestSocket =
|
||||
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
||||
struct Server(tokio::task::JoinHandle<()>);
|
||||
impl Drop for Server {
|
||||
fn drop(&mut self) {
|
||||
self.0.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(socket: &mut TestSocket, value: Value) {
|
||||
socket
|
||||
.send(ClientMessage::Text(value.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn next(socket: &mut TestSocket) -> Value {
|
||||
loop {
|
||||
let message = tokio::time::timeout(Duration::from_secs(5), socket.next())
|
||||
.await
|
||||
.expect("response deadline")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let value: Value = serde_json::from_str(message.to_text().unwrap()).unwrap();
|
||||
if value["type"] != "response.metadata" {
|
||||
return value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
|
||||
let provider = Arc::new(InterruptProvider::default());
|
||||
let admissions = Arc::new(InterruptAdmissions::default());
|
||||
let execution = Arc::new(DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot("sk_ws_interrupt", "openai")),
|
||||
Arc::new(SettlementPorts::default()),
|
||||
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
|
||||
admissions.clone(),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(IgnoredClientApiKeyUsage),
|
||||
));
|
||||
let app = api_router(execution).await;
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let _server = Server(tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap()
|
||||
}));
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer sk_ws_interrupt"),
|
||||
);
|
||||
let (mut socket, _) = connect_async(request).await.unwrap();
|
||||
let create = json!({"type":"response.create","model":"model-a","input":"hello"});
|
||||
let interrupt =
|
||||
|id: &str, mode: &str| json!({"type":"response.interrupt","response_id":id,"mode":mode});
|
||||
send(
|
||||
&mut socket,
|
||||
interrupt("resp_interrupt_0", "discard_partial_items"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(next(&mut socket).await["type"], "error");
|
||||
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
|
||||
send(&mut socket, create.clone()).await;
|
||||
assert_eq!(next(&mut socket).await["type"], "response.created");
|
||||
send(&mut socket, create).await;
|
||||
for invalid in [
|
||||
interrupt("resp_other", "discard_partial_items"),
|
||||
interrupt("resp_interrupt_0", "unknown"),
|
||||
interrupt("", "discard_partial_items"),
|
||||
] {
|
||||
send(&mut socket, invalid).await;
|
||||
let error = next(&mut socket).await;
|
||||
assert_eq!(error["type"], "error");
|
||||
assert_eq!(error["status"], 400);
|
||||
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
send(
|
||||
&mut socket,
|
||||
interrupt("resp_interrupt_0", "discard_partial_items"),
|
||||
)
|
||||
.await;
|
||||
let terminal = next(&mut socket).await;
|
||||
assert_eq!(terminal["type"], "response.incomplete");
|
||||
assert_eq!(
|
||||
terminal["response"]["incomplete_details"]["reason"],
|
||||
"interrupted"
|
||||
);
|
||||
assert_eq!(
|
||||
next(&mut socket).await["response"]["id"],
|
||||
"resp_interrupt_1"
|
||||
);
|
||||
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
|
||||
send(
|
||||
&mut socket,
|
||||
interrupt("resp_interrupt_0", "discard_partial_items"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(next(&mut socket).await["type"], "error");
|
||||
assert_eq!(
|
||||
provider.controls.lock().unwrap()[0].interrupt("resp_interrupt_0"),
|
||||
Err(ResponseInterruptError::Unavailable)
|
||||
);
|
||||
socket.close(None).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
while admissions.active.load(Ordering::SeqCst) != 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("disconnect releases active execution");
|
||||
assert_eq!(
|
||||
provider.controls.lock().unwrap()[1].interrupt("resp_interrupt_1"),
|
||||
Err(ResponseInterruptError::Unavailable)
|
||||
);
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
mod connection;
|
||||
mod forward;
|
||||
mod interrupt;
|
||||
mod protocol;
|
||||
|
||||
use std::sync::{
|
||||
|
||||
@@ -52,6 +52,7 @@ enum AccountSelection {
|
||||
}
|
||||
|
||||
pub(super) struct CoordinationExtensions {
|
||||
response_control: Option<super::response_control::ResponseControl>,
|
||||
continuation: Option<ContinuationBinding>,
|
||||
observation: Option<RequestObservationDispatch>,
|
||||
request_policy: Option<super::policy::RequestPolicyContext>,
|
||||
@@ -65,12 +66,21 @@ pub(super) struct CoordinationExtensions {
|
||||
}
|
||||
|
||||
impl CoordinationExtensions {
|
||||
pub(super) fn with_response_control(
|
||||
mut self,
|
||||
control: Option<super::response_control::ResponseControl>,
|
||||
) -> Self {
|
||||
self.response_control = control;
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn new(
|
||||
continuation: Option<ContinuationBinding>,
|
||||
observation: Option<RequestObservationDispatch>,
|
||||
) -> Self {
|
||||
Self {
|
||||
continuation,
|
||||
response_control: None,
|
||||
observation,
|
||||
request_policy: None,
|
||||
execution_effects: None,
|
||||
@@ -230,6 +240,7 @@ where
|
||||
cancellation: CancellationToken,
|
||||
) -> Result<ResponseExecutionSession<S>, EngineError> {
|
||||
let CoordinationExtensions {
|
||||
response_control,
|
||||
continuation,
|
||||
observation: request_observation,
|
||||
request_policy,
|
||||
@@ -295,6 +306,7 @@ where
|
||||
.fuse(),
|
||||
pending_request: Some(request),
|
||||
request_persisted: false,
|
||||
response_control,
|
||||
operation,
|
||||
plan,
|
||||
request_policy,
|
||||
@@ -409,6 +421,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
deadline_timer: Fuse<Delay>,
|
||||
pending_request: Option<NewModelRequest>,
|
||||
request_persisted: bool,
|
||||
response_control: Option<super::response_control::ResponseControl>,
|
||||
operation: Operation,
|
||||
plan: RoutingPlan,
|
||||
request_policy: Option<super::policy::RequestPolicyContext>,
|
||||
@@ -996,6 +1009,7 @@ where
|
||||
}
|
||||
let context = AttemptContext::new(
|
||||
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
|
||||
.with_response_control(self.response_control.clone())
|
||||
.with_request_profile(self.request_profiles.get(candidate.provider()).cloned())
|
||||
.with_disable_fast(self.plan.disable_fast())
|
||||
.with_pricing(self.plan.pricing())
|
||||
|
||||
@@ -157,6 +157,7 @@ pub struct StartExecution {
|
||||
/// 中间件短路时本值直接释放;调用 `next` 后必须原样交回
|
||||
/// [`ExecutionService::start_prepared`],不能重新认证为另一个 Key。
|
||||
pub struct PreparedRootExecution {
|
||||
response_control: super::response_control::ResponseControl,
|
||||
client: AuthenticatedClient,
|
||||
request_id: ModelRequestId,
|
||||
started_at: SystemTime,
|
||||
@@ -167,6 +168,11 @@ pub struct PreparedRootExecution {
|
||||
}
|
||||
|
||||
impl PreparedRootExecution {
|
||||
#[must_use]
|
||||
pub fn response_control(&self) -> super::response_control::ResponseControl {
|
||||
self.response_control.clone()
|
||||
}
|
||||
|
||||
fn new(client: AuthenticatedClient) -> Result<Self, GatewayError> {
|
||||
let started_at = SystemTime::now();
|
||||
let deadline_at = started_at
|
||||
@@ -178,6 +184,7 @@ impl PreparedRootExecution {
|
||||
let execution_effects_baseline = execution_effects.epoch();
|
||||
Ok(Self {
|
||||
client,
|
||||
response_control: super::response_control::ResponseControl::default(),
|
||||
request_id: new_request_id()?,
|
||||
started_at,
|
||||
deadline_at,
|
||||
@@ -292,6 +299,7 @@ struct PendingStartExecution {
|
||||
}
|
||||
|
||||
struct AuthorizedExecution {
|
||||
response_control: Option<super::response_control::ResponseControl>,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
deadline_at: SystemTime,
|
||||
cancellation: CancellationToken,
|
||||
@@ -894,6 +902,7 @@ impl DefaultExecutionService {
|
||||
) -> Result<StartedExecution, GatewayError> {
|
||||
let PreparedRootExecution {
|
||||
client,
|
||||
response_control,
|
||||
request_id,
|
||||
started_at,
|
||||
deadline_at,
|
||||
@@ -902,6 +911,7 @@ impl DefaultExecutionService {
|
||||
execution_effects_baseline,
|
||||
} = prepared;
|
||||
let authorization = AuthorizedExecution {
|
||||
response_control: Some(response_control),
|
||||
account_scope: Arc::clone(client.policy.account_scope()),
|
||||
deadline_at,
|
||||
cancellation,
|
||||
@@ -1155,6 +1165,7 @@ impl DefaultExecutionService {
|
||||
.map(|graph| Arc::clone(&graph.effects));
|
||||
let extensions =
|
||||
CoordinationExtensions::new(continuation, request_observation.clone())
|
||||
.with_response_control(authorization.response_control.clone())
|
||||
.with_request_policy(request_policy)
|
||||
.with_execution_effects(
|
||||
execution_effects,
|
||||
@@ -1511,6 +1522,7 @@ impl DefaultExecutionService {
|
||||
metadata.user_agent = None;
|
||||
let authorization = AuthorizedExecution {
|
||||
account_scope,
|
||||
response_control: None,
|
||||
deadline_at: parent.deadline_at,
|
||||
cancellation: parent.cancellation.child_token(),
|
||||
extension_scope,
|
||||
@@ -1676,6 +1688,7 @@ impl DefaultExecutionService {
|
||||
},
|
||||
AuthorizedExecution {
|
||||
account_scope,
|
||||
response_control: None,
|
||||
deadline_at: authority.deadline_at,
|
||||
cancellation: authority.cancellation.child_token(),
|
||||
extension_scope: authority.extension_scope.clone(),
|
||||
|
||||
@@ -14,6 +14,7 @@ pub mod observation;
|
||||
pub mod policy;
|
||||
pub mod probe;
|
||||
pub mod provider;
|
||||
pub mod response_control;
|
||||
|
||||
pub use coordinator::{AttemptCoordinator, ResponseExecutionSession};
|
||||
|
||||
@@ -323,6 +324,7 @@ impl ContinuationAttempt {
|
||||
/// Provider 每次执行可见的 request-local context。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RequestAttemptContext {
|
||||
response_control: Option<response_control::ResponseControl>,
|
||||
pricing: Arc<crate::metering::PricingOverrides>,
|
||||
request_profile: Option<crate::account::OpaqueProviderData>,
|
||||
disable_fast: bool,
|
||||
@@ -343,6 +345,15 @@ pub struct RequestAttemptContext {
|
||||
}
|
||||
|
||||
impl RequestAttemptContext {
|
||||
#[must_use]
|
||||
pub fn with_response_control(
|
||||
mut self,
|
||||
control: Option<response_control::ResponseControl>,
|
||||
) -> Self {
|
||||
self.response_control = control;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_pricing(mut self, pricing: Arc<crate::metering::PricingOverrides>) -> Self {
|
||||
self.pricing = pricing;
|
||||
@@ -376,6 +387,7 @@ impl RequestAttemptContext {
|
||||
#[must_use]
|
||||
pub fn new(request_id: ModelRequestId, client_api_key_ref: ClientApiKeyId) -> Self {
|
||||
Self {
|
||||
response_control: None,
|
||||
request_id,
|
||||
client_api_key_ref,
|
||||
request_profile: None,
|
||||
@@ -527,6 +539,11 @@ pub struct AttemptContext {
|
||||
}
|
||||
|
||||
impl AttemptContext {
|
||||
#[must_use]
|
||||
pub fn response_control(&self) -> Option<&response_control::ResponseControl> {
|
||||
self.request.response_control.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn pricing(&self) -> &crate::metering::PricingOverrides {
|
||||
&self.request.pricing
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
//! 当前执行的响应中断信号;只由持有活动响应的 Provider 注册,不查找其他执行或连接。
|
||||
|
||||
use std::fmt;
|
||||
use std::sync::{Arc, Mutex, Weak};
|
||||
|
||||
use crate::lifecycle::CancellationToken;
|
||||
|
||||
/// 单次根执行的控制入口;重试不继承已结束响应的控制目标。
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ResponseControl {
|
||||
active: Arc<Mutex<Weak<InterruptTarget>>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ResponseControl {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("ResponseControl([REDACTED])")
|
||||
}
|
||||
}
|
||||
|
||||
struct InterruptTarget {
|
||||
response_id: String,
|
||||
requested: CancellationToken,
|
||||
}
|
||||
|
||||
/// Provider 持有的活动响应控制权;释放后客户端无法再向该响应发出中断。
|
||||
pub struct ActiveResponseInterrupt(Arc<InterruptTarget>);
|
||||
|
||||
/// 控制帧不能作用于当前执行时的稳定原因,不携带其他响应身份。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ResponseInterruptError {
|
||||
Unavailable,
|
||||
ResponseMismatch,
|
||||
}
|
||||
|
||||
impl ResponseControl {
|
||||
/// 在上游确认响应 ID 后注册;同一执行只允许一个存活的响应 owner。
|
||||
pub fn activate(&self, response_id: String) -> Option<ActiveResponseInterrupt> {
|
||||
let mut active = self
|
||||
.active
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if active.upgrade().is_some() || response_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let target = Arc::new(InterruptTarget {
|
||||
response_id,
|
||||
requested: CancellationToken::new(),
|
||||
});
|
||||
*active = Arc::downgrade(&target);
|
||||
Some(ActiveResponseInterrupt(target))
|
||||
}
|
||||
|
||||
/// 重复中断合并为一次信号;ID 必须匹配本执行当前的上游响应。
|
||||
pub fn interrupt(&self, response_id: &str) -> Result<(), ResponseInterruptError> {
|
||||
let active = self
|
||||
.active
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let target = active
|
||||
.upgrade()
|
||||
.ok_or(ResponseInterruptError::Unavailable)?;
|
||||
if target.response_id != response_id {
|
||||
return Err(ResponseInterruptError::ResponseMismatch);
|
||||
}
|
||||
target.requested.cancel();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl ActiveResponseInterrupt {
|
||||
pub fn response_id(&self) -> &str {
|
||||
&self.0.response_id
|
||||
}
|
||||
|
||||
pub fn requested(&self) -> impl Future<Output = ()> + Send + 'static {
|
||||
// 等待者只持有信号,不能延长活动响应 owner 的生命周期。
|
||||
let requested = self.0.requested.clone();
|
||||
async move { requested.cancelled().await }
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,7 @@ mod middleware;
|
||||
mod policy;
|
||||
mod probe;
|
||||
mod provider;
|
||||
mod response_control;
|
||||
|
||||
use gateway_core::engine::AttemptTrigger;
|
||||
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
use futures::FutureExt;
|
||||
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
|
||||
|
||||
#[test]
|
||||
fn interrupt_requires_the_current_response_owner_and_does_not_extend_its_lifetime() {
|
||||
let control = ResponseControl::default();
|
||||
assert_eq!(
|
||||
control.interrupt("resp_a"),
|
||||
Err(ResponseInterruptError::Unavailable)
|
||||
);
|
||||
assert!(control.activate(String::new()).is_none());
|
||||
let active = control.activate("resp_a".to_owned()).unwrap();
|
||||
let requested = active.requested();
|
||||
assert!(control.activate("resp_b".to_owned()).is_none());
|
||||
assert_eq!(
|
||||
control.interrupt("resp_b"),
|
||||
Err(ResponseInterruptError::ResponseMismatch)
|
||||
);
|
||||
assert!(active.requested().now_or_never().is_none());
|
||||
control.interrupt("resp_a").unwrap();
|
||||
control.interrupt("resp_a").unwrap();
|
||||
assert_eq!(requested.now_or_never(), Some(()));
|
||||
let retained_waiter = active.requested();
|
||||
drop(active);
|
||||
assert_eq!(
|
||||
control.interrupt("resp_a"),
|
||||
Err(ResponseInterruptError::Unavailable)
|
||||
);
|
||||
let next = control.activate("resp_b".to_owned()).unwrap();
|
||||
assert!(next.requested().now_or_never().is_none());
|
||||
drop(retained_waiter);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matching_response_ids_do_not_share_control_between_root_executions() {
|
||||
let first = ResponseControl::default();
|
||||
let second = ResponseControl::default();
|
||||
let active_first = first.activate("resp_same".to_owned()).unwrap();
|
||||
let active_second = second.activate("resp_same".to_owned()).unwrap();
|
||||
first.interrupt("resp_same").unwrap();
|
||||
assert_eq!(active_first.requested().now_or_never(), Some(()));
|
||||
assert!(active_second.requested().now_or_never().is_none());
|
||||
}
|
||||
@@ -137,7 +137,8 @@ fn connection_options(headers: &[MiddlewareHeader]) -> Vec<String> {
|
||||
}
|
||||
|
||||
fn sensitive_or_identity_header(name: &str) -> bool {
|
||||
name.contains("auth")
|
||||
name == "x-openai-fedramp"
|
||||
|| name.contains("auth")
|
||||
|| name.contains("credential")
|
||||
|| name.contains("secret")
|
||||
|| name.contains("token")
|
||||
|
||||
@@ -714,6 +714,11 @@ fn attempt_middleware_context(effects: Arc<ExecutionEffects>) -> MiddlewareConte
|
||||
fn middleware_headers() -> Vec<MiddlewareHeader> {
|
||||
vec![
|
||||
MiddlewareHeader::new("authorization", Bytes::from_static(b"Bearer private")),
|
||||
MiddlewareHeader::new("X-OpenAI-Fedramp", Bytes::from_static(b"true")),
|
||||
MiddlewareHeader::new(
|
||||
"x-openai-account-routing-override",
|
||||
Bytes::from_static(b"us_cr"),
|
||||
),
|
||||
MiddlewareHeader::new("connection", Bytes::from_static(b"x-hop")),
|
||||
MiddlewareHeader::new("x-hop", Bytes::from_static(b"hidden")),
|
||||
MiddlewareHeader::new("x_team", Bytes::from_static(b"team")),
|
||||
|
||||
@@ -863,6 +863,7 @@ impl CodexProvider {
|
||||
})?
|
||||
.with_authentication(lease.authentication())
|
||||
.with_connection_budget(context.connection_budget().clone())
|
||||
.with_response_control(context.response_control().cloned())
|
||||
.with_middleware_headers(middleware_headers),
|
||||
response_origin: self.responses_url.clone(),
|
||||
request: upstream_request,
|
||||
|
||||
@@ -689,6 +689,7 @@ impl OpenAiUpstreamProtocol {
|
||||
/// Codex HTTP/SSE 上游客户端。
|
||||
#[derive(Clone)]
|
||||
pub struct CodexBackendClient {
|
||||
pub(super) response_control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
pub(super) connection_budget: Option<gateway_core::engine::connection::ConnectionBudget>,
|
||||
pub(super) client: Client,
|
||||
pub(super) direct_client: Client,
|
||||
@@ -705,6 +706,14 @@ pub struct CodexBackendClient {
|
||||
}
|
||||
|
||||
impl CodexBackendClient {
|
||||
pub(crate) fn with_response_control(
|
||||
mut self,
|
||||
control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
) -> Self {
|
||||
self.response_control = control;
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_connection_budget(
|
||||
mut self,
|
||||
budget: gateway_core::engine::connection::ConnectionBudget,
|
||||
|
||||
@@ -65,6 +65,7 @@ impl CodexBackendClient {
|
||||
let base_url = base_url.into().trim_end_matches('/').to_string();
|
||||
Self {
|
||||
connection_budget: None,
|
||||
response_control: None,
|
||||
direct_client: client.clone(),
|
||||
client,
|
||||
websocket_origin_key: websocket_origin_key(&base_url),
|
||||
@@ -453,6 +454,7 @@ impl CodexBackendClient {
|
||||
let mut exchange = execute_prepared_response_create_request_stream(
|
||||
&websocket_request,
|
||||
prepared,
|
||||
self.response_control.clone(),
|
||||
context
|
||||
.trace
|
||||
.cloned()
|
||||
|
||||
@@ -30,6 +30,9 @@ pub(super) fn is_managed_identity_header(name: &str) -> bool {
|
||||
| "chatgpt-project-id"
|
||||
| "openai-organization"
|
||||
| "openai-project"
|
||||
// 工作区路由和合规属性属于认证账号,不能继承下游账号的值。
|
||||
| "x-openai-account-routing-override"
|
||||
| "x-openai-fedramp"
|
||||
// 安装身份由当前账号写入 client_metadata,不继承下游安装头。
|
||||
| "x-codex-installation-id"
|
||||
)
|
||||
@@ -217,7 +220,7 @@ impl CodexBackendClient {
|
||||
let provider_headers = headers.keys().cloned().collect::<Vec<_>>();
|
||||
for header in &self.middleware_headers {
|
||||
let name = HeaderName::from_bytes(header.name().as_bytes())?;
|
||||
if provider_headers.contains(&name) {
|
||||
if is_managed_identity_header(name.as_str()) || provider_headers.contains(&name) {
|
||||
return Err(super::client::CodexClientError::MiddlewareHeaderConflict);
|
||||
}
|
||||
headers.append(name, HeaderValue::from_bytes(header.value())?);
|
||||
|
||||
@@ -138,11 +138,13 @@ fn json_value_as_string(value: &Value) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 `response.completed` 旁路提取 response ID,供连接池记录续接能力。
|
||||
/// 从正常完成或官方中断终态旁路提取 response ID,供连接池记录续接能力。
|
||||
///
|
||||
/// 该值不参与客户端 wire 的可交付性判断;无法读取时只是不记录连接内续接状态。
|
||||
pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.completed") {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.completed")
|
||||
&& !websocket_response_is_interrupted(value)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
value
|
||||
@@ -151,6 +153,14 @@ pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn websocket_response_is_interrupted(value: &Value) -> bool {
|
||||
value.get("type").and_then(Value::as_str) == Some("response.incomplete")
|
||||
&& value
|
||||
.pointer("/response/incomplete_details/reason")
|
||||
.and_then(Value::as_str)
|
||||
== Some("interrupted")
|
||||
}
|
||||
|
||||
/// 生成 Responses WebSocket payload 审计快照。
|
||||
pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot {
|
||||
let body = websocket_response_create_payload(request);
|
||||
|
||||
@@ -536,6 +536,9 @@ pub(crate) fn scope_request_to_account(
|
||||
{
|
||||
metadata.remove(*key);
|
||||
}
|
||||
// Guardian 顶层父引用指向原账号的响应;turn metadata 内的同名
|
||||
// 扩展只是客户端关联信息,不能加入各层共用的清理名单。
|
||||
metadata.remove("parent_response_id");
|
||||
}
|
||||
// 官方 Core 在 client_metadata 使用带 x-codex 前缀的键,
|
||||
// turn metadata 内仍使用 installation_id;两处均取当前账号的安装身份。
|
||||
|
||||
@@ -581,6 +581,7 @@ fn continuation_unavailable(
|
||||
pub(crate) async fn execute_prepared_response_create_request_stream(
|
||||
request: &CodexWebSocketRequest,
|
||||
prepared: PreparedWebSocket,
|
||||
response_control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
trace: gateway_core::diagnostics::TraceContext,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let PreparedWebSocket {
|
||||
@@ -638,6 +639,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream(
|
||||
reused,
|
||||
stream_idle_timeout,
|
||||
trace,
|
||||
response_control,
|
||||
);
|
||||
exchange.pool_decision = pool_decision;
|
||||
exchange.connection_local_continuation = connection_local_available;
|
||||
|
||||
@@ -6,6 +6,7 @@ use serde_json::Value;
|
||||
use crate::transport::protocol::websocket::{
|
||||
websocket_event_frame, websocket_event_type, websocket_metadata_headers,
|
||||
websocket_metadata_turn_state, websocket_response_completed_id,
|
||||
websocket_response_is_interrupted,
|
||||
};
|
||||
use crate::transport::response_meta;
|
||||
|
||||
@@ -15,6 +16,7 @@ use super::CodexWebSocketExchangeError;
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(in crate::transport::websocket) enum WebSocketTerminalKind {
|
||||
Completed,
|
||||
Interrupted,
|
||||
Incomplete,
|
||||
Failed,
|
||||
}
|
||||
@@ -29,6 +31,7 @@ pub(super) enum ExchangeAction {
|
||||
}
|
||||
|
||||
pub(super) struct ReducedWebSocketEvent {
|
||||
pub(super) created_response_id: Option<String>,
|
||||
pub(super) action: ExchangeAction,
|
||||
pub(super) diagnostic_event_type: Option<String>,
|
||||
pub(super) turn_state_update: Option<String>,
|
||||
@@ -43,6 +46,7 @@ pub(super) fn reduce_websocket_event(
|
||||
// 不可解析的帧不承载可路由的事件类型,忽略。
|
||||
let Ok(value) = serde_json::from_str::<Value>(raw) else {
|
||||
return Ok(ReducedWebSocketEvent {
|
||||
created_response_id: None,
|
||||
action: ExchangeAction::Ignore,
|
||||
diagnostic_event_type: None,
|
||||
turn_state_update: None,
|
||||
@@ -54,6 +58,7 @@ pub(super) fn reduce_websocket_event(
|
||||
metadata.rate_limit_headers.extend(headers);
|
||||
return Ok(ReducedWebSocketEvent {
|
||||
action: ExchangeAction::RateLimits(parsed),
|
||||
created_response_id: None,
|
||||
diagnostic_event_type,
|
||||
turn_state_update: None,
|
||||
});
|
||||
@@ -75,14 +80,15 @@ pub(super) fn reduce_websocket_event(
|
||||
});
|
||||
|
||||
let event = websocket_event_type(&value);
|
||||
if event == Some("response.completed")
|
||||
&& let Some(response_id) = websocket_response_completed_id(&value)
|
||||
{
|
||||
if let Some(response_id) = websocket_response_completed_id(&value) {
|
||||
continuation.record_completed(response_id);
|
||||
}
|
||||
|
||||
let terminal = match event {
|
||||
Some("response.completed") => Some(WebSocketTerminalKind::Completed),
|
||||
Some("response.incomplete") if websocket_response_is_interrupted(&value) => {
|
||||
Some(WebSocketTerminalKind::Interrupted)
|
||||
}
|
||||
Some("response.incomplete") => Some(WebSocketTerminalKind::Incomplete),
|
||||
Some("response.failed" | "error") => Some(WebSocketTerminalKind::Failed),
|
||||
_ => None,
|
||||
@@ -92,6 +98,10 @@ pub(super) fn reduce_websocket_event(
|
||||
None => ExchangeAction::Ignore,
|
||||
};
|
||||
Ok(ReducedWebSocketEvent {
|
||||
created_response_id: (event == Some("response.created"))
|
||||
.then(|| value.pointer("/response/id").and_then(Value::as_str))
|
||||
.flatten()
|
||||
.map(ToOwned::to_owned),
|
||||
action,
|
||||
diagnostic_event_type,
|
||||
turn_state_update,
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::{sync::Arc, time::Duration};
|
||||
|
||||
use bytes::Bytes;
|
||||
use gateway_core::diagnostics::{TraceContext, diagnostic_json};
|
||||
use gateway_core::engine::response_control::{ActiveResponseInterrupt, ResponseControl};
|
||||
use serde_json::json;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
@@ -65,6 +66,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
reused_connection: bool,
|
||||
stream_idle_timeout: Option<Duration>,
|
||||
trace: TraceContext,
|
||||
response_control: Option<ResponseControl>,
|
||||
) -> CodexWebSocketStreamingExchange {
|
||||
let websocket_connection_id = websocket.connection_id();
|
||||
let response_metadata = metadata.clone();
|
||||
@@ -85,6 +87,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
);
|
||||
let forward = async move {
|
||||
forward_websocket_response_stream(WebSocketStreamForwardState {
|
||||
response_control,
|
||||
websocket,
|
||||
metadata,
|
||||
pool_return,
|
||||
@@ -124,6 +127,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
}
|
||||
|
||||
struct WebSocketStreamForwardState {
|
||||
response_control: Option<ResponseControl>,
|
||||
trace: TraceContext,
|
||||
websocket: PumpedWebSocket,
|
||||
metadata: CodexWebSocketConnectionMetadata,
|
||||
@@ -138,6 +142,7 @@ struct WebSocketStreamForwardState {
|
||||
|
||||
async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let WebSocketStreamForwardState {
|
||||
response_control,
|
||||
trace,
|
||||
mut websocket,
|
||||
mut metadata,
|
||||
@@ -155,6 +160,10 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
.map(|pool_return| std::mem::take(&mut pool_return.continuation))
|
||||
.unwrap_or_default();
|
||||
let mut last_event_type = None;
|
||||
let mut active_interrupt: Option<ActiveResponseInterrupt> = None;
|
||||
let mut interrupt_requested: futures::future::BoxFuture<'static, ()> =
|
||||
Box::pin(std::future::pending());
|
||||
let mut interrupt_sent = false;
|
||||
// 官方 run_websocket_response_stream 按响应消费 metadata 事件;它们不属于握手头。
|
||||
// 复用连接时恢复握手快照,避免上一轮模型信息及普通响应头随每次请求累积。
|
||||
let opening_response_metadata = metadata.response_metadata.clone();
|
||||
@@ -164,6 +173,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
// 客户端断开下游 SSE 流:立即丢弃连接并释放池 slot,
|
||||
// 不再傻等上游 idle 超时(否则同会话后续请求会一直 bypass/busy)。
|
||||
() = tx.closed() => {
|
||||
drop(active_interrupt.take());
|
||||
trace.record("upstream.cancelled", json!({"reason": "receiver_dropped"}));
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
@@ -173,6 +183,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
return;
|
||||
}
|
||||
() = shutdown.cancelled() => {
|
||||
drop(active_interrupt.take());
|
||||
trace.record("upstream.cancelled", json!({"reason": "pool_shutdown"}));
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
@@ -181,6 +192,24 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
).await;
|
||||
return;
|
||||
}
|
||||
() = &mut interrupt_requested, if !interrupt_sent => {
|
||||
let Some(active) = active_interrupt.as_ref() else { continue };
|
||||
let payload = json!({
|
||||
"type": "response.interrupt",
|
||||
"response_id": active.response_id(),
|
||||
"mode": "discard_partial_items",
|
||||
}).to_string();
|
||||
// 复用正在消费此响应的原 socket;信号不经过选号或新请求恢复。
|
||||
if let Err(error) = super::super::handshake::send_websocket_request(&websocket, &payload).await {
|
||||
drop(active_interrupt.take());
|
||||
discard_stream_websocket(websocket, pool_return, StreamWebSocketDiscardReason::UpstreamReceiveFailed).await;
|
||||
let _ = tx.send(Err(error)).await;
|
||||
return;
|
||||
}
|
||||
interrupt_sent = true;
|
||||
trace.record("upstream.interrupt.sent", json!({"mode": "discard_partial_items"}));
|
||||
continue;
|
||||
}
|
||||
message = next_websocket_message(
|
||||
&mut websocket,
|
||||
stream_idle_timeout.filter(|timeout| !timeout.is_zero()).unwrap_or(super::super::pool::DEFAULT_STREAM_IDLE_TIMEOUT),
|
||||
@@ -189,6 +218,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let message = match message {
|
||||
Ok(message) => message,
|
||||
Err(error) => {
|
||||
// 先撤销 owner 再发布失败;下一次 attempt 不等待旧 socket 异步清理。
|
||||
drop(active_interrupt.take());
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
trace.record(
|
||||
"upstream.read.failed",
|
||||
@@ -224,6 +255,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let raw = match message {
|
||||
tungstenite::Message::Text(text) => text.to_string(),
|
||||
tungstenite::Message::Binary(bytes) => {
|
||||
drop(active_interrupt.take());
|
||||
trace.capture("upstream.binary", &bytes);
|
||||
let error = CodexWebSocketExchangeError::UnexpectedBinaryEvent;
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
@@ -239,6 +271,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
return;
|
||||
}
|
||||
tungstenite::Message::Close(frame) => {
|
||||
drop(active_interrupt.take());
|
||||
if let Some(frame) = &frame {
|
||||
trace.dump("upstream.close.reason", frame.reason.as_bytes());
|
||||
}
|
||||
@@ -277,6 +310,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let reduced = match reduce_websocket_event(&raw, &mut metadata, &mut continuation) {
|
||||
Ok(reduced) => reduced,
|
||||
Err(error) => {
|
||||
drop(active_interrupt.take());
|
||||
let observation =
|
||||
(!matches!(error.classified(), CodexWebSocketExchangeError::Upstream(_)))
|
||||
.then(|| connection_observation(&websocket, &error));
|
||||
@@ -297,6 +331,31 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
if let Some(event_type) = reduced.diagnostic_event_type {
|
||||
last_event_type = Some(event_type);
|
||||
}
|
||||
if active_interrupt.is_none()
|
||||
&& let (Some(control), Some(response_id)) =
|
||||
(&response_control, reduced.created_response_id)
|
||||
{
|
||||
active_interrupt = control.activate(response_id);
|
||||
if let Some(active) = &active_interrupt {
|
||||
// 跨消息复用同一个等待 future,避免每帧留下新的取消等待者。
|
||||
interrupt_requested = Box::pin(active.requested());
|
||||
}
|
||||
if active_interrupt.is_none() {
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
pool_return,
|
||||
StreamWebSocketDiscardReason::FailedResponse,
|
||||
)
|
||||
.await;
|
||||
let _ = tx
|
||||
.send(Err(CodexWebSocketExchangeError::PostSendAmbiguous {
|
||||
message: "response control already has an active owner".to_owned(),
|
||||
source: None,
|
||||
}))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
if let Some(turn_state) = reduced.turn_state_update {
|
||||
let mut pending = response_metadata_updates.lock().await;
|
||||
if pending.turn_state.is_none() {
|
||||
@@ -314,7 +373,12 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
ExchangeAction::Forward { frame, terminal } => (frame, terminal),
|
||||
ExchangeAction::Ignore => continue,
|
||||
};
|
||||
if terminal.is_some() {
|
||||
// 先撤销控制权,再交付终态和归还连接,避免迟到的中断影响下一轮。
|
||||
drop(active_interrupt.take());
|
||||
}
|
||||
if tx.send(Ok(Bytes::from(frame))).await.is_err() {
|
||||
drop(active_interrupt.take());
|
||||
trace.record(
|
||||
"upstream.forward.failed",
|
||||
json!({"reason": "receiver_dropped"}),
|
||||
@@ -333,7 +397,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
json!({"kind": format!("{terminal:?}")}),
|
||||
);
|
||||
match terminal {
|
||||
WebSocketTerminalKind::Completed => {
|
||||
WebSocketTerminalKind::Completed | WebSocketTerminalKind::Interrupted => {
|
||||
metadata.response_metadata = opening_response_metadata;
|
||||
finish_stream_websocket(websocket, metadata, continuation, pool_return.take())
|
||||
.await;
|
||||
@@ -392,6 +456,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
last_event_type,
|
||||
},
|
||||
};
|
||||
drop(active_interrupt.take());
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
use super::*;
|
||||
|
||||
const SELECTED_ACCOUNT: &str = "acct_scope_same";
|
||||
|
||||
fn isolation_history() -> Value {
|
||||
json!([
|
||||
{"role":"user", "content":"hello"},
|
||||
{"type":"function_call", "id":"fc_history", "call_id":"call_history", "name":"lookup", "arguments":"{}"},
|
||||
{
|
||||
"type":"function_call_output", "call_id":"call_history", "output":"tool result",
|
||||
"internal_chat_message_metadata_passthrough": {
|
||||
"executed_tool_calls": [{
|
||||
"name":"lookup", "arguments":{},
|
||||
"tool_result_metadata": {
|
||||
"openai/resource_access":{"resource_coverage":"complete", "resources":[]},
|
||||
"parent_response_id":"tool-business-extension"
|
||||
}
|
||||
}]
|
||||
}
|
||||
},
|
||||
{"role":"user", "content":"continue"}
|
||||
])
|
||||
}
|
||||
|
||||
fn isolation_payload() -> ProtocolPayload {
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": isolation_history(),
|
||||
"client_metadata": {
|
||||
"parent_response_id": "resp_parent_old",
|
||||
"guardian_credits_requested": "true",
|
||||
"x-codex-turn-metadata": r#"{"parent_response_id":"client-correlation"}"#,
|
||||
"mcp_attribution": r#"{"status":"complete","sources":[]}"#
|
||||
}
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.with_context(Map::from_iter([(
|
||||
"opaque_request_headers".to_owned(),
|
||||
json!([
|
||||
["X-OpenAI-Account-Routing-Override", STANDARD.encode("us")],
|
||||
[
|
||||
"x-openai-account-routing-override",
|
||||
STANDARD.encode("us_cr")
|
||||
],
|
||||
["X-OpenAI-Fedramp", STANDARD.encode("true")],
|
||||
["x-openai-fedramp", STANDARD.encode("false")],
|
||||
["x-codex-guardian", STANDARD.encode("true")],
|
||||
["x-business-extension", STANDARD.encode("first")],
|
||||
["x-business-extension", STANDARD.encode("second")]
|
||||
]),
|
||||
)]))
|
||||
}
|
||||
|
||||
fn isolation_context(owner: Option<&str>) -> AttemptContext {
|
||||
match owner {
|
||||
Some(owner) => context_with_state_owner("req_account_isolation", owner),
|
||||
None => context("req_account_isolation", CancellationToken::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_isolated_request(
|
||||
headers: &reqwest::header::HeaderMap,
|
||||
body: &Value,
|
||||
owner: Option<&str>,
|
||||
) {
|
||||
for name in ["x-openai-account-routing-override", "x-openai-fedramp"] {
|
||||
assert!(
|
||||
!headers.contains_key(name),
|
||||
"leaked {name} with owner {owner:?}"
|
||||
);
|
||||
}
|
||||
assert_eq!(headers["x-codex-guardian"], "true");
|
||||
assert_eq!(
|
||||
headers
|
||||
.get_all("x-business-extension")
|
||||
.iter()
|
||||
.map(|value| value.as_bytes())
|
||||
.collect::<Vec<_>>(),
|
||||
vec![b"first".as_slice(), b"second".as_slice()],
|
||||
);
|
||||
assert_eq!(
|
||||
body.pointer("/client_metadata/parent_response_id")
|
||||
.and_then(Value::as_str),
|
||||
(owner == Some(SELECTED_ACCOUNT)).then_some("resp_parent_old"),
|
||||
);
|
||||
assert_eq!(
|
||||
body["client_metadata"]["guardian_credits_requested"],
|
||||
"true"
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(
|
||||
body["client_metadata"]["x-codex-turn-metadata"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
)
|
||||
.unwrap(),
|
||||
json!({"parent_response_id": "client-correlation"}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["client_metadata"]["mcp_attribution"],
|
||||
r#"{"status":"complete","sources":[]}"#
|
||||
);
|
||||
assert_eq!(body["input"], isolation_history());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_isolates_account_headers_and_only_scopes_guardian_parent_reference() {
|
||||
for owner in [Some(SELECTED_ACCOUNT), Some("acct_scope_old"), None] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, SELECTED_ACCOUNT).await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(CAPTURE_COMPLETED_SSE),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mut payload = isolation_payload();
|
||||
let mut context = payload.context().clone();
|
||||
context.insert("use_websocket".to_owned(), json!(false));
|
||||
payload = payload.with_context(context);
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
|
||||
),
|
||||
isolation_context(owner),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
while let Some(event) = stream.next().await {
|
||||
event.unwrap();
|
||||
}
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_isolated_request(
|
||||
&requests[0].headers,
|
||||
&captured_request_body(&requests[0]),
|
||||
owner,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_isolates_account_headers_and_only_scopes_guardian_parent_reference() {
|
||||
for owner in [Some(SELECTED_ACCOUNT), Some("acct_scope_old"), None] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, SELECTED_ACCOUNT).await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut headers = None;
|
||||
let mut socket =
|
||||
crate::transport::accept_codex_test_websocket_with(socket, |request, _| {
|
||||
headers = Some(request.headers().clone());
|
||||
})
|
||||
.await;
|
||||
let message = socket.next().await.unwrap().unwrap();
|
||||
let body: Value = serde_json::from_str(message.to_text().unwrap()).unwrap();
|
||||
socket.send(Message::Text(json!({
|
||||
"type": "response.completed",
|
||||
"response": {"id":"resp_isolation", "model":"gpt-5.4", "status":"completed", "output":[]}
|
||||
}).to_string().into())).await.unwrap();
|
||||
(headers.unwrap(), body)
|
||||
});
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(
|
||||
GenerateRequest::from_protocol_payload(isolation_payload()),
|
||||
),
|
||||
),
|
||||
isolation_context(owner),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
while let Some(event) = stream.next().await {
|
||||
event.unwrap();
|
||||
}
|
||||
let (headers, body) = server.await.unwrap();
|
||||
assert_isolated_request(&headers, &body, owner);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn middleware_cannot_reintroduce_account_headers_before_http_send_or_websocket_open() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, SELECTED_ACCOUNT).await;
|
||||
let server = MockServer::start().await;
|
||||
for operation in [http_generate_operation(), generate_operation()] {
|
||||
for name in [
|
||||
"x-openai-account-routing-override",
|
||||
"x-openai-fedramp",
|
||||
"x-codex-installation-id",
|
||||
] {
|
||||
let result = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", operation.clone()),
|
||||
context_with_middleware(
|
||||
"req_managed_account_header",
|
||||
Arc::new(RecordingMiddleware {
|
||||
observed: Arc::default(),
|
||||
replacement: ("service_tier".to_owned(), json!("default")),
|
||||
request_headers: vec![MiddlewareHeader::new(
|
||||
name,
|
||||
Bytes::from_static(b"downstream-value"),
|
||||
)],
|
||||
}),
|
||||
false,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
let error = match result {
|
||||
Err(error) => error,
|
||||
Ok(mut stream) => match stream.next().await {
|
||||
Some(Err(error)) => error,
|
||||
_ => panic!("managed account header must fail before send: {name}"),
|
||||
},
|
||||
};
|
||||
assert_eq!(error.kind(), ProviderErrorKind::Protocol);
|
||||
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
|
||||
}
|
||||
}
|
||||
assert!(server.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
//! 显式指定本机导出文件后才运行;凭据只放入内存,不刷新、不落盘、不打印响应或身份。
|
||||
|
||||
use super::*;
|
||||
use gateway_core::engine::response_control::ResponseControl;
|
||||
use provider_openai::credential::CodexOAuthSecret;
|
||||
use secrecy::SecretString;
|
||||
|
||||
fn live_context(
|
||||
account: &str,
|
||||
owner: Option<&str>,
|
||||
control: Option<ResponseControl>,
|
||||
continuation: Option<ContinuationBinding>,
|
||||
) -> AttemptContext {
|
||||
let trace = control.as_ref().map_or_else(Default::default, |_| {
|
||||
gateway_core::diagnostics::TraceContext::new("req_live_isolation")
|
||||
});
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new("req_live_isolation").unwrap(),
|
||||
ClientApiKeyId::new("key_openai_contract").unwrap(),
|
||||
)
|
||||
.with_response_control(control)
|
||||
.with_trace(trace),
|
||||
NonZeroU32::new(1).unwrap(),
|
||||
SystemTime::now() + Duration::from_secs(90),
|
||||
account_policy(),
|
||||
AccountAttemptContext::diagnostic(
|
||||
BTreeSet::new(),
|
||||
ProviderAccountId::new(account).unwrap(),
|
||||
owner.map(|owner| {
|
||||
ProviderAccountStateOwner::new(
|
||||
ProviderKind::new("openai").unwrap(),
|
||||
ProviderAccountId::new(owner).unwrap(),
|
||||
)
|
||||
}),
|
||||
),
|
||||
continuation,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
}
|
||||
|
||||
async fn live_store() -> Arc<MemoryAccountStore> {
|
||||
let path =
|
||||
std::env::var("CPR_LIVE_ACCOUNTS_FILE").expect("set CPR_LIVE_ACCOUNTS_FILE explicitly");
|
||||
let bytes = std::fs::read(path).expect("read local account export");
|
||||
let document: Value = serde_json::from_slice(&bytes).expect("parse local account export");
|
||||
let accounts: Vec<&Value> = document["documents"]
|
||||
.as_array()
|
||||
.expect("documents")
|
||||
.iter()
|
||||
.filter(|doc| doc["provider"] == "openai")
|
||||
.flat_map(|doc| doc["document"]["accounts"].as_array().expect("accounts"))
|
||||
.take(2)
|
||||
.collect();
|
||||
assert!(accounts.len() == 2, "two OpenAI accounts required");
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
for (account, id) in accounts
|
||||
.into_iter()
|
||||
.zip(["acct_scope_old", "acct_scope_new"])
|
||||
{
|
||||
let text = |key: &str| {
|
||||
account[key]
|
||||
.as_str()
|
||||
.expect("required account field")
|
||||
.to_owned()
|
||||
};
|
||||
let expiration = chrono::DateTime::parse_from_rfc3339(&text("accessTokenExpiresAt"))
|
||||
.expect("expiry timestamp")
|
||||
.with_timezone(&Utc);
|
||||
assert!(
|
||||
expiration > Utc::now(),
|
||||
"live access token expired; test does not rotate credentials"
|
||||
);
|
||||
let mut verified = profile(&text("accountId"));
|
||||
verified.chatgpt_user_id = text("userId");
|
||||
verified.oauth_subject = verified.chatgpt_user_id.clone();
|
||||
verified.poid = None;
|
||||
verified.email = None;
|
||||
verified.access_token_expires_at = Some(expiration);
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: id.to_owned(),
|
||||
name: id.to_owned(),
|
||||
secret: CodexOAuthSecret {
|
||||
access_token: SecretString::from(text("accessToken")),
|
||||
refresh_token: None,
|
||||
id_token: None,
|
||||
},
|
||||
verified_account: verified,
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
store
|
||||
}
|
||||
|
||||
async fn live_response(
|
||||
provider: Arc<CodexProvider>,
|
||||
websocket: bool,
|
||||
account: &str,
|
||||
owner: Option<&str>,
|
||||
input: Value,
|
||||
parent: Option<&str>,
|
||||
) -> (String, Value) {
|
||||
eprintln!(
|
||||
"live request: transport={}, account={}, parent_reference={}",
|
||||
if websocket { "websocket" } else { "http" },
|
||||
account,
|
||||
parent.is_some()
|
||||
);
|
||||
let model = std::env::var("CPR_LIVE_MODEL")
|
||||
.expect("set CPR_LIVE_MODEL to a model available to both accounts");
|
||||
let mut body =
|
||||
json!({"model":model,"instructions":"Reply briefly.","input":input,"store":false});
|
||||
if let Some(parent) = parent {
|
||||
body["client_metadata"] = json!({"parent_response_id":parent});
|
||||
}
|
||||
let payload = ProtocolPayload::json_object("openai", body.as_object().unwrap().clone())
|
||||
.unwrap()
|
||||
.with_context(Map::from_iter([(
|
||||
"use_websocket".to_owned(),
|
||||
json!(websocket),
|
||||
)]));
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request_for_model(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
|
||||
&model,
|
||||
),
|
||||
live_context(account, owner, None, None),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|error| {
|
||||
panic!(
|
||||
"live prepare failed: kind={:?}, status={:?}",
|
||||
error.kind(),
|
||||
error.upstream_status()
|
||||
)
|
||||
});
|
||||
let mut completed = None;
|
||||
let mut actual_transport = None;
|
||||
timeout(Duration::from_secs(90), async {
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.unwrap_or_else(|error| {
|
||||
panic!(
|
||||
"live stream failed: kind={:?}, status={:?}, parent_error={}, model_error={}",
|
||||
error.kind(),
|
||||
error.upstream_status(),
|
||||
error
|
||||
.client_visible_upstream_error()
|
||||
.is_some_and(|detail| detail.message().contains("parent")
|
||||
|| detail.message().contains("Response with id")),
|
||||
error
|
||||
.client_visible_upstream_error()
|
||||
.is_some_and(|detail| detail.message().contains("model")),
|
||||
)
|
||||
});
|
||||
if let Some(observation) = event.response_observation() {
|
||||
actual_transport = Some(observation.transport().as_str().to_owned());
|
||||
}
|
||||
if let Some(wire) = event.wire_event()
|
||||
&& wire.event_type() == Some("response.completed")
|
||||
{
|
||||
let response = &wire.data()["response"];
|
||||
completed = Some((
|
||||
response["id"].as_str().expect("response ID").to_owned(),
|
||||
response["output"].clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("live response deadline");
|
||||
assert_eq!(
|
||||
actual_transport.as_deref(),
|
||||
Some(if websocket { "websocket" } else { "http_sse" })
|
||||
);
|
||||
completed.expect("live request reached response.completed")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires a real account with WebSocket interrupt support and consumes inference quota"]
|
||||
async fn live_websocket_interrupt_and_native_continuation() {
|
||||
let store = live_store().await;
|
||||
let model = std::env::var("CPR_LIVE_MODEL").expect("set CPR_LIVE_MODEL explicitly");
|
||||
let provider =
|
||||
provider_with_base_url_and_retry_budget(&store, OFFICIAL_CODEX_BASE_URL.to_owned(), 0);
|
||||
let account = "acct_scope_old";
|
||||
let session_id = uuid::Uuid::new_v4().to_string();
|
||||
let thread_id = uuid::Uuid::new_v4().to_string();
|
||||
let mut state = Some(
|
||||
ProviderSessionState::new(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("account_id".to_owned(), json!(account)),
|
||||
(
|
||||
"conversation_id".to_owned(),
|
||||
json!(uuid::Uuid::new_v4().to_string()),
|
||||
),
|
||||
("continuation_scope".to_owned(), json!("persisted")),
|
||||
]),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
let mut previous: Option<String> = None;
|
||||
for interrupt in [true, false] {
|
||||
let control = ResponseControl::default();
|
||||
let input = if interrupt {
|
||||
// 输出留出一个网络往返窗口,避免快速完成先于中断到达。
|
||||
"Write 200 numbered lines, with one short word on each line."
|
||||
} else {
|
||||
"Reply with exactly OK."
|
||||
};
|
||||
let payload = ProtocolPayload::json_object("openai", json!({
|
||||
"model":model,"instructions":"Follow the user's request.","input":[{"role":"user","content":input}],
|
||||
"store":false,
|
||||
}).as_object().unwrap().clone()).unwrap().with_context(Map::from_iter([
|
||||
("use_websocket".to_owned(), json!(true)),
|
||||
("session_id".to_owned(), json!(session_id)),
|
||||
("thread_id".to_owned(), json!(thread_id)),
|
||||
]));
|
||||
let request = GenerateRequest::from_protocol_payload(payload)
|
||||
.with_provider_session_state(state.take().expect("session checkpoint"));
|
||||
let continuation = previous.as_ref().map(|id| {
|
||||
ContinuationBinding::Pinned(NativeContinuationPin::new(
|
||||
PreviousResponseId::new(id),
|
||||
PreviousResponseId::new(id),
|
||||
ClientApiKeyId::new("key_openai_contract").unwrap(),
|
||||
ProviderKind::new("openai").unwrap(),
|
||||
ProviderAccountId::new(account).unwrap(),
|
||||
))
|
||||
});
|
||||
let context = live_context(account, Some(account), Some(control.clone()), continuation);
|
||||
let trace = context.trace().clone();
|
||||
assert_eq!(trace.snapshot().unwrap()["wireDumpEnabled"], false);
|
||||
let mut stream = provider
|
||||
.clone()
|
||||
.execute(
|
||||
planned_request_for_model("openai", Operation::Generate(request), &model),
|
||||
context,
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|error| {
|
||||
panic!(
|
||||
"live interrupt prepare failed: {:?} {:?}",
|
||||
error.kind(),
|
||||
error.upstream_status()
|
||||
)
|
||||
});
|
||||
let mut sent = false;
|
||||
let mut terminal = false;
|
||||
let mut interruption_acknowledged = false;
|
||||
let mut output_bytes = 0;
|
||||
let mut actual_transport = None;
|
||||
let started_at = std::time::Instant::now();
|
||||
let result = timeout(Duration::from_secs(90), async {
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.unwrap_or_else(|error| {
|
||||
panic!(
|
||||
"live interrupt stream failed: {:?} {:?}, after_interrupt={}, mentions_session={}, mentions_interrupt={}",
|
||||
error.kind(),
|
||||
error.upstream_status(),
|
||||
sent,
|
||||
error.client_visible_upstream_error().is_some_and(|detail| detail.message().to_ascii_lowercase().contains("session")),
|
||||
error.client_visible_upstream_error().is_some_and(|detail| detail.message().to_ascii_lowercase().contains("interrupt")),
|
||||
)
|
||||
});
|
||||
if let Some(update) = event.session_update() {
|
||||
state = Some(update.clone());
|
||||
}
|
||||
if let Some(observation) = event.response_observation() {
|
||||
let transport = observation.transport().as_str();
|
||||
if actual_transport.as_deref() != Some(transport) {
|
||||
eprintln!("live actual transport={transport}");
|
||||
actual_transport = Some(transport.to_owned());
|
||||
}
|
||||
}
|
||||
if let Some(wire) = event.wire_event() {
|
||||
match wire.event_type() {
|
||||
Some("response.created") if interrupt => {
|
||||
let id = wire
|
||||
.data()
|
||||
.pointer("/response/id")
|
||||
.and_then(Value::as_str)
|
||||
.expect("response ID");
|
||||
sent = control.interrupt(id).is_ok();
|
||||
eprintln!("live interrupt accepted={}, elapsed_ms={}", sent, started_at.elapsed().as_millis());
|
||||
assert!(sent, "current upstream response must support interruption; actual transport={actual_transport:?}");
|
||||
}
|
||||
Some("response.output_text.delta") => {
|
||||
output_bytes += wire.data()["delta"].as_str().map_or(0, str::len);
|
||||
}
|
||||
Some("response.completed" | "response.incomplete") => {
|
||||
let response = &wire.data()["response"];
|
||||
eprintln!("live terminal type={:?}, elapsed_ms={}, output_bytes={}", wire.event_type(), started_at.elapsed().as_millis(), output_bytes);
|
||||
if interrupt {
|
||||
interruption_acknowledged = wire.event_type() == Some("response.incomplete")
|
||||
&& response
|
||||
.pointer("/incomplete_details/reason")
|
||||
.and_then(Value::as_str)
|
||||
== Some("interrupted");
|
||||
} else {
|
||||
assert_eq!(wire.event_type(), Some("response.completed"));
|
||||
}
|
||||
previous =
|
||||
Some(response["id"].as_str().expect("response ID").to_owned());
|
||||
terminal = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
let interrupt_written = trace.snapshot().is_some_and(|snapshot| {
|
||||
snapshot["events"].as_array().is_some_and(|events| {
|
||||
events
|
||||
.iter()
|
||||
.any(|event| event["stage"] == "upstream.interrupt.sent")
|
||||
})
|
||||
});
|
||||
eprintln!("live interrupt written={interrupt_written}, output_bytes={output_bytes}");
|
||||
result.expect("live interrupt response deadline");
|
||||
if interrupt {
|
||||
assert!(
|
||||
interrupt_written,
|
||||
"interrupt was written to the active socket"
|
||||
);
|
||||
assert!(
|
||||
interruption_acknowledged,
|
||||
"upstream acknowledged interruption"
|
||||
);
|
||||
}
|
||||
assert!(terminal && (!interrupt || sent));
|
||||
assert!(state.is_some());
|
||||
eprintln!(
|
||||
"live websocket {}: passed",
|
||||
if interrupt {
|
||||
"interrupt"
|
||||
} else {
|
||||
"native continuation"
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires two real accounts via CPR_LIVE_ACCOUNTS_FILE and consumes minimal inference quota"]
|
||||
async fn live_same_account_parent_and_cross_account_full_replay() {
|
||||
let store = live_store().await;
|
||||
let provider =
|
||||
provider_with_base_url_and_retry_budget(&store, OFFICIAL_CODEX_BASE_URL.to_owned(), 0);
|
||||
for websocket in [false, true] {
|
||||
let first = json!({"role":"user","content":"Reply with exactly OK."});
|
||||
let (parent, output) = live_response(
|
||||
provider.clone(),
|
||||
websocket,
|
||||
"acct_scope_old",
|
||||
None,
|
||||
json!([first]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let mut history = vec![first];
|
||||
history.extend(output.as_array().expect("output items").iter().cloned());
|
||||
history.push(json!({"role":"user","content":"Reply with exactly OK again."}));
|
||||
for account in ["acct_scope_old", "acct_scope_new"] {
|
||||
live_response(
|
||||
provider.clone(),
|
||||
websocket,
|
||||
account,
|
||||
Some("acct_scope_old"),
|
||||
json!(history),
|
||||
Some(&parent),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
eprintln!(
|
||||
"live transport={} initial/same-account-parent/cross-account-full-replay: passed",
|
||||
if websocket { "websocket" } else { "http" }
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,7 @@
|
||||
mod account_isolation;
|
||||
mod capacity;
|
||||
mod live_account_isolation;
|
||||
mod response_interrupt;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::io::{self, Write};
|
||||
@@ -1232,8 +1235,16 @@ fn http_generate_operation() -> Operation {
|
||||
}
|
||||
|
||||
fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest {
|
||||
planned_request_for_model(provider_name, operation, "gpt-5.4")
|
||||
}
|
||||
|
||||
fn planned_request_for_model(
|
||||
provider_name: &str,
|
||||
operation: Operation,
|
||||
model: &str,
|
||||
) -> ProviderRequest {
|
||||
let provider = ProviderKind::new(provider_name).expect("provider");
|
||||
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
|
||||
let upstream_model = UpstreamModelId::new(model).expect("upstream model");
|
||||
let public_model = PublicModelId::new(upstream_model.as_str()).expect("public model");
|
||||
let account_scope = Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
use super::*;
|
||||
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
|
||||
|
||||
async fn next_json(socket: &mut tokio_tungstenite::WebSocketStream<TcpStream>) -> Value {
|
||||
loop {
|
||||
match socket.next().await.unwrap().unwrap() {
|
||||
Message::Text(text) => return serde_json::from_str(&text).unwrap(),
|
||||
Message::Ping(payload) => socket.send(Message::Pong(payload)).await.unwrap(),
|
||||
Message::Pong(_) => {}
|
||||
other => panic!("unexpected synthetic upstream frame: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn interrupt_context(control: ResponseControl, previous: Option<&str>) -> AttemptContext {
|
||||
let account = ProviderAccountId::new("acct_provider_contract").unwrap();
|
||||
let provider = ProviderKind::new("openai").unwrap();
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new("req_interrupt").unwrap(),
|
||||
ClientApiKeyId::new("key_openai_contract").unwrap(),
|
||||
)
|
||||
.with_response_control(Some(control)),
|
||||
NonZeroU32::new(1).unwrap(),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(
|
||||
BTreeSet::new(),
|
||||
None,
|
||||
previous.map(|_| ProviderAccountStateOwner::new(provider.clone(), account.clone())),
|
||||
)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
previous.map(|id| {
|
||||
ContinuationBinding::Pinned(NativeContinuationPin::new(
|
||||
PreviousResponseId::new(id),
|
||||
PreviousResponseId::new(id),
|
||||
ClientApiKeyId::new("key_openai_contract").unwrap(),
|
||||
provider,
|
||||
account,
|
||||
))
|
||||
}),
|
||||
CancellationToken::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_interrupt_uses_the_active_socket_and_releases_control_before_reuse() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut socket = accept_codex_test_websocket(socket).await;
|
||||
for id in ["resp_interrupt_a", "resp_interrupt_b"] {
|
||||
let create = next_json(&mut socket).await;
|
||||
assert_eq!(create["type"], "response.create");
|
||||
if id == "resp_interrupt_b" {
|
||||
assert_eq!(create["previous_response_id"], "resp_interrupt_a");
|
||||
}
|
||||
socket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {"id":id, "model":"gpt-5.4", "status":"in_progress"}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
socket.send(Message::Text(json!({"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"partial"}).to_string().into())).await.unwrap();
|
||||
let interrupt = next_json(&mut socket).await;
|
||||
assert_eq!(
|
||||
interrupt,
|
||||
json!({"type":"response.interrupt", "response_id":id,"mode":"discard_partial_items"})
|
||||
);
|
||||
socket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type":"response.incomplete",
|
||||
"response":{"id":id,"model":"gpt-5.4","status":"incomplete","incomplete_details":{"reason":"interrupted"},"output":[]}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
let provider = provider_with_base_url(&store, base_url);
|
||||
let mut session_state = None;
|
||||
for id in ["resp_interrupt_a", "resp_interrupt_b"] {
|
||||
let control = ResponseControl::default();
|
||||
let previous = (id == "resp_interrupt_b").then_some("resp_interrupt_a");
|
||||
let mut request = generate_with_persisted_session_context(
|
||||
"acct_provider_contract",
|
||||
"interrupt-conversation",
|
||||
"interrupt-session",
|
||||
"interrupt-thread",
|
||||
);
|
||||
if let Some(state) = session_state.take() {
|
||||
request = request.with_provider_session_state(state);
|
||||
}
|
||||
let mut stream = provider
|
||||
.clone()
|
||||
.execute(
|
||||
planned_request("openai", Operation::Generate(request)),
|
||||
interrupt_context(control.clone(), previous),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut interrupted = false;
|
||||
timeout(Duration::from_secs(5), async {
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.unwrap();
|
||||
if let Some(state) = event.session_update() {
|
||||
session_state = Some(state.clone());
|
||||
}
|
||||
if !interrupted && control.interrupt(id).is_ok() {
|
||||
assert_eq!(
|
||||
control.interrupt("resp_foreign"),
|
||||
Err(ResponseInterruptError::ResponseMismatch)
|
||||
);
|
||||
control.interrupt(id).unwrap();
|
||||
interrupted = true;
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("interrupt completes the active response");
|
||||
assert!(interrupted);
|
||||
assert!(session_state.is_some());
|
||||
assert_eq!(
|
||||
control.interrupt(id),
|
||||
Err(ResponseInterruptError::Unavailable)
|
||||
);
|
||||
}
|
||||
timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
@@ -262,6 +262,18 @@ fn websocket_completed_id_should_read_the_id_without_validating_the_response_sha
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_interrupted_id_is_reusable_but_other_incomplete_reasons_are_not() {
|
||||
for (reason, expected) in [
|
||||
("interrupted", Some("resp_interrupted".to_owned())),
|
||||
("max_output_tokens", None),
|
||||
("content_filter", None),
|
||||
] {
|
||||
let event = json!({"type":"response.incomplete","response":{"id":"resp_interrupted","incomplete_details":{"reason":reason}}});
|
||||
assert_eq!(websocket_response_completed_id(&event), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_completed_id_should_preserve_an_empty_opaque_id() {
|
||||
let event = json!({
|
||||
|
||||
+13
-1
@@ -217,6 +217,8 @@ API Key 与 OAuth 共用模拟客户端画像(`User-Agent`、`originator`、`v
|
||||
包括会话、线程、Lite 和其他业务扩展头。
|
||||
上游认证只来自选中的账号;API Key 不携带 OAuth Cookie、ChatGPT 账号身份或下游的
|
||||
`X-OpenAI-Actor-Authorization` 托管认证声明。
|
||||
下游的 `x-openai-account-routing-override`、`x-openai-fedramp` 也不透传,
|
||||
工作区路由与合规属性不能从原账号继承;请求中间件不能重新注入这些托管身份头。
|
||||
|
||||
Responses 也不透传 `x-stainless-*`、`Origin`、`Referer`、`sec-ch-ua*` 和 `sec-fetch-*`
|
||||
携带的下游 SDK/浏览器环境或页面来源。过滤规则适用于所有下游客户端,与 User-Agent 无关;
|
||||
@@ -250,10 +252,20 @@ Codex/OAuth 上游的历史回填按字段形状兼容,不以 User-Agent 品
|
||||
|
||||
请求头过滤不提供客户端匿名化;系统提示词、工具定义、工具结果、工作目录及其他业务 metadata
|
||||
保持原有语义,可能包含客户端环境信息。
|
||||
`client_metadata.parent_response_id` 是 Guardian 的账号内响应引用,只有归属可信且仍为同一账号时保留;
|
||||
切号或归属未知时移除。`x-codex-guardian`、`guardian_credits_requested` 和序列化
|
||||
`x-codex-turn-metadata` 内普通扩展的同名 `parent_response_id` 保持原样。
|
||||
|
||||
Responses WebSocket 仅接受文本 `response.create`,同一连接串行执行。当前响应期间收到的后续业务帧
|
||||
Responses WebSocket 接受文本 `response.create` 和 `response.interrupt`,创建请求在同一连接串行执行。当前响应期间收到的后续业务帧
|
||||
留在有界接收队列中,待当前响应完成终结和写出后再逐条校验、准入与执行,不因请求提前到达而断开。
|
||||
接收队列容量为 32 个事件,超载仍关闭连接;Ping/Pong、客户端关闭和服务关闭不等待队列中的请求执行。
|
||||
活动响应期间会即时处理 `{"type":"response.interrupt","response_id":"当前响应 ID","mode":"discard_partial_items"}`。
|
||||
中断只能发送到该执行占用的原上游 WS,不重新选号或创建推理 attempt;重复中断合并为一次发送。
|
||||
ID 不匹配、没有可中断响应、实际走 HTTP 或 Provider 不支持控制时返回 `400` 协议错误,原执行继续;
|
||||
客户端需要终止这类执行时可关闭连接,后续按既有续接合同恢复。
|
||||
中断后继续转发上游事件与终态;只有 `response.incomplete` 的 `incomplete_details.reason` 为
|
||||
`interrupted` 时,才按官方中断语义保留原连接续接能力。部分输出是否被丢弃以上游终态为准。
|
||||
控制帧发送成功不等于上游已确认中断;若上游仍返回 `response.completed`,按正常完成处理。
|
||||
|
||||
OAuth 账号默认 `prefer_websocket`,客户端使用 HTTP/SSE 时仍可能选择上游 WebSocket;
|
||||
可将账号上游传输方式设为 `http`,固定使用 HTTP/SSE。API Key 账号默认使用 HTTP/SSE,
|
||||
|
||||
@@ -527,6 +527,10 @@ Admin 的手动重置复用同一账本与 Key 行锁,在一个事务中清零
|
||||
- SSE 持有并发名额直到终态;WebSocket 每个 `response.create` 独立准入,空闲连接不占名额。
|
||||
Core 在每次请求开始时重新鉴权并冻结当前策略,既有连接也应用已发布的限额与授权范围变更。
|
||||
换号和内部重试复用同一名额,Key 与账号并发上限同时生效;预算拒绝或启动失败释放名额。
|
||||
- 响应中断入口由 Core 的根执行持有,API 只控制当前连接正在交付的执行。Provider 在收到上游
|
||||
`response.created` 后注册活动响应 owner,在终态、失败或断连时释放;控制句柄本身不延长 owner
|
||||
生命周期。中断信号交给占用原账号、原 socket 的 exchange,不经过全局响应查询、新选号或独立 attempt。
|
||||
API 暂存的下一轮请求和连接泵接收队列共用容量名额,控制帧处理不取消中间件正在读取响应体的 future。
|
||||
|
||||
完成、失败、取消和断连使用同一结算路径,在结束时以网关请求 ID 幂等累计已取得费用。
|
||||
执行会话持有结算与准入释放的同一个清理 future;协议等待被取消后,后续驱动或宿主 detached finalizer
|
||||
|
||||
Reference in New Issue
Block a user