mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(gateway): restore upstream response transparency
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Generated
+2
@@ -1121,6 +1121,7 @@ dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"gateway-admin",
|
||||
@@ -1142,6 +1143,7 @@ name = "gateway-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
|
||||
@@ -14,6 +14,7 @@ path = "tests/main.rs"
|
||||
[dependencies]
|
||||
axum = { version = "0.8.9", default-features = false, features = ["json", "query", "ws"] }
|
||||
base64 = "0.22.1"
|
||||
bytes = "1.11.1"
|
||||
chrono = { version = "0.4.45", features = ["serde"] }
|
||||
futures = "0.3.32"
|
||||
gateway-admin = { path = "../gateway-admin" }
|
||||
|
||||
@@ -19,7 +19,7 @@ use gateway_core::engine::execution::{ClientTransport, ExecutionSession, Started
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
|
||||
use gateway_core::lifecycle::ConnectionGuard;
|
||||
use gateway_protocol::openai::sse::{DONE_SSE_FRAME, response_failed_sse_event};
|
||||
use gateway_protocol::openai::sse::{DONE_SSE_FRAME, response_failed_sse_event_with_id};
|
||||
|
||||
use crate::ApiState;
|
||||
use crate::openai::{
|
||||
@@ -468,13 +468,13 @@ impl ResponsesStreamState {
|
||||
fn new(
|
||||
session: Box<dyn ExecutionSession>,
|
||||
encoder: OpenAiResponsesEncoder,
|
||||
initial_frames: Vec<String>,
|
||||
initial_frames: Vec<Bytes>,
|
||||
connection_guard: Option<Box<dyn ConnectionGuard>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
session: Some(session),
|
||||
encoder,
|
||||
pending: initial_frames.into_iter().map(Bytes::from).collect(),
|
||||
pending: initial_frames.into_iter().collect(),
|
||||
output_finished: false,
|
||||
execution_terminal: false,
|
||||
_connection_guard: connection_guard,
|
||||
@@ -526,7 +526,13 @@ impl ResponsesStreamState {
|
||||
.session
|
||||
.as_ref()
|
||||
.is_some_and(|session| session.is_finalized());
|
||||
self.finish_with_gateway_error(gateway_error_from_engine(&error));
|
||||
if self.encoder.has_wire_failure() {
|
||||
self.pending
|
||||
.push_back(Bytes::from_static(DONE_SSE_FRAME.as_bytes()));
|
||||
self.output_finished = true;
|
||||
} else {
|
||||
self.finish_with_gateway_error(gateway_error_from_engine(&error));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -537,7 +543,7 @@ impl ResponsesStreamState {
|
||||
if self.encoder.is_completed() {
|
||||
self.finish_completed(frames).await;
|
||||
} else {
|
||||
self.pending.extend(frames.into_iter().map(Bytes::from));
|
||||
self.pending.extend(frames);
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
@@ -547,7 +553,7 @@ impl ResponsesStreamState {
|
||||
}
|
||||
}
|
||||
|
||||
async fn finish_completed(&mut self, frames: Vec<String>) {
|
||||
async fn finish_completed(&mut self, frames: Vec<Bytes>) {
|
||||
let next = match self.session.as_mut() {
|
||||
Some(session) => session.next_event().await,
|
||||
None => {
|
||||
@@ -565,7 +571,7 @@ impl ResponsesStreamState {
|
||||
.as_ref()
|
||||
.is_some_and(|session| session.is_finalized()) =>
|
||||
{
|
||||
self.pending.extend(frames.into_iter().map(Bytes::from));
|
||||
self.pending.extend(frames);
|
||||
self.execution_terminal = true;
|
||||
self.pending
|
||||
.push_back(Bytes::from_static(DONE_SSE_FRAME.as_bytes()));
|
||||
@@ -608,7 +614,8 @@ impl ResponsesStreamState {
|
||||
fn finish_with_gateway_error(&mut self, error: GatewayError) {
|
||||
let (_, default_type, default_code) = gateway_error_contract(error.kind());
|
||||
self.pending
|
||||
.push_back(Bytes::from(response_failed_sse_event(
|
||||
.push_back(Bytes::from(response_failed_sse_event_with_id(
|
||||
self.encoder.response_id(),
|
||||
error.client_error_type().unwrap_or(default_type),
|
||||
error.client_error_code().unwrap_or(default_code),
|
||||
error.client_message(),
|
||||
@@ -621,7 +628,8 @@ impl ResponsesStreamState {
|
||||
fn finish_with_encode_error(&mut self, error: ResponseEncodeError) {
|
||||
let protocol = error.protocol_body().error;
|
||||
self.pending
|
||||
.push_back(Bytes::from(response_failed_sse_event(
|
||||
.push_back(Bytes::from(response_failed_sse_event_with_id(
|
||||
self.encoder.response_id(),
|
||||
protocol.kind,
|
||||
protocol.code,
|
||||
&protocol.message,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use bytes::Bytes;
|
||||
use gateway_core::accounting::Usage;
|
||||
use gateway_core::event::{
|
||||
CompactionOutput, ContentItem, ContentKind, EventSequenceValidator, FinishReason, GatewayEvent,
|
||||
@@ -21,14 +22,6 @@ enum ResponseRepresentation {
|
||||
OpenAiWire,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PendingWireEvent {
|
||||
event_type: Option<String>,
|
||||
data: Value,
|
||||
sse_id: Option<String>,
|
||||
sse_retry: Option<u64>,
|
||||
}
|
||||
|
||||
/// Provider event 封套到 OpenAI Responses wire 的唯一表达边界。
|
||||
///
|
||||
/// Provider 已提供 OpenAI wire 时只下发该 wire;否则由 [`ResponsesCollector`]
|
||||
@@ -37,10 +30,9 @@ struct PendingWireEvent {
|
||||
pub struct OpenAiResponsesEncoder {
|
||||
canonical: ResponsesCollector,
|
||||
representation: ResponseRepresentation,
|
||||
upstream_response_id: Option<String>,
|
||||
gateway_response_id: Option<String>,
|
||||
response_id: Option<String>,
|
||||
wire_terminal: Option<Value>,
|
||||
pending_wire: Vec<PendingWireEvent>,
|
||||
wire_failure: bool,
|
||||
canonical_completed: bool,
|
||||
}
|
||||
|
||||
@@ -51,10 +43,9 @@ impl OpenAiResponsesEncoder {
|
||||
Self {
|
||||
canonical: ResponsesCollector::new(created_at),
|
||||
representation: ResponseRepresentation::Undecided,
|
||||
upstream_response_id: None,
|
||||
gateway_response_id: None,
|
||||
response_id: None,
|
||||
wire_terminal: None,
|
||||
pending_wire: Vec::new(),
|
||||
wire_failure: false,
|
||||
canonical_completed: false,
|
||||
}
|
||||
}
|
||||
@@ -64,8 +55,29 @@ impl OpenAiResponsesEncoder {
|
||||
/// # Errors
|
||||
///
|
||||
/// canonical 顺序无效、响应身份变化或表达模式混用时返回错误。
|
||||
pub fn push_sse(&mut self, event: &ProviderEvent) -> Result<Vec<String>, ResponseEncodeError> {
|
||||
self.push(event, WireEncoding::Sse)
|
||||
pub fn push_sse(&mut self, event: &ProviderEvent) -> Result<Vec<Bytes>, ResponseEncodeError> {
|
||||
let openai_wire = self.select_representation(event)?;
|
||||
if self.representation == ResponseRepresentation::OpenAiWire {
|
||||
let Some(wire) = openai_wire else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
self.observe_wire_terminal(wire);
|
||||
if let Some(raw_sse_frame) = wire.raw_sse_frame() {
|
||||
return Ok(vec![raw_sse_frame.clone()]);
|
||||
}
|
||||
return Ok(vec![Bytes::from(encode_sse_event_with_metadata(
|
||||
wire.event_type().unwrap_or_default(),
|
||||
&wire.data().to_string(),
|
||||
wire.sse_id(),
|
||||
wire.sse_retry(),
|
||||
))]);
|
||||
}
|
||||
|
||||
Ok(self
|
||||
.push_canonical(event)?
|
||||
.into_iter()
|
||||
.map(Bytes::from)
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// 消费一个 Provider event,并返回 WebSocket JSON messages。
|
||||
@@ -77,7 +89,28 @@ impl OpenAiResponsesEncoder {
|
||||
&mut self,
|
||||
event: &ProviderEvent,
|
||||
) -> Result<Vec<String>, ResponseEncodeError> {
|
||||
self.push(event, WireEncoding::WebSocket)
|
||||
let openai_wire = self.select_representation(event)?;
|
||||
if self.representation == ResponseRepresentation::OpenAiWire {
|
||||
let Some(wire) = openai_wire else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if !wire.has_json_data() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.observe_wire_terminal(wire);
|
||||
return Ok(vec![wire.data().to_string()]);
|
||||
}
|
||||
|
||||
self.push_canonical(event)?
|
||||
.into_iter()
|
||||
.map(|frame| {
|
||||
frame
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("data: "))
|
||||
.map(ToOwned::to_owned)
|
||||
.ok_or(ResponseEncodeError::InvalidEventEncoding)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 返回是否已经观察到 canonical 或 wire 终态。
|
||||
@@ -87,10 +120,16 @@ impl OpenAiResponsesEncoder {
|
||||
|| self.canonical.final_response.is_some()
|
||||
}
|
||||
|
||||
/// 返回 Core 冻结并已写入客户端事件的网关响应 ID。
|
||||
/// 返回是否已把 Provider 原生失败 event 交付给客户端。
|
||||
#[must_use]
|
||||
pub fn gateway_response_id(&self) -> Option<&str> {
|
||||
self.gateway_response_id.as_deref()
|
||||
pub const fn has_wire_failure(&self) -> bool {
|
||||
self.wire_failure
|
||||
}
|
||||
|
||||
/// 返回 Core 已观察到的客户端可见 Provider 原生响应 ID。
|
||||
#[must_use]
|
||||
pub fn response_id(&self) -> Option<&str> {
|
||||
self.response_id.as_deref()
|
||||
}
|
||||
|
||||
/// 校验完整响应并返回非流式 JSON。
|
||||
@@ -110,11 +149,10 @@ impl OpenAiResponsesEncoder {
|
||||
}
|
||||
}
|
||||
|
||||
fn push(
|
||||
fn select_representation<'event>(
|
||||
&mut self,
|
||||
event: &ProviderEvent,
|
||||
encoding: WireEncoding,
|
||||
) -> Result<Vec<String>, ResponseEncodeError> {
|
||||
event: &'event ProviderEvent,
|
||||
) -> Result<Option<&'event gateway_core::event::ProtocolWireEvent>, ResponseEncodeError> {
|
||||
self.observe_identity(event)?;
|
||||
let openai_wire = event
|
||||
.wire_event()
|
||||
@@ -130,70 +168,18 @@ impl OpenAiResponsesEncoder {
|
||||
self.representation = ResponseRepresentation::Canonical;
|
||||
}
|
||||
|
||||
if self.representation == ResponseRepresentation::OpenAiWire {
|
||||
let Some(wire) = openai_wire else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
self.pending_wire.push(PendingWireEvent {
|
||||
event_type: wire.event_type().map(ToOwned::to_owned),
|
||||
data: wire.data().clone(),
|
||||
sse_id: wire.sse_id().map(ToOwned::to_owned),
|
||||
sse_retry: wire.sse_retry(),
|
||||
});
|
||||
if self.gateway_response_id.is_none() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let pending = std::mem::take(&mut self.pending_wire);
|
||||
return pending
|
||||
.into_iter()
|
||||
.map(|pending| {
|
||||
let PendingWireEvent {
|
||||
event_type,
|
||||
mut data,
|
||||
sse_id,
|
||||
sse_retry,
|
||||
} = pending;
|
||||
self.rewrite_wire_identity(&mut data);
|
||||
let effective_type = event_type
|
||||
.as_deref()
|
||||
.or_else(|| data.get("type").and_then(Value::as_str));
|
||||
if matches!(
|
||||
effective_type,
|
||||
Some("response.completed" | "response.incomplete")
|
||||
) {
|
||||
self.wire_terminal = data.get("response").cloned();
|
||||
}
|
||||
Ok(match encoding {
|
||||
WireEncoding::Sse => encode_sse_event_with_metadata(
|
||||
event_type.as_deref().unwrap_or_default(),
|
||||
&data.to_string(),
|
||||
sse_id.as_deref(),
|
||||
sse_retry,
|
||||
),
|
||||
WireEncoding::WebSocket => data.to_string(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
Ok(openai_wire)
|
||||
}
|
||||
|
||||
fn push_canonical(
|
||||
&mut self,
|
||||
event: &ProviderEvent,
|
||||
) -> Result<Vec<String>, ResponseEncodeError> {
|
||||
let mut canonical_frames = Vec::new();
|
||||
for fact in event.canonical_facts() {
|
||||
canonical_frames.extend(self.canonical.push(fact)?);
|
||||
}
|
||||
|
||||
match encoding {
|
||||
WireEncoding::Sse => Ok(canonical_frames),
|
||||
WireEncoding::WebSocket => canonical_frames
|
||||
.into_iter()
|
||||
.map(|frame| {
|
||||
frame
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("data: "))
|
||||
.map(ToOwned::to_owned)
|
||||
.ok_or(ResponseEncodeError::InvalidEventEncoding)
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
Ok(canonical_frames)
|
||||
}
|
||||
|
||||
fn observe_identity(&mut self, event: &ProviderEvent) -> Result<(), ResponseEncodeError> {
|
||||
@@ -206,67 +192,33 @@ impl OpenAiResponsesEncoder {
|
||||
}
|
||||
_ => continue,
|
||||
};
|
||||
let Some(upstream) = metadata.upstream_response_id() else {
|
||||
continue;
|
||||
};
|
||||
if self
|
||||
.upstream_response_id
|
||||
.response_id
|
||||
.as_deref()
|
||||
.is_some_and(|current| current != upstream.as_str())
|
||||
|| self
|
||||
.gateway_response_id
|
||||
.as_deref()
|
||||
.is_some_and(|current| current != metadata.response_id())
|
||||
.is_some_and(|current| current != metadata.response_id())
|
||||
{
|
||||
return Err(ResponseEncodeError::WireIdentityChanged);
|
||||
}
|
||||
self.upstream_response_id = Some(upstream.as_str().to_owned());
|
||||
self.gateway_response_id = Some(metadata.response_id().to_owned());
|
||||
self.response_id = Some(metadata.response_id().to_owned());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn rewrite_wire_identity(&self, data: &mut Value) {
|
||||
let (Some(upstream), Some(gateway)) = (
|
||||
self.upstream_response_id.as_deref(),
|
||||
self.gateway_response_id.as_deref(),
|
||||
) else {
|
||||
fn observe_wire_terminal(&mut self, wire: &gateway_core::event::ProtocolWireEvent) {
|
||||
if !wire.has_json_data() {
|
||||
return;
|
||||
};
|
||||
rewrite_response_identity(data, upstream, gateway);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum WireEncoding {
|
||||
Sse,
|
||||
WebSocket,
|
||||
}
|
||||
|
||||
fn rewrite_response_identity(value: &mut Value, upstream: &str, gateway: &str) {
|
||||
match value {
|
||||
Value::Array(values) => {
|
||||
for value in values {
|
||||
rewrite_response_identity(value, upstream, gateway);
|
||||
}
|
||||
}
|
||||
Value::Object(object) => {
|
||||
if let Some(Value::String(response_id)) = object.get_mut("response_id")
|
||||
&& response_id == upstream
|
||||
{
|
||||
*response_id = gateway.to_owned();
|
||||
}
|
||||
if let Some(Value::Object(response)) = object.get_mut("response")
|
||||
&& let Some(Value::String(response_id)) = response.get_mut("id")
|
||||
&& response_id == upstream
|
||||
{
|
||||
*response_id = gateway.to_owned();
|
||||
}
|
||||
for value in object.values_mut() {
|
||||
rewrite_response_identity(value, upstream, gateway);
|
||||
}
|
||||
let effective_type = wire
|
||||
.event_type()
|
||||
.or_else(|| wire.data().get("type").and_then(Value::as_str));
|
||||
if matches!(
|
||||
effective_type,
|
||||
Some("response.completed" | "response.incomplete")
|
||||
) {
|
||||
self.wire_terminal = wire.data().get("response").cloned();
|
||||
} else if matches!(effective_type, Some("response.failed" | "error")) {
|
||||
self.wire_failure = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -428,6 +428,10 @@ async fn forward_execution(
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
if encoder.has_wire_failure() {
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
@@ -446,7 +450,7 @@ fn commit_connection_replay(
|
||||
encoder: &OpenAiResponsesEncoder,
|
||||
provider_state: Option<ProviderSessionState>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let response_id = encoder.gateway_response_id().ok_or_else(|| {
|
||||
let response_id = encoder.response_id().ok_or_else(|| {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway completed a response without an identity",
|
||||
|
||||
@@ -10,6 +10,7 @@ use axum::{
|
||||
extract::connect_info::ConnectInfo,
|
||||
http::{HeaderMap, Request, StatusCode, header::AUTHORIZATION},
|
||||
};
|
||||
use bytes::Bytes;
|
||||
use futures::future::{BoxFuture, pending};
|
||||
use gateway_core::engine::execution::{
|
||||
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
|
||||
@@ -642,17 +643,19 @@ async fn streaming_response_forwards_only_allowlisted_provider_headers() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_commit_batch_encodes_pre_identity_wire_before_committing() {
|
||||
async fn streaming_commit_batch_preserves_pre_identity_wire_before_committing() {
|
||||
let trace = Arc::new(Trace::default());
|
||||
let upstream_id = SafeUpstreamValue::new("resp_upstream").expect("upstream response ID");
|
||||
let metadata =
|
||||
ResponseMeta::new("resp_gateway", "public-model").with_upstream_response_id(upstream_id);
|
||||
let metadata = ResponseMeta::new("resp_upstream", "public-model");
|
||||
let raw_future = Bytes::from_static(
|
||||
b"id: evt_before_identity\r\nevent: response.future_metadata\r\nretry: 3000\r\ndata: { \"type\": \"response.future_metadata\", \"opaque\": true }\r\n\r\n",
|
||||
);
|
||||
let events = vec![
|
||||
ProviderEvent::wire(
|
||||
ProtocolWireEvent::json_with_sse_metadata(
|
||||
ProtocolWireEvent::json_with_raw_sse_metadata(
|
||||
"openai",
|
||||
Some("response.future_metadata".to_owned()),
|
||||
json!({"type":"response.future_metadata","opaque":true}),
|
||||
raw_future.clone(),
|
||||
Some("evt_before_identity".to_owned()),
|
||||
Some(3_000),
|
||||
)
|
||||
@@ -700,10 +703,8 @@ async fn streaming_commit_batch_encodes_pre_identity_wire_before_committing() {
|
||||
let created = body.find("response.created").expect("created event");
|
||||
let completed = body.find("response.completed").expect("completed event");
|
||||
assert!(future < created && created < completed);
|
||||
assert!(body.contains("id: evt_before_identity\n"));
|
||||
assert!(body.contains("retry: 3000\n"));
|
||||
assert!(body.contains("resp_gateway"));
|
||||
assert!(!body.contains("resp_upstream"));
|
||||
assert!(body.as_bytes().starts_with(raw_future.as_ref()));
|
||||
assert!(body.contains("resp_upstream"));
|
||||
assert!(body.ends_with("data: [DONE]\n\n"));
|
||||
assert_eq!(trace.client_statuses(), vec![200]);
|
||||
assert_eq!(trace.snapshot(), vec!["next_event", "commit", "next_end"]);
|
||||
@@ -921,6 +922,90 @@ async fn streaming_error_should_emit_client_visible_upstream_details() {
|
||||
assert!(body.ends_with("data: [DONE]\n\n"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_upstream_wire_failure_should_not_be_rewritten_as_a_gateway_error() {
|
||||
let trace = Arc::new(Trace::default());
|
||||
let raw_failure = Bytes::from_static(
|
||||
b"event: response.failed\r\ndata: { \"type\": \"response.failed\", \"response\": { \"id\": \"resp_test\", \"status\": \"failed\", \"error\": { \"code\": \"rate_limit_exceeded\", \"message\": \"upstream raw failure marker\" } } }\r\n\r\n",
|
||||
);
|
||||
let wire_failure = ProviderEvent::wire(
|
||||
ProtocolWireEvent::json_with_raw_sse_metadata(
|
||||
"openai",
|
||||
Some("response.failed".to_owned()),
|
||||
json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": "resp_test",
|
||||
"status": "failed",
|
||||
"error": {
|
||||
"code": "rate_limit_exceeded",
|
||||
"message": "upstream raw failure marker"
|
||||
}
|
||||
}
|
||||
}),
|
||||
raw_failure.clone(),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("valid upstream failure wire"),
|
||||
);
|
||||
let started_wire = ProviderEvent::canonical_with_wire(
|
||||
vec![started()],
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.created".to_owned()),
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_test",
|
||||
"model": "public-model",
|
||||
"status": "in_progress"
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("valid upstream started wire"),
|
||||
);
|
||||
let session = FakeSession::streaming(
|
||||
Arc::clone(&trace),
|
||||
vec![
|
||||
NextStep::Event(CoordinatedEvent::single(
|
||||
started_wire,
|
||||
CommitRequirement::CommitBeforeDelivery,
|
||||
)),
|
||||
NextStep::Event(CoordinatedEvent::single(
|
||||
wire_failure,
|
||||
CommitRequirement::AlreadyCommitted,
|
||||
)),
|
||||
NextStep::Error(EngineError::Provider(ProviderError::new(
|
||||
ProviderErrorKind::RateLimited,
|
||||
UpstreamSendState::Sent,
|
||||
))),
|
||||
],
|
||||
);
|
||||
|
||||
let response = stream_execution_response(Box::new(session), 1, None).await;
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("read SSE body");
|
||||
|
||||
assert!(
|
||||
body.windows(raw_failure.len())
|
||||
.any(|frame| frame == raw_failure.as_ref())
|
||||
);
|
||||
assert_eq!(
|
||||
String::from_utf8_lossy(&body)
|
||||
.matches("upstream raw failure marker")
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert!(!String::from_utf8_lossy(&body).contains("resp_proxy_"));
|
||||
assert!(String::from_utf8_lossy(&body).ends_with("data: [DONE]\n\n"));
|
||||
assert_eq!(
|
||||
trace.snapshot(),
|
||||
vec!["next_event", "commit", "next_event", "next_error"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_second_commit_request_should_fail_and_finalize_once() {
|
||||
let trace = Arc::new(Trace::default());
|
||||
|
||||
@@ -2,6 +2,7 @@ mod http;
|
||||
mod websocket;
|
||||
|
||||
use axum::http::HeaderMap;
|
||||
use bytes::Bytes;
|
||||
use gateway_api::openai::responses::{
|
||||
ContinuationIntent, DecodedResponsesRequest, OpenAiRequestHeaders, OpenAiResponsesEncoder,
|
||||
RequestDecodeError, ResponseCreateFrameError, ResponseEncodeError, ResponsesCollector,
|
||||
@@ -9,7 +10,6 @@ use gateway_api::openai::responses::{
|
||||
};
|
||||
use gateway_core::{
|
||||
accounting::{CalculatedCost, ProviderReportedCost, Usage},
|
||||
error::SafeUpstreamValue,
|
||||
event::{
|
||||
CompactionOutput, CompactionSummary, ContentItem, ContentKind, EventSequenceError,
|
||||
FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, ReasoningDelta, ResponseMeta,
|
||||
@@ -804,11 +804,17 @@ fn websocket_events_should_reuse_the_exact_sse_event_json() {
|
||||
let websocket_events = websocket_encoder
|
||||
.push_websocket(&provider_event)
|
||||
.expect("WebSocket event encoding");
|
||||
let parsed = parse_sse_events(&sse_frames.join(""))
|
||||
.expect("generated SSE frames remain parseable")
|
||||
.into_iter()
|
||||
.map(|event| event.data)
|
||||
let sse_body = sse_frames
|
||||
.iter()
|
||||
.flat_map(|frame| frame.as_ref().iter().copied())
|
||||
.collect::<Vec<_>>();
|
||||
let parsed = parse_sse_events(
|
||||
std::str::from_utf8(&sse_body).expect("generated SSE frames remain UTF-8"),
|
||||
)
|
||||
.expect("generated SSE frames remain parseable")
|
||||
.into_iter()
|
||||
.map(|event| event.data)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(websocket_events, parsed);
|
||||
}
|
||||
@@ -1070,17 +1076,14 @@ fn collector_should_reject_unrepresentable_media_output() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_response_id() {
|
||||
fn openai_wire_encoder_should_preserve_unknown_media_events_and_upstream_response_id() {
|
||||
let upstream_id = "resp_upstream_private";
|
||||
let gateway_id = "resp_gateway_public";
|
||||
let started = ResponseMeta::new(gateway_id, "gpt-test").with_upstream_response_id(
|
||||
SafeUpstreamValue::new(upstream_id).expect("safe upstream response ID"),
|
||||
let started = ResponseMeta::new(upstream_id, "gpt-test");
|
||||
let completed =
|
||||
ResponseMeta::new(upstream_id, "gpt-test").with_finish_reason(FinishReason::Stop);
|
||||
let raw_partial_image = Bytes::from_static(
|
||||
b"id: evt_partial_image\r\nevent: response.image_generation_call.partial_image\r\nretry: 2000\r\ndata: { \"type\": \"response.image_generation_call.partial_image\", \"response_id\": \"resp_upstream_private\", \"partial_image_b64\": \"opaque-image-fragment\" }\r\n\r\n",
|
||||
);
|
||||
let completed = ResponseMeta::new(gateway_id, "gpt-test")
|
||||
.with_finish_reason(FinishReason::Stop)
|
||||
.with_upstream_response_id(
|
||||
SafeUpstreamValue::new(upstream_id).expect("safe upstream response ID"),
|
||||
);
|
||||
let events = [
|
||||
openai_wire_event(
|
||||
vec![GatewayEvent::Started(started)],
|
||||
@@ -1091,7 +1094,7 @@ fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_re
|
||||
}),
|
||||
),
|
||||
ProviderEvent::wire(
|
||||
ProtocolWireEvent::json_with_sse_metadata(
|
||||
ProtocolWireEvent::json_with_raw_sse_metadata(
|
||||
"openai",
|
||||
Some("response.image_generation_call.partial_image".to_owned()),
|
||||
json!({
|
||||
@@ -1099,6 +1102,7 @@ fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_re
|
||||
"response_id": upstream_id,
|
||||
"partial_image_b64": "opaque-image-fragment"
|
||||
}),
|
||||
raw_partial_image.clone(),
|
||||
Some("evt_partial_image".to_owned()),
|
||||
Some(2_000),
|
||||
)
|
||||
@@ -1123,7 +1127,12 @@ fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_re
|
||||
.flat_map(|event| encoder.push_sse(event).expect("wire event should encode"))
|
||||
.collect::<Vec<_>>();
|
||||
let response = encoder.finish().expect("wire response should finish");
|
||||
let parsed = parse_sse_events(&frames.join("")).expect("wire SSE should parse");
|
||||
let body = frames
|
||||
.iter()
|
||||
.flat_map(|frame| frame.as_ref().iter().copied())
|
||||
.collect::<Vec<_>>();
|
||||
let body_text = std::str::from_utf8(&body).expect("wire SSE should be UTF-8");
|
||||
let parsed = parse_sse_events(body_text).expect("wire SSE should parse");
|
||||
|
||||
assert_eq!(parsed.len(), 3);
|
||||
assert_eq!(parsed[1].id.as_deref(), Some("evt_partial_image"));
|
||||
@@ -1134,8 +1143,12 @@ fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_re
|
||||
.get("partial_image_b64"),
|
||||
Some(&json!("opaque-image-fragment"))
|
||||
);
|
||||
assert_eq!(response.get("id"), Some(&json!(gateway_id)));
|
||||
assert!(!frames.join("").contains(upstream_id));
|
||||
assert_eq!(response.get("id"), Some(&json!(upstream_id)));
|
||||
assert!(body_text.contains(upstream_id));
|
||||
assert!(
|
||||
body.windows(raw_partial_image.len())
|
||||
.any(|frame| frame == raw_partial_image.as_ref())
|
||||
);
|
||||
}
|
||||
|
||||
fn openai_wire_event(canonical: Vec<GatewayEvent>, event_type: &str, data: Value) -> ProviderEvent {
|
||||
|
||||
@@ -13,6 +13,7 @@ path = "tests/main.rs"
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1.89"
|
||||
bytes = "1.11.1"
|
||||
chrono = "0.4.45"
|
||||
futures = "0.3.32"
|
||||
futures-timer = "3.0.3"
|
||||
|
||||
@@ -5,10 +5,6 @@ use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use futures::{FutureExt, StreamExt, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::accounting::{CostEstimate, CostSource, Usage};
|
||||
use crate::engine::continuation::ContinuationBinding;
|
||||
use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream};
|
||||
@@ -25,6 +21,8 @@ use crate::event::{
|
||||
};
|
||||
use crate::operation::{Operation, RetrySafety};
|
||||
use crate::routing::RoutingPlan;
|
||||
use futures::{FutureExt, StreamExt, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
/// Request 级协调器;不会创建或写入 `request_attempts`。
|
||||
pub struct AttemptCoordinator<S: ?Sized> {
|
||||
@@ -60,7 +58,6 @@ where
|
||||
let client_api_key_ref = request.client_api_key_ref.clone();
|
||||
let request_started_at = request.started_at;
|
||||
let deadline = request.deadline_at;
|
||||
let gateway_response_id = format!("resp_{}", Uuid::now_v7().simple());
|
||||
let account_state_owner = continuation
|
||||
.as_ref()
|
||||
.and_then(ContinuationBinding::pinned)
|
||||
@@ -103,7 +100,6 @@ where
|
||||
image_generation_requested,
|
||||
cost: CostEstimate::unavailable(),
|
||||
timings: ModelRequestTimings::default(),
|
||||
gateway_response_id,
|
||||
client_response_id: None,
|
||||
upstream_response_id: None,
|
||||
last_account_exhaustion: None,
|
||||
@@ -182,7 +178,6 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
image_generation_requested: bool,
|
||||
cost: CostEstimate,
|
||||
timings: ModelRequestTimings,
|
||||
gateway_response_id: String,
|
||||
client_response_id: Option<String>,
|
||||
upstream_response_id: Option<String>,
|
||||
last_account_exhaustion: Option<ProviderError>,
|
||||
@@ -408,8 +403,8 @@ where
|
||||
continue;
|
||||
}
|
||||
let mut identity_error = None;
|
||||
for fact in event.canonical_facts_mut() {
|
||||
if let Err(error) = self.freeze_response_identity(fact) {
|
||||
for fact in event.canonical_facts() {
|
||||
if let Err(error) = self.observe_response_identity(fact) {
|
||||
identity_error = Some(error);
|
||||
break;
|
||||
}
|
||||
@@ -735,27 +730,32 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn freeze_response_identity(&mut self, event: &mut GatewayEvent) -> Result<(), ProviderError> {
|
||||
let upstream_response_id = event
|
||||
.freeze_gateway_response_id(&self.gateway_response_id)
|
||||
.map_err(|_| {
|
||||
ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
|
||||
})?;
|
||||
let Some(upstream_response_id) = upstream_response_id else {
|
||||
return Ok(());
|
||||
fn observe_response_identity(&mut self, event: &GatewayEvent) -> Result<(), ProviderError> {
|
||||
let metadata = match event {
|
||||
GatewayEvent::Started(metadata) | GatewayEvent::Completed(metadata) => metadata,
|
||||
_ => return Ok(()),
|
||||
};
|
||||
let response_id = crate::error::SafeUpstreamValue::new(metadata.response_id().to_owned())
|
||||
.map_err(|_| {
|
||||
ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
|
||||
})?;
|
||||
let response_id = response_id.as_str();
|
||||
if response_id.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
if self
|
||||
.upstream_response_id
|
||||
.as_deref()
|
||||
.is_some_and(|expected| expected != upstream_response_id)
|
||||
.is_some_and(|expected| expected != response_id)
|
||||
{
|
||||
return Err(ProviderError::new(
|
||||
ProviderErrorKind::Protocol,
|
||||
UpstreamSendState::Sent,
|
||||
));
|
||||
}
|
||||
self.client_response_id = Some(self.gateway_response_id.clone());
|
||||
self.upstream_response_id = Some(upstream_response_id);
|
||||
let response_id = response_id.to_owned();
|
||||
self.client_response_id = Some(response_id.clone());
|
||||
self.upstream_response_id = Some(response_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
|
||||
use bytes::Bytes;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
@@ -17,7 +18,6 @@ pub struct ResponseMeta {
|
||||
response_id: String,
|
||||
model: String,
|
||||
finish_reason: Option<FinishReason>,
|
||||
upstream_response_id: Option<SafeUpstreamValue>,
|
||||
}
|
||||
|
||||
impl ResponseMeta {
|
||||
@@ -28,7 +28,6 @@ impl ResponseMeta {
|
||||
response_id: response_id.into(),
|
||||
model: model.into(),
|
||||
finish_reason: None,
|
||||
upstream_response_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,7 +38,7 @@ impl ResponseMeta {
|
||||
self
|
||||
}
|
||||
|
||||
/// 返回网关响应 ID。
|
||||
/// 返回客户端可见的 Provider 原生响应 ID。
|
||||
#[must_use]
|
||||
pub fn response_id(&self) -> &str {
|
||||
&self.response_id
|
||||
@@ -56,19 +55,6 @@ impl ResponseMeta {
|
||||
pub const fn finish_reason(&self) -> Option<FinishReason> {
|
||||
self.finish_reason
|
||||
}
|
||||
|
||||
/// 附加 adapter 已分类为非 bearer 的上游 response ID。
|
||||
#[must_use]
|
||||
pub fn with_upstream_response_id(mut self, response_id: SafeUpstreamValue) -> Self {
|
||||
self.upstream_response_id = Some(response_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// 返回仅供 attempt 遥测持久化的安全上游 response ID。
|
||||
#[must_use]
|
||||
pub const fn upstream_response_id(&self) -> Option<&SafeUpstreamValue> {
|
||||
self.upstream_response_id.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 规范化终止原因。
|
||||
@@ -273,6 +259,8 @@ pub struct ProtocolWireEvent {
|
||||
protocol: String,
|
||||
event_type: Option<String>,
|
||||
data: Value,
|
||||
has_json_data: bool,
|
||||
raw_sse_frame: Option<Bytes>,
|
||||
sse_id: Option<String>,
|
||||
sse_retry: Option<u64>,
|
||||
}
|
||||
@@ -318,11 +306,59 @@ impl ProtocolWireEvent {
|
||||
protocol,
|
||||
event_type,
|
||||
data,
|
||||
has_json_data: true,
|
||||
raw_sse_frame: None,
|
||||
sse_id,
|
||||
sse_retry,
|
||||
})
|
||||
}
|
||||
|
||||
/// 创建携带未经改写 SSE 原始帧的协议 JSON event。
|
||||
///
|
||||
/// `data` 只供 Core 旁路观测;客户端 SSE 输出必须优先使用 `raw_sse_frame`。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 与 [`Self::json_with_sse_metadata`] 相同。
|
||||
pub fn json_with_raw_sse_metadata(
|
||||
protocol: impl Into<String>,
|
||||
event_type: Option<String>,
|
||||
data: Value,
|
||||
raw_sse_frame: Bytes,
|
||||
sse_id: Option<String>,
|
||||
sse_retry: Option<u64>,
|
||||
) -> Result<Self, IdentifierError> {
|
||||
let mut event =
|
||||
Self::json_with_sse_metadata(protocol, event_type, data, sse_id, sse_retry)?;
|
||||
event.raw_sse_frame = Some(raw_sse_frame);
|
||||
Ok(event)
|
||||
}
|
||||
|
||||
/// 创建只有原始 SSE 字节、没有可解析 JSON 的协议 event。
|
||||
///
|
||||
/// 例如 keep-alive 注释和未知的非 JSON SSE 帧仍须原样透传,但不会参与
|
||||
/// WebSocket JSON 输出或旁路观测。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 协议名不满足 wire 安全约束时返回错误。
|
||||
pub fn raw_sse(
|
||||
protocol: impl Into<String>,
|
||||
raw_sse_frame: Bytes,
|
||||
) -> Result<Self, IdentifierError> {
|
||||
let protocol = protocol.into();
|
||||
validate_text(&protocol, 64, true, None)?;
|
||||
Ok(Self {
|
||||
protocol,
|
||||
event_type: None,
|
||||
data: Value::Null,
|
||||
has_json_data: false,
|
||||
raw_sse_frame: Some(raw_sse_frame),
|
||||
sse_id: None,
|
||||
sse_retry: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// 返回客户端协议名称。
|
||||
#[must_use]
|
||||
pub fn protocol(&self) -> &str {
|
||||
@@ -353,6 +389,18 @@ impl ProtocolWireEvent {
|
||||
&self.data
|
||||
}
|
||||
|
||||
/// 返回该 event 是否携带可供旁路观测的 JSON 数据。
|
||||
#[must_use]
|
||||
pub const fn has_json_data(&self) -> bool {
|
||||
self.has_json_data
|
||||
}
|
||||
|
||||
/// 返回可直接写给 SSE 客户端的未经改写原始帧。
|
||||
#[must_use]
|
||||
pub const fn raw_sse_frame(&self) -> Option<&Bytes> {
|
||||
self.raw_sse_frame.as_ref()
|
||||
}
|
||||
|
||||
/// 拆出协议原生 JSON 数据。
|
||||
#[must_use]
|
||||
pub fn into_data(self) -> Value {
|
||||
@@ -366,6 +414,8 @@ impl fmt::Debug for ProtocolWireEvent {
|
||||
.debug_struct("ProtocolWireEvent")
|
||||
.field("protocol", &self.protocol)
|
||||
.field("event_type", &self.event_type)
|
||||
.field("has_json_data", &self.has_json_data)
|
||||
.field("has_raw_sse_frame", &self.raw_sse_frame.is_some())
|
||||
.field("has_sse_id", &self.sse_id.is_some())
|
||||
.field("sse_retry", &self.sse_retry)
|
||||
.field("data", &"<not included in Debug>")
|
||||
@@ -760,24 +810,6 @@ impl GatewayEvent {
|
||||
| Self::Completed(_)
|
||||
)
|
||||
}
|
||||
|
||||
/// 冻结客户端可见的网关 response ID,并返回 Provider 原生 response ID。
|
||||
///
|
||||
/// 只有 Coordinator 会调用该边界;Provider 产生的 ID 永远不会直接下发。
|
||||
pub(crate) fn freeze_gateway_response_id(
|
||||
&mut self,
|
||||
gateway_response_id: &str,
|
||||
) -> Result<Option<String>, crate::error::IdentifierError> {
|
||||
let metadata = match self {
|
||||
Self::Started(metadata) | Self::Completed(metadata) => metadata,
|
||||
_ => return Ok(None),
|
||||
};
|
||||
let upstream_response_id = SafeUpstreamValue::new(metadata.response_id.clone())?;
|
||||
let upstream_response_id_text = upstream_response_id.as_str().to_owned();
|
||||
metadata.response_id = gateway_response_id.to_owned();
|
||||
metadata.upstream_response_id = Some(upstream_response_id);
|
||||
Ok(Some(upstream_response_id_text))
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonical event 顺序错误。
|
||||
|
||||
@@ -343,11 +343,10 @@ fn complete_stream(total_tokens: Option<u64>) -> Vec<Result<GatewayEvent, Provid
|
||||
..Usage::new()
|
||||
})));
|
||||
}
|
||||
events.push(Ok(GatewayEvent::Completed(
|
||||
ResponseMeta::new("response-1", "gpt-5").with_upstream_response_id(
|
||||
SafeUpstreamValue::new("upstream-response").expect("safe response id"),
|
||||
),
|
||||
)));
|
||||
events.push(Ok(GatewayEvent::Completed(ResponseMeta::new(
|
||||
"response-1",
|
||||
"gpt-5",
|
||||
))));
|
||||
events
|
||||
}
|
||||
|
||||
@@ -505,9 +504,8 @@ fn success_updates_one_model_request_and_persists_usage() {
|
||||
GatewayEvent::Completed(metadata) => metadata.response_id(),
|
||||
event => panic!("unexpected final event: {event:?}"),
|
||||
};
|
||||
assert!(started_id.starts_with("resp_"));
|
||||
assert_eq!(started_id, "response-1");
|
||||
assert_eq!(started_id, completed_id);
|
||||
assert_ne!(started_id, "response-1");
|
||||
assert!(!session.is_finalized());
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit response");
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! SSE 事件解析与编码。
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use serde_json::json;
|
||||
use thiserror::Error;
|
||||
|
||||
@@ -16,6 +18,50 @@ pub struct SseEvent {
|
||||
pub retry: Option<u64>,
|
||||
}
|
||||
|
||||
/// 保留完整原始字节与旁路解析结果的单个 SSE 帧。
|
||||
///
|
||||
/// 这用于需要原样转发的 transport;`events` 只供旁路观测,不能反过来重建
|
||||
/// `raw`。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct SseFrame {
|
||||
raw: Vec<u8>,
|
||||
events: Vec<SseEvent>,
|
||||
}
|
||||
|
||||
impl SseFrame {
|
||||
fn new(raw: Vec<u8>, events: Vec<SseEvent>) -> Self {
|
||||
Self { raw, events }
|
||||
}
|
||||
|
||||
/// 返回未经改写的完整 SSE 帧字节。
|
||||
#[must_use]
|
||||
pub fn raw(&self) -> &[u8] {
|
||||
&self.raw
|
||||
}
|
||||
|
||||
/// 返回从同一原始帧旁路解析出的事件。
|
||||
#[must_use]
|
||||
pub fn events(&self) -> &[SseEvent] {
|
||||
&self.events
|
||||
}
|
||||
|
||||
/// 拆出原始帧与旁路解析事件。
|
||||
#[must_use]
|
||||
pub fn into_parts(self) -> (Vec<u8>, Vec<SseEvent>) {
|
||||
(self.raw, self.events)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for SseFrame {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("SseFrame")
|
||||
.field("raw_len", &self.raw.len())
|
||||
.field("events", &self.events)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// 单事件缓冲上限。
|
||||
pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
@@ -54,6 +100,38 @@ impl SseEventDecoder {
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// 追加一个字节块,保留每个完整 SSE 帧的原始字节及其旁路解析结果。
|
||||
///
|
||||
/// 与 [`Self::push`] 相比,此方法仅应由需要原样转发的 transport 使用;普通
|
||||
/// 解析路径无需为原始帧额外分配内存。旁路解析失败时仍保留原始帧,并以空
|
||||
/// event 列表表示不可观测,避免观测失败截断客户端字节流。
|
||||
pub fn push_frames(&mut self, chunk: &[u8]) -> Result<Vec<SseFrame>, SseError> {
|
||||
self.pending.extend_from_slice(chunk);
|
||||
let mut frames = Vec::new();
|
||||
let mut consumed = 0usize;
|
||||
|
||||
while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) {
|
||||
let end = consumed.saturating_add(frame_len);
|
||||
let raw = self.pending[consumed..end].to_vec();
|
||||
let events = std::str::from_utf8(&raw)
|
||||
.ok()
|
||||
.and_then(|frame| parse_sse_events(frame).ok())
|
||||
.unwrap_or_default();
|
||||
frames.push(SseFrame::new(raw, events));
|
||||
consumed = end;
|
||||
}
|
||||
|
||||
if consumed != 0 {
|
||||
self.pending.drain(..consumed);
|
||||
}
|
||||
if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES {
|
||||
return Err(SseError::BufferExceeded {
|
||||
max_bytes: MAX_SSE_EVENT_BUFFER_BYTES,
|
||||
});
|
||||
}
|
||||
Ok(frames)
|
||||
}
|
||||
|
||||
/// 流结束时解析尚未带空行分隔符的最后一帧。
|
||||
pub fn finish(&mut self) -> Result<Vec<SseEvent>, SseError> {
|
||||
if self.pending.is_empty() {
|
||||
@@ -64,6 +142,21 @@ impl SseEventDecoder {
|
||||
.map_err(|error| SseError::ParseError(error.to_string()))?;
|
||||
parse_sse_events(frame)
|
||||
}
|
||||
|
||||
/// 在流结束时返回未带空行分隔符的最后一帧及其旁路解析结果。
|
||||
///
|
||||
/// 与 [`Self::push_frames`] 一致,旁路解析失败不会丢弃原始帧。
|
||||
pub fn finish_frames(&mut self) -> Result<Vec<SseFrame>, SseError> {
|
||||
if self.pending.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let raw = std::mem::take(&mut self.pending);
|
||||
let events = std::str::from_utf8(&raw)
|
||||
.ok()
|
||||
.and_then(|frame| parse_sse_events(frame).ok())
|
||||
.unwrap_or_default();
|
||||
Ok(vec![SseFrame::new(raw, events)])
|
||||
}
|
||||
}
|
||||
|
||||
/// 编码 OpenAI Responses `response.failed` SSE 事件。
|
||||
@@ -116,6 +209,17 @@ pub fn sse_frame_end(bytes: &[u8]) -> Option<usize> {
|
||||
sse_frame_separator_bytes(bytes).map(|(position, separator_len)| position + separator_len)
|
||||
}
|
||||
|
||||
/// 判断完整 SSE 帧是否是单个 `[DONE]` 控制帧。
|
||||
#[must_use]
|
||||
pub fn sse_frame_is_done(frame: &str) -> bool {
|
||||
let mut data = frame.lines().filter_map(|raw_line| {
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
let (field, value) = split_sse_field(line);
|
||||
(field == "data").then_some(value)
|
||||
});
|
||||
matches!((data.next(), data.next()), (Some("[DONE]"), None))
|
||||
}
|
||||
|
||||
/// SSE 解析错误。
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum SseError {
|
||||
|
||||
@@ -1,20 +1,9 @@
|
||||
use gateway_protocol::openai::sse::{
|
||||
DONE_SSE_FRAME, MAX_SSE_EVENT_BUFFER_BYTES, SseError, SseEvent, SseEventDecoder,
|
||||
encode_sse_event, encode_sse_event_with_metadata, parse_sse_events, response_failed_sse_event,
|
||||
response_failed_sse_event_with_id, sse_frame_end,
|
||||
response_failed_sse_event_with_id, sse_frame_end, sse_frame_is_done,
|
||||
};
|
||||
|
||||
fn sse_frame_is_done(frame: &str) -> bool {
|
||||
let mut data = frame.lines().filter_map(|raw_line| {
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
let (field, value) = line.split_once(':').map_or((line, ""), |(field, value)| {
|
||||
(field, value.strip_prefix(' ').unwrap_or(value))
|
||||
});
|
||||
(field == "data").then_some(value)
|
||||
});
|
||||
matches!((data.next(), data.next()), (Some("[DONE]"), None))
|
||||
}
|
||||
|
||||
fn sse_body_has_done(body: &str) -> bool {
|
||||
body.trim_end_matches(['\r', '\n'])
|
||||
.ends_with(DONE_SSE_FRAME.trim_end_matches(['\r', '\n']))
|
||||
@@ -118,6 +107,69 @@ fn incremental_decoder_should_decode_across_arbitrary_chunk_boundaries() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_frame_decoder_should_preserve_exact_frame_across_chunk_boundaries() {
|
||||
let raw = b": keep-alive\r\nid: evt_1\r\nevent: response.future\r\nretry: 250\r\ndata: { \"type\": \"response.future\", \"opaque\": true }\r\n\r\n";
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
|
||||
assert!(
|
||||
decoder
|
||||
.push_frames(&raw[..37])
|
||||
.expect("partial raw frame")
|
||||
.is_empty()
|
||||
);
|
||||
let frames = decoder.push_frames(&raw[37..]).expect("complete raw frame");
|
||||
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].raw(), raw);
|
||||
assert_eq!(
|
||||
frames[0].events()[0].event.as_deref(),
|
||||
Some("response.future")
|
||||
);
|
||||
assert_eq!(
|
||||
frames[0].events()[0].data,
|
||||
r#"{ "type": "response.future", "opaque": true }"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_frame_decoder_should_preserve_unparseable_frames_for_transparent_transport() {
|
||||
let invalid_retry = b"retry: later\ndata: opaque\n\n";
|
||||
let non_utf8 = b"data: \xff\n\n";
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
|
||||
let frames = decoder
|
||||
.push_frames(invalid_retry)
|
||||
.expect("raw frame remains deliverable");
|
||||
assert_eq!(frames[0].raw(), invalid_retry);
|
||||
assert!(frames[0].events().is_empty());
|
||||
|
||||
let frames = decoder
|
||||
.push_frames(non_utf8)
|
||||
.expect("non-UTF-8 raw frame remains deliverable");
|
||||
assert_eq!(frames[0].raw(), non_utf8);
|
||||
assert!(frames[0].events().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_frame_decoder_finish_should_preserve_unparseable_trailing_frame() {
|
||||
let raw = b"retry: later\ndata: opaque";
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
|
||||
assert!(
|
||||
decoder
|
||||
.push_frames(raw)
|
||||
.expect("unfinished raw frame remains buffered")
|
||||
.is_empty()
|
||||
);
|
||||
let frames = decoder
|
||||
.finish_frames()
|
||||
.expect("trailing raw frame remains deliverable");
|
||||
|
||||
assert_eq!(frames[0].raw(), raw);
|
||||
assert!(frames[0].events().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_decoder_finish_should_decode_unterminated_final_frame() {
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
|
||||
@@ -797,6 +797,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let failure_rate_limit_headers = response.rate_limit_headers.clone();
|
||||
let rate_limit_updates = response.rate_limit_header_updates;
|
||||
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str());
|
||||
if response.transport == CodexBackendTransport::HttpSse {
|
||||
decoder = decoder.with_raw_sse_passthrough();
|
||||
}
|
||||
loop {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::{
|
||||
fmt,
|
||||
};
|
||||
|
||||
use bytes::Bytes;
|
||||
use gateway_core::accounting::Usage;
|
||||
use gateway_core::engine::UpstreamSendState;
|
||||
use gateway_core::error::{ProviderError, ProviderErrorKind};
|
||||
@@ -13,7 +14,7 @@ use gateway_core::event::{
|
||||
ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta,
|
||||
};
|
||||
use gateway_protocol::openai::events::{TokenUsage, extract_usage};
|
||||
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder};
|
||||
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder, SseFrame, sse_frame_is_done};
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
@@ -38,6 +39,7 @@ pub struct CodexCanonicalDecoder {
|
||||
reasoning_output_seen: BTreeSet<u32>,
|
||||
usage_emitted: bool,
|
||||
semantic_output_seen: bool,
|
||||
raw_sse_passthrough: bool,
|
||||
}
|
||||
|
||||
/// 上游 Responses 事件的两类失败:协议损坏,或上游明确报告业务失败。
|
||||
@@ -130,10 +132,32 @@ impl CodexCanonicalDecoder {
|
||||
reasoning_output_seen: BTreeSet::new(),
|
||||
usage_emitted: false,
|
||||
semantic_output_seen: false,
|
||||
raw_sse_passthrough: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 让 HTTP SSE 上游帧以原始字节随 wire event 向下游传递。
|
||||
///
|
||||
/// JSON 投影仍只供 canonical facts、计费与亲和旁路读取;不会用于重建 SSE。
|
||||
#[must_use]
|
||||
pub fn with_raw_sse_passthrough(mut self) -> Self {
|
||||
self.raw_sse_passthrough = true;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
|
||||
if self.raw_sse_passthrough {
|
||||
let frames = match self.decoder.push_frames(chunk) {
|
||||
Ok(frames) => frames,
|
||||
Err(error) => {
|
||||
return self.failure(
|
||||
Vec::new(),
|
||||
CodexCanonicalError::Protocol(protocol_error(error)),
|
||||
);
|
||||
}
|
||||
};
|
||||
return self.decode_frames(frames);
|
||||
}
|
||||
let events = match self.decoder.push(chunk) {
|
||||
Ok(events) => events,
|
||||
Err(error) => {
|
||||
@@ -147,6 +171,18 @@ impl CodexCanonicalDecoder {
|
||||
}
|
||||
|
||||
pub fn finish(&mut self) -> CodexCanonicalOutcome {
|
||||
if self.raw_sse_passthrough {
|
||||
let frames = match self.decoder.finish_frames() {
|
||||
Ok(frames) => frames,
|
||||
Err(error) => {
|
||||
return self.failure(
|
||||
Vec::new(),
|
||||
CodexCanonicalError::Protocol(protocol_error(error)),
|
||||
);
|
||||
}
|
||||
};
|
||||
return self.decode_frames(frames);
|
||||
}
|
||||
let events = match self.decoder.finish() {
|
||||
Ok(events) => events,
|
||||
Err(error) => {
|
||||
@@ -162,51 +198,108 @@ impl CodexCanonicalDecoder {
|
||||
fn decode(&mut self, events: Vec<SseEvent>) -> CodexCanonicalOutcome {
|
||||
let mut output = Vec::new();
|
||||
for event in events {
|
||||
if event.data.trim() == "[DONE]" {
|
||||
if let Err(error) = self.decode_one(event, None, &mut output) {
|
||||
return self.failure(output, error);
|
||||
}
|
||||
}
|
||||
CodexCanonicalOutcome::Events(output)
|
||||
}
|
||||
|
||||
fn decode_frames(&mut self, frames: Vec<SseFrame>) -> CodexCanonicalOutcome {
|
||||
let mut output = Vec::new();
|
||||
for frame in frames {
|
||||
let done = std::str::from_utf8(frame.raw()).is_ok_and(sse_frame_is_done);
|
||||
let (raw, events) = frame.into_parts();
|
||||
if done {
|
||||
continue;
|
||||
}
|
||||
let value = match serde_json::from_str::<Value>(&event.data) {
|
||||
Ok(value) => value,
|
||||
// P1 恢复 raw bytes 后会连该帧原样交付;当前 event carrier 只能承载
|
||||
// JSON,所以先跳过这条无法投影的帧,继续读取后续上游事件。
|
||||
Err(_) => continue,
|
||||
};
|
||||
let event_type = event
|
||||
.event
|
||||
.as_deref()
|
||||
.or_else(|| value.get("type").and_then(Value::as_str));
|
||||
if matches!(event_type, Some("response.failed" | "error")) {
|
||||
return self.failure(
|
||||
output,
|
||||
CodexCanonicalError::Upstream(ResponsesSseFailure::from_event(
|
||||
event_type.unwrap_or_default(),
|
||||
&value,
|
||||
)),
|
||||
);
|
||||
let raw = Bytes::from(raw);
|
||||
if events.is_empty() {
|
||||
if let Ok(wire) = ProtocolWireEvent::raw_sse("openai", raw) {
|
||||
output.push(ProviderEvent::wire(wire));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
let mut canonical = Vec::new();
|
||||
if let Some(event_type) = event_type {
|
||||
let _ = self.decode_event(event_type, &value, &mut canonical);
|
||||
for (index, event) in events.into_iter().enumerate() {
|
||||
let raw_sse_frame = (index == 0).then(|| raw.clone());
|
||||
if let Err(error) = self.decode_one(event, raw_sse_frame, &mut output) {
|
||||
return self.failure(output, error);
|
||||
}
|
||||
}
|
||||
let semantic_output = response_event_signals(event_type, &value).semantic_output;
|
||||
let wire = match ProtocolWireEvent::json_with_sse_metadata(
|
||||
}
|
||||
CodexCanonicalOutcome::Events(output)
|
||||
}
|
||||
|
||||
fn decode_one(
|
||||
&mut self,
|
||||
event: SseEvent,
|
||||
raw_sse_frame: Option<Bytes>,
|
||||
output: &mut Vec<ProviderEvent>,
|
||||
) -> Result<(), CodexCanonicalError> {
|
||||
if event.data.trim() == "[DONE]" {
|
||||
return Ok(());
|
||||
}
|
||||
let value = match serde_json::from_str::<Value>(&event.data) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
if let Some(raw_sse_frame) = raw_sse_frame
|
||||
&& let Ok(wire) = ProtocolWireEvent::raw_sse("openai", raw_sse_frame)
|
||||
{
|
||||
output.push(ProviderEvent::wire(wire));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let event_type = event
|
||||
.event
|
||||
.as_deref()
|
||||
.or_else(|| value.get("type").and_then(Value::as_str));
|
||||
if matches!(event_type, Some("response.failed" | "error")) {
|
||||
let failure = ResponsesSseFailure::from_event(event_type.unwrap_or_default(), &value);
|
||||
if let Ok(wire) = Self::wire_for_event(event, value, raw_sse_frame) {
|
||||
output.push(ProviderEvent::wire(wire));
|
||||
}
|
||||
return Err(CodexCanonicalError::Upstream(failure));
|
||||
}
|
||||
let mut canonical = Vec::new();
|
||||
if let Some(event_type) = event_type {
|
||||
let _ = self.decode_event(event_type, &value, &mut canonical);
|
||||
}
|
||||
let semantic_output = response_event_signals(event_type, &value).semantic_output;
|
||||
let Ok(wire) = Self::wire_for_event(event, value, raw_sse_frame) else {
|
||||
return Ok(());
|
||||
};
|
||||
output.push(if canonical.is_empty() {
|
||||
ProviderEvent::wire(wire)
|
||||
} else {
|
||||
ProviderEvent::canonical_with_wire(canonical, wire)
|
||||
});
|
||||
self.semantic_output_seen |= semantic_output;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn wire_for_event(
|
||||
event: SseEvent,
|
||||
value: Value,
|
||||
raw_sse_frame: Option<Bytes>,
|
||||
) -> Result<ProtocolWireEvent, gateway_core::error::IdentifierError> {
|
||||
match raw_sse_frame {
|
||||
Some(raw_sse_frame) => ProtocolWireEvent::json_with_raw_sse_metadata(
|
||||
"openai",
|
||||
event.event,
|
||||
value,
|
||||
raw_sse_frame,
|
||||
event.id,
|
||||
event.retry,
|
||||
),
|
||||
None => ProtocolWireEvent::json_with_sse_metadata(
|
||||
"openai",
|
||||
event.event,
|
||||
value,
|
||||
event.id,
|
||||
event.retry,
|
||||
) {
|
||||
Ok(wire) => wire,
|
||||
Err(_) => continue,
|
||||
};
|
||||
output.push(if canonical.is_empty() {
|
||||
ProviderEvent::wire(wire)
|
||||
} else {
|
||||
ProviderEvent::canonical_with_wire(canonical, wire)
|
||||
});
|
||||
self.semantic_output_seen |= semantic_output;
|
||||
),
|
||||
}
|
||||
CodexCanonicalOutcome::Events(output)
|
||||
}
|
||||
|
||||
fn failure(
|
||||
|
||||
@@ -29,6 +29,57 @@ fn decoder_should_preserve_existing_metadata_fixture_as_openai_wire() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_sse_passthrough_should_keep_original_bytes_alongside_canonical_facts() {
|
||||
let heartbeat = ": keep-alive\r\n\r\n";
|
||||
let created = concat!(
|
||||
"id: evt_created\r\n",
|
||||
"event: response.created\r\n",
|
||||
"retry: 250\r\n",
|
||||
"data: { \"type\": \"response.created\", \"response\": { \"id\": \"resp_raw\", \"model\": \"gpt-test\" } }\r\n\r\n",
|
||||
);
|
||||
let body = format!("{heartbeat}{created}");
|
||||
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_raw_sse_passthrough()
|
||||
.push(body.as_bytes())
|
||||
.expect("raw SSE should remain deliverable");
|
||||
let heartbeat_wire = events[0].wire_event().expect("heartbeat wire");
|
||||
let created_wire = events[1].wire_event().expect("created wire");
|
||||
|
||||
assert!(!heartbeat_wire.has_json_data());
|
||||
assert_eq!(
|
||||
heartbeat_wire.raw_sse_frame().map(AsRef::as_ref),
|
||||
Some(heartbeat.as_bytes())
|
||||
);
|
||||
assert!(created_wire.has_json_data());
|
||||
assert_eq!(
|
||||
created_wire.raw_sse_frame().map(AsRef::as_ref),
|
||||
Some(created.as_bytes())
|
||||
);
|
||||
assert!(matches!(
|
||||
events[1].canonical_facts(),
|
||||
[GatewayEvent::Started(metadata)] if metadata.response_id() == "resp_raw"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_sse_passthrough_should_forward_unparseable_frames_without_failure() {
|
||||
let raw = b"retry: later\ndata: opaque\n\n";
|
||||
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_raw_sse_passthrough()
|
||||
.push(raw)
|
||||
.expect("unparseable frame must not abort transparent transport");
|
||||
let wire = events[0].wire_event().expect("raw wire event");
|
||||
|
||||
assert!(!wire.has_json_data());
|
||||
assert_eq!(
|
||||
wire.raw_sse_frame().map(AsRef::as_ref),
|
||||
Some(raw.as_slice())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
|
||||
let body = concat!(
|
||||
@@ -328,6 +379,29 @@ fn decoder_should_classify_official_cyber_policy_as_an_invalid_request() {
|
||||
assert_failed_event("cyber_policy", "policy-secret-marker");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_sse_failure_should_keep_the_original_frame_before_reporting_the_typed_failure() {
|
||||
let raw = concat!(
|
||||
"event: response.failed\r\n",
|
||||
"data: { \"type\": \"response.failed\", \"response\": { \"id\": \"resp_raw_failed\", \"status\": \"failed\", \"error\": { \"code\": \"rate_limit_exceeded\", \"message\": \"raw failure marker\" } } }\r\n\r\n",
|
||||
);
|
||||
|
||||
let failure = CodexCanonicalDecoder::new("fallback")
|
||||
.with_raw_sse_passthrough()
|
||||
.push(raw.as_bytes())
|
||||
.expect_err("response.failed remains a typed lifecycle failure");
|
||||
let wire = failure.events()[0]
|
||||
.wire_event()
|
||||
.expect("upstream failure wire remains deliverable");
|
||||
|
||||
assert_eq!(wire.event_type(), Some("response.failed"));
|
||||
assert_eq!(
|
||||
wire.raw_sse_frame().map(|frame| frame.as_ref()),
|
||||
Some(raw.as_bytes())
|
||||
);
|
||||
assert!(!format!("{failure:?}").contains("raw failure marker"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_preserve_same_chunk_output_before_typed_failure() {
|
||||
let marker = "same-chunk-secret-marker";
|
||||
@@ -365,7 +439,8 @@ fn decoder_should_preserve_same_chunk_output_before_typed_failure() {
|
||||
vec![
|
||||
"response.created",
|
||||
"response.content_part.added",
|
||||
"response.output_text.delta"
|
||||
"response.output_text.delta",
|
||||
"response.failed"
|
||||
]
|
||||
);
|
||||
assert!(!format!("{failure:?}").contains(marker));
|
||||
|
||||
Reference in New Issue
Block a user