mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(stream-bridge): 支持上游 sync 响应转 SSE 流式输出,记录 client/upstream 流模式差异
- 新增 sync_to_stream 桥接模块,将非 SSE 上游响应转换为 SSE 格式回传给流式客户端 - stream_pump 检测非 SSE 响应头后缓冲整包并通过桥接逻辑重写为 SSE 帧 - proxy handler 同步支持 sync→stream 聚合与转换(覆盖 openai/claude/gemini 四种格式) - sync_products 补全 openai:cli 的完整流式事件聚合(text delta、reasoning、tool call 等) - usage runtime 写入 client_requested_stream / upstream_is_stream 到 request_metadata - SQL 查询层将两个布尔字段从 request_metadata jsonb 中提取并回传给前端 - 前端 status.ts 新增 resolveUsageStreamLabelSegments,优先读取 client_requested_stream - RequestDetailDrawer 在流式转换场景下显示"客户端→上游"两段 Badge
This commit is contained in:
@@ -9,9 +9,14 @@ use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
pub(crate) mod stream;
|
||||
#[path = "sync_finalize.rs"]
|
||||
pub(crate) mod sync;
|
||||
#[path = "sync_to_stream.rs"]
|
||||
pub(crate) mod sync_to_stream;
|
||||
|
||||
pub(crate) use stream::LocalStreamRewriter;
|
||||
pub(crate) use sync::LocalCoreSyncFinalizeOutcome;
|
||||
pub(crate) use sync_to_stream::{
|
||||
maybe_bridge_standard_sync_json_to_stream, SyncToStreamBridgeOutcome,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_sync_finalize_outcome(
|
||||
trace_id: &str,
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
convert_claude_cli_response_to_openai_cli, convert_gemini_cli_response_to_openai_cli,
|
||||
convert_openai_chat_response_to_openai_cli, ClaudeClientEmitter, GeminiClientEmitter,
|
||||
OpenAIChatClientEmitter, OpenAICliClientEmitter, OpenAICliProviderState,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) struct SyncToStreamBridgeOutcome {
|
||||
pub(crate) sse_body: Vec<u8>,
|
||||
pub(crate) terminal_summary: Option<ExecutionStreamTerminalSummary>,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_bridge_standard_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
|
||||
let provider_api_format = normalize_api_format(provider_api_format);
|
||||
let client_api_format = normalize_api_format(client_api_format);
|
||||
if !is_standard_api_format(provider_api_format.as_str())
|
||||
|| !is_standard_api_format(client_api_format.as_str())
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let bridge_context = build_bridge_report_context(
|
||||
report_context,
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
);
|
||||
let Some(openai_cli_response) = convert_provider_sync_response_to_openai_cli(
|
||||
provider_body_json,
|
||||
provider_api_format.as_str(),
|
||||
&bridge_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let terminal_summary = build_terminal_summary_from_openai_cli_response(&openai_cli_response);
|
||||
let canonical_frames =
|
||||
build_canonical_frames_from_openai_cli_response(&openai_cli_response, &bridge_context)?;
|
||||
let sse_body =
|
||||
emit_client_stream_from_canonical_frames(canonical_frames, client_api_format.as_str())?;
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary,
|
||||
}))
|
||||
}
|
||||
|
||||
fn normalize_api_format(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn is_standard_api_format(value: &str) -> bool {
|
||||
matches!(
|
||||
value,
|
||||
"openai:chat"
|
||||
| "openai:cli"
|
||||
| "openai:compact"
|
||||
| "claude:chat"
|
||||
| "claude:cli"
|
||||
| "gemini:chat"
|
||||
| "gemini:cli"
|
||||
)
|
||||
}
|
||||
|
||||
fn build_bridge_report_context(
|
||||
report_context: Option<&Value>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
) -> Value {
|
||||
let mut context = report_context
|
||||
.cloned()
|
||||
.filter(Value::is_object)
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let object = context
|
||||
.as_object_mut()
|
||||
.expect("bridge report context should stay object");
|
||||
object
|
||||
.entry("provider_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(provider_api_format.to_string()));
|
||||
object
|
||||
.entry("client_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(client_api_format.to_string()));
|
||||
context
|
||||
}
|
||||
|
||||
fn convert_provider_sync_response_to_openai_cli(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
report_context: &Value,
|
||||
) -> Option<Value> {
|
||||
match provider_api_format {
|
||||
"openai:cli" | "openai:compact" => Some(provider_body_json.clone()),
|
||||
"openai:chat" => {
|
||||
convert_openai_chat_response_to_openai_cli(provider_body_json, report_context, false)
|
||||
}
|
||||
"claude:chat" | "claude:cli" => {
|
||||
convert_claude_cli_response_to_openai_cli(provider_body_json, report_context)
|
||||
}
|
||||
"gemini:chat" | "gemini:cli" => {
|
||||
convert_gemini_cli_response_to_openai_cli(provider_body_json, report_context)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_canonical_frames_from_openai_cli_response(
|
||||
openai_cli_response: &Value,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<crate::ai_pipeline::CanonicalStreamFrame>, GatewayError> {
|
||||
let mut state = OpenAICliProviderState::default();
|
||||
let line = format!(
|
||||
"data: {}\n",
|
||||
serde_json::to_string(&json!({
|
||||
"type": "response.completed",
|
||||
"response": openai_cli_response,
|
||||
}))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
);
|
||||
let mut frames = state
|
||||
.push_line(report_context, line.into_bytes())
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
frames.extend(
|
||||
state
|
||||
.finish(report_context)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
Ok(frames)
|
||||
}
|
||||
|
||||
fn emit_client_stream_from_canonical_frames(
|
||||
canonical_frames: Vec<crate::ai_pipeline::CanonicalStreamFrame>,
|
||||
client_api_format: &str,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
match client_api_format {
|
||||
"openai:chat" => {
|
||||
let mut emitter = OpenAIChatClientEmitter::default();
|
||||
emit_with_openai_chat_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"openai:cli" | "openai:compact" => {
|
||||
let mut emitter = OpenAICliClientEmitter::default();
|
||||
emit_with_openai_cli_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"claude:chat" | "claude:cli" => {
|
||||
let mut emitter = ClaudeClientEmitter::default();
|
||||
emit_with_claude_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"gemini:chat" | "gemini:cli" => {
|
||||
let mut emitter = GeminiClientEmitter::default();
|
||||
emit_with_gemini_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_with_openai_chat_emitter(
|
||||
emitter: &mut OpenAIChatClientEmitter,
|
||||
canonical_frames: Vec<crate::ai_pipeline::CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_openai_cli_emitter(
|
||||
emitter: &mut OpenAICliClientEmitter,
|
||||
canonical_frames: Vec<crate::ai_pipeline::CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_claude_emitter(
|
||||
emitter: &mut ClaudeClientEmitter,
|
||||
canonical_frames: Vec<crate::ai_pipeline::CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_gemini_emitter(
|
||||
emitter: &mut GeminiClientEmitter,
|
||||
canonical_frames: Vec<crate::ai_pipeline::CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn build_terminal_summary_from_openai_cli_response(
|
||||
openai_cli_response: &Value,
|
||||
) -> Option<ExecutionStreamTerminalSummary> {
|
||||
let response = openai_cli_response.as_object()?;
|
||||
let response_id = response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let model = response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let finish_reason = response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.map(|output| resolve_openai_cli_finish_reason(output))
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
let standardized_usage = response
|
||||
.get("usage")
|
||||
.and_then(standardized_usage_from_openai_usage);
|
||||
Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage,
|
||||
finish_reason,
|
||||
response_id,
|
||||
model,
|
||||
observed_finish: true,
|
||||
parser_error: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_openai_cli_finish_reason(output: &[Value]) -> String {
|
||||
let has_tool_calls = output.iter().filter_map(Value::as_object).any(|item| {
|
||||
item.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value == "function_call")
|
||||
});
|
||||
if has_tool_calls {
|
||||
"tool_calls".to_string()
|
||||
} else {
|
||||
"stop".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsage> {
|
||||
let usage = value.as_object()?;
|
||||
let input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_tokens = usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_creation_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let cache_read_tokens = usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage.get("total_tokens").and_then(Value::as_i64).unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
);
|
||||
let mut standardized_usage = StandardizedUsage::new();
|
||||
standardized_usage.input_tokens = input_tokens;
|
||||
standardized_usage.output_tokens = output_tokens;
|
||||
standardized_usage.cache_creation_tokens = cache_creation_tokens;
|
||||
standardized_usage.cache_read_tokens = cache_read_tokens;
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("total_tokens".to_string(), json!(total_tokens));
|
||||
Some(standardized_usage.normalize_cache_creation_breakdown())
|
||||
}
|
||||
@@ -1,11 +1,112 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_pipeline::maybe_bridge_standard_sync_json_to_stream;
|
||||
|
||||
use super::maybe_build_local_stream_rewriter;
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_sync_bridge_converts_openai_chat_sync_json_to_openai_chat_sse() {
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"id": "chatcmpl_sync_123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5.4",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello from sync bridge"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 1,
|
||||
"completion_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}),
|
||||
"openai:chat",
|
||||
"openai:chat",
|
||||
None,
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output_text = utf8(outcome.sse_body);
|
||||
assert!(output_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(output_text.contains("\"role\":\"assistant\""));
|
||||
assert!(output_text.contains("\"content\":\"Hello from sync bridge\""));
|
||||
assert!(output_text.contains("\"finish_reason\":\"stop\""));
|
||||
assert!(output_text.contains("data: [DONE]"));
|
||||
let summary = outcome
|
||||
.terminal_summary
|
||||
.expect("terminal summary should exist");
|
||||
assert_eq!(summary.response_id.as_deref(), Some("resp_sync_123"));
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-5.4"));
|
||||
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
|
||||
assert_eq!(
|
||||
summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.dimensions.get("total_tokens"))
|
||||
.cloned(),
|
||||
Some(json!(3))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_sync_bridge_converts_claude_sync_json_to_openai_cli_sse() {
|
||||
let report_context = json!({
|
||||
"mapped_model": "claude-sonnet-4-5",
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"id": "msg_sync_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "Hello from Claude sync"
|
||||
}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3
|
||||
}
|
||||
}),
|
||||
"claude:chat",
|
||||
"openai:cli",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output_text = utf8(outcome.sse_body);
|
||||
assert!(output_text.contains("event: response.created"));
|
||||
assert!(output_text.contains("event: response.output_text.delta"));
|
||||
assert!(output_text.contains("event: response.completed"));
|
||||
assert!(output_text.contains("\"text\":\"Hello from Claude sync\""));
|
||||
let summary = outcome
|
||||
.terminal_summary
|
||||
.expect("terminal summary should exist");
|
||||
assert_eq!(summary.response_id.as_deref(), Some("msg_sync_123"));
|
||||
assert_eq!(summary.model.as_deref(), Some("claude-sonnet-4-5"));
|
||||
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
|
||||
assert_eq!(
|
||||
summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.dimensions.get("total_tokens"))
|
||||
.cloned(),
|
||||
Some(json!(5))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn antigravity_stream_rewriter_unwraps_and_injects_tool_ids() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -178,7 +178,57 @@ fn aggregates_openai_cli_stream_completed_event_to_final_response() {
|
||||
"object": "response",
|
||||
"model": "gpt-5",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "resp_123_msg",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3,
|
||||
},
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregates_openai_cli_stream_tool_call_events_to_final_response() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_tool_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||
"event: response.output_item.added\n",
|
||||
"data: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"type\":\"function_call\",\"id\":\"call_123\",\"call_id\":\"call_123\",\"name\":\"get_weather\",\"arguments\":\"\",\"status\":\"in_progress\"}}\n\n",
|
||||
"event: response.function_call_arguments.delta\n",
|
||||
"data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":1,\"item_id\":\"call_123\",\"call_id\":\"call_123\",\"delta\":\"{\\\"city\\\":\\\"SF\\\"}\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_tool_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n",
|
||||
);
|
||||
|
||||
let result =
|
||||
aggregate_openai_cli_stream_sync_response(body.as_bytes()).expect("result should exist");
|
||||
|
||||
assert_eq!(
|
||||
result,
|
||||
json!({
|
||||
"id": "resp_tool_123",
|
||||
"object": "response",
|
||||
"model": "gpt-5",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "function_call",
|
||||
"id": "call_123",
|
||||
"call_id": "call_123",
|
||||
"name": "get_weather",
|
||||
"arguments": "{\"city\":\"SF\"}",
|
||||
"status": "in_progress",
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
@@ -1328,6 +1378,56 @@ fn local_finalize_handles_openai_compact_openai_family_stream_response_even_when
|
||||
assert_eq!(provider_body["status"], "completed");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_finalize_preserves_provider_stream_body_for_same_format_stream_aggregated_to_sync() {
|
||||
let body = concat!(
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello sync\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cli_samefmt_stream_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"id\":\"msg_123\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello sync\",\"annotations\":[]}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n",
|
||||
);
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: "trace-openai-cli-samefmt-stream-123".to_string(),
|
||||
report_kind: "openai_cli_sync_finalize".to_string(),
|
||||
report_context: Some(json!({
|
||||
"client_api_format": "openai:cli",
|
||||
"provider_api_format": "openai:cli",
|
||||
"model": "gpt-5.4",
|
||||
"mapped_model": "gpt-5.4",
|
||||
"needs_conversion": false,
|
||||
"has_envelope": false,
|
||||
"upstream_is_stream": true,
|
||||
})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string())]),
|
||||
body_json: None,
|
||||
client_body_json: None,
|
||||
body_base64: Some(base64::engine::general_purpose::STANDARD.encode(body.as_bytes())),
|
||||
telemetry: None,
|
||||
};
|
||||
|
||||
let outcome = maybe_build_local_core_sync_finalize_response(
|
||||
"trace-openai-cli-samefmt-stream-123",
|
||||
&test_decision(),
|
||||
&payload,
|
||||
)
|
||||
.expect("local finalize should succeed")
|
||||
.expect("local finalize should match");
|
||||
|
||||
let report = outcome
|
||||
.background_report
|
||||
.expect("same-format stream finalize should downgrade to success report");
|
||||
let client_body = report.client_body_json.expect("client body should exist");
|
||||
assert!(report.body_json.is_none());
|
||||
assert_eq!(
|
||||
report.body_base64,
|
||||
Some(base64::engine::general_purpose::STANDARD.encode(body.as_bytes()))
|
||||
);
|
||||
assert_eq!(client_body["object"], "response");
|
||||
assert_eq!(client_body["status"], "completed");
|
||||
assert_eq!(client_body["output"][0]["content"][0]["text"], "Hello sync");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_finalize_handles_openai_chat_cross_format_sync_response_from_claude() {
|
||||
let payload = GatewaySyncReportRequest {
|
||||
|
||||
@@ -19,8 +19,9 @@ pub(crate) use self::adaptation::{
|
||||
};
|
||||
pub(crate) use self::finalize::common::LocalCoreSyncFinalizeOutcome;
|
||||
pub(crate) use self::finalize::internal::{
|
||||
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,
|
||||
maybe_compile_sync_finalize_response,
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_stream_response_rewriter,
|
||||
maybe_build_sync_finalize_outcome, maybe_compile_sync_finalize_response,
|
||||
SyncToStreamBridgeOutcome,
|
||||
};
|
||||
pub(crate) use self::planner::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
|
||||
@@ -101,7 +101,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.is_kiro || resolved.is_antigravity,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -27,6 +27,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) original_request_body_json: Option<&'a Value>,
|
||||
pub(crate) original_request_body_base64: Option<&'a str>,
|
||||
pub(crate) client_requested_stream: bool,
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
pub(crate) has_envelope: bool,
|
||||
pub(crate) needs_conversion: bool,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
@@ -122,6 +123,10 @@ pub(crate) fn build_local_execution_report_context(
|
||||
"client_requested_stream".to_string(),
|
||||
Value::Bool(parts.client_requested_stream),
|
||||
);
|
||||
object.insert(
|
||||
"upstream_is_stream".to_string(),
|
||||
Value::Bool(parts.upstream_is_stream),
|
||||
);
|
||||
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
|
||||
object.insert(
|
||||
"needs_conversion".to_string(),
|
||||
|
||||
@@ -87,6 +87,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -79,6 +79,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: body_base64,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -70,6 +70,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -75,7 +75,11 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: false,
|
||||
needs_conversion: true,
|
||||
extra_fields,
|
||||
|
||||
@@ -92,7 +92,11 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: upstream_is_stream,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream,
|
||||
has_envelope: false,
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
|
||||
@@ -96,7 +96,11 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.is_antigravity,
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
|
||||
@@ -20,9 +20,9 @@ pub(crate) use crate::ai_pipeline::{
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter,
|
||||
maybe_build_sync_finalize_outcome, maybe_compile_sync_finalize_response,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
|
||||
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,
|
||||
maybe_compile_sync_finalize_response, LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
pub(crate) use aether_ai_pipeline::api::{
|
||||
build_core_error_body_for_client_format, core_error_background_report_kind,
|
||||
@@ -97,3 +97,19 @@ pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_pipeline::api::supports_stream_scheduler_decision_kind(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_openai_chat_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_pipeline::api::aggregate_openai_chat_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_openai_cli_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_pipeline::api::aggregate_openai_cli_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_claude_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_pipeline::api::aggregate_claude_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_gemini_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_pipeline::api::aggregate_gemini_stream_sync_response(body)
|
||||
}
|
||||
|
||||
@@ -39,8 +39,8 @@ use self::execution_failures::{
|
||||
handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport,
|
||||
};
|
||||
use crate::ai_pipeline_api::{
|
||||
maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter,
|
||||
normalize_provider_private_report_context,
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
|
||||
maybe_build_stream_response_rewriter, normalize_provider_private_report_context,
|
||||
};
|
||||
use crate::api::response::{
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
@@ -52,7 +52,8 @@ use crate::execution_runtime::build_direct_execution_frame_stream;
|
||||
#[cfg(test)]
|
||||
use crate::execution_runtime::remote_compat::post_stream_plan_to_remote_execution_runtime;
|
||||
use crate::execution_runtime::submission::{
|
||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
||||
resolve_core_error_background_report_kind, strip_utf8_bom_and_ws,
|
||||
submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::execution_runtime::transport::{
|
||||
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
|
||||
@@ -481,6 +482,11 @@ fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
|
||||
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
|
||||
}
|
||||
|
||||
fn parse_prefetched_sync_json_body(body: &[u8]) -> Option<Value> {
|
||||
let stripped = strip_utf8_bom_and_ws(body);
|
||||
serde_json::from_slice::<Value>(stripped).ok()
|
||||
}
|
||||
|
||||
fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Bytes, std::io::Error> {
|
||||
let payload = failure
|
||||
.to_json_string()
|
||||
@@ -913,6 +919,7 @@ async fn execute_stream_from_frame_stream(
|
||||
let mut prefetched_inspection_body = Vec::new();
|
||||
let mut prefetched_telemetry: Option<ExecutionTelemetry> = None;
|
||||
let mut reached_eof = false;
|
||||
let mut sync_json_stream_bridge_active = false;
|
||||
if skip_direct_finalize_prefetch {
|
||||
debug!(
|
||||
event_name = "execution_runtime_stream_prefetch_skipped",
|
||||
@@ -1055,6 +1062,59 @@ async fn execute_stream_from_frame_stream(
|
||||
StreamPrefetchInspection::NonError => {}
|
||||
}
|
||||
|
||||
if !response_headers_indicate_sse(&headers) && (200..300).contains(&status_code)
|
||||
{
|
||||
if let Some(body_json) =
|
||||
parse_prefetched_sync_json_body(&prefetched_inspection_body)
|
||||
{
|
||||
match maybe_bridge_standard_sync_json_to_stream(
|
||||
&body_json,
|
||||
plan.provider_api_format.as_str(),
|
||||
plan.client_api_format.as_str(),
|
||||
report_context.as_ref(),
|
||||
) {
|
||||
Ok(Some(outcome)) => {
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert(
|
||||
"content-type".to_string(),
|
||||
"text/event-stream".to_string(),
|
||||
);
|
||||
stream_terminal_summary = outcome.terminal_summary;
|
||||
prefetched_body.extend_from_slice(&outcome.sse_body);
|
||||
prefetched_chunks.push(Bytes::from(outcome.sse_body));
|
||||
sync_json_stream_bridge_active = true;
|
||||
break;
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_sync_json_stream_bridge_error",
|
||||
format!(
|
||||
"failed to bridge execution runtime sync json to stream: {err:?}"
|
||||
),
|
||||
502,
|
||||
);
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context,
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
headers,
|
||||
prefetched_telemetry,
|
||||
&provider_prefetched_body,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let normalized_chunk = if let Some(normalizer) =
|
||||
private_stream_normalizer.as_mut()
|
||||
{
|
||||
@@ -1134,7 +1194,9 @@ async fn execute_stream_from_frame_stream(
|
||||
prefetched_telemetry = Some(frame_telemetry);
|
||||
}
|
||||
StreamFramePayload::Eof { summary } => {
|
||||
stream_terminal_summary = summary;
|
||||
if summary.is_some() {
|
||||
stream_terminal_summary = summary;
|
||||
}
|
||||
reached_eof = true;
|
||||
break;
|
||||
}
|
||||
@@ -1214,6 +1276,7 @@ async fn execute_stream_from_frame_stream(
|
||||
let provider_prefetched_body_for_report = provider_prefetched_body;
|
||||
let prefetched_body_for_report = prefetched_body;
|
||||
let prefetched_chunks_for_body = prefetched_chunks;
|
||||
let sync_json_stream_bridge_active_for_report = sync_json_stream_bridge_active;
|
||||
let initial_telemetry = prefetched_telemetry;
|
||||
let initial_reached_eof = reached_eof;
|
||||
let direct_stream_finalize_kind_owned = direct_stream_finalize_kind;
|
||||
@@ -1256,10 +1319,16 @@ async fn execute_stream_from_frame_stream(
|
||||
let mut buffered_body = Vec::new();
|
||||
let mut provider_body_truncated = false;
|
||||
let mut client_body_truncated = false;
|
||||
let mut private_stream_normalizer =
|
||||
maybe_build_provider_private_stream_normalizer(report_context_owned.as_ref());
|
||||
let mut local_stream_rewriter =
|
||||
maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref());
|
||||
let mut private_stream_normalizer = if sync_json_stream_bridge_active_for_report {
|
||||
None
|
||||
} else {
|
||||
maybe_build_provider_private_stream_normalizer(report_context_owned.as_ref())
|
||||
};
|
||||
let mut local_stream_rewriter = if sync_json_stream_bridge_active_for_report {
|
||||
None
|
||||
} else {
|
||||
maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref())
|
||||
};
|
||||
append_stream_capture_bytes(
|
||||
&mut provider_buffered_body,
|
||||
&provider_prefetched_body_for_report,
|
||||
@@ -1360,6 +1429,9 @@ async fn execute_stream_from_frame_stream(
|
||||
};
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
if sync_json_stream_bridge_active_for_report {
|
||||
continue;
|
||||
}
|
||||
let chunk =
|
||||
match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) {
|
||||
Ok(chunk) => chunk,
|
||||
@@ -1488,7 +1560,9 @@ async fn execute_stream_from_frame_stream(
|
||||
telemetry = Some(frame_telemetry);
|
||||
}
|
||||
StreamFramePayload::Eof { summary } => {
|
||||
stream_terminal_summary = summary;
|
||||
if summary.is_some() {
|
||||
stream_terminal_summary = summary;
|
||||
}
|
||||
break;
|
||||
}
|
||||
StreamFramePayload::Error { error } => {
|
||||
@@ -1874,8 +1948,11 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||
use axum::body::to_bytes;
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::extract::ws::Message;
|
||||
use axum::extract::Request;
|
||||
use axum::routing::any;
|
||||
use axum::{http::header, http::HeaderValue, Router};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::watch;
|
||||
|
||||
@@ -1954,6 +2031,113 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_execution_runtime_stream_bridges_sync_json_body_from_remote_runtime_to_sse() {
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("local addr should resolve");
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(|_request: Request| async move {
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"id\\\":\\\"resp-remote-runtime-sync-json-123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[{\\\"type\\\":\\\"message\\\",\\\"id\\\":\\\"msg-remote-runtime-sync-json-123\\\",\\\"role\\\":\\\"assistant\\\",\\\"content\\\":[{\\\"type\\\":\\\"output_text\\\",\\\"text\\\":\\\"Hello from remote runtime sync json\\\",\\\"annotations\\\":[]}]}],\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = axum::http::Response::new(Body::from(frames));
|
||||
response.headers_mut().insert(
|
||||
header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_execution_runtime_override_base_url(format!("http://{addr}"));
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-remote-runtime-sync-json-stream".into(),
|
||||
candidate_id: Some("cand-remote-runtime-sync-json-stream".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://chatgpt.com/backend-api/codex/responses".into(),
|
||||
headers: BTreeMap::from([
|
||||
("content-type".into(), "application/json".into()),
|
||||
("accept".into(), "text/event-stream".into()),
|
||||
]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"stream": true
|
||||
})),
|
||||
stream: true,
|
||||
client_api_format: "openai:cli".into(),
|
||||
provider_api_format: "openai:cli".into(),
|
||||
model_name: Some("gpt-5.4".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
};
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/responses",
|
||||
Some("ai_public".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
Some("openai:cli".to_string()),
|
||||
)
|
||||
.with_execution_runtime_candidate(true);
|
||||
|
||||
let response = execute_execution_runtime_stream(
|
||||
&state,
|
||||
plan,
|
||||
"trace-remote-runtime-sync-json-stream",
|
||||
&decision,
|
||||
"openai_cli_stream",
|
||||
None,
|
||||
Some(json!({
|
||||
"provider_api_format": "openai:cli",
|
||||
"client_api_format": "openai:cli",
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.expect("execution should succeed")
|
||||
.expect("execution should return a client response");
|
||||
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body should read");
|
||||
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
|
||||
assert!(text.contains("event: response.output_text.delta"));
|
||||
assert!(text.contains("Hello from remote runtime sync json"));
|
||||
assert!(text.contains("event: response.completed"));
|
||||
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data(
|
||||
) {
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::io::Error as IoError;
|
||||
use std::time::Instant;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionStreamTerminalSummary,
|
||||
@@ -12,12 +14,13 @@ use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline_api::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
StreamingStandardTerminalObserver,
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
|
||||
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
|
||||
};
|
||||
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
|
||||
use crate::execution_runtime::transport::DirectUpstreamResponse;
|
||||
use crate::execution_runtime::DirectUpstreamStreamExecution;
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) fn build_direct_execution_frame_stream(
|
||||
execution: DirectUpstreamStreamExecution,
|
||||
@@ -57,6 +60,113 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
|
||||
let mut observer_buffered = Vec::new();
|
||||
|
||||
if !response_headers_indicate_sse(&headers) {
|
||||
let original_headers = headers.clone();
|
||||
match buffer_non_sse_upstream_body(response, started_at).await {
|
||||
Ok(buffered) => {
|
||||
let mut response_headers = original_headers;
|
||||
let mut response_body = Bytes::from(buffered.body_bytes);
|
||||
let mut summary = None;
|
||||
match maybe_bridge_non_sse_sync_json_to_stream(
|
||||
status_code,
|
||||
&response_headers,
|
||||
response_body.as_ref(),
|
||||
provider_api_format.as_str(),
|
||||
&observer_context,
|
||||
) {
|
||||
Ok(Some(outcome)) => {
|
||||
response_headers = rewrite_headers_for_bridged_sse_response(
|
||||
&response_headers,
|
||||
outcome.sse_body.len(),
|
||||
);
|
||||
response_body = Bytes::from(outcome.sse_body);
|
||||
summary = outcome.terminal_summary;
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
yield Err(IoError::other(format!("{err:?}")));
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
match encode_headers_frame(status_code, response_headers) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
if !response_body.is_empty() {
|
||||
match encode_telemetry_frame(buffered.ttfb_ms, buffered.ttfb_ms, 0) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_data_frame(&response_body) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
match encode_telemetry_frame(
|
||||
buffered.ttfb_ms,
|
||||
Some(started_at.elapsed().as_millis() as u64),
|
||||
buffered.upstream_bytes,
|
||||
) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_stream_frame_ndjson(&StreamFrame::eof_with_summary(summary)) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => yield Err(err),
|
||||
}
|
||||
}
|
||||
Err(BufferedUpstreamBodyError {
|
||||
message,
|
||||
ttfb_ms,
|
||||
upstream_bytes,
|
||||
}) => {
|
||||
match encode_headers_frame(status_code, original_headers) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_error_frame(status_code, message) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_telemetry_frame(
|
||||
ttfb_ms,
|
||||
Some(started_at.elapsed().as_millis() as u64),
|
||||
upstream_bytes,
|
||||
) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_stream_frame_ndjson(&StreamFrame::eof_with_summary(None)) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => yield Err(err),
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
match encode_headers_frame(status_code, headers) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
@@ -206,7 +316,7 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
|
||||
fn encode_headers_frame(
|
||||
status_code: u16,
|
||||
headers: std::collections::BTreeMap<String, String>,
|
||||
headers: BTreeMap<String, String>,
|
||||
) -> Result<Bytes, IoError> {
|
||||
encode_stream_frame_ndjson(&StreamFrame {
|
||||
frame_type: StreamFrameType::Headers,
|
||||
@@ -260,6 +370,178 @@ fn encode_error_frame(status_code: u16, message: String) -> Result<Bytes, IoErro
|
||||
})
|
||||
}
|
||||
|
||||
struct BufferedUpstreamBody {
|
||||
body_bytes: Vec<u8>,
|
||||
ttfb_ms: Option<u64>,
|
||||
upstream_bytes: u64,
|
||||
}
|
||||
|
||||
struct BufferedUpstreamBodyError {
|
||||
message: String,
|
||||
ttfb_ms: Option<u64>,
|
||||
upstream_bytes: u64,
|
||||
}
|
||||
|
||||
fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
|
||||
headers
|
||||
.get("content-type")
|
||||
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
|
||||
}
|
||||
|
||||
async fn buffer_non_sse_upstream_body(
|
||||
response: DirectUpstreamResponse,
|
||||
started_at: Instant,
|
||||
) -> Result<BufferedUpstreamBody, BufferedUpstreamBodyError> {
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut upstream_bytes = 0u64;
|
||||
let mut ttfb_ms = None;
|
||||
|
||||
match response {
|
||||
DirectUpstreamResponse::Reqwest(response) => {
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
while let Some(item) = bytes_stream.next().await {
|
||||
match item {
|
||||
Ok(chunk) => {
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Err(err) => {
|
||||
let message = format_error_chain(&err);
|
||||
warn!(
|
||||
event_name = "stream_pump_body_read_error",
|
||||
log_type = "ops",
|
||||
upstream_bytes,
|
||||
error = %message,
|
||||
"upstream body stream read error"
|
||||
);
|
||||
return Err(BufferedUpstreamBodyError {
|
||||
message,
|
||||
ttfb_ms,
|
||||
upstream_bytes,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
|
||||
match response.next_chunk().await {
|
||||
Ok(Some(chunk)) => {
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(message) => {
|
||||
warn!(
|
||||
event_name = "stream_pump_body_read_error",
|
||||
log_type = "ops",
|
||||
upstream_bytes,
|
||||
error = %message,
|
||||
"upstream body stream read error"
|
||||
);
|
||||
return Err(BufferedUpstreamBodyError {
|
||||
message,
|
||||
ttfb_ms,
|
||||
upstream_bytes,
|
||||
});
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
Ok(BufferedUpstreamBody {
|
||||
body_bytes,
|
||||
ttfb_ms,
|
||||
upstream_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
fn maybe_bridge_non_sse_sync_json_to_stream(
|
||||
status_code: u16,
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &[u8],
|
||||
provider_api_format: &str,
|
||||
report_context: &Value,
|
||||
) -> Result<Option<crate::ai_pipeline::SyncToStreamBridgeOutcome>, GatewayError> {
|
||||
if !(200..300).contains(&status_code) || body_bytes.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let decoded_body_bytes = decode_non_sse_response_body_bytes(headers, body_bytes)
|
||||
.unwrap_or_else(|| body_bytes.to_vec());
|
||||
if !response_body_is_json(headers, &decoded_body_bytes) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let client_api_format = report_context
|
||||
.get("client_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(provider_api_format);
|
||||
maybe_bridge_standard_sync_json_to_stream(
|
||||
&body_json,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
Some(report_context),
|
||||
)
|
||||
}
|
||||
|
||||
fn rewrite_headers_for_bridged_sse_response(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_len: usize,
|
||||
) -> BTreeMap<String, String> {
|
||||
let mut rewritten = headers.clone();
|
||||
rewritten.remove("content-encoding");
|
||||
rewritten.insert("content-type".to_string(), "text/event-stream".to_string());
|
||||
rewritten.insert("content-length".to_string(), body_len.to_string());
|
||||
rewritten
|
||||
}
|
||||
|
||||
fn decode_non_sse_response_body_bytes(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &[u8],
|
||||
) -> Option<Vec<u8>> {
|
||||
let encoding = headers
|
||||
.get("content-encoding")
|
||||
.map(String::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase());
|
||||
match encoding.as_deref() {
|
||||
Some("gzip") => {
|
||||
let mut decoder = flate2::read::GzDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
std::io::Read::read_to_end(&mut decoder, &mut out).ok()?;
|
||||
Some(out)
|
||||
}
|
||||
Some("deflate") => {
|
||||
let mut decoder = flate2::read::DeflateDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
std::io::Read::read_to_end(&mut decoder, &mut out).ok()?;
|
||||
Some(out)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
if headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.is_some_and(|value| value.contains("json"))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
serde_json::from_slice::<Value>(body_bytes).is_ok()
|
||||
}
|
||||
|
||||
fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
|
||||
let mut message = err.to_string();
|
||||
let mut source = err.source();
|
||||
@@ -361,6 +643,7 @@ mod tests {
|
||||
use axum::extract::ws::Message;
|
||||
use axum::routing::post;
|
||||
use axum::{http::header, http::HeaderValue, Router};
|
||||
use base64::Engine as _;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::watch;
|
||||
@@ -568,6 +851,142 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_execution_frame_stream_bridges_sync_json_body_to_sse_for_standard_stream_request(
|
||||
) {
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("local addr should resolve");
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new().route(
|
||||
"/responses",
|
||||
post(|| async {
|
||||
let body = serde_json::json!({
|
||||
"id": "resp_sync_bridge_123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg_sync_bridge_123",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello from buffered JSON stream",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
});
|
||||
let mut response = axum::http::Response::new(Body::from(
|
||||
serde_json::to_vec(&body).expect("json should encode"),
|
||||
));
|
||||
response.headers_mut().insert(
|
||||
header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
|
||||
let runtime = DirectSyncExecutionRuntime::new();
|
||||
let execution = runtime
|
||||
.execute_stream(&ExecutionPlan {
|
||||
request_id: "req-sync-bridge".to_string(),
|
||||
candidate_id: Some("cand-sync-bridge".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!("http://{addr}/responses"),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(serde_json::json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"stream": true
|
||||
})),
|
||||
stream: true,
|
||||
client_api_format: "openai:cli".to_string(),
|
||||
provider_api_format: "openai:cli".to_string(),
|
||||
model_name: Some("gpt-5.4".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("stream execution should succeed");
|
||||
|
||||
let frames = build_direct_execution_frame_stream(execution)
|
||||
.map(|item| item.expect("frame should encode"))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|bytes| String::from_utf8(bytes.to_vec()).expect("frame should be utf8"))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
server.abort();
|
||||
|
||||
let header_frame: Value =
|
||||
serde_json::from_str(&frames[0]).expect("headers frame should parse");
|
||||
assert_eq!(
|
||||
header_frame
|
||||
.get("payload")
|
||||
.and_then(|payload| payload.get("headers"))
|
||||
.and_then(|headers| headers.get("content-type"))
|
||||
.and_then(Value::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
|
||||
let data_frame = frames
|
||||
.iter()
|
||||
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
|
||||
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("data"))
|
||||
.expect("data frame should exist");
|
||||
let bridged_body = base64::engine::general_purpose::STANDARD
|
||||
.decode(
|
||||
data_frame
|
||||
.get("payload")
|
||||
.and_then(|payload| payload.get("chunk_b64"))
|
||||
.and_then(Value::as_str)
|
||||
.expect("chunk_b64 should exist"),
|
||||
)
|
||||
.expect("data frame should decode");
|
||||
let bridged_text = String::from_utf8(bridged_body).expect("bridged body should be utf8");
|
||||
assert!(bridged_text.contains("event: response.output_text.delta"));
|
||||
assert!(bridged_text.contains("\"delta\":\"Hello from buffered JSON stream\""));
|
||||
assert!(bridged_text.contains("event: response.completed"));
|
||||
|
||||
let eof_frame = frames
|
||||
.iter()
|
||||
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
|
||||
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
|
||||
.expect("eof frame should exist");
|
||||
assert_eq!(
|
||||
eof_frame
|
||||
.get("payload")
|
||||
.and_then(|payload| payload.get("summary"))
|
||||
.and_then(|summary| summary.get("response_id"))
|
||||
.and_then(Value::as_str),
|
||||
Some("resp_sync_bridge_123")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_execution_frame_stream_preserves_local_tunnel_stream_error_message() {
|
||||
let state = AppState::new().expect("app state should build");
|
||||
|
||||
@@ -327,6 +327,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
|
||||
fn build_stream_summary_report_context(plan: &ExecutionPlan) -> Value {
|
||||
json!({
|
||||
"provider_api_format": plan.provider_api_format,
|
||||
"client_api_format": plan.client_api_format,
|
||||
"model": plan.model_name,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,10 +5,15 @@ use self::local::{
|
||||
};
|
||||
use super::internal::resolve_local_proxy_execution_path;
|
||||
pub(crate) use super::public::matches_model_mapping_for_models;
|
||||
use crate::ai_pipeline_api;
|
||||
use crate::ai_pipeline_api::{
|
||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_cli_stream_sync_response,
|
||||
maybe_bridge_standard_sync_json_to_stream,
|
||||
};
|
||||
use crate::api::response::{
|
||||
build_client_response, build_local_auth_rejection_response, build_local_http_error_response,
|
||||
build_local_overloaded_response, build_local_user_rpm_limited_response,
|
||||
build_client_response, build_client_response_from_parts, build_local_auth_rejection_response,
|
||||
build_local_http_error_response, build_local_overloaded_response,
|
||||
build_local_user_rpm_limited_response,
|
||||
};
|
||||
use crate::constants::{
|
||||
DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
|
||||
@@ -53,7 +58,7 @@ use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::extract::{ConnectInfo, Request, State};
|
||||
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::Instant;
|
||||
use std::{collections::BTreeMap, time::Instant};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
|
||||
@@ -406,8 +411,13 @@ async fn maybe_forward_public_request_to_tunnel_owner(
|
||||
message: format!("owner gateway affinity forward failed: {err}"),
|
||||
})?;
|
||||
|
||||
let mut response =
|
||||
build_client_response(upstream_response, &request_context.trace_id, Some(decision))?;
|
||||
let mut response = build_sync_aware_affinity_forward_response(
|
||||
request_context,
|
||||
buffered_body,
|
||||
decision,
|
||||
upstream_response,
|
||||
)
|
||||
.await?;
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER),
|
||||
HeaderValue::from_str(owner.gateway_instance_id.as_str())
|
||||
@@ -416,6 +426,219 @@ async fn maybe_forward_public_request_to_tunnel_owner(
|
||||
Ok(Some(response))
|
||||
}
|
||||
|
||||
fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool {
|
||||
headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
|
||||
}
|
||||
|
||||
fn collect_upstream_response_headers(
|
||||
headers: &reqwest::header::HeaderMap,
|
||||
) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
(
|
||||
name.as_str().to_string(),
|
||||
value.to_str().unwrap_or_default().to_string(),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn aggregate_sync_sse_response_for_client(
|
||||
decision: &GatewayControlDecision,
|
||||
public_path: &str,
|
||||
body: &[u8],
|
||||
) -> Option<serde_json::Value> {
|
||||
let api_format = decision
|
||||
.auth_endpoint_signature
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
match api_format {
|
||||
Some(value) if value.eq_ignore_ascii_case("openai:chat") => {
|
||||
aggregate_openai_chat_stream_sync_response(body)
|
||||
}
|
||||
Some(value)
|
||||
if value.eq_ignore_ascii_case("openai:cli")
|
||||
|| value.eq_ignore_ascii_case("openai:compact") =>
|
||||
{
|
||||
aggregate_openai_cli_stream_sync_response(body)
|
||||
}
|
||||
Some(value)
|
||||
if value.eq_ignore_ascii_case("claude:chat")
|
||||
|| value.eq_ignore_ascii_case("claude:cli") =>
|
||||
{
|
||||
aggregate_claude_stream_sync_response(body)
|
||||
}
|
||||
Some(value)
|
||||
if value.eq_ignore_ascii_case("gemini:chat")
|
||||
|| value.eq_ignore_ascii_case("gemini:cli") =>
|
||||
{
|
||||
aggregate_gemini_stream_sync_response(body)
|
||||
}
|
||||
_ if public_path == "/v1/chat/completions" => {
|
||||
aggregate_openai_chat_stream_sync_response(body)
|
||||
}
|
||||
_ if public_path == "/v1/responses" || public_path == "/v1/responses/compact" => {
|
||||
aggregate_openai_cli_stream_sync_response(body)
|
||||
}
|
||||
_ if public_path == "/v1/messages" => aggregate_claude_stream_sync_response(body),
|
||||
_ if decision.route_family.as_deref() == Some("gemini")
|
||||
&& (public_path.contains(":generateContent")
|
||||
|| public_path.contains(":streamGenerateContent")) =>
|
||||
{
|
||||
aggregate_gemini_stream_sync_response(body)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_sync_json_proxy_response(
|
||||
status_code: u16,
|
||||
upstream_headers: &BTreeMap<String, String>,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let mut headers = upstream_headers.clone();
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
let body_bytes =
|
||||
serde_json::to_vec(body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
fn resolve_affinity_forward_client_api_format(
|
||||
decision: &GatewayControlDecision,
|
||||
public_path: &str,
|
||||
) -> Option<&'static str> {
|
||||
let api_format = decision
|
||||
.auth_endpoint_signature
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
match api_format {
|
||||
Some(value) if value.eq_ignore_ascii_case("openai:chat") => Some("openai:chat"),
|
||||
Some(value) if value.eq_ignore_ascii_case("openai:cli") => Some("openai:cli"),
|
||||
Some(value) if value.eq_ignore_ascii_case("openai:compact") => Some("openai:compact"),
|
||||
Some(value) if value.eq_ignore_ascii_case("claude:chat") => Some("claude:chat"),
|
||||
Some(value) if value.eq_ignore_ascii_case("claude:cli") => Some("claude:cli"),
|
||||
Some(value) if value.eq_ignore_ascii_case("gemini:chat") => Some("gemini:chat"),
|
||||
Some(value) if value.eq_ignore_ascii_case("gemini:cli") => Some("gemini:cli"),
|
||||
_ if public_path == "/v1/chat/completions" => Some("openai:chat"),
|
||||
_ if public_path == "/v1/responses" => Some("openai:cli"),
|
||||
_ if public_path == "/v1/responses/compact" => Some("openai:compact"),
|
||||
_ if public_path == "/v1/messages" => Some("claude:chat"),
|
||||
_ if decision.route_family.as_deref() == Some("gemini")
|
||||
&& (public_path.contains(":generateContent")
|
||||
|| public_path.contains(":streamGenerateContent")) =>
|
||||
{
|
||||
Some("gemini:chat")
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_stream_sse_proxy_response(
|
||||
status_code: u16,
|
||||
upstream_headers: &BTreeMap<String, String>,
|
||||
sse_body: &[u8],
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let mut headers = upstream_headers.clone();
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "text/event-stream".to_string());
|
||||
headers.insert("content-length".to_string(), sse_body.len().to_string());
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from(sse_body.to_vec()),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
async fn build_sync_aware_affinity_forward_response(
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
buffered_body: Option<&Bytes>,
|
||||
decision: &GatewayControlDecision,
|
||||
upstream_response: reqwest::Response,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let Some(buffered_body) = buffered_body else {
|
||||
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
|
||||
};
|
||||
let stream_request = request_wants_stream(request_context, buffered_body);
|
||||
let upstream_is_sse = upstream_response_is_sse(upstream_response.headers());
|
||||
if (!stream_request && !upstream_is_sse) || (stream_request && upstream_is_sse) {
|
||||
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
|
||||
}
|
||||
|
||||
let status_code = upstream_response.status().as_u16();
|
||||
let headers = collect_upstream_response_headers(upstream_response.headers());
|
||||
let body_bytes = upstream_response
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if stream_request {
|
||||
if (200..300).contains(&status_code) {
|
||||
if let Some(client_api_format) = resolve_affinity_forward_client_api_format(
|
||||
decision,
|
||||
request_context.request_path.as_str(),
|
||||
) {
|
||||
if let Ok(body_json) = serde_json::from_slice::<serde_json::Value>(&body_bytes) {
|
||||
if let Some(outcome) = maybe_bridge_standard_sync_json_to_stream(
|
||||
&body_json,
|
||||
client_api_format,
|
||||
client_api_format,
|
||||
None,
|
||||
)? {
|
||||
return build_stream_sse_proxy_response(
|
||||
status_code,
|
||||
&headers,
|
||||
&outcome.sse_body,
|
||||
&request_context.trace_id,
|
||||
decision,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if let Some(body_json) = aggregate_sync_sse_response_for_client(
|
||||
decision,
|
||||
request_context.request_path.as_str(),
|
||||
&body_bytes,
|
||||
) {
|
||||
return build_sync_json_proxy_response(
|
||||
status_code,
|
||||
&headers,
|
||||
&body_json,
|
||||
&request_context.trace_id,
|
||||
decision,
|
||||
);
|
||||
}
|
||||
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
&request_context.trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn proxy_request(
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,
|
||||
|
||||
@@ -232,6 +232,73 @@ fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUs
|
||||
!body.contains_key("stream") && users_me_usage_api_format_defaults_to_non_stream(item)
|
||||
}
|
||||
|
||||
fn users_me_usage_headers_stream_flag(headers: Option<&serde_json::Value>) -> Option<bool> {
|
||||
let object = headers.and_then(serde_json::Value::as_object)?;
|
||||
let raw = object
|
||||
.iter()
|
||||
.find(|(key, _)| key.eq_ignore_ascii_case("content-type"))
|
||||
.and_then(|(_, value)| match value {
|
||||
serde_json::Value::String(text) => Some(text.as_str()),
|
||||
serde_json::Value::Array(values) => values.iter().find_map(serde_json::Value::as_str),
|
||||
_ => None,
|
||||
})?
|
||||
.trim();
|
||||
if raw.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let normalized = raw.to_ascii_lowercase();
|
||||
Some(
|
||||
normalized.contains("event-stream")
|
||||
|| normalized.contains("eventstream")
|
||||
|| normalized.contains("x-ndjson"),
|
||||
)
|
||||
}
|
||||
|
||||
fn users_me_usage_body_is_sse_capture(value: Option<&serde_json::Value>) -> bool {
|
||||
let Some(object) = value.and_then(serde_json::Value::as_object) else {
|
||||
return false;
|
||||
};
|
||||
object
|
||||
.get("chunks")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.is_some()
|
||||
&& object
|
||||
.get("metadata")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|metadata| metadata.get("stream"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn users_me_usage_infer_client_stream_from_captured_bodies(
|
||||
item: &StoredRequestUsageAudit,
|
||||
) -> Option<bool> {
|
||||
let provider_stream = users_me_usage_body_is_sse_capture(item.response_body.as_ref());
|
||||
let client_stream = users_me_usage_body_is_sse_capture(item.client_response_body.as_ref());
|
||||
if client_stream {
|
||||
Some(true)
|
||||
} else if provider_stream && item.client_response_body.is_some() {
|
||||
Some(false)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn users_me_usage_infer_upstream_stream_from_captured_bodies(
|
||||
item: &StoredRequestUsageAudit,
|
||||
) -> Option<bool> {
|
||||
let provider_stream = users_me_usage_body_is_sse_capture(item.response_body.as_ref());
|
||||
let client_stream = users_me_usage_body_is_sse_capture(item.client_response_body.as_ref());
|
||||
if provider_stream {
|
||||
Some(true)
|
||||
} else if client_stream && item.response_body.is_some() {
|
||||
Some(false)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn users_me_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
item.request_metadata
|
||||
.as_ref()
|
||||
@@ -239,7 +306,20 @@ fn users_me_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
.and_then(|metadata| metadata.get("client_requested_stream"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.or_else(|| users_me_usage_request_body_stream_flag(item))
|
||||
.or_else(|| users_me_usage_headers_stream_flag(item.client_response_headers.as_ref()))
|
||||
.or_else(|| users_me_usage_request_body_implies_default_non_stream(item).then_some(false))
|
||||
.or_else(|| users_me_usage_infer_client_stream_from_captured_bodies(item))
|
||||
.unwrap_or(item.is_stream)
|
||||
}
|
||||
|
||||
fn users_me_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
item.request_metadata
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|metadata| metadata.get("upstream_is_stream"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.or_else(|| users_me_usage_headers_stream_flag(item.response_headers.as_ref()))
|
||||
.or_else(|| users_me_usage_infer_upstream_stream_from_captured_bodies(item))
|
||||
.unwrap_or(item.is_stream)
|
||||
}
|
||||
|
||||
@@ -256,6 +336,7 @@ fn build_users_me_usage_record_payload(
|
||||
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
|
||||
let rate_multiplier = item.settlement_rate_multiplier();
|
||||
let client_is_stream = users_me_usage_client_is_stream(item);
|
||||
let upstream_is_stream = users_me_usage_upstream_is_stream(item);
|
||||
let mut payload = json!({
|
||||
"id": item.id,
|
||||
"model": item.model,
|
||||
@@ -271,7 +352,7 @@ fn build_users_me_usage_record_payload(
|
||||
"response_time_ms": item.response_time_ms,
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"is_stream": item.is_stream,
|
||||
"upstream_is_stream": item.is_stream,
|
||||
"upstream_is_stream": upstream_is_stream,
|
||||
"client_requested_stream": client_is_stream,
|
||||
"client_is_stream": client_is_stream,
|
||||
"status": item.status,
|
||||
@@ -307,6 +388,7 @@ fn build_users_me_usage_record_payload(
|
||||
fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_json::Value {
|
||||
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
|
||||
let client_is_stream = users_me_usage_client_is_stream(item);
|
||||
let upstream_is_stream = users_me_usage_upstream_is_stream(item);
|
||||
let mut payload = json!({
|
||||
"id": item.id,
|
||||
"status": item.status,
|
||||
@@ -325,7 +407,7 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
|
||||
"api_format": item.api_format,
|
||||
"endpoint_api_format": item.endpoint_api_format,
|
||||
"is_stream": item.is_stream,
|
||||
"upstream_is_stream": item.is_stream,
|
||||
"upstream_is_stream": upstream_is_stream,
|
||||
"client_requested_stream": client_is_stream,
|
||||
"client_is_stream": client_is_stream,
|
||||
"has_format_conversion": item.has_format_conversion,
|
||||
@@ -1151,7 +1233,7 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
build_users_me_usage_active_payload, build_users_me_usage_record_payload,
|
||||
users_me_usage_client_is_stream,
|
||||
users_me_usage_client_is_stream, users_me_usage_upstream_is_stream,
|
||||
};
|
||||
|
||||
fn sample_usage(status: &str) -> StoredRequestUsageAudit {
|
||||
@@ -1303,4 +1385,105 @@ mod tests {
|
||||
assert_eq!(active_payload["client_requested_stream"], false);
|
||||
assert_eq!(active_payload["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_upstream_stream_prefers_request_metadata_flag() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: false,
|
||||
request_metadata: Some(json!({
|
||||
"client_requested_stream": false,
|
||||
"upstream_is_stream": true
|
||||
})),
|
||||
..sample_usage("completed")
|
||||
};
|
||||
|
||||
let record_payload =
|
||||
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
|
||||
assert_eq!(record_payload["is_stream"], false);
|
||||
assert_eq!(record_payload["upstream_is_stream"], true);
|
||||
assert_eq!(record_payload["client_requested_stream"], false);
|
||||
assert_eq!(record_payload["client_is_stream"], false);
|
||||
|
||||
let active_payload = build_users_me_usage_active_payload(&item);
|
||||
assert_eq!(active_payload["is_stream"], false);
|
||||
assert_eq!(active_payload["upstream_is_stream"], true);
|
||||
assert_eq!(active_payload["client_requested_stream"], false);
|
||||
assert_eq!(active_payload["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_stream_modes_fall_back_to_captured_response_bodies_when_request_metadata_is_missing(
|
||||
) {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: true,
|
||||
response_body: Some(json!({
|
||||
"chunks": [
|
||||
{"type": "response.created"},
|
||||
{"type": "response.output_text.delta", "delta": "Hello"}
|
||||
],
|
||||
"metadata": {
|
||||
"stream": true,
|
||||
"stored_chunks": 2,
|
||||
"total_chunks": 2
|
||||
}
|
||||
})),
|
||||
client_response_body: Some(json!({
|
||||
"id": "resp-1",
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"output": []
|
||||
})),
|
||||
..sample_usage("completed")
|
||||
};
|
||||
|
||||
assert!(!users_me_usage_client_is_stream(&item));
|
||||
assert!(users_me_usage_upstream_is_stream(&item));
|
||||
|
||||
let record_payload =
|
||||
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
|
||||
assert_eq!(record_payload["is_stream"], true);
|
||||
assert_eq!(record_payload["upstream_is_stream"], true);
|
||||
assert_eq!(record_payload["client_requested_stream"], false);
|
||||
assert_eq!(record_payload["client_is_stream"], false);
|
||||
|
||||
let active_payload = build_users_me_usage_active_payload(&item);
|
||||
assert_eq!(active_payload["is_stream"], true);
|
||||
assert_eq!(active_payload["upstream_is_stream"], true);
|
||||
assert_eq!(active_payload["client_requested_stream"], false);
|
||||
assert_eq!(active_payload["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_stream_modes_fall_back_to_captured_response_headers_when_bodies_are_detached() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: true,
|
||||
response_headers: Some(json!({
|
||||
"content-type": "text/event-stream; charset=utf-8"
|
||||
})),
|
||||
client_response_headers: Some(json!({
|
||||
"content-type": "application/json"
|
||||
})),
|
||||
response_body_ref: Some("usage://request/req-1/response_body".to_string()),
|
||||
client_response_body_ref: Some(
|
||||
"usage://request/req-1/client_response_body".to_string(),
|
||||
),
|
||||
..sample_usage("completed")
|
||||
};
|
||||
|
||||
assert!(!users_me_usage_client_is_stream(&item));
|
||||
assert!(users_me_usage_upstream_is_stream(&item));
|
||||
|
||||
let record_payload =
|
||||
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
|
||||
assert_eq!(record_payload["is_stream"], true);
|
||||
assert_eq!(record_payload["upstream_is_stream"], true);
|
||||
assert_eq!(record_payload["client_requested_stream"], false);
|
||||
assert_eq!(record_payload["client_is_stream"], false);
|
||||
|
||||
let active_payload = build_users_me_usage_active_payload(&item);
|
||||
assert_eq!(active_payload["is_stream"], true);
|
||||
assert_eq!(active_payload["upstream_is_stream"], true);
|
||||
assert_eq!(active_payload["client_requested_stream"], false);
|
||||
assert_eq!(active_payload["client_is_stream"], false);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,7 +395,17 @@ async fn gateway_executes_openai_compact_openai_family_upstream_stream_via_local
|
||||
"object": "response",
|
||||
"model": "gpt-5",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "resp_compact_openai_family_123_msg",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello Compact",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3,
|
||||
|
||||
@@ -387,7 +387,17 @@ async fn gateway_executes_openai_cli_sync_upstream_stream_via_local_finalize_res
|
||||
"object": "response",
|
||||
"model": "gpt-5-upstream",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "resp_stream_001_msg",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
|
||||
@@ -23,6 +23,7 @@ use aether_data_contracts::repository::candidates::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use base64::Engine as _;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -3380,21 +3381,19 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref
|
||||
"request_id": "trace-codex-cli-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
"content-type": "text/event-stream"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "resp-codex-local-123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}
|
||||
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-codex-local-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello from Codex\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-codex-local-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
|
||||
)
|
||||
)
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 37
|
||||
@@ -3459,6 +3458,10 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(response_json["id"], "resp-codex-local-123");
|
||||
assert_eq!(
|
||||
response_json["output"][0]["content"][0]["text"],
|
||||
"Hello from Codex"
|
||||
);
|
||||
|
||||
let seen_refresh_request = seen_refresh
|
||||
.lock()
|
||||
|
||||
@@ -57,6 +57,37 @@ fn sample_auth_snapshot(
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_cli_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
allowed_model: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
Some(serde_json::json!([allowed_model])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
Some(serde_json::json!([allowed_model])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
provider_id.to_string(),
|
||||
@@ -91,6 +122,29 @@ fn sample_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalo
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_codex_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
endpoint_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_key(key_id: &str, provider_id: &str, node_id: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
key_id.to_string(),
|
||||
@@ -119,6 +173,34 @@ fn sample_key(key_id: &str, provider_id: &str, node_id: &str) -> StoredProviderC
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
fn sample_codex_key(key_id: &str, provider_id: &str, node_id: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
key_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["openai:cli"])),
|
||||
"plain-upstream-key".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(json!({
|
||||
"enabled": true,
|
||||
"mode": "tunnel",
|
||||
"node_id": node_id,
|
||||
})),
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
fn tunnel_attachment_key(node_id: &str) -> String {
|
||||
format!("tunnel.attachments.{node_id}")
|
||||
}
|
||||
@@ -536,3 +618,471 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
|
||||
owner_handle.abort();
|
||||
fallback_probe_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_aggregates_sync_sse_from_remote_tunnel_owner_before_returning_to_client() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenOwnerRequest {
|
||||
path: String,
|
||||
body: String,
|
||||
trace_id: String,
|
||||
gateway_marker: String,
|
||||
trusted_user_id: String,
|
||||
trusted_api_key_id: String,
|
||||
trusted_access_allowed: String,
|
||||
forwarded_by: String,
|
||||
owner_instance_id: String,
|
||||
}
|
||||
|
||||
let seen_owner = Arc::new(Mutex::new(None::<SeenOwnerRequest>));
|
||||
let seen_owner_clone = Arc::clone(&seen_owner);
|
||||
let owner = Router::new().route(
|
||||
"/v1/responses",
|
||||
any(move |request: Request| {
|
||||
let seen_owner_inner = Arc::clone(&seen_owner_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest {
|
||||
path: parts
|
||||
.uri
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str())
|
||||
.unwrap_or("/")
|
||||
.to_string(),
|
||||
body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"),
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
gateway_marker: parts
|
||||
.headers
|
||||
.get(GATEWAY_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_user_id: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_USER_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_api_key_id: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_API_KEY_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_access_allowed: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
forwarded_by: parts
|
||||
.headers
|
||||
.get(TUNNEL_AFFINITY_FORWARDED_BY_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
owner_instance_id: parts
|
||||
.headers
|
||||
.get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-codex-affinity-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello from Codex\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-codex-affinity-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
|
||||
)))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::CACHE_CONTROL,
|
||||
HeaderValue::from_static("no-cache"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::HeaderName::from_static(GATEWAY_HEADER),
|
||||
HeaderValue::from_static("gateway-b-owner"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-cli-affinity")),
|
||||
sample_cli_auth_snapshot("api-key-affinity-cli-1", "user-affinity-cli-1", "gpt-5.4"),
|
||||
)]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-cli-owner")],
|
||||
vec![sample_codex_endpoint(
|
||||
"endpoint-cli-owner",
|
||||
"provider-cli-owner",
|
||||
)],
|
||||
vec![sample_codex_key(
|
||||
"key-cli-owner",
|
||||
"provider-cli-owner",
|
||||
"node-cli-owner",
|
||||
)],
|
||||
));
|
||||
let observed_at_unix_secs = current_unix_secs();
|
||||
let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
"development-key",
|
||||
)
|
||||
.with_auth_api_key_reader(auth_repository)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
tunnel_attachment_key("node-cli-owner"),
|
||||
serde_json::to_value(crate::tunnel::TunnelAttachmentRecord {
|
||||
gateway_instance_id: "gateway-b".to_string(),
|
||||
relay_base_url: owner_url.clone(),
|
||||
conn_count: 1,
|
||||
observed_at_unix_secs,
|
||||
})
|
||||
.expect("attachment should serialize"),
|
||||
)]);
|
||||
|
||||
let mut state = AppState::new().expect("gateway state should build");
|
||||
state = state
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
state.scheduler_affinity_cache.insert(
|
||||
"scheduler_affinity:api-key-affinity-cli-1:openai:cli:gpt-5.4".to_string(),
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
provider_id: "provider-cli-owner".to_string(),
|
||||
endpoint_id: "endpoint-cli-owner".to_string(),
|
||||
key_id: "key-cli-owner".to_string(),
|
||||
},
|
||||
Duration::from_secs(300),
|
||||
100,
|
||||
);
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-cli-affinity",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-tunnel-affinity-cli-sync-1")
|
||||
.json(&json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"stream": false
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(GATEWAY_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-b-owner")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("tunnel_affinity_forward")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-b")
|
||||
);
|
||||
assert!(response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|value| value.contains("application/json")));
|
||||
let body: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(body["id"], "resp-codex-affinity-123");
|
||||
assert_eq!(body["object"], "response");
|
||||
assert_eq!(body["status"], "completed");
|
||||
assert_eq!(body["output"][0]["content"][0]["text"], "Hello from Codex");
|
||||
assert_eq!(body["usage"]["total_tokens"], 3);
|
||||
|
||||
let owner_request = seen_owner
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("owner request should be captured");
|
||||
assert_eq!(owner_request.path, "/v1/responses");
|
||||
let owner_body: serde_json::Value =
|
||||
serde_json::from_str(&owner_request.body).expect("owner body should parse");
|
||||
assert_eq!(owner_body["model"], "gpt-5.4");
|
||||
assert_eq!(owner_body["stream"], false);
|
||||
assert_eq!(owner_request.trace_id, "trace-tunnel-affinity-cli-sync-1");
|
||||
assert_eq!(owner_request.gateway_marker, "rust-phase3b-affinity");
|
||||
assert_eq!(owner_request.trusted_user_id, "user-affinity-cli-1");
|
||||
assert_eq!(owner_request.trusted_api_key_id, "api-key-affinity-cli-1");
|
||||
assert_eq!(owner_request.trusted_access_allowed, "true");
|
||||
assert_eq!(owner_request.forwarded_by, "gateway-a");
|
||||
assert_eq!(owner_request.owner_instance_id, "gateway-b");
|
||||
|
||||
gateway_handle.abort();
|
||||
owner_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning_to_client() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenOwnerRequest {
|
||||
path: String,
|
||||
body: String,
|
||||
trace_id: String,
|
||||
gateway_marker: String,
|
||||
trusted_user_id: String,
|
||||
trusted_api_key_id: String,
|
||||
trusted_access_allowed: String,
|
||||
forwarded_by: String,
|
||||
owner_instance_id: String,
|
||||
}
|
||||
|
||||
let seen_owner = Arc::new(Mutex::new(None::<SeenOwnerRequest>));
|
||||
let seen_owner_clone = Arc::clone(&seen_owner);
|
||||
let owner = Router::new().route(
|
||||
"/v1/responses",
|
||||
any(move |request: Request| {
|
||||
let seen_owner_inner = Arc::clone(&seen_owner_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest {
|
||||
path: parts
|
||||
.uri
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str())
|
||||
.unwrap_or("/")
|
||||
.to_string(),
|
||||
body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"),
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
gateway_marker: parts
|
||||
.headers
|
||||
.get(GATEWAY_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_user_id: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_USER_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_api_key_id: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_API_KEY_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_access_allowed: parts
|
||||
.headers
|
||||
.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
forwarded_by: parts
|
||||
.headers
|
||||
.get(TUNNEL_AFFINITY_FORWARDED_BY_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
owner_instance_id: parts
|
||||
.headers
|
||||
.get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&json!({
|
||||
"id": "resp-codex-affinity-stream-123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg-codex-affinity-stream-123",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello from affinity sync json",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("body should encode"),
|
||||
))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::HeaderName::from_static(GATEWAY_HEADER),
|
||||
HeaderValue::from_static("gateway-b-owner"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-cli-affinity")),
|
||||
sample_cli_auth_snapshot("api-key-affinity-cli-1", "user-affinity-cli-1", "gpt-5.4"),
|
||||
)]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-cli-owner")],
|
||||
vec![sample_codex_endpoint(
|
||||
"endpoint-cli-owner",
|
||||
"provider-cli-owner",
|
||||
)],
|
||||
vec![sample_codex_key(
|
||||
"key-cli-owner",
|
||||
"provider-cli-owner",
|
||||
"node-cli-owner",
|
||||
)],
|
||||
));
|
||||
let observed_at_unix_secs = current_unix_secs();
|
||||
let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
"development-key",
|
||||
)
|
||||
.with_auth_api_key_reader(auth_repository)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
tunnel_attachment_key("node-cli-owner"),
|
||||
serde_json::to_value(crate::tunnel::TunnelAttachmentRecord {
|
||||
gateway_instance_id: "gateway-b".to_string(),
|
||||
relay_base_url: owner_url.clone(),
|
||||
conn_count: 1,
|
||||
observed_at_unix_secs,
|
||||
})
|
||||
.expect("attachment should serialize"),
|
||||
)]);
|
||||
|
||||
let mut state = AppState::new().expect("gateway state should build");
|
||||
state = state
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
state.scheduler_affinity_cache.insert(
|
||||
"scheduler_affinity:api-key-affinity-cli-1:openai:cli:gpt-5.4".to_string(),
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
provider_id: "provider-cli-owner".to_string(),
|
||||
endpoint_id: "endpoint-cli-owner".to_string(),
|
||||
key_id: "key-cli-owner".to_string(),
|
||||
},
|
||||
Duration::from_secs(300),
|
||||
100,
|
||||
);
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-cli-affinity",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-tunnel-affinity-cli-stream-1")
|
||||
.json(&json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"stream": true
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(GATEWAY_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-b-owner")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("tunnel_affinity_forward")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-b")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("event: response.output_text.delta"));
|
||||
assert!(body.contains("Hello from affinity sync json"));
|
||||
assert!(body.contains("event: response.completed"));
|
||||
|
||||
let owner_request = seen_owner
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("owner request should be captured");
|
||||
assert_eq!(owner_request.path, "/v1/responses");
|
||||
let owner_body: serde_json::Value =
|
||||
serde_json::from_str(&owner_request.body).expect("owner body should parse");
|
||||
assert_eq!(owner_body["model"], "gpt-5.4");
|
||||
assert_eq!(owner_body["stream"], true);
|
||||
assert_eq!(owner_request.trace_id, "trace-tunnel-affinity-cli-stream-1");
|
||||
assert_eq!(owner_request.gateway_marker, "rust-phase3b-affinity");
|
||||
assert_eq!(owner_request.trusted_user_id, "user-affinity-cli-1");
|
||||
assert_eq!(owner_request.trusted_api_key_id, "api-key-affinity-cli-1");
|
||||
assert_eq!(owner_request.trusted_access_allowed, "true");
|
||||
assert_eq!(owner_request.forwarded_by, "gateway-a");
|
||||
assert_eq!(owner_request.owner_instance_id, "gateway-b");
|
||||
|
||||
gateway_handle.abort();
|
||||
owner_handle.abort();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user