diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs index 685418a2f..48542658a 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs @@ -12,7 +12,7 @@ use gateway_core::engine::middleware::{ MiddlewareBody, MiddlewareError, MiddlewareFrame, MiddlewareFraming, MiddlewareHeader, MiddlewareResponse, }; -use gateway_core::engine::response_control::ResponseControl; +use gateway_core::engine::response_control::{ResponseControl, ResponseControlUnavailable}; use gateway_core::error::{GatewayError, GatewayErrorKind}; use gateway_core::event::ProviderResponseHeader; use gateway_core::operation::ProviderSessionState; @@ -25,7 +25,7 @@ use super::{ super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody}, connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext}, protocol::{ - decode_response_interrupt, error_event, initial_engine_error_event, response_metadata_event, + error_event, initial_engine_error_event, is_response_create, response_metadata_event, }, }; @@ -137,7 +137,7 @@ pub(super) async fn forward_response( replay: &mut ConnectionReplaySnapshot, capture: ReplayCaptureHandle, validation: ResponseValidationFacts, - response_control: ResponseControl, + response_control: &ResponseControl, ) -> ForwardOutcome { let (protocol, status, headers, mut body, _) = response.into_parts(); if protocol != "openai" || StatusCode::from_u16(status).is_err() { @@ -150,8 +150,14 @@ pub(super) async fn forward_response( ProviderResponseHeader::new(name, value) }) .collect::>(); - let first = match next_body_input(connection, body.as_mut(), &response_control, &request_id) - .await + let first = match next_body_input( + connection, + body.as_mut(), + response_control, + &request_id, + false, + ) + .await { BodyInput::Frame(Ok(Some(frame))) => frame, BodyInput::Frame(Ok(None)) => { @@ -197,8 +203,14 @@ pub(super) async fn forward_response( let frame = match current.take() { Some(frame) => frame, None => { - match next_body_input(connection, body.as_mut(), &response_control, &request_id) - .await + match next_body_input( + connection, + body.as_mut(), + response_control, + &request_id, + terminal_seen, + ) + .await { BodyInput::Frame(Ok(Some(frame))) => frame, BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => { @@ -302,6 +314,7 @@ async fn next_body_input( body: &mut dyn MiddlewareBody, response_control: &ResponseControl, request_id: &Arc, + terminal_seen: bool, ) -> BodyInput { // 处理控制帧时继续持有同一个读取 future，不能取消正在加工正文的中间件 let frame = body.next_frame(); @@ -312,19 +325,11 @@ async fn next_body_input( 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 - { + if is_response_create(payload) { + connection.defer(event); + continue; + } + if send_control(connection, response_control, payload, request_id).await == ForwardOutcome::Disconnect { return BodyInput::Disconnect; } } @@ -337,6 +342,11 @@ async fn next_body_input( ConnectionEvent::Exited(_) => return BodyInput::Disconnect, } } + event = response_control.receive(), if terminal_seen => { + if forward_control_event(connection, response_control, event).await == ForwardOutcome::Disconnect { + return BodyInput::Disconnect; + } + } frame = &mut frame => return BodyInput::Frame(frame), } } @@ -720,3 +730,46 @@ async fn send_metadata( .await .is_ok() } + +pub(super) async fn send_control( + connection: &mut ResponsesWebSocketConnection, + control: &ResponseControl, + payload: &str, + request_id: &Arc, +) -> ForwardOutcome { + let result = tokio::select! { + _ = connection.wait_for_exit() => return ForwardOutcome::Disconnect, + result = control.send(payload) => result, + }; + if result.is_err() { + return send_gateway_error( + connection, + &GatewayError::new( + GatewayErrorKind::InvalidRequest, + "No upstream WebSocket connection is available for control messages", + ), + request_id, + ) + .await; + } + ForwardOutcome::Continue +} + +pub(super) async fn forward_control_event( + connection: &mut ResponsesWebSocketConnection, + control: &ResponseControl, + event: Result, +) -> ForwardOutcome { + let Ok(payload) = event else { + control.clear(); + return ForwardOutcome::Continue; + }; + if connection + .send_text(payload, WriteContext::connection(FramePhase::Data)) + .await + .is_err() + { + return ForwardOutcome::Disconnect; + } + ForwardOutcome::Continue +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs index 342479b25..88164269b 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs @@ -24,6 +24,7 @@ use gateway_core::{ engine::{ execution::{AuthenticatedClient, ClientTransport}, middleware::{FrozenMiddlewarePlan, MiddlewareError}, + response_control::ResponseControl, }, lifecycle::{ConnectionGuard, ConnectionLifecycle}, operation::OperationKind, @@ -46,11 +47,12 @@ use super::{ }; use connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext}; use forward::{ - ConnectionReplaySnapshot, ForwardOutcome, execution_response, forward_response, - new_replay_capture, send_gateway_error, send_middleware_error, send_protocol_error, + ConnectionReplaySnapshot, ForwardOutcome, execution_response, forward_control_event, + forward_response, new_replay_capture, send_control, send_gateway_error, send_middleware_error, + send_protocol_error, }; -use protocol::connection_limit_event; pub use protocol::{ResponseCreateFrameError, decode_response_create_with_context}; +use protocol::{connection_limit_event, is_response_create}; const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only"; const CONNECTION_LIMIT_CLOSE_REASON: &str = "Responses websocket connection limit reached"; @@ -191,10 +193,20 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke ); let mut request_count = 0_u64; let mut replay = ConnectionReplaySnapshot::default(); + let mut response_control = ResponseControl::default(); loop { - let Some(event) = connection.next_event().await else { - break; + let event = tokio::select! { + event = connection.next_event() => { + let Some(event) = event else { break }; + event + } + event = response_control.receive() => { + if forward_control_event(&mut connection, &response_control, event).await == ForwardOutcome::Disconnect { + break; + } + continue; + } }; let payload = match event { ConnectionEvent::Text(payload) => payload, @@ -212,8 +224,22 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke expire_connection(&mut connection).await; break; } - request_count = request_count.saturating_add(1); let correlation_id = Arc::::from(service.next_request_id()); + if !is_response_create(&payload) { + if send_control( + &mut connection, + &response_control, + &payload, + &correlation_id, + ) + .await + == ForwardOutcome::Disconnect + { + break; + } + continue; + } + request_count = request_count.saturating_add(1); // 初步解码只保留路由事实，避免整份正文跨越准入等待和响应交付 let model_hint = match decode_response_create_with_context(&payload, &request_headers) { Ok(decoded) => Some(decoded.metadata().requested_model().to_owned()), @@ -251,6 +277,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke expire_connection(&mut connection).await; break; } + response_control.clear(); let execution = service.execution(); let preparation = async { // 握手的显式改写可继承，Key 策略与宿主设置仍在每轮执行前刷新 @@ -291,7 +318,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke } }; let request_id = Arc::::from(prepared.request_id().to_string()); - let response_control = prepared.response_control(); + response_control = prepared.response_control(); let capture = new_replay_capture(); let validation = ResponseValidationFacts::default(); let input = RequestInput { @@ -381,7 +408,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke &mut replay, capture, validation, - response_control, + &response_control, ) .await == ForwardOutcome::Disconnect diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs b/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs index edf08bbe0..1ca95c858 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs @@ -31,42 +31,22 @@ pub fn decode_response_create_with_context( decode_response_create_inner(payload, request_headers) } -/// 中断只引用活动响应；普通创建帧仍交给原有解码与串行准入路径 -pub(super) fn decode_response_interrupt( - payload: &str, -) -> Result, ResponseCreateFrameError> { +/// 只识别需要本地准入的创建帧，其余文本不在网关解释 +pub(super) fn is_response_create(payload: &str) -> bool { #[derive(serde::Deserialize)] struct FrameType { #[serde(rename = "type")] - message_type: Option, + message_type: String, } - // 非控制帧仍在原来的串行请求边界校验，不能抢先拒绝排队的普通请求 // 只投影类型，避免为排队的大型 response.create 再构造完整 JSON 树 - let Ok(frame) = serde_json::from_str::(payload) else { - return Ok(None); - }; - if frame.message_type.as_deref() != Some("response.interrupt") { - return Ok(None); + match serde_json::from_str::(payload) { + Ok(frame) => frame.message_type == "response.create", + // struct 投影拒绝重复 type；沿用创建解码器的 JSON object 语义， + // 避免这类创建请求绕过准入进入透传路径 + Err(_) => serde_json::from_str::(payload).is_ok_and(|frame| { + frame.get("type").and_then(Value::as_str) == Some("response.create") + }), } - 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( diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs index f6c069a2c..c295ec644 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs @@ -106,10 +106,18 @@ async fn next_error(socket: &mut TestSocket) -> Value { } async fn assert_no_duplicate_failure(socket: &mut TestSocket) { - // 后续入站帧只会在本轮收敛后处理，比固定等待窗口更可靠地检查没有重复失败终态 - socket.send(Message::Text("{".into())).await.unwrap(); + // 无效创建帧仍走串行队列，以本地业务校验作为本轮收敛后的屏障 + socket + .send(Message::Text( + json!({"type":"response.create","model":"model-a","stream":false}) + .to_string() + .into(), + )) + .await + .unwrap(); let error = next_event(socket).await; - assert_eq!(error["error"]["code"], "invalid_json"); + assert_eq!(error["error"]["code"], "invalid_value"); + assert_eq!(error["error"]["param"], "stream"); assert_eq!(error["status"], 400); assert_eq!(error["headers"]["x-request-id"], error["request_id"]); assert_eq!( diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/interrupt.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/interrupt.rs index 03ac069d5..ed8f47391 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/interrupt.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/interrupt.rs @@ -1,12 +1,69 @@ -//! 验证 Responses WebSocket 中断向当前执行传递并释放相关资源 +//! 验证未知控制消息、迟到中断、结算期间收发和创建请求的串行准入 use super::*; -use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError}; +use gateway_core::engine::response_control::{ + ResponseControlTransport, ResponseControlUnavailable, +}; + +struct ControlTransport { + id: String, + active: std::sync::atomic::AtomicBool, + requests: Mutex>, + sender: tokio::sync::mpsc::Sender, + receiver: tokio::sync::Mutex>, +} + +impl ControlTransport { + fn new(id: String) -> Self { + let (sender, receiver) = tokio::sync::mpsc::channel(32); + Self { + id, + active: std::sync::atomic::AtomicBool::new(true), + requests: Mutex::default(), + sender, + receiver: tokio::sync::Mutex::new(receiver), + } + } +} + +#[async_trait] +impl ResponseControlTransport for ControlTransport { + async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> { + self.requests.lock().unwrap().push(payload.to_owned()); + let original: Value = + serde_json::from_str(payload).unwrap_or_else(|_| Value::String(payload.to_owned())); + let event = if original["type"] == "future.reject" { + upstream_control_error() + } else if original["type"] == "response.interrupt" + && original["response_id"] == self.id + && original["mode"] == "discard_partial_items" + && self.active.load(Ordering::SeqCst) + { + json!({"type":"response.incomplete", "response":{"id":self.id,"model":"model-a","status":"incomplete","output":[],"incomplete_details":{"reason":"interrupted"}}}) + } else { + json!({"type":"future.control.ack", "original":original}) + }; + self.sender + .send(event) + .await + .map_err(|_| ResponseControlUnavailable) + } + async fn receive(&self) -> Result { + self.receiver + .lock() + .await + .recv() + .await + .map(|event| event.to_string()) + .ok_or(ResponseControlUnavailable) + } +} #[derive(Default)] struct InterruptProvider { calls: AtomicUsize, - controls: Mutex>, + transport: Mutex>>, + complete: tokio::sync::Notify, } #[async_trait] @@ -30,12 +87,10 @@ impl Provider for InterruptProvider { ) -> Result { 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 transport = Arc::new(ControlTransport::new(id)); + let owner: Arc = transport.clone(); + context.response_control().unwrap().bind(&owner); + *self.transport.lock().unwrap() = Some(transport.clone()); let candidate = request.candidate(); let metadata = ProviderCallMetadata::new( candidate.provider().clone(), @@ -44,41 +99,47 @@ impl Provider for InterruptProvider { 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, transport, self), + |(phase, transport, provider)| async move { + let (wire, events, next_phase) = match phase { 0 => ( - "response.created", - "in_progress", - GatewayEvent::Started(ResponseMeta::new(&id, "model-a")), + json!({"type":"response.created","response":{"id":transport.id,"model":"model-a","status":"in_progress","output":[]}}), + vec![GatewayEvent::Started(ResponseMeta::new( + &transport.id, + "model-a", + ))], + 1, ), 1 => { - active.as_ref().unwrap().requested().await; - drop(active.take()); - ( - "response.incomplete", - "incomplete", - GatewayEvent::Completed(ResponseMeta::new(&id, "model-a")), - ) + let wire = tokio::select! { + reply = transport.receive() => serde_json::from_str::(&reply.unwrap()).unwrap(), + () = provider.complete.notified() => json!({"type":"response.completed","response":{"id":transport.id,"model":"model-a","status":"completed","output":[]}}), + }; + if matches!( + wire["type"].as_str(), + Some("response.completed" | "response.incomplete") + ) { + transport.active.store(false, Ordering::SeqCst); + ( + wire, + vec![GatewayEvent::Completed(ResponseMeta::new( + &transport.id, + "model-a", + ))], + 2, + ) + } else { + (wire, Vec::new(), 1) + } } _ => 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_type = wire["type"].as_str().unwrap().to_owned(); let event = ProviderEvent::canonical_with_wire( - vec![event], - ProtocolWireEvent::json( - "openai", - Some(event_type.to_owned()), - json!({ - "type":event_type, "response":response - }), - ) - .unwrap(), + events, + ProtocolWireEvent::json("openai", Some(event_type), wire).unwrap(), ); - Some((Ok(event), (phase + 1, active, id))) + Some((Ok(event), (next_phase, transport, provider))) }, ); Ok(ProviderStream::new(metadata, body, ())) @@ -88,6 +149,8 @@ impl Provider for InterruptProvider { #[derive(Default)] struct InterruptAdmissions { active: AtomicUsize, + release_started: tokio::sync::Notify, + release_gate: Mutex>>, } impl ClientAdmissionPort for InterruptAdmissions { @@ -109,6 +172,11 @@ impl ClientAdmissionPort for InterruptAdmissions { _: &'a ModelRequestId, ) -> BoxFuture<'a, Result> { Box::pin(async { + let gate = self.release_gate.lock().unwrap().take(); + self.release_started.notify_one(); + if let Some(gate) = gate { + gate.await.unwrap(); + } self.active.fetch_sub(1, Ordering::SeqCst); Ok(true) }) @@ -151,15 +219,15 @@ async fn next(socket: &mut TestSocket) -> 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()); +async fn connect( + provider: Arc, + admissions: Arc, +) -> (TestSocket, Server) { let execution = Arc::new(DefaultExecutionService::new( RuntimeSnapshotHandle::new(snapshot("sk_ws_interrupt", "openai")), Arc::new(SettlementPorts::default()), - ProviderRegistry::new([provider.clone() as Arc]).unwrap(), - admissions.clone(), + ProviderRegistry::new([provider as Arc]).unwrap(), + admissions, Arc::new(UnusedContinuation), Arc::new(IgnoredClientApiKeyUsage), Arc::new(crate::support::RecordingDiagnostics::default()), @@ -167,7 +235,7 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() { 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 { + let server = Server(tokio::spawn(async move { axum::serve(listener, app).await.unwrap() })); let mut request = format!("ws://{address}/v1/responses") @@ -177,29 +245,136 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() { AUTHORIZATION, HeaderValue::from_static("Bearer sk_ws_interrupt"), ); - let (mut socket, _) = connect_async(request).await.unwrap(); + let (socket, _) = connect_async(request).await.unwrap(); + (socket, server) +} + +fn upstream_control_error() -> Value { + json!({"type":"error", "status":400, "error":{"type":"invalid_request_error","code":"future_rejected","message":"Synthetic upstream rejection","param":"type"},"vendor_extension":true}) +} + +fn interrupt(id: &str, mode: &str) -> Value { + json!({"type":"response.interrupt","response_id":id,"mode":mode,"future_extension":true}) +} + +#[tokio::test] +async fn late_interrupt_after_terminal_is_sent_upstream_and_keeps_the_connection_reusable() { + let provider = Arc::new(InterruptProvider::default()); + let admissions = Arc::new(InterruptAdmissions::default()); + let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await; + let create = json!({"type":"response.create","model":"model-a","input":"hello"}); + for (number, terminal) in [(0, "response.completed"), (1, "response.incomplete")] { + let id = format!("resp_interrupt_{number}"); + send(&mut socket, create.clone()).await; + assert_eq!(next(&mut socket).await["type"], "response.created"); + if number == 0 { + provider.complete.notify_one(); + } else { + send(&mut socket, interrupt(&id, "discard_partial_items")).await; + } + assert_eq!(next(&mut socket).await["type"], terminal); + tokio::time::timeout(Duration::from_secs(5), async { + while admissions.active.load(Ordering::SeqCst) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + for control in [ + interrupt(&id, "discard_partial_items"), + interrupt(&id, "future_mode"), + json!({"type":"future.control","extra":[1,2]}), + ] { + send(&mut socket, control.clone()).await; + assert_eq!( + next(&mut socket).await, + json!({"type":"future.control.ack","original":control}) + ); + } + send(&mut socket, json!({"type":"future.reject"})).await; + assert_eq!(next(&mut socket).await, upstream_control_error()); + assert_eq!(provider.calls.load(Ordering::SeqCst), number + 1); + assert_eq!(admissions.active.load(Ordering::SeqCst), 0); + } + socket.close(None).await.unwrap(); +} + +#[tokio::test] +async fn late_interrupt_during_terminal_settlement_does_not_abort_or_reorder_requests() { + let provider = Arc::new(InterruptProvider::default()); + let (release, wait_for_release) = tokio::sync::oneshot::channel(); + let admissions = Arc::new(InterruptAdmissions { + release_gate: Mutex::new(Some(wait_for_release)), + ..Default::default() + }); + let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await; 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"); + provider.complete.notify_one(); + assert_eq!(next(&mut socket).await["type"], "response.completed"); + tokio::time::timeout( + Duration::from_secs(5), + admissions.release_started.notified(), + ) + .await + .unwrap(); send(&mut socket, create).await; - for invalid in [ - interrupt("resp_other", "discard_partial_items"), - interrupt("resp_interrupt_0", "unknown"), - interrupt("", "discard_partial_items"), + let late = interrupt("resp_interrupt_0", "discard_partial_items"); + send(&mut socket, late.clone()).await; + assert_eq!( + next(&mut socket).await, + json!({"type":"future.control.ack", "original":late}) + ); + assert_eq!(provider.calls.load(Ordering::SeqCst), 1); + assert_eq!(admissions.active.load(Ordering::SeqCst), 1); + release.send(()).unwrap(); + assert_eq!( + next(&mut socket).await["response"]["id"], + "resp_interrupt_1" + ); + socket.close(None).await.unwrap(); +} + +#[tokio::test] +async fn unknown_controls_preserve_bytes_and_do_not_admit_queued_creates() { + let provider = Arc::new(InterruptProvider::default()); + let admissions = Arc::new(InterruptAdmissions::default()); + let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await; + let create = json!({"type":"response.create","model":"model-a","input":"hello"}); + send(&mut socket, json!({"type":"future.control"})).await; + let error = next(&mut socket).await; + assert_eq!(error["type"], "error"); + assert_ne!(error["error"]["param"], "type"); + assert_eq!(provider.calls.load(Ordering::SeqCst), 0); + send(&mut socket, create).await; + assert_eq!(next(&mut socket).await["type"], "response.created"); + // 重复 type 仍是创建请求，不能借投影解码失败绕过串行准入 + socket.send(ClientMessage::Text(r#"{"type":"response.create","type":"response.create","model":"model-a","input":"next"}"#.into())).await.unwrap(); + for raw in [ + r#"{ "type": "future.control", "mode": "unknown", "extension": [1, 2] }"#, + r#"{"extension":"no local type contract"}"#, + r#"{"type":42,"extension":true}"#, + "not-json", ] { - send(&mut socket, invalid).await; - let error = next(&mut socket).await; - assert_eq!(error["type"], "error"); - assert_eq!(error["status"], 400); + socket.send(ClientMessage::Text(raw.into())).await.unwrap(); + let expected = + serde_json::from_str::(raw).unwrap_or_else(|_| Value::String(raw.to_owned())); + assert_eq!(next(&mut socket).await["original"], expected); + assert_eq!( + provider + .transport + .lock() + .unwrap() + .as_ref() + .unwrap() + .requests + .lock() + .unwrap() + .last() + .unwrap(), + raw + ); assert_eq!(provider.calls.load(Ordering::SeqCst), 1); } send( @@ -207,27 +382,12 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() { 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["type"], "response.incomplete"); 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 { @@ -236,8 +396,4 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() { }) .await .expect("disconnect releases active execution"); - assert_eq!( - provider.controls.lock().unwrap()[1].interrupt("resp_interrupt_1"), - Err(ResponseInterruptError::Unavailable) - ); } diff --git a/backend/crates/gateway-core/src/engine/response_control.rs b/backend/crates/gateway-core/src/engine/response_control.rs index 14815d82a..3b04cb08e 100644 --- a/backend/crates/gateway-core/src/engine/response_control.rs +++ b/backend/crates/gateway-core/src/engine/response_control.rs @@ -1,14 +1,27 @@ -//! 当前执行的响应中断信号；只由持有活动响应的 Provider 注册，不查找其他执行或连接 +//! 原上游连接的控制消息端口；不解释协议类型、不创建执行或延长连接生命周期 use std::fmt; use std::sync::{Arc, Mutex, Weak}; -use crate::lifecycle::CancellationToken; +use async_trait::async_trait; -/// 单次根执行的控制入口；重试不继承已结束响应的控制目标 +/// Provider 实现的连接控制通道，消息内容由上游协议解释 +#[async_trait] +pub trait ResponseControlTransport: Send + Sync { + async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable>; + + /// 仅在响应正文已结束后读取，调用方必须先取消此读取再开始下一轮执行 + async fn receive(&self) -> Result; +} + +/// 控制消息没有可用的原上游连接，不能通过重新选号或重试恢复 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResponseControlUnavailable; + +/// 单次根执行绑定的原连接；终态后仍可使用，连接销毁或重新绑定时失效 #[derive(Clone, Default)] pub struct ResponseControl { - active: Arc>>, + target: Arc>>>, } impl fmt::Debug for ResponseControl { @@ -17,64 +30,42 @@ impl fmt::Debug for ResponseControl { } } -struct InterruptTarget { - response_id: String, - requested: CancellationToken, -} - -/// Provider 持有的活动响应控制权；释放后客户端无法再向该响应发出中断 -pub struct ActiveResponseInterrupt(Arc); - -/// 控制帧不能作用于当前执行时的稳定原因，不携带其他响应身份 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ResponseInterruptError { - Unavailable, - ResponseMismatch, -} - impl ResponseControl { - /// 在上游确认响应 ID 后注册；同一执行只允许一个存活的响应 owner - pub fn activate(&self, response_id: String) -> Option { - let mut active = self - .active + /// Provider 在请求成功写入上游后绑定；连接 owner 持有强引用并负责撤销 + pub fn bind(&self, transport: &Arc) { + *self + .target .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)) + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::downgrade(transport)); } - /// 重复中断合并为一次信号；ID 必须匹配本执行当前的上游响应 - pub fn interrupt(&self, response_id: &str) -> Result<(), ResponseInterruptError> { - let active = self - .active + pub fn clear(&self) { + *self + .target .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); + .unwrap_or_else(std::sync::PoisonError::into_inner) = None; + } + + fn transport(&self) -> Option> { + self.target + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .and_then(Weak::upgrade) + } + + pub async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> { + self.transport() + .ok_or(ResponseControlUnavailable)? + .send(payload) + .await + } + + /// 未绑定时保持等待，供连接空闲循环与客户端输入共同 select + pub async fn receive(&self) -> Result { + match self.transport() { + Some(transport) => transport.receive().await, + None => std::future::pending().await, } - target.requested.cancel(); - Ok(()) - } -} - -impl ActiveResponseInterrupt { - pub fn response_id(&self) -> &str { - &self.0.response_id - } - - pub fn requested(&self) -> impl Future + Send + 'static { - // 等待者只持有信号，不能延长活动响应 owner 的生命周期 - let requested = self.0.requested.clone(); - async move { requested.cancelled().await } } } diff --git a/backend/crates/gateway-core/tests/engine/response_control.rs b/backend/crates/gateway-core/tests/engine/response_control.rs index 4dc8c352b..d4b3cc250 100644 --- a/backend/crates/gateway-core/tests/engine/response_control.rs +++ b/backend/crates/gateway-core/tests/engine/response_control.rs @@ -1,45 +1,68 @@ -//! 验证响应中断的当前 owner 校验及独立执行间的控制隔离 +//! 验证原连接控制端口的透明传递、owner 生命周期与执行隔离 -use futures::FutureExt; -use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError}; +use std::sync::{Arc, Mutex}; -#[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); +use async_trait::async_trait; +use futures::{FutureExt, executor::block_on}; +use gateway_core::engine::response_control::{ + ResponseControl, ResponseControlTransport, ResponseControlUnavailable, +}; + +#[derive(Default)] +struct Transport(Mutex>); + +#[async_trait] +impl ResponseControlTransport for Transport { + async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> { + self.0.lock().unwrap().push(payload.to_owned()); + Ok(()) + } + async fn receive(&self) -> Result { + Ok("upstream reply".to_owned()) + } } #[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()); +fn control_requires_a_live_connection_owner_and_preserves_unknown_payloads() { + block_on(async { + let control = ResponseControl::default(); + assert_eq!( + control.send("unknown").await, + Err(ResponseControlUnavailable) + ); + assert!(control.receive().now_or_never().is_none()); + let transport = Arc::new(Transport::default()); + let owner: Arc = transport.clone(); + control.bind(&owner); + let payload = "{ \"type\": \"future.control\", \"extension\": [1, 2] }"; + control.send(payload).await.unwrap(); + control.send(payload).await.unwrap(); + assert_eq!(*transport.0.lock().unwrap(), [payload, payload]); + assert_eq!(control.receive().await.unwrap(), "upstream reply"); + drop(owner); + drop(transport); + assert_eq!(control.send(payload).await, Err(ResponseControlUnavailable)); + }); +} + +#[test] +fn replacing_or_clearing_a_binding_does_not_share_control_between_executions() { + block_on(async { + let first = ResponseControl::default(); + let second = ResponseControl::default(); + let old: Arc = Arc::new(Transport::default()); + let current = Arc::new(Transport::default()); + let owner: Arc = current.clone(); + first.bind(&old); + second.bind(&owner); + drop(old); + assert_eq!(first.send("old").await, Err(ResponseControlUnavailable)); + second.send("current").await.unwrap(); + assert_eq!(*current.0.lock().unwrap(), ["current"]); + second.clear(); + assert_eq!( + second.send("cleared").await, + Err(ResponseControlUnavailable) + ); + }); } diff --git a/backend/crates/providers/openai/src/transport/websocket/control.rs b/backend/crates/providers/openai/src/transport/websocket/control.rs new file mode 100644 index 000000000..94f492568 --- /dev/null +++ b/backend/crates/providers/openai/src/transport/websocket/control.rs @@ -0,0 +1,60 @@ +//! WebSocket 有界消息通道与原连接文本透传；空闲读取与响应正文消费互斥 + +use std::{sync::Arc, time::Duration}; + +use async_trait::async_trait; +use gateway_core::engine::response_control::{ + ResponseControlTransport, ResponseControlUnavailable, +}; +use tokio::sync::{Mutex, mpsc, oneshot}; +use tokio_tungstenite::tungstenite::{self, Message}; + +pub(super) const WEBSOCKET_SEND_TIMEOUT: Duration = Duration::from_secs(5 * 60); + +pub(super) enum PumpCommand { + Send { + message: Message, + ack: oneshot::Sender>, + }, +} + +pub(super) type SharedMessages = Arc>>>; + +pub(super) struct PumpControl { + pub(super) tx_command: mpsc::Sender, + pub(super) rx_message: SharedMessages, +} + +pub(super) async fn send_message( + commands: &mpsc::Sender, + message: Message, +) -> Result<(), tungstenite::Error> { + let (ack, rx_ack) = oneshot::channel(); + commands + .send(PumpCommand::Send { message, ack }) + .await + .map_err(|_| tungstenite::Error::ConnectionClosed)?; + rx_ack + .await + .unwrap_or(Err(tungstenite::Error::ConnectionClosed)) +} + +#[async_trait] +impl ResponseControlTransport for PumpControl { + async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> { + tokio::time::timeout( + WEBSOCKET_SEND_TIMEOUT, + send_message(&self.tx_command, Message::Text(payload.to_owned().into())), + ) + .await + .map_err(|_| ResponseControlUnavailable)? + .map_err(|_| ResponseControlUnavailable) + } + + async fn receive(&self) -> Result { + match self.rx_message.lock().await.recv().await { + Some(Ok(Message::Text(text))) => Ok(text.to_string()), + _ => Err(ResponseControlUnavailable), + } + } +} diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs index 1ae816ee7..3ed01715f 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs @@ -31,7 +31,6 @@ pub(super) enum ExchangeAction { } pub(super) struct ReducedWebSocketEvent { - pub(super) created_response_id: Option, pub(super) action: ExchangeAction, pub(super) diagnostic_event_type: Option, pub(super) turn_state_update: Option, @@ -46,7 +45,6 @@ pub(super) fn reduce_websocket_event( // 不可解析的帧不承载可路由的事件类型，忽略 let Ok(value) = serde_json::from_str::(raw) else { return Ok(ReducedWebSocketEvent { - created_response_id: None, action: ExchangeAction::Ignore, diagnostic_event_type: None, turn_state_update: None, @@ -58,7 +56,6 @@ 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, }); @@ -98,10 +95,6 @@ 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, diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index eb129746d..859bdb756 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -4,7 +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 gateway_core::engine::response_control::ResponseControl; use serde_json::json; use tokio::sync::{Mutex, mpsc}; use tokio_util::sync::CancellationToken; @@ -60,7 +60,7 @@ impl StreamWebSocketDiscardReason { } pub(in crate::transport::websocket) fn stream_websocket_response( - websocket: PumpedWebSocket, + mut websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, pool_return: Option, reused_connection: bool, @@ -68,6 +68,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response( trace: TraceContext, response_control: Option, ) -> CodexWebSocketStreamingExchange { + websocket.bind_control(response_control.as_ref()); let websocket_connection_id = websocket.connection_id(); let response_metadata = metadata.clone(); let rate_limit_updates = Arc::new(Mutex::new(Vec::new())); @@ -87,7 +88,6 @@ pub(in crate::transport::websocket) fn stream_websocket_response( ); let forward = async move { forward_websocket_response_stream(WebSocketStreamForwardState { - response_control, websocket, metadata, pool_return, @@ -127,7 +127,6 @@ pub(in crate::transport::websocket) fn stream_websocket_response( } struct WebSocketStreamForwardState { - response_control: Option, trace: TraceContext, websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, @@ -142,7 +141,6 @@ struct WebSocketStreamForwardState { async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { let WebSocketStreamForwardState { - response_control, trace, mut websocket, mut metadata, @@ -160,10 +158,6 @@ 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 = 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(); @@ -173,7 +167,6 @@ 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, @@ -183,7 +176,6 @@ 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, @@ -192,24 +184,6 @@ 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), @@ -218,8 +192,6 @@ 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", @@ -255,7 +227,6 @@ 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); @@ -271,7 +242,6 @@ 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()); } @@ -310,7 +280,6 @@ 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)); @@ -331,31 +300,6 @@ 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() { @@ -373,12 +317,7 @@ 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"}), @@ -456,7 +395,6 @@ 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, diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index dffc5e680..bb8bc55a3 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -33,9 +33,10 @@ use super::{ pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket, RawWsStream}, }; +use super::control::WEBSOCKET_SEND_TIMEOUT; + const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; const WEBSOCKET_CONNECT_TIMEOUT: Duration = Duration::from_secs(15); -const WEBSOCKET_SEND_TIMEOUT: Duration = Duration::from_secs(5 * 60); impl CodexWebSocketConnection { /// 构造 Responses WebSocket 连接描述 diff --git a/backend/crates/providers/openai/src/transport/websocket/mod.rs b/backend/crates/providers/openai/src/transport/websocket/mod.rs index 9455f4da9..89f71c2bc 100644 --- a/backend/crates/providers/openai/src/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/mod.rs @@ -2,6 +2,7 @@ mod audit; mod breaker; +mod control; mod coordinator; mod error; mod exchange; diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index 8b3b031e2..5a21f29c2 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -13,6 +13,8 @@ //! - 接收：`next` 从 message channel 取出 pump 转发的入站帧（`Ping`/`Pong` 已被 pump 吞掉） //! - 入站缓冲满时暂停读取 socket 和主动探活，仅继续处理发送/关闭命令；消费恢复后按原顺序继续转发 +use gateway_core::engine::response_control::{ResponseControl, ResponseControlTransport}; + use std::sync::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, @@ -21,14 +23,17 @@ use std::time::Duration; use futures::{Sink, SinkExt, Stream, StreamExt}; use tokio::{ - sync::{mpsc, oneshot}, + sync::mpsc, task::JoinHandle, time::{Instant, MissedTickBehavior}, }; use tokio_tungstenite::tungstenite::Message; use uuid::Uuid; -use super::CodexWebSocketCloseError; +use super::{ + CodexWebSocketCloseError, + control::{PumpCommand, PumpControl, SharedMessages, send_message}, +}; /// 底层 tungstenite WebSocket 流 pub(crate) trait WebSocketIo: @@ -89,13 +94,6 @@ impl PumpKeepalive { } } -enum PumpCommand { - Send { - message: Message, - ack: oneshot::Sender>, - }, -} - const PUMP_MESSAGE_BUFFER: usize = 64; #[derive(Debug, Clone)] @@ -284,7 +282,8 @@ impl PumpLifecycleState { pub(crate) struct PumpedWebSocket { connection_id: Uuid, tx_command: mpsc::Sender, - rx_message: mpsc::Receiver>, + rx_message: SharedMessages, + control: Option>, closed: Arc, lifecycle: Arc>, pump: Option>, @@ -335,7 +334,8 @@ impl PumpedWebSocket { Self { connection_id, tx_command, - rx_message, + rx_message: Arc::new(tokio::sync::Mutex::new(rx_message)), + control: None, closed, lifecycle, pump: Some(pump), @@ -344,25 +344,26 @@ impl PumpedWebSocket { /// 通过 pump 发送一帧，返回底层 `send` 的结果 pub(crate) async fn send(&self, message: Message) -> Result<(), tungstenite::Error> { - let (ack, rx_ack) = oneshot::channel(); - if self - .tx_command - .send(PumpCommand::Send { message, ack }) - .await - .is_err() - { - return Err(tungstenite::Error::ConnectionClosed); - } - rx_ack - .await - .unwrap_or(Err(tungstenite::Error::ConnectionClosed)) + send_message(&self.tx_command, message).await + } + + /// 每轮发送后替换控制 owner，旧执行的弱引用不能控制下一轮连接 + pub(crate) fn bind_control(&mut self, control: Option<&ResponseControl>) { + self.control = control.map(|control| { + let transport: Arc = Arc::new(PumpControl { + tx_command: self.tx_command.clone(), + rx_message: Arc::clone(&self.rx_message), + }); + control.bind(&transport); + transport + }); } /// 取出下一帧入站消息（`Ping`/`Pong` 已被 pump 处理，不会到达这里） /// /// 返回 `None` 表示连接已结束（pump 已退出且缓冲已排空） pub(crate) async fn next(&mut self) -> Option> { - self.rx_message.recv().await + self.rx_message.lock().await.recv().await } /// 连接是否已被后台 pump 判定关闭/失活 diff --git a/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs b/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs index f3923ede6..ebbd1f481 100644 --- a/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs +++ b/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs @@ -1,7 +1,7 @@ -//! 验证响应中断使用当前连接并在复用前释放控制权 +//! 验证未知控制帧和迟到中断使用原连接，并在复用时撤销旧执行控制权 use super::*; -use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError}; +use gateway_core::engine::response_control::{ResponseControl, ResponseControlUnavailable}; async fn next_json(socket: &mut tokio_tungstenite::WebSocketStream) -> Value { loop { @@ -46,11 +46,12 @@ fn interrupt_context(control: ResponseControl, previous: Option<&str>) -> Attemp } #[tokio::test] -async fn response_interrupt_uses_the_active_socket_and_releases_control_before_reuse() { +async fn response_controls_use_the_original_socket_through_terminal_and_revoke_on_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 (finish, finished) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { let (socket, _) = listener.accept().await.unwrap(); let mut socket = accept_codex_test_websocket(socket).await; @@ -72,6 +73,19 @@ async fn response_interrupt_uses_the_active_socket_and_releases_control_before_r .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 unknown = next_json(&mut socket).await; + assert_eq!( + unknown, + json!({"type":"future.control", "extension":{"mode":"future"}}) + ); + socket + .send(Message::Text( + json!({"type":"future.ack", "extension":true}) + .to_string() + .into(), + )) + .await + .unwrap(); let interrupt = next_json(&mut socket).await; assert_eq!( interrupt, @@ -88,10 +102,25 @@ async fn response_interrupt_uses_the_active_socket_and_releases_control_before_r )) .await .unwrap(); + let late = next_json(&mut socket).await; + assert_eq!( + late, + json!({"type":"response.interrupt", "response_id":id,"mode":"future_mode", "extra":1}) + ); + socket + .send(Message::Text( + json!({"type":"future.idle_ack", "original":late}) + .to_string() + .into(), + )) + .await + .unwrap(); } + let _ = finished.await; }); let provider = provider_with_base_url(&store, base_url); let mut session_state = None; + let mut old_control: Option = 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"); @@ -112,34 +141,62 @@ async fn response_interrupt_uses_the_active_socket_and_releases_control_before_r ) .await .unwrap(); - let mut interrupted = false; + let mut unknown_reply = false; + let mut sent = 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; + if !sent { + if let Some(old) = old_control.take() { + assert_eq!(old.send("old control").await, Err(ResponseControlUnavailable)); + } + control.send(&json!({"type":"future.control", "extension":{"mode":"future"}}).to_string()).await.unwrap(); + control.send(&json!({"type":"response.interrupt", "response_id":id,"mode":"discard_partial_items"}).to_string()).await.unwrap(); + sent = true; } + unknown_reply |= event.wire_event().and_then(|wire| wire.event_type()) == Some("future.ack"); } }) .await - .expect("interrupt completes the active response"); - assert!(interrupted); + .expect("controls complete the active response"); + assert!(unknown_reply, "unknown upstream events remain visible"); assert!(session_state.is_some()); + let late = + json!({"type":"response.interrupt", "response_id":id,"mode":"future_mode", "extra":1}); + control.send(&late.to_string()).await.unwrap(); + let reply = timeout(Duration::from_secs(5), control.receive()) + .await + .unwrap() + .unwrap(); assert_eq!( - control.interrupt(id), - Err(ResponseInterruptError::Unavailable) + serde_json::from_str::(&reply).unwrap(), + json!({"type":"future.idle_ack", "original":late}) ); + // 取消空闲读取后，下一轮正文必须能继续消费同一个有界接收通道 + assert!( + timeout(Duration::from_millis(10), control.receive()) + .await + .is_err() + ); + old_control = Some(control); } + finish.send(()).unwrap(); timeout(Duration::from_secs(5), server) .await .unwrap() .unwrap(); + let control = old_control.unwrap(); + assert_eq!( + timeout(Duration::from_secs(5), control.receive()) + .await + .unwrap(), + Err(ResponseControlUnavailable) + ); + assert_eq!( + control.send("closed connection").await, + Err(ResponseControlUnavailable) + ); } diff --git a/docs/api.md b/docs/api.md index fc58c7309..93c5bfda2 100644 --- a/docs/api.md +++ b/docs/api.md @@ -309,13 +309,16 @@ turn metadata 的 `workspaces` 也会保留仓库绝对路径、Git 远端地址 已升级的 Responses 连接在入站解析前和出站写入前经过 `websocket` 中间件，完整消息与主动发送接口见 [SDK 洋葱中间件](../backend/crates/gateway-plugin/sdk/docs/capabilities.md#洋葱中间件)。下述规则描述默认协议处理 -Responses WebSocket 接受文本 `response.create` 和 `response.interrupt`，创建请求在同一连接串行执行。当前响应期间收到的后续业务帧 +Responses WebSocket 接受文本消息，`response.create` 在同一连接串行执行。当前响应期间收到的后续创建请求 留在有界接收队列中，待当前响应完成终结和写出后再逐条校验、准入与执行，不因请求提前到达而断开。 -接收队列容量为 32 个事件，超载仍关闭连接；Ping/Pong、客户端关闭和服务关闭不等待队列中的请求执行。 -活动响应期间会即时处理 `{"type":"response.interrupt","response_id":"当前响应 ID","mode":"discard_partial_items"}`。 -中断只能发送到该执行占用的原上游 WS，不重新选号或创建推理 attempt；重复中断合并为一次发送。 -ID 不匹配、没有可中断响应、实际走 HTTP 或 Provider 不支持控制时返回 `400` 协议错误，原执行继续； -客户端需要终止这类执行时可关闭连接，后续按既有续接合同恢复。 +接收队列容量为 32 个事件，超载仍关闭连接；Ping/Pong、客户端关闭和服务关闭不等待队列中的请求执行 + +其余文本消息原样发送到当前执行绑定的原上游 WebSocket，包括 +`response.interrupt` 和未知类型；类型、响应 ID、模式及扩展字段由上游校验。 +控制帧不重新选号、不创建推理 attempt，也不合并重复消息。响应结束后原连接仍可接收控制帧并返回上游事件； +下一轮执行使用新绑定的控制入口。尚未建立原连接、原连接已失效、实际走 HTTP 或 Provider 未提供控制通道时返回 `400`， +不能通过控制帧新建上游连接 + 中断后继续转发上游事件与终态；只有 `response.incomplete` 的 `incomplete_details.reason` 为 `interrupted` 时，才按官方中断语义保留原连接续接能力。部分输出是否被丢弃以上游终态为准。 控制帧发送成功不等于上游已确认中断；若上游仍返回 `response.completed`，按正常完成处理 diff --git a/docs/architecture.md b/docs/architecture.md index 571ed7c3b..d46f22e7d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -708,9 +708,10 @@ Admin 的手动重置复用同一账本与 Key 行锁，在一个事务中清零 - SSE 持有并发名额直到终态；WebSocket 每个 `response.create` 独立准入，空闲连接不占名额。 Core 在每次请求开始时重新鉴权并冻结当前策略，既有连接也应用已发布的限额与授权范围变更。 换号和内部重试复用同一名额，Key 与账号并发上限同时生效；预算拒绝或启动失败释放名额 -- 响应中断入口由 Core 的根执行持有，API 只控制当前连接正在交付的执行。Provider 在收到上游 - `response.created` 后注册活动响应 owner，在终态、失败或断连时释放；控制句柄本身不延长 owner - 生命周期。中断信号交给占用原账号、原 socket 的 exchange，不经过全局响应查询、新选号或独立 attempt。 +- Core 的根执行携带不解释协议的连接控制端口，Provider 在请求写入上游后绑定原 socket。 + API 在响应期间发送控制帧，终态后还可读取原连接事件；空闲读取必须在下一轮执行前取消， + 与 Provider 正文消费共用单一有界接收通道。连接 owner 持有控制目标，复用时替换，销毁时释放； + 控制句柄使用弱引用，不延长连接生命周期，也不经过全局响应查询、新选号或独立 attempt。 API 暂存的下一轮请求和连接泵接收队列共用容量名额，控制帧处理不取消中间件正在读取响应体的 future 完成、失败、取消和断连使用同一结算路径，在结束时以网关请求 ID 幂等累计已取得费用。